Update app.py

app.py

@@ -14,14 +14,19 @@ import time
 import wave
 import contextlib
 import asyncio
-import uuid # For unique filenames
-import shutil # For cleaning up temp dirs
+import uuid # For unique identifiers
+import shutil # For directory operations
+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
 
 # Video and audio processing
-from moviepy.editor import ImageClip, AudioFileClip, CompositeVideoClip, concatenate_videoclips
+from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
+from moviepy.config import change_settings # Potential for setting imagemagick path if needed
 
 # Type hints
 import typing_extensions as typing
@@ -30,464 +35,731 @@ import typing_extensions as typing
 import nest_asyncio
 nest_asyncio.apply() # Apply patch for asyncio in environments like Streamlit/Jupyter
 
+# --- Logging Setup ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 # --- Configuration ---
-st.set_page_config(page_title="ChronoWeave", layout="wide")
-st.title("🌀 ChronoWeave: Branching Narrative Generator")
+st.set_page_config(page_title="ChronoWeave", layout="wide", initial_sidebar_state="expanded")
+st.title("🌀 ChronoWeave: Advanced Branching Narrative Generator")
 st.markdown("""
-Generate multiple, branching story timelines from a single theme using AI.
-Based on the work of Yousif Ahmed. Copyright 2025 Google LLC.
+Generate multiple, branching story timelines from a single theme using AI, complete with images and narration.
+*Based on the work of Yousif Ahmed. Copyright 2025 Google LLC.*
 """)
 
 # --- Constants ---
-MODEL = "models/gemini-1.5-flash" # Or other suitable text model supporting JSON
-# Using v1alpha for the Live API for audio output.
-AUDIO_MODEL_VERSION = 'v1alpha' # Must be alpha for audio modality
-IMAGE_MODEL_ID = "imagen-3" # Or your preferred Imagen model "imagen-3.0-generate-002"
+# 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
+AUDIO_MODEL_ID = f"models/gemini-1.5-flash" # Model used via the v1alpha endpoint
+AUDIO_SAMPLING_RATE = 24000 # Standard for TTS models like Google's
+# Image Model Config
+IMAGE_MODEL_ID = "imagen-3" # Or specific version like "imagen-3.0-generate-002"
+DEFAULT_ASPECT_RATIO = "1:1"
+# Video Config
+VIDEO_FPS = 24
+VIDEO_CODEC = "libx264" # Widely compatible H.264
+AUDIO_CODEC = "aac" # Common audio codec for MP4
+# File Management
+TEMP_DIR_BASE = ".chrono_temp" # Base name for temporary directories
 
 # --- API Key Handling ---
+# This section correctly handles the missing secret. The error in the traceback means
+# the secret wasn't set in the *deployment environment*.
+GOOGLE_API_KEY = None
 try:
-    # Preferred way to handle secrets in Streamlit sharing/HF Spaces
+    # Preferred way: Use Streamlit secrets when deployed
     GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
-    os.environ['GOOGLE_API_KEY'] = GOOGLE_API_KEY
+    logger.info("Google API Key loaded from Streamlit secrets.")
 except KeyError:
-    st.error("🚨 Google API Key not found! Please add it as a Secret named 'GOOGLE_API_KEY' in your Hugging Face Space settings.", icon="🚨")
-    st.stop() # Halt execution if no key
+    # 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.",
+            icon="🚨"
+        )
+        st.stop() # Halt execution
 
-# --- Initialize Google Client ---
+# --- Initialize Google Clients ---
 try:
-    # Initialize the client with the API key
     genai.configure(api_key=GOOGLE_API_KEY)
 
-    # Create separate clients or configure one for different API versions if needed
-    # Client for Text/Imagen (standard API)
-    client_standard = genai.GenerativeModel(MODEL)
-    # Client for Live Audio (v1alpha) - requires different client init
+    # Client for Text/Imagen Generation (using standard API endpoint)
+    client_standard = genai.GenerativeModel(TEXT_MODEL_ID)
+    logger.info(f"Initialized standard GenerativeModel for {TEXT_MODEL_ID}.")
+
+    # Client specifically for the Live API (Audio) using the v1alpha endpoint
+    # Note: Ensure the 'audiomodality' package or relevant parts of google-cloud-aiplatform are available if needed,
+    # depending on the exact library version and API stability.
+    # As of late 2023/early 2024, using a separate client instance targeting the specific endpoint is often necessary.
     client_live = genai.Client(
-         client_options={'api_endpoint': f'{AUDIO_MODEL_VERSION}.generativelanguage.googleapis.com'}
+         client_options={'api_endpoint': f'{AUDIO_API_VERSION}.generativelanguage.googleapis.com'}
     )
-    # Note: As of recent updates, genai.configure might handle this better,
-    # but separating clients or explicitly setting endpoints can be more robust.
-    # Adjust based on the library version and observed behavior.
-
+    live_model = client_live.get_model(AUDIO_MODEL_ID) # Get the model handle via the live client
+    logger.info(f"Initialized live client for audio generation ({AUDIO_MODEL_ID} via {AUDIO_API_VERSION}).")
 
 except Exception as e:
-    st.error(f"🚨 Failed to initialize Google AI Client: {e}", icon="🚨")
+    logger.exception("Failed to initialize Google AI Clients.")
+    st.error(f"🚨 Failed to initialize Google AI Clients: {e}", icon="🚨")
     st.stop()
 
 
-# --- Define Structured Output Schemas ---
-class StorySegment(typing.TypedDict):
-    scene_id: int
-    image_prompt: str
-    audio_text: str
-    character_description: str
-    timeline_visual_modifier: typing.Optional[str]
-
-class Timeline(typing.TypedDict):
-    timeline_id: int
-    divergence_reason: str
-    segments: list[StorySegment]
-
-class ChronoWeaveResponse(typing.TypedDict):
-    core_theme: str
-    timelines: list[Timeline]
-    total_scenes_per_timeline: int
+# --- Define Pydantic Schemas for Robust Validation ---
+class StorySegment(BaseModel):
+    scene_id: int = Field(..., ge=0, description="Scene number within the timeline, starting from 0.")
+    image_prompt: str = Field(..., min_length=10, max_length=150, description="Concise visual description for image generation (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):
+        if any(word in v.lower() for word in ["person", "people", "human", "man", "woman", "boy", "girl", "child"]):
+            # Instead of raising error, we'll try to guide the LLM better in the main prompt
+            # raise ValueError("Image prompt must not contain descriptions of humans.")
+            logger.warning(f"Image prompt '{v[:50]}...' may contain human descriptions. Relying on API-level controls.")
+        return v
+
+class Timeline(BaseModel):
+    timeline_id: int = Field(..., ge=0, description="Unique identifier for this timeline.")
+    divergence_reason: str = Field(..., min_length=5, description="Clear reason why this timeline branched off.")
+    segments: List[StorySegment] = Field(..., min_items=1, description="List of scenes composing this timeline.")
+
+class ChronoWeaveResponse(BaseModel):
+    core_theme: str = Field(..., min_length=5, description="The central theme provided by the user.")
+    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):
+        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
 
 # --- Helper Functions ---
 
 @contextlib.contextmanager
-def wave_file(filename, channels=1, rate=24000, sample_width=2):
-    """Context manager to write WAV files."""
-    with wave.open(filename, "wb") as wf:
+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."""
+    wf = None
+    try:
+        wf = wave.open(filename, "wb")
         wf.setnchannels(channels)
-        wf.setsampwidth(sample_width)
+        wf.setsampwidth(sample_width) # 2 bytes for 16-bit audio
         wf.setframerate(rate)
         yield wf
-
-async def generate_audio_live_async(api_text, output_filename):
-    """Generates audio using Gemini Live API (async version)."""
+    except Exception as e:
+        logger.error(f"Error opening/configuring wave file {filename}: {e}")
+        raise # Re-raise the exception
+    finally:
+        if wf:
+            wf.close()
+            # logger.debug(f"Closed wave file: {filename}")
+
+
+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) with improved error handling.
+    Returns the path to the generated audio file or None on failure.
+    """
     collected_audio = bytearray()
-    st.write(f"🎙️ Generating audio for: '{api_text[:50]}...'") # Log start
+    task_id = os.path.basename(output_filename).split('.')[0] # Extract T#_S# for logging
+    logger.info(f"🎙️ [{task_id}] Requesting audio for: '{api_text[:60]}...'")
 
     try:
-        # Use the 'client_live' specifically configured for v1alpha
-        live_model = client_live.get_model(f"models/gemini-1.5-flash") # Specify model within the live client context
-
+        # Use the 'live_model' obtained from the 'client_live' instance.
         config = {
-            "response_modalities": ["AUDIO"]
+            "response_modalities": ["AUDIO"],
+            "audio_config": { # Optional: Specify voice, etc.
+                "audio_encoding": "LINEAR16", # Required format for WAV output
+                "sample_rate_hertz": AUDIO_SAMPLING_RATE,
+                # "voice": voice if voice else "aura-asteria-en" # Example voice - check availability
+            }
         }
-        # Connect to the Live API using the live client.
+
+        # Add a strong negative prompt to avoid conversational filler
+        # This is prepended here, but could also be part of the main Gemini prompt structure
+        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'. "
+            "Speak only the sentence itself:\n\n"
+            f'"{api_text}"'
+        )
+
+        # Use the live client's model to connect
         async with live_model.connect(config=config) as session:
-            await session.send_request([api_text]) # Simpler send for single prompt
+            # Send the refined request
+            await session.send_request([directive_prompt])
+
             async for response in session.stream_content():
-                 if response.audio_chunk:
+                if response.audio_chunk and response.audio_chunk.data:
                     collected_audio.extend(response.audio_chunk.data)
+                # Handle potential errors within the stream if the API provides them
+                if response.error:
+                     logger.error(f"   ❌ [{task_id}] Error during audio stream: {response.error}")
+                     st.error(f"Audio stream error for scene {task_id}: {response.error}", icon="🔊")
+                     return None # Stop processing this audio request
 
         if not collected_audio:
-             st.warning(f"⚠️ No audio data received for: '{api_text[:50]}...'")
-             return None # Indicate failure
-
-        audio_bytes = bytes(collected_audio)
-        # Write the collected audio bytes into a WAV file.
-        with wave_file(output_filename) as wf:
-            wf.writeframes(audio_bytes)
-        st.write(f"   ✅ Audio saved: {os.path.basename(output_filename)}")
+            logger.warning(f"⚠️ [{task_id}] No audio data received for: '{api_text[:60]}...'")
+            st.warning(f"No audio data generated for scene {task_id}.", icon="🔊")
+            return None
+
+        # Write the collected audio bytes into a WAV file using the context manager.
+        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="🔇")
+         return None
     except Exception as e:
-        st.error(f"   ❌ Audio generation failed for '{api_text[:50]}...': {e}", icon="🚨")
+        logger.exception(f"   ❌ [{task_id}] Audio generation failed unexpectedly for '{api_text[:60]}...': {e}")
+        st.error(f"Audio generation failed for scene {task_id}: {e}", icon="🔊")
         return None
 
 
-def generate_story_sequence_chrono(theme: str, num_scenes: int, num_timelines: int, divergence_prompt: str = "") -> ChronoWeaveResponse | None:
-    """Generates branching story sequences using Gemini structured output."""
-    st.write(f"📚 Generating {num_timelines} timeline(s) for theme: '{theme}'...")
-    divergence_instruction = f"Introduce divergence between timelines. {divergence_prompt}" if divergence_prompt else "Introduce natural points of divergence between timelines after the first scene or two."
-
-    prompt = f'''
-    As an expert narrative designer, create a branching story based on the theme: "{theme}".
-    Generate exactly {num_timelines} distinct timelines, each containing exactly {num_scenes} scenes.
-    Each scene should be approximately 5-10 seconds long when narrated.
-
-    {divergence_instruction} Clearly state the reason for divergence for each timeline after the first.
-
-    For each scene in each timeline, provide:
-    - scene_id: An integer starting from 0 for the scene number within its timeline.
-    - image_prompt: A concise (15-25 words) description for an image generation model. Focus on visual details, characters (animals/objects only, NO PEOPLE), background, and action. Maintain a consistent 'kids animation style' (e.g., simple, rounded shapes, bright colors) across all scenes and timelines unless specified by a timeline_visual_modifier.
-    - audio_text: A single, engaging sentence of narration or dialogue for the scene (max 25 words).
-    - character_description: Brief description of recurring characters (names, key features) mentioned in *this specific scene's image prompt*. Keep consistent within a timeline. (Max 30 words).
-    - timeline_visual_modifier: (Optional, string or null) A *brief* hint if this timeline should have a slightly different visual feel from this scene onwards (e.g., "slightly darker lighting", "more cluttered background", "character looks worried"). Keep it subtle. Use null if no specific modifier.
-
-    Constraint: Ensure the output strictly adheres to the following JSON schema. Do not include preamble or explanations outside the JSON structure. Respond ONLY with the JSON object.
-
-    JSON Schema:
-    {{
-      "type": "object",
-      "properties": {{
-        "core_theme": {{"type": "string"}},
-        "timelines": {{
-          "type": "array",
-          "items": {{
-            "type": "object",
-            "properties": {{
-              "timeline_id": {{"type": "integer"}},
-              "divergence_reason": {{"type": "string"}},
-              "segments": {{
-                "type": "array",
-                "items": {{
-                  "type": "object",
-                  "properties": {{
-                    "scene_id": {{"type": "integer"}},
-                    "image_prompt": {{"type": "string"}},
-                    "audio_text": {{"type": "string"}},
-                    "character_description": {{"type": "string"}},
-                    "timeline_visual_modifier": {{"type": ["string", "null"]}}
-                  }},
-                  "required": ["scene_id", "image_prompt", "audio_text", "character_description", "timeline_visual_modifier"]
-                }}
-              }}
-            }},
-            "required": ["timeline_id", "divergence_reason", "segments"]
-          }}
-        }},
-        "total_scenes_per_timeline": {{"type": "integer"}}
-      }},
-      "required": ["core_theme", "timelines", "total_scenes_per_timeline"]
-    }}
-    '''
+def generate_story_sequence_chrono(
+    theme: str,
+    num_scenes: int,
+    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.
+    """
+    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}")
+
+    divergence_instruction = (
+        f"Introduce clear points of divergence between timelines, starting potentially after the first scene. "
+        f"If provided, use this hint for divergence: '{divergence_prompt}'. "
+        f"Clearly state the divergence reason for each timeline (except potentially the first)."
+    )
+
+    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}".
+
+    **Instructions:**
+    1.  Generate exactly **{num_timelines}** distinct timelines.
+    2.  Each timeline must contain exactly **{num_scenes}** sequential scenes.
+    3.  **Crucially, DO NOT include any humans, people, or humanoid figures** in the descriptions or actions. Focus strictly on animals, fantasy creatures, animated objects, or natural elements.
+    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 (15-35 words), focusing on the non-human character(s), setting, action, and visual style. Explicitly mention the main character(s) for consistency.
+    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.
+
+    **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_json(indent=2))}
+    ```
+    """
 
     try:
         response = client_standard.generate_content(
             contents=prompt,
             generation_config=genai.types.GenerationConfig(
-                 response_mime_type="application/json",
-                 # Optional: Add temperature, etc. if needed
+                response_mime_type="application/json",
+                temperature=0.7 # Add some creativity
             )
-            # The schema can also be passed via generation_config in some versions/models
-            # config={
-            #     'response_mime_type': 'application/json',
-            #     'response_schema': ChronoWeaveResponse # Pass the TypedDict directly
-            # }
         )
 
-        # Debugging: Print raw response text
-        # st.text_area("Raw Gemini Response:", response.text, height=200)
+        # Debugging: Log raw response
+        # logger.debug(f"Raw Gemini Response Text:\n{response.text}")
+
+        # 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.text_area("Problematic AI Response:", response.text, height=200)
+            return None
 
-        story_data = json.loads(response.text) # Assuming response.text contains the JSON string
-        st.success("✅ Story structure generated successfully!")
-        # Basic validation (can be more thorough)
-        if 'timelines' in story_data and isinstance(story_data['timelines'], list):
-             # Further validation could check segment structure, etc.
-             return story_data # Return the parsed dictionary
-        else:
-             st.error("🚨 Generated story data is missing the 'timelines' list.", icon="🚨")
-             return None
+        # Validate the parsed data using Pydantic
+        try:
+            validated_data = ChronoWeaveResponse.parse_obj(raw_data)
+            logger.info("✅ Story structure generated and validated successfully!")
+            st.success("✅ Story structure generated and validated!")
+            return validated_data
+        except ValidationError as val_err:
+            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
+            return None
 
-    except json.JSONDecodeError as e:
-        st.error(f"🚨 Failed to decode JSON response from Gemini: {e}", icon="🚨")
-        st.text_area("Problematic Response Text:", response.text if 'response' in locals() else "No response object.", height=150)
-        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="🚫")
+         return None
     except Exception as e:
-        st.error(f"🚨 Error generating story sequence: {e}", icon="🚨")
-        # Log the prompt potentially? Be careful with sensitive data if applicable.
-        # st.text_area("Failed Prompt:", prompt, height=200)
+        logger.exception("Error during story sequence generation:")
+        st.error(f"🚨 An unexpected error occurred during story generation: {e}", icon="💥")
+        # Optional: Show the prompt that failed (be mindful of length/PII)
+        # st.text_area("Failed Prompt (excerpt):", prompt[:500]+"...", height=150)
         return None
 
 
-def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1") -> Image.Image | None:
-    """Generates an image using Imagen."""
-    st.write(f"🖼️ Generating image for: '{prompt[:60]}...'")
+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.
+    """
+    logger.info(f"🖼️ [{task_id}] Requesting image for: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
+
+    # Refined prompt incorporating negative constraints and style guidance
+    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"Prompt: {prompt}"
+    )
+
     try:
-        # Use the standard client's dedicated image generation method
+        # Use the standard client's generate_content method.
+        # How aspect ratio and negative prompts are passed can vary slightly with API versions.
+        # This uses the model's understanding of the text prompt.
+        # For more explicit control, future API versions might use parameters in GenerationConfig or Tools.
         response = client_standard.generate_content(
-            f"Generate an image with the following prompt, ensuring a child-friendly animation style and NO human figures: {prompt}",
+            full_prompt,
              generation_config=genai.types.GenerationConfig(
-                 candidate_count=1, # Generate one image
-                 # Imagen specific parameters are often passed differently or rely on model defaults
-                 # Check documentation for precise Imagen control via the unified API
+                 candidate_count=1,
+                 # The following are conceptual parameters for Imagen via the unified API.
+                 # Check the latest google-generativeai library documentation for the exact syntax.
+                 # stop_sequences=["human", "person"], # May not be directly supported this way
+                 # custom_params={"aspect_ratio": aspect_ratio, "negative_prompt": "human, person, people, child, realistic, photo"}
+                 # As of now, embedding these in the text prompt is the most reliable way.
              ),
-             # If the model/API version requires specific image parameters:
-             # tools=[genai.ImageParams(model=IMAGE_MODEL_ID, number_of_images=1, aspect_ratio=aspect_ratio, person_generation="DONT_ALLOW")]
+             # safety_settings={'HARM_CATEGORY_DANGEROUS_CONTENT': 'BLOCK_NONE'} # Adjust safety cautiously if needed
         )
 
-        # Accessing image data might vary slightly depending on API response structure
-        # This assumes response.parts contains the image data if successful
-        if response.parts and response.parts[0].inline_data:
+        # Check for valid response and image data
+        if response.parts and response.parts[0].inline_data and response.parts[0].inline_data.data:
             image_bytes = response.parts[0].inline_data.data
-            image = Image.open(BytesIO(image_bytes))
-            st.write("   ✅ Image generated.")
-            return image
+            try:
+                image = Image.open(BytesIO(image_bytes))
+                logger.info(f"   ✅ [{task_id}] Image generated successfully.")
+                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:
-            # Check for safety blocks or other reasons for failure
-            if response.prompt_feedback.block_reason:
-                 st.warning(f"   ⚠️ Image generation blocked for prompt '{prompt[:60]}...'. Reason: {response.prompt_feedback.block_reason}", icon="⚠️")
+            # Check for blocking or other issues
+            block_reason = getattr(response.prompt_feedback, 'block_reason', None)
+            safety_ratings = getattr(response.prompt_feedback, 'safety_ratings', [])
+            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="🚫")
+            elif safety_ratings:
+                 filtered_ratings = [f"{r.category}: {r.probability}" for r in safety_ratings if r.probability != 'NEGLIGIBLE']
+                 if filtered_ratings:
+                    logger.warning(f"   ⚠️ [{task_id}] Image generated but flagged by safety filters: {', '.join(filtered_ratings)}. Prompt: '{prompt[:70]}...'")
+                    st.warning(f"Image generation for scene {task_id} flagged by safety filters: {', '.join(filtered_ratings)}", icon="⚠️")
+                    # Proceeding, but warning the user. Consider returning None if strict safety is needed.
+                    # return None # Uncomment this line to block flagged images
+                    # If we proceed, we need to extract the image data despite the warning
+                    if response.parts and response.parts[0].inline_data and response.parts[0].inline_data.data:
+                         image_bytes = response.parts[0].inline_data.data
+                         try:
+                              image = Image.open(BytesIO(image_bytes))
+                              logger.info(f"   ✅ [{task_id}] Image generated (with safety flags).")
+                              return image
+                         except Exception as img_err:
+                              logger.error(f"   ❌ [{task_id}] Failed to decode flagged image data: {img_err}")
+                              st.warning(f"Failed to decode flagged image data for scene {task_id}.", icon="🖼️")
+                              return None
+                    else:
+                         # Should not happen if flagged but still contains data, but handle defensively
+                         logger.warning(f"   ⚠️ [{task_id}] Image flagged but no image data found. Prompt: '{prompt[:70]}...'")
+                         st.warning(f"No image data received for scene {task_id} (safety flagged).", icon="🖼️")
+                         return None
             else:
-                 st.warning(f"   ⚠️ No image data received for prompt '{prompt[:60]}...'.", icon="⚠️")
-            # Debugging: st.write(response)
+                 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="🖼️")
+                 # You might want to inspect the full 'response' object here for clues
+                 # logger.debug(f"Full Imagen response object: {response}")
             return None
 
+    except genai.types.generation_types.BlockedPromptException as bpe:
+         logger.error(f"   ❌ [{task_id}] Image generation blocked for prompt '{prompt[:70]}...': {bpe}")
+         st.error(f"Image generation blocked for scene {task_id} due to safety settings.", icon="🚫")
+         return None
     except Exception as e:
-        st.error(f"   ❌ Image generation failed for '{prompt[:60]}...': {e}", icon="🚨")
+        logger.exception(f"   ❌ [{task_id}] Image generation failed unexpectedly for '{prompt[:70]}...': {e}")
+        st.error(f"Image generation failed for scene {task_id}: {e}", icon="🖼️")
         return None
 
 
 # --- Streamlit UI Elements ---
-st.sidebar.header("Configuration")
+st.sidebar.header("⚙️ Configuration")
 
-# API Key display/check (already handled above, but sidebar is a good place)
+# API Key Status (already handled, just display status)
 if GOOGLE_API_KEY:
-    st.sidebar.success("Google API Key Loaded!", icon="✅")
+    st.sidebar.success("Google API Key Loaded", icon="✅")
 else:
+    # This part should technically not be reached due to st.stop() earlier
     st.sidebar.error("Google API Key Missing!", icon="🚨")
 
-theme = st.sidebar.text_input("Story Theme:", "A curious squirrel finds a shiny object")
-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 it started raining?")
-aspect_ratio = st.sidebar.selectbox("Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0)
+# 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.")
+
+# 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.")
+# Add audio voice selection if API supports it and voices are known
+# available_voices = ["aura-asteria-en", "aura-luna-en", "aura-stella-en"] # Example
+# audio_voice = st.sidebar.selectbox("🗣️ Narration Voice:", available_voices, index=0)
+audio_voice = None # Placeholder if voice selection isn't implemented/stable
 
-generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY))
+generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY), use_container_width=True)
 
 st.sidebar.markdown("---")
-st.sidebar.info("Note: Generation can take several minutes depending on settings.")
+st.sidebar.info("⏳ Generation can take several minutes, especially with more scenes or timelines.", icon="⏳")
+st.sidebar.markdown(f"<small>Models: Text={TEXT_MODEL_ID}, Image={IMAGE_MODEL_ID}, Audio={AUDIO_MODEL_ID} ({AUDIO_API_VERSION})</small>", unsafe_allow_html=True)
+
 
 # --- Main Logic ---
 if generate_button:
     if not theme:
-        st.error("Please enter a story theme.", icon="👈")
+        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())
-        temp_dir = os.path.join(".", f"chrono_temp_{run_id}") # Create in current dir
-        os.makedirs(temp_dir, exist_ok=True)
-        st.write(f"Working directory: {temp_dir}")
-
-        final_video_paths = {} # To store {timeline_id: video_path}
-
-        with st.spinner("Generating narrative structure..."):
-            chrono_data = generate_story_sequence_chrono(theme, num_scenes, num_timelines, divergence_prompt)
-
-        if chrono_data and 'timelines' in chrono_data:
-            st.success(f"Found {len(chrono_data['timelines'])} timelines. Processing each...")
-
-            all_timelines_successful = True # Flag to track if all timelines worked
-
-            # Use st.status for detailed progress
+        run_id = str(uuid.uuid4()).split('-')[0] # Short unique ID
+        temp_dir = os.path.join(TEMP_DIR_BASE, f"run_{run_id}")
+        try:
+            os.makedirs(temp_dir, exist_ok=True)
+            logger.info(f"Created temporary directory: {temp_dir}")
+        except OSError as e:
+            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]}
+
+        # --- 1. Generate Narrative Structure ---
+        chrono_response: Optional[ChronoWeaveResponse] = None
+        with st.spinner("Generating narrative structure... 🤔"):
+            chrono_response = generate_story_sequence_chrono(theme, num_scenes, num_timelines, divergence_prompt)
+
+        if chrono_response:
+            st.success(f"Narrative structure received for {len(chrono_response.timelines)} timelines.")
+            logger.info(f"Successfully generated structure for {len(chrono_response.timelines)} timelines.")
+
+            # --- 2. Process Each Timeline ---
+            overall_start_time = time.time()
+            all_timelines_successful = True # Assume success initially
+
+            # Use st.status for collapsible progress updates
             with st.status("Generating assets and composing videos...", expanded=True) as status:
 
-                for timeline in chrono_data['timelines']:
-                    timeline_id = timeline['timeline_id']
-                    divergence = timeline['divergence_reason']
-                    segments = timeline['segments']
-                    st.subheader(f"Timeline {timeline_id}: {divergence}")
-
-                    temp_image_files = []
-                    temp_audio_files = []
-                    video_clips = []
-                    timeline_successful = True # Flag for this specific timeline
-
-                    for i, segment in enumerate(segments):
-                        status.update(label=f"Processing Timeline {timeline_id}, Scene {i+1}/{num_scenes}...")
-                        scene_id = segment['scene_id']
-                        image_prompt = segment['image_prompt']
-                        audio_text = segment['audio_text']
-                        char_desc = segment['character_description']
-                        vis_mod = segment['timeline_visual_modifier']
-
-                        st.write(f"--- Scene {i+1} (T{timeline_id}) ---")
-                        st.write(f"*   **Image Prompt:** {image_prompt}" + (f" (Modifier: {vis_mod})" if vis_mod else ""))
-                        st.write(f"*   **Audio Text:** {audio_text}")
-                        # st.write(f"*   Character Desc: {char_desc}") # Can be verbose
-
-                        # --- Image Generation ---
-                        combined_prompt = f"{image_prompt} {char_desc}"
-                        if vis_mod:
-                            combined_prompt += f" Style hint: {vis_mod}"
-
-                        generated_image = generate_image_imagen(combined_prompt, aspect_ratio)
+                for timeline_index, timeline in enumerate(chrono_response.timelines):
+                    timeline_id = timeline.timeline_id
+                    divergence = timeline.divergence_reason
+                    segments = timeline.segments
+                    timeline_label = f"Timeline {timeline_id}" # Consistent label
+                    st.subheader(f"Processing {timeline_label}: {divergence}")
+                    logger.info(f"--- Processing {timeline_label} (Index: {timeline_index}) ---")
+                    generation_errors[timeline_id] = [] # Initialize error list for this timeline
+
+                    # Store paths for this timeline's assets
+                    temp_image_files = {} # {scene_id: path}
+                    temp_audio_files = {} # {scene_id: path}
+                    video_clips = [] # List of moviepy clips for concatenation
+                    timeline_start_time = time.time()
+                    scene_success_count = 0
+
+
+                    for scene_index, segment in enumerate(segments):
+                        scene_id = segment.scene_id
+                        task_id = f"T{timeline_id}_S{scene_id}" # Unique ID for logging/filenames
+                        status_message = f"Processing {timeline_label}, Scene {scene_id + 1}/{len(segments)}..."
+                        status.update(label=status_message)
+                        st.markdown(f"--- **Scene {scene_id + 1} ({task_id})** ---")
+                        logger.info(status_message)
+
+                        scene_has_error = False
+
+                        # Log scene details
+                        st.write(f"   *Image Prompt:* {segment.image_prompt}" + (f" *(Modifier: {segment.timeline_visual_modifier})*" if segment.timeline_visual_modifier else ""))
+                        st.write(f"   *Audio Text:* {segment.audio_text}")
+                        # st.write(f"*   Character Desc: {segment.character_description}") # Optional verbosity
+
+                        # --- 2a. Image Generation ---
+                        with st.spinner(f"[{task_id}] Generating image... 🎨"):
+                            combined_prompt = f"{segment.image_prompt}. {segment.character_description}"
+                            if segment.timeline_visual_modifier:
+                                combined_prompt += f" Visual style hint: {segment.timeline_visual_modifier}."
+
+                            generated_image: Optional[Image.Image] = generate_image_imagen(combined_prompt, aspect_ratio, task_id)
 
                         if generated_image:
-                            image_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_image.png")
-                            generated_image.save(image_path)
-                            temp_image_files.append(image_path)
-                            st.image(generated_image, width=200) # Show thumbnail
-                        else:
-                            st.warning(f"Skipping scene {i+1} in timeline {timeline_id} due to image generation failure.")
-                            timeline_successful = False
-                            continue # Skip to next segment if image fails
-
-                        # --- Audio Generation ---
-                        # Add negative prompt to prevent conversational filler
-                        audio_negative_prompt = "Narrate the following sentence directly, with expression, without any introduction or closing remarks like 'Okay' or 'Here is the narration'. Just read the sentence:"
-                        full_audio_prompt = f"{audio_negative_prompt}\n{audio_text}"
-                        audio_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_audio.wav")
-
-                        # Run the async audio generation function
-                        try:
-                           generated_audio_path = asyncio.run(generate_audio_live_async(full_audio_prompt, audio_path))
-                        except Exception as e:
-                            st.error(f"Asyncio error during audio gen: {e}")
-                            generated_audio_path = None
-
-
-                        if generated_audio_path:
-                            temp_audio_files.append(generated_audio_path)
-                            # st.audio(generated_audio_path) # Optional: Preview audio
+                            image_path = os.path.join(temp_dir, f"{task_id}_image.png")
+                            try:
+                                generated_image.save(image_path)
+                                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}")
+                                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"Skipping video clip for scene {i+1} in timeline {timeline_id} due to audio generation failure.")
-                            # Clean up the image file for this failed scene segment
-                            if os.path.exists(image_path):
-                                os.remove(image_path)
-                                temp_image_files.remove(image_path)
-                            timeline_successful = False
-                            continue # Skip making video clip if audio fails
-
-
-                        # --- Create Video Clip ---
-                        try:
-                            st.write("   🎬 Creating video clip...")
-                            audio_clip = AudioFileClip(generated_audio_path)
-                            # Ensure PIL Image is used if needed, or numpy array directly
-                            np_image = np.array(Image.open(image_path))
-                            # Create ImageClip, ensure duration matches audio
-                            image_clip = ImageClip(np_image).set_duration(audio_clip.duration)
-
-                            # Handle potential size mismatch if needed (resize image_clip or set size explicitly)
-                            # image_clip = image_clip.resize(width=...)
-
-                            composite_clip = image_clip.set_audio(audio_clip) # Simpler composition
-                            video_clips.append(composite_clip)
-                            st.write("      ✅ Clip created.")
-                        except Exception as e:
-                             st.error(f"      ❌ Failed to create video clip for scene {i+1} (T{timeline_id}): {e}", icon="🚨")
-                             timeline_successful = False
-                             # Don't break the whole timeline, just skip this clip maybe? Or mark timeline as failed.
-
-
-                    # --- Assemble Timeline Video ---
-                    if video_clips and timeline_successful: # Only assemble if clips were made and no major errors
-                        status.update(label=f"Composing final video for Timeline {timeline_id}...")
-                        st.write(f"🎞️ Assembling final video for Timeline {timeline_id}...")
+                            st.warning(f"Image generation failed for scene {task_id}. Skipping scene.", icon="🖼️")
+                            scene_has_error = True
+                            generation_errors[timeline_id].append(f"Scene {scene_id+1}: Image generation failed.")
+                            # No image, so skip audio and video for this scene
+                            continue
+
+                        # --- 2b. Audio Generation ---
+                        generated_audio_path: Optional[str] = None
+                        if not scene_has_error: # Only generate audio if image succeeded
+                            with st.spinner(f"[{task_id}] Generating audio... 🔊"):
+                                audio_path_temp = os.path.join(temp_dir, f"{task_id}_audio.wav")
+                                try:
+                                    # Run the async function using asyncio.run() which works with nest_asyncio
+                                    generated_audio_path = asyncio.run(
+                                        generate_audio_live_async(segment.audio_text, audio_path_temp, audio_voice)
+                                    )
+                                except RuntimeError as e:
+                                    # Catch potential issues if asyncio loop is misconfigured despite nest_asyncio
+                                    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="⚡")
+                                    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="💥")
+                                    scene_has_error = True
+                                    generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation error.")
+
+
+                            if generated_audio_path:
+                                temp_audio_files[scene_id] = generated_audio_path
+                                # Optional: Preview audio
+                                try:
+                                    with open(generated_audio_path, 'rb') as ap:
+                                        st.audio(ap.read(), format='audio/wav')
+                                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="🔊")
+                                scene_has_error = True
+                                generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation failed.")
+                                # Clean up the image for this failed segment if audio fails
+                                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}")
+                                continue # Skip video clip creation
+
+                        # --- 2c. Create Video Clip ---
+                        if not scene_has_error and scene_id in temp_image_files and scene_id in temp_audio_files:
+                            st.write(f"   🎬 Creating video clip for Scene {scene_id+1}...")
+                            try:
+                                # Ensure files exist before creating clips
+                                img_path = temp_image_files[scene_id]
+                                aud_path = temp_audio_files[scene_id]
+                                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}")
+
+                                audio_clip = AudioFileClip(aud_path)
+                                # Use numpy array for ImageClip to avoid potential PIL issues with moviepy versions
+                                np_image = np.array(Image.open(img_path))
+                                image_clip = ImageClip(np_image).set_duration(audio_clip.duration)
+
+                                # Composite the clip
+                                composite_clip = image_clip.set_audio(audio_clip)
+                                video_clips.append(composite_clip)
+                                logger.info(f"      ✅ [{task_id}] Video clip created (Duration: {audio_clip.duration:.2f}s).")
+                                st.write(f"      ✅ Clip created (Duration: {audio_clip.duration:.2f}s).")
+                                scene_success_count += 1
+
+                            except Exception as e:
+                                logger.exception(f"      ❌ [{task_id}] Failed to create video clip for scene {scene_id+1}: {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.")
+                                # Attempt cleanup of related files if clip creation fails
+                                if scene_id in temp_image_files and os.path.exists(temp_image_files[scene_id]): os.remove(temp_image_files[scene_id])
+                                if scene_id in temp_audio_files and os.path.exists(temp_audio_files[scene_id]): os.remove(temp_audio_files[scene_id])
+
+                        # --- End of Scene Loop ---
+
+                    # --- 2d. Assemble Timeline Video ---
+                    timeline_duration = time.time() - timeline_start_time
+                    if video_clips and scene_success_count == len(segments): # Only assemble if all scenes were successful
+                        status.update(label=f"Composing final video for {timeline_label}...")
+                        st.write(f"🎞️ Assembling final video for {timeline_label}...")
+                        logger.info(f"🎞️ Assembling final video for {timeline_label} ({len(video_clips)} clips)...")
+                        output_filename = os.path.join(temp_dir, f"timeline_{timeline_id}_final.mp4")
+                        final_timeline_video = None # Define before try block
                         try:
                             final_timeline_video = concatenate_videoclips(video_clips, method="compose")
-                            output_filename = os.path.join(temp_dir, f"timeline_{timeline_id}_final_video.mp4")
-                            # Use 'libx264' for broader compatibility, specify audio codec
-                            final_timeline_video.write_videofile(output_filename, fps=24, codec='libx264', audio_codec='aac')
+                            # Write video file with specified codecs and fps
+                            final_timeline_video.write_videofile(
+                                output_filename,
+                                fps=VIDEO_FPS,
+                                codec=VIDEO_CODEC,
+                                audio_codec=AUDIO_CODEC,
+                                logger=None # Suppress moviepy console spam
+                                # threads=4 # Optional: specify threads
+                            )
                             final_video_paths[timeline_id] = output_filename
-                            st.success(f"   ✅ Video for Timeline {timeline_id} saved: {os.path.basename(output_filename)}")
-
-                            # Close clips to release resources
-                            for clip in video_clips:
-                                if hasattr(clip, 'close'): clip.close()
-                                if hasattr(clip, 'audio') and hasattr(clip.audio, 'close'): clip.audio.close()
-                            if hasattr(final_timeline_video, 'close'): final_timeline_video.close()
-
+                            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:
-                            st.error(f"   ❌ Failed to write final video for Timeline {timeline_id}: {e}", icon="🚨")
+                            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
-                    elif not video_clips:
-                         st.warning(f"No video clips were successfully generated for Timeline {timeline_id}. Skipping final video assembly.")
-                         all_timelines_successful = False
-                    else:
-                         st.warning(f"Timeline {timeline_id} encountered errors. Skipping final video assembly.")
-                         all_timelines_successful = False
-
+                            generation_errors[timeline_id].append(f"Timeline {timeline_id}: Final video assembly failed.")
+                        finally:
+                            # Crucially, close all clips to release file handles
+                            logger.debug(f"[{timeline_label}] Closing {len(video_clips)} video clips...")
+                            for clip in video_clips:
+                                try:
+                                    clip.close()
+                                    if clip.audio:
+                                        clip.audio.close()
+                                except Exception as e_close:
+                                    logger.warning(f"   ⚠️ [{timeline_label}] Error closing clip: {e_close}")
+                            if final_timeline_video:
+                                try:
+                                     final_timeline_video.close()
+                                except Exception as e_close_final:
+                                     logger.warning(f"   ⚠️ [{timeline_label}] Error closing final video object: {e_close_final}")
+                            logger.debug(f"[{timeline_label}] Clips closed.")
 
-                    # Intermediate cleanup for the timeline (optional, helps manage files)
-                    # for file in temp_audio_files:
-                    #     if os.path.exists(file): os.remove(file)
-                    # for file in temp_image_files:
-                    #     if os.path.exists(file): os.remove(file)
+                    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="🚫")
+                        all_timelines_successful = False
+                    else: # Some scenes failed
+                        logger.warning(f"[{timeline_label}] Encountered errors in {len(generation_errors[timeline_id])} scene(s). Skipping final video assembly.")
+                        st.warning(f"{timeline_label} had errors in {len(generation_errors[timeline_id])} scene(s). Final video not assembled.", icon="⚠️")
+                        all_timelines_successful = False
+
+                    # Log errors for the timeline if any occurred
+                    if generation_errors[timeline_id]:
+                         logger.error(f"Errors occurred in {timeline_label}: {generation_errors[timeline_id]}")
+
+                    # Intermediate cleanup (optional, can free up disk space during long runs)
+                    # logger.debug(f"[{timeline_label}] Cleaning up intermediate files...")
+                    # for scene_id, fpath in temp_image_files.items():
+                    #     if os.path.exists(fpath): os.remove(fpath)
+                    # for scene_id, fpath in temp_audio_files.items():
+                    #     if os.path.exists(fpath): os.remove(fpath)
+
+                # --- End of Timelines Loop ---
 
                 # Final status update
+                overall_duration = time.time() - overall_start_time
                 if all_timelines_successful and final_video_paths:
-                    status.update(label="ChronoWeave Generation Complete!", state="complete", expanded=False)
+                    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)
                 elif final_video_paths:
-                     status.update(label="ChronoWeave Generation Partially Complete (some errors occurred).", state="warning", expanded=False)
+                    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) # Keep expanded if errors
+                    logger.warning(status_msg)
                 else:
-                    status.update(label="ChronoWeave Generation Failed.", state="error", expanded=False)
-
+                    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)
 
-            # --- Display Results ---
-            st.header("Generated Timelines")
+            # --- 3. Display Results ---
+            st.header("🎬 Generated Timelines")
             if final_video_paths:
                 sorted_timeline_ids = sorted(final_video_paths.keys())
-                for timeline_id in sorted_timeline_ids:
+                cols = st.columns(len(sorted_timeline_ids)) # Create columns for side-by-side display
+
+                for idx, timeline_id in enumerate(sorted_timeline_ids):
                     video_path = final_video_paths[timeline_id]
-                    # Find matching timeline divergence reason
-                    reason = "Unknown"
-                    for t in chrono_data.get('timelines', []):
-                         if t.get('timeline_id') == timeline_id:
-                             reason = t.get('divergence_reason', 'N/A')
-                             break
-                    st.subheader(f"Timeline {timeline_id}: {reason}")
-                    try:
-                        video_file = open(video_path, 'rb')
-                        video_bytes = video_file.read()
-                        st.video(video_bytes)
-                        video_file.close()
-                    except FileNotFoundError:
-                         st.error(f"Could not find video file: {video_path}", icon="🚨")
-                    except Exception as e:
-                         st.error(f"Could not display video {video_path}: {e}", icon="🚨")
+                    # Find matching timeline data for context
+                    timeline_data = next((t for t in chrono_response.timelines if t.timeline_id == timeline_id), None)
+                    reason = timeline_data.divergence_reason if timeline_data else "Unknown Divergence"
+                    col = cols[idx]
+                    with col:
+                        st.subheader(f"Timeline {timeline_id}")
+                        st.caption(f"Divergence: {reason}")
+                        try:
+                            # Read video bytes for display
+                            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}")
+                            # Add download button
+                            st.download_button(
+                                label=f"Download T{timeline_id} Video",
+                                data=video_bytes,
+                                file_name=f"chronoweave_timeline_{timeline_id}.mp4",
+                                mime="video/mp4"
+                            )
+                            # Display errors for this timeline if any occurred
+                            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}")
+
+                        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="🚨")
+                        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="🚨")
             else:
-                st.warning("No final videos were successfully generated.")
-
-            # --- Cleanup ---
-            st.write("Cleaning up temporary files...")
+                st.warning("No final videos were successfully generated in this run.")
+                # Display global errors if no videos were made
+                all_errors = [msg for err_list in generation_errors.values() for msg in err_list]
+                if all_errors:
+                    st.subheader("Summary of Generation Issues")
+                    for error_msg in all_errors:
+                        st.error(f"- {error_msg}")
+
+            # --- 4. Cleanup ---
+            st.info(f"Cleaning up temporary directory: {temp_dir}")
             try:
                 shutil.rmtree(temp_dir)
-                st.write("   ✅ Temporary files removed.")
+                logger.info(f"✅ Temporary directory removed: {temp_dir}")
+                st.success("✅ Temporary files cleaned up.")
             except Exception as e:
-                st.warning(f"   ⚠️ Could not remove temporary directory {temp_dir}: {e}", icon="⚠️")
-
+                logger.error(f"⚠️ Could not remove temporary directory {temp_dir}: {e}")
+                st.warning(f"Could not automatically remove temporary files: {temp_dir}. Please remove it manually if needed.", icon="⚠️")
 
-        elif not chrono_data:
-            st.error("Story generation failed. Cannot proceed.", icon="🛑")
+        elif not chrono_response:
+            # Error message already shown by generate_story_sequence_chrono
+            logger.error("Story generation failed, cannot proceed.")
         else:
-            # This case might happen if chrono_data is returned but is malformed (e.g., no 'timelines' key)
-            st.error("Story data seems malformed. Cannot proceed.", icon="🛑")
-            # st.json(chrono_data) # Display the problematic data
+            # This case implies chrono_response exists but somehow failed validation logic (should be caught earlier)
+            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.")
 
 else:
-    st.info("Configure settings in the sidebar and click 'Generate ChronoWeave'")
+    st.info("Configure settings in the sidebar and click '✨ Generate ChronoWeave ✨' to start.")