Update app.py

app.py

@@ -31,13 +31,14 @@ from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
 # Type hints
 import typing_extensions as typing
 
-# Async support for Streamlit/Google API
+# Async support
 import nest_asyncio
 nest_asyncio.apply()
 
 # Import Vertex AI SDK for image generation (Preview API)
 import vertexai
 from vertexai.preview.vision_models import ImageGenerationModel
+from google.oauth2 import service_account
 
 # --- Logging Setup ---
 logging.basicConfig(
@@ -58,7 +59,7 @@ Generate multiple, branching story timelines from a single theme using AI, compl
 TEXT_MODEL_ID = "models/gemini-1.5-flash"
 AUDIO_MODEL_ID = "models/gemini-1.5-flash"
 AUDIO_SAMPLING_RATE = 24000
-# Pretrained Imagen model used with Vertex AI preview
+# Pretrained Imagen model identifier for Vertex AI preview
 IMAGE_MODEL_ID = "imagen-3.0-generate-002"
 DEFAULT_ASPECT_RATIO = "1:1"
 VIDEO_FPS = 24
@@ -79,15 +80,24 @@ except KeyError:
         st.error("🚨 **Google API Key Not Found!** Please configure it.", icon="🚨")
         st.stop()
 
-# Vertex AI configuration: Get project and location
+# Vertex AI configuration: load PROJECT_ID and LOCATION from secrets or environment.
 PROJECT_ID = st.secrets.get("PROJECT_ID") or os.environ.get("PROJECT_ID")
 LOCATION = st.secrets.get("LOCATION") or os.environ.get("LOCATION", "us-central1")
 if not PROJECT_ID:
     st.error("🚨 **PROJECT_ID not set!** Please add PROJECT_ID to your secrets.", icon="🚨")
     st.stop()
 
-# Initialize Vertex AI
-vertexai.init(project=PROJECT_ID, location=LOCATION)
+# Load service account JSON from environment (secret name: SERVICE_ACCOUNT_JSON)
+try:
+    service_account_info = json.loads(os.environ["SERVICE_ACCOUNT_JSON"])
+    credentials = service_account.Credentials.from_service_account_info(service_account_info)
+    logger.info("Service account credentials loaded successfully.")
+except Exception as e:
+    st.error(f"🚨 Failed to load service account JSON: {e}", icon="🚨")
+    st.stop()
+
+# Initialize Vertex AI with credentials
+vertexai.init(project=PROJECT_ID, location=LOCATION, credentials=credentials)
 
 # --- Initialize Google Clients for text/audio ---
 try:
@@ -106,7 +116,7 @@ except Exception as e:
     st.error(f"🚨 Failed Init: {e}", icon="🚨")
     st.stop()
 
-# --- Define Pydantic Schemas (Using V2 Syntax) ---
+# --- Define Pydantic Schemas ---
 class StorySegment(BaseModel):
     scene_id: int = Field(..., ge=0)
     image_prompt: str = Field(..., min_length=10, max_length=250)
@@ -140,7 +150,6 @@ class ChronoWeaveResponse(BaseModel):
         return self
 
 # --- Helper Functions ---
-
 @contextlib.contextmanager
 def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLING_RATE, sample_width: int = 2):
     """Context manager to safely write WAV files."""
@@ -265,12 +274,11 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
     """
     Generates an image using Vertex AI's Imagen model via the Vertex AI preview API.
     
-    This function loads the pretrained model "imagen-3.0-generate-002" and attempts to generate an image.
+    This function loads the pretrained Imagen model "imagen-3.0-generate-002" and attempts to generate an image.
     If authentication fails, it provides guidance on how to resolve the issue.
     """
     logger.info(f"🖼️ [{task_id}] Requesting image: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
     try:
-        # Load the pretrained Imagen model from Vertex AI.
         generation_model = ImageGenerationModel.from_pretrained(IMAGE_MODEL_ID)
         images = generation_model.generate_images(
             prompt=prompt,
@@ -293,8 +301,7 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
                 "If running in Colab, try:\n"
                 "    from google.colab import auth\n"
                 "    auth.authenticate_user()\n"
-                "If using a service account or other environment, please refer to "
-                "https://cloud.google.com/docs/authentication for guidance."
+                "If using a service account or other environment, please refer to https://cloud.google.com/docs/authentication for guidance."
             )
         else:
             error_msg = f"Image generation for {task_id} failed: {e}"
@@ -560,7 +567,7 @@ if generate_button:
                 logger.info(f"✅ Temp dir removed: {temp_dir}")
                 st.success("✅ Temp files cleaned.")
             except Exception as e:
-                logger.error(f"⚠️ Failed remove temp dir {temp_dir}: {e}")
+                logger.error(f"⚠️ Failed to remove temp dir {temp_dir}: {e}")
                 st.warning(f"Could not remove temp files: {temp_dir}.", icon="⚠️")
 
         elif not chrono_response: