Update app.py

app.py

@@ -26,7 +26,7 @@ from typing import List, Optional, Literal
 
 # Video and audio processing
 from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
-from moviepy.config import change_settings # Potential for setting imagemagick path if needed
+# from moviepy.config import change_settings # Potential for setting imagemagick path if needed
 
 # Type hints
 import typing_extensions as typing
@@ -51,8 +51,8 @@ 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
-AUDIO_MODEL_ID = f"models/gemini-1.5-flash" # Model used via the v1alpha endpoint
+AUDIO_API_VERSION = 'v1alpha' # Required for audio modality (though endpoint set implicitly now)
+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
 IMAGE_MODEL_ID = "imagen-3" # Or specific version like "imagen-3.0-generate-002"
@@ -65,8 +65,6 @@ AUDIO_CODEC = "aac" # Common audio codec for MP4
 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: Use Streamlit secrets when deployed
@@ -83,29 +81,35 @@ except KeyError:
             "🚨 **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.",
+            "2.  **Local Development:** Set the `GOOGLE_API_KEY` environment variable or create a `.streamlit/secrets.toml` file.",
             icon="🚨"
         )
         st.stop() # Halt execution
 
 # --- 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.")
 
-    # Client for Text/Imagen Generation (using standard API endpoint)
+    # Model/Client Handle for Text/Imagen Generation
     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_API_VERSION}.generativelanguage.googleapis.com'}
-    )
-    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}).")
-
+    # Model Handle for Audio Generation
+    # Use the standard GenerativeModel initialization.
+    # The necessary methods (like .connect) are part of this object.
+    live_model = genai.GenerativeModel(AUDIO_MODEL_ID) # Use GenerativeModel here
+    logger.info(f"Initialized GenerativeModel handle for audio ({AUDIO_MODEL_ID}).")
+    # We no longer use or need 'client_live' or explicit endpoint setting here.
+    # The audio config is handled within the generate_audio_live_async function.
+
+except AttributeError as ae:
+     # Keep this specific error catch just in case library structure is very old/unexpected
+     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.stop()
 except Exception as e:
     logger.exception("Failed to initialize Google AI Clients.")
     st.error(f"🚨 Failed to initialize Google AI Clients: {e}", icon="🚨")
@@ -123,9 +127,7 @@ class StorySegment(BaseModel):
     @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.")
+            logger.warning(f"Image prompt '{v[:50]}...' may contain human descriptions. Relying on API-level controls & prompt instructions.")
         return v
 
 class Timeline(BaseModel):
@@ -164,13 +166,15 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
         raise # Re-raise the exception
     finally:
         if wf:
-            wf.close()
-            # logger.debug(f"Closed wave file: {filename}")
+            try:
+                wf.close()
+            except Exception as e_close:
+                logger.error(f"Error closing wave file {filename}: {e_close}")
 
 
 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.
+    Generates audio using Gemini Live API (async version) via the GenerativeModel.
     Returns the path to the generated audio file or None on failure.
     """
     collected_audio = bytearray()
@@ -178,18 +182,17 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
     logger.info(f"🎙️ [{task_id}] Requesting audio for: '{api_text[:60]}...'")
 
     try:
-        # Use the 'live_model' obtained from the 'client_live' instance.
+        # Use the 'live_model' (a GenerativeModel instance) initialized earlier.
         config = {
             "response_modalities": ["AUDIO"],
-            "audio_config": { # Optional: Specify voice, etc.
+            "audio_config": {
                 "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
+                # "voice": voice if voice else "aura-asteria-en" # Optional: Specify voice if needed and available
             }
         }
 
-        # Add a strong negative prompt to avoid conversational filler
-        # This is prepended here, but could also be part of the main Gemini prompt structure
+        # 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'. "
@@ -197,16 +200,14 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
             f'"{api_text}"'
         )
 
-        # Use the live client's model to connect
+        # Connect and stream using the GenerativeModel instance
         async with live_model.connect(config=config) as session:
-            # Send the refined request
             await session.send_request([directive_prompt])
-
             async for response in session.stream_content():
                 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:
+                if hasattr(response, 'error') and 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
@@ -227,6 +228,7 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
          st.error(f"Audio generation blocked for scene {task_id} due to safety settings.", icon="🔇")
          return None
     except Exception as e:
+        # Catch other potential errors during connect/send/stream
         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
@@ -270,11 +272,12 @@ def generate_story_sequence_chrono(
 
     **JSON Schema:**
     ```json
-    {json.dumps(ChronoWeaveResponse.schema_json(indent=2))}
+    {json.dumps(ChronoWeaveResponse.schema(), indent=2)}
     ```
-    """
+    """ # Using .schema() which is the Pydantic v1 way, adjust if using v2 (.model_json_schema())
 
     try:
+        # Use the standard client (GenerativeModel instance) for text generation
         response = client_standard.generate_content(
             contents=prompt,
             generation_config=genai.types.GenerationConfig(
@@ -288,6 +291,7 @@ def generate_story_sequence_chrono(
 
         # Attempt to parse the JSON
         try:
+            # Use response.text which should contain the JSON string
             raw_data = json.loads(response.text)
         except json.JSONDecodeError as json_err:
             logger.error(f"Failed to decode JSON response: {json_err}")
@@ -295,9 +299,17 @@ def generate_story_sequence_chrono(
             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
+        except Exception as e:
+             logger.error(f"Error accessing or decoding response text: {e}")
+             st.error(f"🚨 Error processing AI response: {e}", icon="📄")
+             # Log the response object itself if possible
+             # logger.debug(f"Response object: {response}")
+             return None
+
 
         # Validate the parsed data using Pydantic
         try:
+            # Use parse_obj for Pydantic v1, or YourModel.model_validate(raw_data) for v2
             validated_data = ChronoWeaveResponse.parse_obj(raw_data)
             logger.info("✅ Story structure generated and validated successfully!")
             st.success("✅ Story structure generated and validated!")
@@ -332,33 +344,48 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
     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}. " # Explicitly state aspect ratio in prompt too
         f"Prompt: {prompt}"
     )
 
     try:
         # 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(
             full_prompt,
              generation_config=genai.types.GenerationConfig(
                  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.
+                 # Add other config like temperature if desired
              ),
-             # safety_settings={'HARM_CATEGORY_DANGEROUS_CONTENT': 'BLOCK_NONE'} # Adjust safety cautiously if needed
+             # Safety settings can be adjusted here if necessary and permitted
+             # safety_settings={'HARM_CATEGORY_DANGEROUS_CONTENT': 'BLOCK_NONE'} # Use cautiously
         )
 
         # 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
+        # Accessing image data might depend slightly on the exact API response structure
+        # common pattern is response.candidates[0].content.parts[0].inline_data.data
+        # or directly response.parts if simpler structure
+        image_bytes = None
+        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
+             part = response.candidates[0].content.parts[0]
+             if hasattr(part, 'inline_data') and part.inline_data and hasattr(part.inline_data,'data'):
+                  image_bytes = part.inline_data.data
+             elif hasattr(part, 'file_data') and part.file_data: # Handle potential file URIs if API changes
+                  logger.warning(f"   ⚠️ [{task_id}] Received file URI instead of inline data. Handling not implemented.")
+                  # Potentially download from part.file_data.file_uri here
+                  return None # Or implement download
+
+        if image_bytes:
             try:
                 image = Image.open(BytesIO(image_bytes))
                 logger.info(f"   ✅ [{task_id}] Image generated successfully.")
+                # Check safety feedback even on success
+                safety_ratings = getattr(response.candidates[0], 'safety_ratings', [])
+                if safety_ratings:
+                    filtered_ratings = [f"{r.category.name}: {r.probability.name}" for r in safety_ratings if 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="⚠️")
+
                 return image
             except Exception as img_err:
                 logger.error(f"   ❌ [{task_id}] Failed to decode generated image data: {img_err}")
@@ -366,43 +393,24 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
                 return None
         else:
             # Check for blocking or other issues
-            block_reason = getattr(response.prompt_feedback, 'block_reason', None)
-            safety_ratings = getattr(response.prompt_feedback, 'safety_ratings', [])
+            block_reason = None
+            prompt_feedback = getattr(response, 'prompt_feedback', None)
+            if prompt_feedback:
+                block_reason = getattr(prompt_feedback, 'block_reason', None)
+
             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:
                  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
+                 # Log the full response for debugging
                  # 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}")
+         # This might be caught by the block_reason check above, but good to have explicit catch
+         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="🚫")
          return None
     except Exception as e:
@@ -414,12 +422,11 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
 # --- Streamlit UI Elements ---
 st.sidebar.header("⚙️ Configuration")
 
-# API Key Status (already handled, just display status)
+# API Key Status
 if GOOGLE_API_KEY:
     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="🚨")
+    st.sidebar.error("Google API Key Missing!", icon="🚨") # Should not be reached if st.stop() works
 
 # Story Parameters
 theme = st.sidebar.text_input("📖 Story Theme:", "A curious squirrel finds a mysterious, glowing acorn")
@@ -433,13 +440,13 @@ aspect_ratio = st.sidebar.selectbox("🖼️ Image Aspect Ratio:", ["1:1", "16:9
 # 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
+audio_voice = None # Placeholder
 
 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.markdown(f"<small>Models: Text={TEXT_MODEL_ID}, Image={IMAGE_MODEL_ID}, Audio={AUDIO_MODEL_ID} ({AUDIO_API_VERSION})</small>", unsafe_allow_html=True)
+st.sidebar.markdown(f"<small>Models: Text={TEXT_MODEL_ID}, Image={IMAGE_MODEL_ID}, Audio={AUDIO_MODEL_ID}</small>", unsafe_allow_html=True)
 
 
 # --- Main Logic ---
@@ -483,19 +490,18 @@ if generate_button:
                     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
+                    generation_errors[timeline_id] = [] # Initialize error list
 
-                    # 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
+                    video_clips = [] # List of moviepy clips
                     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
+                        task_id = f"T{timeline_id}_S{scene_id}" # Unique ID
                         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})** ---")
@@ -506,15 +512,14 @@ if generate_button:
                         # 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 ---
+                        generated_image: Optional[Image.Image] = None # Define before spinner
                         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)
+                            generated_image = generate_image_imagen(combined_prompt, aspect_ratio, task_id)
 
                         if generated_image:
                             image_path = os.path.join(temp_dir, f"{task_id}_image.png")
@@ -531,21 +536,19 @@ if generate_button:
                             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
+                            continue # Skip audio/video for this scene
 
                         # --- 2b. Audio Generation ---
                         generated_audio_path: Optional[str] = None
-                        if not scene_has_error: # Only generate audio if image succeeded
+                        if not scene_has_error:
                             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
+                                    # Run the async function using asyncio.run()
                                     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
@@ -556,10 +559,8 @@ if generate_button:
                                     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')
@@ -569,7 +570,7 @@ if generate_button:
                                 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
+                                # Clean up image 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])
@@ -582,54 +583,59 @@ if generate_button:
                         # --- 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}...")
+                            img_path = temp_image_files[scene_id]
+                            aud_path = temp_audio_files[scene_id]
+                            audio_clip_instance = None # Define before try
+                            image_clip_instance = None # Define before try
+                            composite_clip = None # Define before try
                             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
+                                audio_clip_instance = AudioFileClip(aud_path)
                                 np_image = np.array(Image.open(img_path))
-                                image_clip = ImageClip(np_image).set_duration(audio_clip.duration)
+                                image_clip_instance = ImageClip(np_image).set_duration(audio_clip_instance.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).")
+                                composite_clip = image_clip_instance.set_audio(audio_clip_instance)
+                                video_clips.append(composite_clip) # Add the clip to be concatenated later
+                                logger.info(f"      ✅ [{task_id}] Video clip created (Duration: {audio_clip_instance.duration:.2f}s).")
+                                st.write(f"      ✅ Clip created (Duration: {audio_clip_instance.duration:.2f}s).")
                                 scene_success_count += 1
+                                # Don't close individual clips here yet, needed for concatenation
 
                             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])
+                                # Cleanup resources if clip creation failed for *this* scene
+                                if audio_clip_instance: audio_clip_instance.close()
+                                if image_clip_instance: image_clip_instance.close()
+                                # Attempt cleanup of related files
+                                if os.path.exists(img_path): os.remove(img_path)
+                                if os.path.exists(aud_path): os.remove(aud_path)
 
                         # --- 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
+                    # Only assemble if clips were created and no *fatal* errors occurred during scene processing
+                    # (We check scene_success_count against expected number)
+                    if video_clips and scene_success_count == len(segments):
                         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:
+                            # Concatenate the collected clips
                             final_timeline_video = concatenate_videoclips(video_clips, method="compose")
-                            # 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
                             logger.info(f"   ✅ [{timeline_label}] Final video saved: {os.path.basename(output_filename)}")
@@ -641,41 +647,36 @@ if generate_button:
                             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:
+                            # Now close all individual clips and the final concatenated clip
+                            logger.debug(f"[{timeline_label}] Closing {len(video_clips)} source clips...")
+                            for i, clip in enumerate(video_clips):
                                 try:
-                                    clip.close()
-                                    if clip.audio:
-                                        clip.audio.close()
+                                    if clip: # Check if clip object exists
+                                      if clip.audio: clip.audio.close()
+                                      clip.close()
                                 except Exception as e_close:
-                                    logger.warning(f"   ⚠️ [{timeline_label}] Error closing clip: {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}")
-                            logger.debug(f"[{timeline_label}] Clips closed.")
 
                     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="⚠️")
+                    else: # Some scenes failed, so scene_success_count < len(segments)
+                        error_count = len(segments) - scene_success_count
+                        logger.warning(f"[{timeline_label}] Encountered 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
 
-                    # Log errors for the timeline if any occurred
+                    # Log accumulated 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)
+                         logger.error(f"Summary of errors in {timeline_label}: {generation_errors[timeline_id]}")
 
                 # --- End of Timelines Loop ---
 
@@ -685,11 +686,11 @@ if generate_button:
                     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:
+                elif final_video_paths: # Some videos made, but errors occurred
                     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
+                    status.update(label=status_msg, state="warning", expanded=True)
                     logger.warning(status_msg)
-                else:
+                else: # No videos made
                     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)
@@ -698,29 +699,30 @@ if generate_button:
             st.header("🎬 Generated Timelines")
             if final_video_paths:
                 sorted_timeline_ids = sorted(final_video_paths.keys())
-                cols = st.columns(len(sorted_timeline_ids)) # Create columns for side-by-side display
+                # Adjust column count based on number of videos, max 3-4 wide?
+                num_cols = min(len(sorted_timeline_ids), 3)
+                cols = st.columns(num_cols)
 
                 for idx, timeline_id in enumerate(sorted_timeline_ids):
+                    col = cols[idx % num_cols] # Cycle through columns
                     video_path = final_video_paths[timeline_id]
-                    # 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"
+                                mime="video/mp4",
+                                key=f"download_btn_{timeline_id}" # Unique key for download button
                             )
                             # Display errors for this timeline if any occurred
                             if generation_errors.get(timeline_id):
@@ -736,15 +738,19 @@ if generate_button:
                              st.error(f"Error displaying video for Timeline {timeline_id}: {e}", icon="🚨")
             else:
                 st.warning("No final videos were successfully generated in this run.")
-                # Display global errors if no videos were made
+                # Display summary of all 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}")
+                    with st.expander("View All Errors", expanded=True):
+                        for tid, errors in generation_errors.items():
+                            if errors:
+                                st.error(f"Timeline {tid}:")
+                                for msg in errors:
+                                    st.error(f"  - {msg}")
 
             # --- 4. Cleanup ---
-            st.info(f"Cleaning up temporary directory: {temp_dir}")
+            st.info(f"Attempting to clean up temporary directory: {temp_dir}")
             try:
                 shutil.rmtree(temp_dir)
                 logger.info(f"✅ Temporary directory removed: {temp_dir}")
@@ -754,10 +760,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
+            # Error message likely already shown by generate_story_sequence_chrono
             logger.error("Story generation failed, cannot proceed.")
         else:
-            # This case implies chrono_response exists but somehow failed validation logic (should be caught earlier)
+            # Fallback for unexpected state
             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.")