Update app.py

app.py

@@ -16,7 +16,7 @@ import contextlib
 import asyncio
 import uuid  # For unique identifiers
 import shutil  # For directory operations
-import logging  # For improved logging
+import logging  # For logging
 
 # Image handling
 from PIL import Image
@@ -27,7 +27,6 @@ from typing import List, Optional, Dict, Any
 
 # Video and audio processing
 from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
-# from moviepy.config import change_settings  # Uncomment if you need to change settings
 
 # Type hints
 import typing_extensions as typing
@@ -36,8 +35,9 @@ import typing_extensions as typing
 import nest_asyncio
 nest_asyncio.apply()
 
-# Import Vertex AI SDK
-from google.cloud import aiplatform
+# Import Vertex AI SDK for image generation (Preview API)
+import vertexai
+from vertexai.preview.vision_models import ImageGenerationModel
 
 # --- Logging Setup ---
 logging.basicConfig(
@@ -58,14 +58,15 @@ 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
-IMAGE_MODEL_ID = "imagen-3"  # Now used with Vertex AI
+# IMAGE_MODEL_ID is now used with the preview API
+IMAGE_MODEL_ID = "imagen-3.0-generate-002"
 DEFAULT_ASPECT_RATIO = "1:1"
 VIDEO_FPS = 24
 VIDEO_CODEC = "libx264"
 AUDIO_CODEC = "aac"
 TEMP_DIR_BASE = ".chrono_temp"
 
-# --- API Key and Vertex AI Config Handling ---
+# --- Secrets and Environment Variables ---
 GOOGLE_API_KEY = None
 try:
     GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
@@ -78,17 +79,16 @@ except KeyError:
         st.error("🚨 **Google API Key Not Found!** Please configure it.", icon="🚨")
         st.stop()
 
-# --- Vertex AI Configuration ---
-# Set up environment variables for Vertex AI; ensure these are in your Streamlit secrets or environment.
+# For Vertex AI, we also need PROJECT_ID and LOCATION.
 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")
-IMAGE_ENDPOINT_ID = st.secrets.get("IMAGE_ENDPOINT_ID") or os.environ.get("IMAGE_ENDPOINT_ID")
-
-if not PROJECT_ID or not IMAGE_ENDPOINT_ID:
-    st.error("🚨 **Vertex AI is not configured properly!** "
-             "Please set PROJECT_ID and IMAGE_ENDPOINT_ID in your secrets.", icon="🚨")
+if not PROJECT_ID:
+    st.error("🚨 **PROJECT_ID not set!** Please add PROJECT_ID to your secrets.", icon="🚨")
     st.stop()
 
+# Initialize Vertex AI (used for image generation)
+vertexai.init(project=PROJECT_ID, location=LOCATION)
+
 # --- Initialize Google Clients for text/audio ---
 try:
     genai.configure(api_key=GOOGLE_API_KEY)
@@ -117,7 +117,7 @@ class StorySegment(BaseModel):
     @field_validator('image_prompt')
     @classmethod
     def image_prompt_no_humans(cls, v: str) -> str:
-        if any(w in v.lower() for w in ["person", "people", "human", "man", "woman", "boy", "girl", "child"]):
+        if any(word in v.lower() for word in ["person", "people", "human", "man", "woman", "boy", "girl", "child"]):
             logger.warning(f"Prompt '{v[:50]}...' may contain humans.")
         return v
 
@@ -263,36 +263,31 @@ JSON Schema: ```json
 
 def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str = "IMG") -> Optional[Image.Image]:
     """
-    Generates an image using Vertex AI's Imagen model via the Google Cloud AI Platform SDK.
+    Generates an image using Vertex AI's Imagen model via the Vertex AI preview API.
     
-    Ensure that the following environment variables or Streamlit secrets are set:
-      - PROJECT_ID: Your Google Cloud project ID.
-      - LOCATION: The Vertex AI region (e.g., "us-central1").
-      - IMAGE_ENDPOINT_ID: The resource ID of your deployed Imagen endpoint.
+    It calls the ImageGenerationModel from vertexai.preview.vision_models with the pretrained model "imagen-3.0-generate-002" and returns a PIL Image.
     """
     logger.info(f"🖼️ [{task_id}] Requesting image: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
     try:
-        # Initialize Vertex AI with your project and location.
-        aiplatform.init(project=PROJECT_ID, location=LOCATION)
-        # Retrieve your deployed endpoint.
-        endpoint = aiplatform.Endpoint(IMAGE_ENDPOINT_ID)
-        # Create a prediction instance. (The instance structure depends on your model.)
-        instance = {"prompt": prompt, "aspect_ratio": aspect_ratio}
-        prediction_response = endpoint.predict(instances=[instance])
-        # Assume the prediction returns a base64-encoded image string under the key "image".
-        import base64
-        image_base64 = prediction_response.predictions[0].get("image")
-        if not image_base64:
-            logger.error(f"❌ [{task_id}] No image returned in prediction.")
-            st.error(f"Image prediction failed for {task_id}: No image returned.", icon="🖼️")
-            return None
-        image_data = base64.b64decode(image_base64)
-        image = Image.open(BytesIO(image_data))
+        # Load the pretrained Imagen model
+        generation_model = ImageGenerationModel.from_pretrained(IMAGE_MODEL_ID)
+        # Generate the image (here we generate one image)
+        images = generation_model.generate_images(
+            prompt=prompt,
+            number_of_images=1,
+            aspect_ratio=aspect_ratio,
+            negative_prompt="",
+            person_generation="",
+            safety_filter_level="",
+            add_watermark=True,
+        )
+        # Return the generated PIL image (using the internal _pil_image attribute)
+        image = images[0]._pil_image
         logger.info(f"✅ [{task_id}] Image generated successfully.")
         return image
     except Exception as e:
         logger.exception(f"❌ [{task_id}] Image generation failed: {e}")
-        st.error(f"Image generation failed for {task_id}: {e}", icon="🖼️")
+        st.error(f"Image generation for {task_id} failed: {e}", icon="🖼️")
         return None
 
 # --- Streamlit UI Elements ---
@@ -325,7 +320,7 @@ if generate_button:
             os.makedirs(temp_dir, exist_ok=True)
             logger.info(f"Created temp dir: {temp_dir}")
         except OSError as e:
-            st.error(f"🚨 Failed create temp dir {temp_dir}: {e}", icon="📂")
+            st.error(f"🚨 Failed to create temp dir {temp_dir}: {e}", icon="📂")
             st.stop()
         final_video_paths, generation_errors = {}, {}