Update app.py

app.py

@@ -21,8 +21,9 @@ import logging # For better logging
 # Image handling
 from PIL import Image
 # Pydantic for data validation
-from pydantic import BaseModel, Field, ValidationError, validator
-from typing import List, Optional, Literal
+# Updated imports for Pydantic v2 syntax
+from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator
+from typing import List, Optional, Literal, Dict, Any
 
 # Video and audio processing
 from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
@@ -51,7 +52,7 @@ Generate multiple, branching story timelines from a single theme using AI, compl
 # Text/JSON Model
 TEXT_MODEL_ID = "models/gemini-1.5-flash" # Or "gemini-1.5-pro" for potentially higher quality/cost
 # Audio Model Config
-AUDIO_API_VERSION = 'v1alpha' # Required for audio modality (though endpoint set implicitly now)
+AUDIO_API_VERSION = 'v1alpha' # May not be strictly needed for endpoint if library handles it
 AUDIO_MODEL_ID = f"models/gemini-1.5-flash" # Model used for audio tasks
 AUDIO_SAMPLING_RATE = 24000 # Standard for TTS models like Google's
 # Image Model Config
@@ -67,44 +68,33 @@ TEMP_DIR_BASE = ".chrono_temp" # Base name for temporary directories
 # --- API Key Handling ---
 GOOGLE_API_KEY = None
 try:
-    # Preferred way: Use Streamlit secrets when deployed
     GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
     logger.info("Google API Key loaded from Streamlit secrets.")
 except KeyError:
-    # Fallback: Check environment variable (useful for local development)
     GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
     if GOOGLE_API_KEY:
         logger.info("Google API Key loaded from environment variable.")
     else:
-        # Error if neither is found
         st.error(
-            "🚨 **Google API Key Not Found!**\n"
-            "Please configure your Google API Key:\n"
-            "1.  **Streamlit Cloud/Hugging Face Spaces:** Add it as a Secret named `GOOGLE_API_KEY` in your app's settings.\n"
-            "2.  **Local Development:** Set the `GOOGLE_API_KEY` environment variable or create a `.streamlit/secrets.toml` file.",
+            "🚨 **Google API Key Not Found!** Please configure it via Streamlit secrets or environment variable.",
             icon="🚨"
         )
-        st.stop() # Halt execution
+        st.stop()
 
 # --- Initialize Google Clients ---
-# CORRECTED SECTION: Uses genai.GenerativeModel for both models
 try:
-    # Configure globally
     genai.configure(api_key=GOOGLE_API_KEY)
     logger.info("Configured google-generativeai with API key.")
 
-    # Model/Client Handle for Text/Imagen Generation
     client_standard = genai.GenerativeModel(TEXT_MODEL_ID)
     logger.info(f"Initialized standard GenerativeModel for {TEXT_MODEL_ID}.")
 
-    # Model Handle for Audio Generation
-    # Use the standard GenerativeModel initialization.
-    live_model = genai.GenerativeModel(AUDIO_MODEL_ID) # Use GenerativeModel here
+    live_model = genai.GenerativeModel(AUDIO_MODEL_ID)
     logger.info(f"Initialized GenerativeModel handle for audio ({AUDIO_MODEL_ID}).")
 
 except AttributeError as ae:
      logger.exception("AttributeError during Google AI Client Initialization.")
-     st.error(f"🚨 Failed to initialize Google AI Clients due to an unexpected library structure error: {ae}. Please ensure 'google-generativeai' is up-to-date.", icon="🚨")
+     st.error(f"🚨 Initialization Error: {ae}. Ensure 'google-generativeai' is up-to-date.", icon="🚨")
      st.stop()
 except Exception as e:
     logger.exception("Failed to initialize Google AI Clients.")
@@ -112,18 +102,21 @@ except Exception as e:
     st.stop()
 
 
-# --- Define Pydantic Schemas for Robust Validation ---
+# --- Define Pydantic Schemas (Using V2 Syntax) ---
 class StorySegment(BaseModel):
     scene_id: int = Field(..., ge=0, description="Scene number within the timeline, starting from 0.")
-    # Increased max_length for image_prompt
-    image_prompt: str = Field(..., min_length=10, max_length=250, # <-- Increased from 150
-                              description="Concise visual description for image generation (target 15-35 words). Focus on non-human characters, setting, action, style.")
-    audio_text: str = Field(..., min_length=5, max_length=150, description="Single sentence of narration/dialogue for the scene (max 30 words).")
-    character_description: str = Field(..., max_length=100, description="Brief description of key non-human characters/objects in *this* scene's prompt for consistency.")
-    timeline_visual_modifier: Optional[str] = Field(None, max_length=50, description="Optional subtle visual style hint (e.g., 'slightly darker', 'more vibrant colors').")
-
-    @validator('image_prompt')
-    def image_prompt_no_humans(cls, v):
+    image_prompt: str = Field(..., min_length=10, max_length=250, # Keep increased limit
+                              description="Concise visual description (target 15-35 words). Focus on non-human characters, setting, action, style.")
+    audio_text: str = Field(..., min_length=5, max_length=150, description="Single sentence of narration/dialogue (max 30 words).")
+    # Increased max_length for character_description
+    character_description: str = Field(..., max_length=250, # <-- Increased from 100
+                                      description="Brief description of key non-human characters/objects in *this* scene's prompt (target under 20 words).")
+    timeline_visual_modifier: Optional[str] = Field(None, max_length=50, description="Optional subtle visual style hint.")
+
+    # Pydantic v2 style field validator
+    @field_validator('image_prompt')
+    @classmethod
+    def image_prompt_no_humans(cls, v: str) -> str:
         if any(word in v.lower() for word in ["person", "people", "human", "man", "woman", "boy", "girl", "child"]):
             logger.warning(f"Image prompt '{v[:50]}...' may contain human descriptions. Relying on API-level controls & prompt instructions.")
         return v
@@ -138,15 +131,14 @@ class ChronoWeaveResponse(BaseModel):
     timelines: List[Timeline] = Field(..., min_items=1, description="List of generated timelines.")
     total_scenes_per_timeline: int = Field(..., gt=0, description="The requested number of scenes per timeline.")
 
-    @validator('timelines')
-    def check_timeline_segment_count(cls, timelines, values):
-        # Pydantic v1 style validation. For v2, use model_validator(mode='before') or similar
-        if 'total_scenes_per_timeline' in values:
-            expected_scenes = values['total_scenes_per_timeline']
-            for i, timeline in enumerate(timelines):
-                if len(timeline.segments) != expected_scenes:
-                    raise ValueError(f"Timeline {i} (ID: {timeline.timeline_id}) has {len(timeline.segments)} segments, but expected {expected_scenes}.")
-        return timelines
+    # Pydantic v2 style model validator
+    @model_validator(mode='after') # Use 'after' to access validated fields
+    def check_timeline_segment_count(self) -> 'ChronoWeaveResponse':
+        expected_scenes = self.total_scenes_per_timeline
+        for i, timeline in enumerate(self.timelines):
+            if len(timeline.segments) != expected_scenes:
+                raise ValueError(f"Timeline {i} (ID: {timeline.timeline_id}) has {len(timeline.segments)} segments, but expected {expected_scenes}.")
+        return self
 
 # --- Helper Functions ---
 
@@ -162,7 +154,7 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
         yield wf
     except Exception as e:
         logger.error(f"Error opening/configuring wave file {filename}: {e}")
-        raise # Re-raise the exception
+        raise
     finally:
         if wf:
             try:
@@ -172,34 +164,26 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
 
 
 async def generate_audio_live_async(api_text: str, output_filename: str, voice: Optional[str] = None) -> Optional[str]:
-    """
-    Generates audio using Gemini Live API (async version) via the GenerativeModel.
-    Returns the path to the generated audio file or None on failure.
-    """
+    """Generates audio using Gemini Live API (async version) via the GenerativeModel."""
     collected_audio = bytearray()
-    task_id = os.path.basename(output_filename).split('.')[0] # Extract T#_S# for logging
+    task_id = os.path.basename(output_filename).split('.')[0]
     logger.info(f"🎙️ [{task_id}] Requesting audio for: '{api_text[:60]}...'")
 
     try:
-        # Use the 'live_model' (a GenerativeModel instance) initialized earlier.
         config = {
             "response_modalities": ["AUDIO"],
             "audio_config": {
-                "audio_encoding": "LINEAR16", # Required format for WAV output
+                "audio_encoding": "LINEAR16",
                 "sample_rate_hertz": AUDIO_SAMPLING_RATE,
-                # "voice": voice if voice else "aura-asteria-en" # Optional
             }
         }
-
-        # Prepend directive to discourage conversational filler
         directive_prompt = (
             "Narrate the following sentence directly and engagingly. "
-            "Do not add any introductory or concluding remarks like 'Okay', 'Sure', or 'Here is the narration'. "
+            "Do not add any introductory or concluding remarks. "
             "Speak only the sentence itself:\n\n"
             f'"{api_text}"'
         )
 
-        # Connect and stream using the GenerativeModel instance
         async with live_model.connect(config=config) as session:
             await session.send_request([directive_prompt])
             async for response in session.stream_content():
@@ -215,18 +199,17 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
             st.warning(f"No audio data generated for scene {task_id}.", icon="🔊")
             return None
 
-        # Write the collected audio bytes into a WAV file.
         with wave_file_writer(output_filename, rate=AUDIO_SAMPLING_RATE) as wf:
             wf.writeframes(bytes(collected_audio))
         logger.info(f"   ✅ [{task_id}] Audio saved: {os.path.basename(output_filename)} ({len(collected_audio)} bytes)")
         return output_filename
 
     except genai.types.generation_types.BlockedPromptException as bpe:
-         logger.error(f"   ❌ [{task_id}] Audio generation blocked for prompt '{api_text[:60]}...': {bpe}")
-         st.error(f"Audio generation blocked for scene {task_id} due to safety settings.", icon="🔇")
+         logger.error(f"   ❌ [{task_id}] Audio generation blocked: {bpe}")
+         st.error(f"Audio generation blocked for scene {task_id}.", icon="🔇")
          return None
     except Exception as e:
-        logger.exception(f"   ❌ [{task_id}] Audio generation failed unexpectedly for '{api_text[:60]}...': {e}")
+        logger.exception(f"   ❌ [{task_id}] Audio generation failed unexpectedly: {e}")
         st.error(f"Audio generation failed for scene {task_id}: {e}", icon="🔊")
         return None
 
@@ -237,10 +220,7 @@ def generate_story_sequence_chrono(
     num_timelines: int,
     divergence_prompt: str = ""
 ) -> Optional[ChronoWeaveResponse]:
-    """
-    Generates branching story sequences using Gemini structured output and validates with Pydantic.
-    Returns a validated Pydantic object or None on failure.
-    """
+    """Generates branching story sequences using Gemini structured output and validates with Pydantic."""
     st.info(f"📚 Generating {num_timelines} timeline(s) x {num_scenes} scenes for theme: '{theme}'...")
     logger.info(f"Requesting story structure: Theme='{theme}', Timelines={num_timelines}, Scenes={num_scenes}")
 
@@ -250,7 +230,7 @@ def generate_story_sequence_chrono(
         f"Clearly state the divergence reason for each timeline (except potentially the first)."
     )
 
-    # Updated prompt with stricter image_prompt length guidance
+    # Updated prompt with stricter guidance on description lengths
     prompt = f"""
     Act as an expert narrative designer specializing in short, visual, branching stories for children.
     Create a story based on the core theme: "{theme}".
@@ -262,20 +242,19 @@ def generate_story_sequence_chrono(
     4.  {divergence_instruction}
     5.  Maintain a consistent visual style across all scenes and timelines: **'Simple, friendly kids animation style with bright colors and rounded shapes'**, unless a `timeline_visual_modifier` subtly alters it.
     6.  Each scene's narration (`audio_text`) should be a single, concise sentence (approx. 5-10 seconds spoken length, max 30 words).
-    7.  Image prompts (`image_prompt`) should be descriptive **and concise (target 15-35 words MAXIMUM)**, focusing only on the non-human character(s), setting, action, and essential visual style elements for *this specific scene*. Explicitly mention the main character(s) for consistency. **Do NOT repeat the general 'Simple, friendly kids animation style...' description in every image prompt unless it is essential for a specific visual change**; rely on the overall style instruction and the optional `timeline_visual_modifier`.
-    8.  `character_description` should briefly describe recurring non-human characters mentioned *in the specific scene's image prompt* (name, key visual features). Keep consistent within a timeline.
+    7.  Image prompts (`image_prompt`) should be descriptive **and concise (target 15-35 words MAXIMUM)**, focusing only on the non-human character(s), setting, action, and essential visual style elements for *this specific scene*. Explicitly mention the main character(s) for consistency. **Do NOT repeat the general style description** unless essential for a specific visual change.
+    8.  `character_description` should **very briefly** describe recurring non-human characters mentioned *in the specific scene's image prompt* (name, key visual features). **Keep descriptions extremely concise (e.g., "Nutsy: fluffy squirrel"). Aim for under 20 words total per scene.** Keep consistent within a timeline.
 
     **Output Format:**
     Respond ONLY with a valid JSON object adhering strictly to the provided schema. Do not include any text before or after the JSON object.
 
     **JSON Schema:**
     ```json
-    {json.dumps(ChronoWeaveResponse.schema(), indent=2)}
+    {json.dumps(ChronoWeaveResponse.model_json_schema(), indent=2)}
     ```
-    """ # Using .schema() for Pydantic v1. Use .model_json_schema() for v2.
+    """ # Using .model_json_schema() for Pydantic v2.
 
     try:
-        # Use the standard client (GenerativeModel instance) for text generation
         response = client_standard.generate_content(
             contents=prompt,
             generation_config=genai.types.GenerationConfig(
@@ -284,13 +263,12 @@ def generate_story_sequence_chrono(
             )
         )
 
-        # Attempt to parse the JSON
         try:
             raw_data = json.loads(response.text)
         except json.JSONDecodeError as json_err:
             logger.error(f"Failed to decode JSON response: {json_err}")
             logger.error(f"Problematic Response Text:\n{response.text}")
-            st.error(f"🚨 Failed to parse the story structure from the AI. Error: {json_err}", icon="📄")
+            st.error(f"🚨 Failed to parse story structure from AI: {json_err}", icon="📄")
             st.text_area("Problematic AI Response:", response.text, height=200)
             return None
         except Exception as e:
@@ -298,10 +276,9 @@ def generate_story_sequence_chrono(
              st.error(f"🚨 Error processing AI response: {e}", icon="📄")
              return None
 
-        # Validate the parsed data using Pydantic
         try:
-            # Use parse_obj for Pydantic v1, or model_validate(raw_data) for v2
-            validated_data = ChronoWeaveResponse.parse_obj(raw_data)
+            # Use model_validate for Pydantic v2
+            validated_data = ChronoWeaveResponse.model_validate(raw_data)
             logger.info("✅ Story structure generated and validated successfully!")
             st.success("✅ Story structure generated and validated!")
             return validated_data
@@ -309,33 +286,28 @@ def generate_story_sequence_chrono(
             logger.error(f"JSON structure validation failed: {val_err}")
             logger.error(f"Received Data:\n{json.dumps(raw_data, indent=2)}")
             st.error(f"🚨 The generated story structure is invalid: {val_err}", icon="🧬")
-            st.json(raw_data) # Show the invalid structure that failed validation
+            st.json(raw_data)
             return None
 
     except genai.types.generation_types.BlockedPromptException as bpe:
          logger.error(f"Story generation prompt blocked: {bpe}")
-         st.error("🚨 The story generation prompt was blocked, likely due to safety filters. Try rephrasing the theme.", icon="🚫")
+         st.error("🚨 Story generation prompt blocked (safety filters). Try rephrasing.", icon="🚫")
          return None
     except Exception as e:
         logger.exception("Error during story sequence generation:")
-        st.error(f"🚨 An unexpected error occurred during story generation: {e}", icon="💥")
+        st.error(f"🚨 Unexpected error during story generation: {e}", icon="💥")
         return None
 
 
 def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str = "IMG") -> Optional[Image.Image]:
-    """
-    Generates an image using Imagen via the standard client with specific controls.
-    Returns a PIL Image object or None on failure.
-    """
+    """Generates an image using Imagen via the standard client."""
     logger.info(f"🖼️ [{task_id}] Requesting image for: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
 
-    # Refined prompt - relies on the story generator to provide concise prompts now
-    # Still includes base style and negative constraints as reinforcement
     full_prompt = (
         f"Generate an image in a child-friendly, simple animation style with bright colors and rounded shapes. "
         f"Ensure absolutely NO humans or human-like figures are present. Focus on animals or objects. "
         f"Aspect ratio should be {aspect_ratio}. "
-        f"Scene Description: {prompt}" # Use the potentially shorter prompt from story gen
+        f"Scene Description: {prompt}"
     )
 
     try:
@@ -349,25 +321,24 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
         image_bytes = None
         safety_ratings = []
         block_reason = None
+        finish_reason = None # Store finish reason if available
 
-        # Check response structure - adjust based on actual API behavior
+        # Consolidate response checking
         if hasattr(response, 'candidates') and response.candidates:
             candidate = response.candidates[0]
+            if hasattr(candidate, 'finish_reason'):
+                finish_reason = candidate.finish_reason.name if hasattr(candidate.finish_reason, 'name') else str(candidate.finish_reason)
             if hasattr(candidate, 'content') and candidate.content and hasattr(candidate.content, 'parts') and candidate.content.parts:
                 part = candidate.content.parts[0]
                 if hasattr(part, 'inline_data') and part.inline_data and hasattr(part.inline_data, 'data'):
                     image_bytes = part.inline_data.data
             if hasattr(candidate, 'safety_ratings'):
                 safety_ratings = candidate.safety_ratings
-            # Finish reason might also indicate issues (e.g., SAFETY)
-            # if hasattr(candidate, 'finish_reason') and candidate.finish_reason != 'STOP': ...
 
-        # Check prompt feedback for blocking outside of candidates
         if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
-            if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason != 'BLOCK_REASON_UNSPECIFIED':
-                 block_reason = response.prompt_feedback.block_reason.name # Get the name of the enum
+            if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason.name != 'BLOCK_REASON_UNSPECIFIED':
+                 block_reason = response.prompt_feedback.block_reason.name
             if hasattr(response.prompt_feedback, 'safety_ratings'):
-                 # Combine prompt feedback ratings with candidate ratings if necessary
                  safety_ratings.extend(response.prompt_feedback.safety_ratings)
 
 
@@ -375,39 +346,40 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
             try:
                 image = Image.open(BytesIO(image_bytes))
                 logger.info(f"   ✅ [{task_id}] Image generated successfully.")
-                # Log safety flags if present
                 filtered_ratings = [f"{r.category.name}: {r.probability.name}" for r in safety_ratings if hasattr(r,'probability') and r.probability.name != 'NEGLIGIBLE']
                 if filtered_ratings:
-                    logger.warning(f"   ⚠️ [{task_id}] Image generated but flagged by safety filters: {', '.join(filtered_ratings)}.")
-                    st.warning(f"Image for scene {task_id} flagged by safety filters: {', '.join(filtered_ratings)}", icon="⚠️")
+                    logger.warning(f"   ⚠️ [{task_id}] Image flagged by safety filters: {', '.join(filtered_ratings)}.")
+                    st.warning(f"Image for scene {task_id} flagged: {', '.join(filtered_ratings)}", icon="⚠️")
                 return image
             except Exception as img_err:
                 logger.error(f"   ❌ [{task_id}] Failed to decode generated image data: {img_err}")
                 st.warning(f"Failed to decode image data for scene {task_id}.", icon="🖼️")
                 return None
         else:
-            # If no image bytes, determine why
+            # Enhanced logging for failure reasons
+            fail_reason = "Unknown reason."
             if block_reason:
-                 logger.warning(f"   ⚠️ [{task_id}] Image generation blocked. Reason: {block_reason}. Prompt: '{prompt[:70]}...'")
-                 st.warning(f"Image generation blocked for scene {task_id}. Reason: {block_reason}", icon="🚫")
+                fail_reason = f"Blocked (Reason: {block_reason})."
+            elif finish_reason and finish_reason not in ['STOP', 'FINISH_REASON_UNSPECIFIED']:
+                 fail_reason = f"Finished early (Reason: {finish_reason})."
             else:
-                 # Check for safety flags even if no block reason explicitly given
                  filtered_ratings = [f"{r.category.name}: {r.probability.name}" for r in safety_ratings if hasattr(r,'probability') and r.probability.name != 'NEGLIGIBLE']
                  if filtered_ratings:
-                      logger.warning(f"   ⚠️ [{task_id}] Image generation failed, safety filters triggered: {', '.join(filtered_ratings)}. Prompt: '{prompt[:70]}...'")
-                      st.warning(f"Image generation failed for scene {task_id}, safety filters triggered: {', '.join(filtered_ratings)}", icon="⚠️")
-                 else:
-                      logger.warning(f"   ⚠️ [{task_id}] No image data received, unknown reason. Prompt: '{prompt[:70]}...'")
-                      st.warning(f"No image data received for scene {task_id}, reason unclear.", icon="🖼️")
-                      # logger.debug(f"Full Imagen response object: {response}")
+                     fail_reason = f"Safety filters triggered: {', '.join(filtered_ratings)}."
+
+            logger.warning(f"   ⚠️ [{task_id}] No image data received. Reason: {fail_reason} Prompt: '{prompt[:70]}...'")
+            st.warning(f"No image data received for scene {task_id}. Reason: {fail_reason}", icon="🖼️")
+            # Uncomment to log full response only on unknown failures
+            # if fail_reason == "Unknown reason.":
+            #      logger.debug(f"   ⚠️ [{task_id}] Full API response object: {response}")
             return None
 
     except genai.types.generation_types.BlockedPromptException as bpe:
          logger.error(f"   ❌ [{task_id}] Image generation blocked (exception): {bpe}")
-         st.error(f"Image generation blocked for scene {task_id} due to safety settings.", icon="🚫")
+         st.error(f"Image generation blocked for scene {task_id} (exception).", icon="🚫")
          return None
     except Exception as e:
-        logger.exception(f"   ❌ [{task_id}] Image generation failed unexpectedly for '{prompt[:70]}...': {e}")
+        logger.exception(f"   ❌ [{task_id}] Image generation failed unexpectedly: {e}")
         st.error(f"Image generation failed for scene {task_id}: {e}", icon="🖼️")
         return None
 
@@ -415,27 +387,24 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
 # --- Streamlit UI Elements ---
 st.sidebar.header("⚙️ Configuration")
 
-# API Key Status
 if GOOGLE_API_KEY:
     st.sidebar.success("Google API Key Loaded", icon="✅")
 else:
     st.sidebar.error("Google API Key Missing!", icon="🚨")
 
-# Story Parameters
 theme = st.sidebar.text_input("📖 Story Theme:", "A curious squirrel finds a mysterious, glowing acorn")
-num_scenes = st.sidebar.slider("🎬 Scenes per Timeline:", min_value=2, max_value=7, value=3, help="Number of scenes (image+narration) in each timeline.")
-num_timelines = st.sidebar.slider("🌿 Number of Timelines:", min_value=1, max_value=4, value=2, help="Number of parallel storylines to generate.")
-divergence_prompt = st.sidebar.text_input("↔️ Divergence Hint (Optional):", placeholder="e.g., What if a bird tried to steal it?", help="A suggestion for how the timelines might differ.")
+num_scenes = st.sidebar.slider("🎬 Scenes per Timeline:", min_value=2, max_value=7, value=3)
+num_timelines = st.sidebar.slider("🌿 Number of Timelines:", min_value=1, max_value=4, value=2)
+divergence_prompt = st.sidebar.text_input("↔️ Divergence Hint (Optional):", placeholder="e.g., What if a bird tried to steal it?")
 
-# Generation Settings
 st.sidebar.subheader("🎨 Visual & Audio Settings")
-aspect_ratio = st.sidebar.selectbox("🖼️ Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0, help="Aspect ratio for generated images.")
-audio_voice = None # Placeholder
+aspect_ratio = st.sidebar.selectbox("🖼️ Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0)
+audio_voice = None
 
 generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY), use_container_width=True)
 
 st.sidebar.markdown("---")
-st.sidebar.info("⏳ Generation can take several minutes, especially with more scenes or timelines.", icon="⏳")
+st.sidebar.info("⏳ Generation can take several minutes.", icon="⏳")
 st.sidebar.markdown(f"<small>Models: Text={TEXT_MODEL_ID}, Image={IMAGE_MODEL_ID}, Audio={AUDIO_MODEL_ID}</small>", unsafe_allow_html=True)
 
 
@@ -444,8 +413,7 @@ if generate_button:
     if not theme:
         st.error("Please enter a story theme in the sidebar.", icon="👈")
     else:
-        # Create a unique temporary directory for this run
-        run_id = str(uuid.uuid4()).split('-')[0] # Short unique ID
+        run_id = str(uuid.uuid4()).split('-')[0]
         temp_dir = os.path.join(TEMP_DIR_BASE, f"run_{run_id}")
         try:
             os.makedirs(temp_dir, exist_ok=True)
@@ -454,8 +422,8 @@ if generate_button:
             st.error(f"🚨 Failed to create temporary directory {temp_dir}: {e}", icon="📂")
             st.stop()
 
-        final_video_paths = {} # Stores {timeline_id: video_path}
-        generation_errors = {} # Stores {timeline_id: [error_messages]}
+        final_video_paths = {}
+        generation_errors = {}
 
         # --- 1. Generate Narrative Structure ---
         chrono_response: Optional[ChronoWeaveResponse] = None
@@ -463,15 +431,11 @@ if generate_button:
             chrono_response = generate_story_sequence_chrono(theme, num_scenes, num_timelines, divergence_prompt)
 
         if chrono_response:
-            # Structure generated and validated successfully by the function
-            # st.success(...) is now inside generate_story_sequence_chrono on success
-
             # --- 2. Process Each Timeline ---
             overall_start_time = time.time()
-            all_timelines_successful = True # Assume success initially
+            all_timelines_successful = True
 
             with st.status("Generating assets and composing videos...", expanded=True) as status:
-
                 for timeline_index, timeline in enumerate(chrono_response.timelines):
                     timeline_id = timeline.timeline_id
                     divergence = timeline.divergence_reason
@@ -494,18 +458,16 @@ if generate_button:
                         status.update(label=status_message)
                         st.markdown(f"--- **Scene {scene_id + 1} ({task_id})** ---")
                         logger.info(status_message)
-
                         scene_has_error = False
 
-                        st.write(f"   *Image Prompt:* {segment.image_prompt}" + (f" *(Modifier: {segment.timeline_visual_modifier})*" if segment.timeline_visual_modifier else ""))
+                        st.write(f"   *Image Prompt:* {segment.image_prompt}" + (f" *(Mod: {segment.timeline_visual_modifier})*" if segment.timeline_visual_modifier else ""))
                         st.write(f"   *Audio Text:* {segment.audio_text}")
 
                         # --- 2a. Image Generation ---
                         generated_image: Optional[Image.Image] = None
                         with st.spinner(f"[{task_id}] Generating image... 🎨"):
-                            # Combine prompt using the (hopefully shorter) prompt from the structure
-                            combined_prompt = segment.image_prompt # Use directly
-                            if segment.character_description: # Add character desc if present
+                            combined_prompt = segment.image_prompt
+                            if segment.character_description:
                                 combined_prompt += f" Featuring: {segment.character_description}"
                             if segment.timeline_visual_modifier:
                                 combined_prompt += f" Style hint: {segment.timeline_visual_modifier}."
@@ -518,12 +480,12 @@ if generate_button:
                                 temp_image_files[scene_id] = image_path
                                 st.image(generated_image, width=180, caption=f"Scene {scene_id+1} Image")
                             except Exception as e:
-                                logger.error(f"   ❌ [{task_id}] Failed to save image {image_path}: {e}")
+                                logger.error(f"   ❌ [{task_id}] Failed to save image: {e}")
                                 st.error(f"Failed to save image for scene {task_id}.", icon="💾")
                                 scene_has_error = True
                                 generation_errors[timeline_id].append(f"Scene {scene_id+1}: Image save failed.")
                         else:
-                            st.warning(f"Image generation failed for scene {task_id}. Skipping scene.", icon="🖼️")
+                            # Error/warning already logged by generate_image_imagen
                             scene_has_error = True
                             generation_errors[timeline_id].append(f"Scene {scene_id+1}: Image generation failed.")
                             continue
@@ -538,13 +500,13 @@ if generate_button:
                                         generate_audio_live_async(segment.audio_text, audio_path_temp, audio_voice)
                                     )
                                 except RuntimeError as e:
-                                    logger.error(f"   ❌ [{task_id}] Asyncio runtime error during audio gen: {e}")
-                                    st.error(f"Asyncio error during audio generation for {task_id}: {e}", icon="⚡")
+                                    logger.error(f"   ❌ [{task_id}] Asyncio runtime error: {e}")
+                                    st.error(f"Asyncio error during audio gen for {task_id}: {e}", icon="⚡")
                                     scene_has_error = True
                                     generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio async error.")
                                 except Exception as e:
-                                    logger.exception(f"   ❌ [{task_id}] Unexpected error during audio generation call for {task_id}: {e}")
-                                    st.error(f"Unexpected error in audio generation for {task_id}: {e}", icon="💥")
+                                    logger.exception(f"   ❌ [{task_id}] Unexpected audio error: {e}")
+                                    st.error(f"Unexpected audio error for {task_id}: {e}", icon="💥")
                                     scene_has_error = True
                                     generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation error.")
 
@@ -556,16 +518,12 @@ if generate_button:
                                 except Exception as e:
                                      logger.warning(f"   ⚠️ [{task_id}] Could not display audio preview: {e}")
                             else:
-                                st.warning(f"Audio generation failed for {task_id}. Skipping video clip.", icon="🔊")
+                                # Error/warning already logged by generate_audio_live_async
                                 scene_has_error = True
                                 generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation failed.")
                                 if scene_id in temp_image_files and os.path.exists(temp_image_files[scene_id]):
-                                    try:
-                                        os.remove(temp_image_files[scene_id])
-                                        logger.info(f"   🗑️ [{task_id}] Removed image file due to audio failure.")
-                                        del temp_image_files[scene_id]
-                                    except OSError as e:
-                                         logger.warning(f"   ⚠️ [{task_id}] Could not remove image file {temp_image_files[scene_id]} after audio failure: {e}")
+                                    try: os.remove(temp_image_files[scene_id]); logger.info(f"   🗑️ [{task_id}] Removed image file due to audio failure."); del temp_image_files[scene_id]
+                                    except OSError as e: logger.warning(f"   ⚠️ [{task_id}] Failed to remove image after audio failure: {e}")
                                 continue
 
                         # --- 2c. Create Video Clip ---
@@ -573,17 +531,14 @@ if generate_button:
                             st.write(f"   🎬 Creating video clip for Scene {scene_id+1}...")
                             img_path = temp_image_files[scene_id]
                             aud_path = temp_audio_files[scene_id]
-                            audio_clip_instance = None
-                            image_clip_instance = None
-                            composite_clip = None
+                            audio_clip_instance, image_clip_instance, composite_clip = None, None, None
                             try:
-                                if not os.path.exists(img_path): raise FileNotFoundError(f"Image file not found: {img_path}")
-                                if not os.path.exists(aud_path): raise FileNotFoundError(f"Audio file not found: {aud_path}")
+                                if not os.path.exists(img_path): raise FileNotFoundError(f"Image file missing: {img_path}")
+                                if not os.path.exists(aud_path): raise FileNotFoundError(f"Audio file missing: {aud_path}")
 
                                 audio_clip_instance = AudioFileClip(aud_path)
                                 np_image = np.array(Image.open(img_path))
                                 image_clip_instance = ImageClip(np_image).set_duration(audio_clip_instance.duration)
-
                                 composite_clip = image_clip_instance.set_audio(audio_clip_instance)
                                 video_clips.append(composite_clip)
                                 logger.info(f"      ✅ [{task_id}] Video clip created (Duration: {audio_clip_instance.duration:.2f}s).")
@@ -591,16 +546,16 @@ if generate_button:
                                 scene_success_count += 1
 
                             except Exception as e:
-                                logger.exception(f"      ❌ [{task_id}] Failed to create video clip for scene {scene_id+1}: {e}")
+                                logger.exception(f"      ❌ [{task_id}] Failed to create video clip: {e}")
                                 st.error(f"Failed to create video clip for {task_id}: {e}", icon="🎬")
                                 scene_has_error = True
                                 generation_errors[timeline_id].append(f"Scene {scene_id+1}: Video clip creation failed.")
                                 if audio_clip_instance: audio_clip_instance.close()
                                 if image_clip_instance: image_clip_instance.close()
-                                if os.path.exists(img_path): os.remove(img_path)
-                                if os.path.exists(aud_path): os.remove(aud_path)
-
-                        # --- End of Scene Loop ---
+                                try:
+                                     if os.path.exists(img_path): os.remove(img_path)
+                                     if os.path.exists(aud_path): os.remove(aud_path)
+                                except OSError as e_rem: logger.warning(f"   ⚠️ [{task_id}] Failed to remove files after clip error: {e_rem}")
 
                     # --- 2d. Assemble Timeline Video ---
                     timeline_duration = time.time() - timeline_start_time
@@ -612,46 +567,36 @@ if generate_button:
                         final_timeline_video = None
                         try:
                             final_timeline_video = concatenate_videoclips(video_clips, method="compose")
-                            final_timeline_video.write_videofile(
-                                output_filename,
-                                fps=VIDEO_FPS,
-                                codec=VIDEO_CODEC,
-                                audio_codec=AUDIO_CODEC,
-                                logger=None
-                            )
+                            final_timeline_video.write_videofile(output_filename, fps=VIDEO_FPS, codec=VIDEO_CODEC, audio_codec=AUDIO_CODEC, logger=None)
                             final_video_paths[timeline_id] = output_filename
                             logger.info(f"   ✅ [{timeline_label}] Final video saved: {os.path.basename(output_filename)}")
                             st.success(f"✅ Video for {timeline_label} completed in {timeline_duration:.2f}s.")
-
                         except Exception as e:
                             logger.exception(f"   ❌ [{timeline_label}] Failed to write final video: {e}")
                             st.error(f"Failed to assemble video for {timeline_label}: {e}", icon="📼")
                             all_timelines_successful = False
                             generation_errors[timeline_id].append(f"Timeline {timeline_id}: Final video assembly failed.")
                         finally:
-                            logger.debug(f"[{timeline_label}] Closing {len(video_clips)} source clips...")
+                            logger.debug(f"[{timeline_label}] Closing {len(video_clips)} source clips and final video...")
                             for i, clip in enumerate(video_clips):
                                 try:
                                     if clip:
                                       if clip.audio: clip.audio.close()
                                       clip.close()
-                                except Exception as e_close:
-                                    logger.warning(f"   ⚠️ [{timeline_label}] Error closing source clip {i}: {e_close}")
+                                except Exception as e_close: logger.warning(f"   ⚠️ [{timeline_label}] Error closing source clip {i}: {e_close}")
                             if final_timeline_video:
                                 try:
                                      if final_timeline_video.audio: final_timeline_video.audio.close()
                                      final_timeline_video.close()
-                                     logger.debug(f"[{timeline_label}] Closed final video object.")
-                                except Exception as e_close_final:
-                                     logger.warning(f"   ⚠️ [{timeline_label}] Error closing final video object: {e_close_final}")
+                                except Exception as e_close_final: logger.warning(f"   ⚠️ [{timeline_label}] Error closing final video object: {e_close_final}")
 
                     elif not video_clips:
-                        logger.warning(f"[{timeline_label}] No video clips successfully generated. Skipping final assembly.")
-                        st.warning(f"No scenes were successfully processed for {timeline_label}. Video cannot be created.", icon="🚫")
+                        logger.warning(f"[{timeline_label}] No video clips generated. Skipping final assembly.")
+                        st.warning(f"No scenes processed for {timeline_label}. Video cannot be created.", icon="🚫")
                         all_timelines_successful = False
                     else: # Some scenes failed
                         error_count = len(segments) - scene_success_count
-                        logger.warning(f"[{timeline_label}] Encountered errors in {error_count} scene(s). Skipping final video assembly.")
+                        logger.warning(f"[{timeline_label}] Errors in {error_count} scene(s). Skipping final video assembly.")
                         st.warning(f"{timeline_label} had errors in {error_count} scene(s). Final video not assembled.", icon="⚠️")
                         all_timelines_successful = False
 
@@ -659,20 +604,16 @@ if generate_button:
                          logger.error(f"Summary of errors in {timeline_label}: {generation_errors[timeline_id]}")
 
                 # --- End of Timelines Loop ---
-
                 overall_duration = time.time() - overall_start_time
                 if all_timelines_successful and final_video_paths:
                     status_msg = f"ChronoWeave Generation Complete! ({len(final_video_paths)} videos in {overall_duration:.2f}s)"
-                    status.update(label=status_msg, state="complete", expanded=False)
-                    logger.info(status_msg)
+                    status.update(label=status_msg, state="complete", expanded=False); logger.info(status_msg)
                 elif final_video_paths:
-                    status_msg = f"ChronoWeave Partially Complete ({len(final_video_paths)} videos, some errors occurred). Total time: {overall_duration:.2f}s"
-                    status.update(label=status_msg, state="warning", expanded=True)
-                    logger.warning(status_msg)
+                    status_msg = f"ChronoWeave Partially Complete ({len(final_video_paths)} videos, errors occurred). Time: {overall_duration:.2f}s"
+                    status.update(label=status_msg, state="warning", expanded=True); logger.warning(status_msg)
                 else:
-                    status_msg = f"ChronoWeave Generation Failed. No videos produced. Total time: {overall_duration:.2f}s"
-                    status.update(label=status_msg, state="error", expanded=True)
-                    logger.error(status_msg)
+                    status_msg = f"ChronoWeave Generation Failed. No videos produced. Time: {overall_duration:.2f}s"
+                    status.update(label=status_msg, state="error", expanded=True); logger.error(status_msg)
 
             # --- 3. Display Results ---
             st.header("🎬 Generated Timelines")
@@ -691,30 +632,21 @@ if generate_button:
                         st.subheader(f"Timeline {timeline_id}")
                         st.caption(f"Divergence: {reason}")
                         try:
-                            with open(video_path, 'rb') as video_file:
-                                video_bytes = video_file.read()
+                            with open(video_path, 'rb') as video_file: video_bytes = video_file.read()
                             st.video(video_bytes)
                             logger.info(f"Displaying video for Timeline {timeline_id}")
-                            st.download_button(
-                                label=f"Download T{timeline_id} Video",
-                                data=video_bytes,
-                                file_name=f"chronoweave_timeline_{timeline_id}.mp4",
-                                mime="video/mp4",
-                                key=f"download_btn_{timeline_id}"
-                            )
+                            st.download_button(label=f"Download T{timeline_id} Video", data=video_bytes, file_name=f"chronoweave_timeline_{timeline_id}.mp4", mime="video/mp4", key=f"dl_{timeline_id}")
                             if generation_errors.get(timeline_id):
-                                with st.expander(f"⚠️ View {len(generation_errors[timeline_id])} Generation Issues"):
-                                     for error_msg in generation_errors[timeline_id]:
-                                         st.warning(f"- {error_msg}")
-
+                                with st.expander(f"⚠️ View {len(generation_errors[timeline_id])} Issues"):
+                                     for error_msg in generation_errors[timeline_id]: st.warning(f"- {error_msg}")
                         except FileNotFoundError:
-                             logger.error(f"Could not find video file for display: {video_path}")
-                             st.error(f"Error: Video file not found for Timeline {timeline_id}.", icon="🚨")
+                             logger.error(f"Video file not found for display: {video_path}")
+                             st.error(f"Error: Video file missing for T{timeline_id}.", icon="🚨")
                         except Exception as e:
                              logger.exception(f"Could not display video {video_path}: {e}")
-                             st.error(f"Error displaying video for Timeline {timeline_id}: {e}", icon="🚨")
+                             st.error(f"Error displaying video T{timeline_id}: {e}", icon="🚨")
             else:
-                st.warning("No final videos were successfully generated in this run.")
+                st.warning("No final videos were successfully generated.")
                 all_errors = [msg for err_list in generation_errors.values() for msg in err_list]
                 if all_errors:
                     st.subheader("Summary of Generation Issues")
@@ -722,8 +654,7 @@ if generate_button:
                         for tid, errors in generation_errors.items():
                             if errors:
                                 st.error(f"Timeline {tid}:")
-                                for msg in errors:
-                                    st.error(f"  - {msg}")
+                                for msg in errors: st.error(f"  - {msg}")
 
             # --- 4. Cleanup ---
             st.info(f"Attempting to clean up temporary directory: {temp_dir}")
@@ -736,11 +667,10 @@ if generate_button:
                 st.warning(f"Could not automatically remove temporary files: {temp_dir}. Please remove it manually if needed.", icon="⚠️")
 
         elif not chrono_response:
-            # Error message already shown by generate_story_sequence_chrono or validation
             logger.error("Story generation or validation failed, cannot proceed.")
         else:
-            st.error("An unexpected issue occurred after story generation. Cannot proceed.", icon="🛑")
-            logger.error("Chrono_response existed but was falsy in the main logic block.")
+            st.error("An unexpected issue occurred after story generation.", icon="🛑")
+            logger.error("Chrono_response existed but was falsy in main logic.")
 
 else:
     st.info("Configure settings in the sidebar and click '✨ Generate ChronoWeave ✨' to start.")