Update app.py

app.py

@@ -99,14 +99,10 @@ try:
 
     # 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()
@@ -119,7 +115,9 @@ except Exception as e:
 # --- 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.")
+    # Increased max_length for image_prompt
+    image_prompt: str = Field(..., min_length=10, max_length=250, # <-- Increased from 150
+                              description="Concise visual description for image generation (target 15-35 words). Focus on non-human characters, setting, action, style.")
     audio_text: str = Field(..., min_length=5, max_length=150, description="Single sentence of narration/dialogue for the scene (max 30 words).")
     character_description: str = Field(..., max_length=100, description="Brief description of key non-human characters/objects in *this* scene's prompt for consistency.")
     timeline_visual_modifier: Optional[str] = Field(None, max_length=50, description="Optional subtle visual style hint (e.g., 'slightly darker', 'more vibrant colors').")
@@ -142,6 +140,7 @@ class ChronoWeaveResponse(BaseModel):
 
     @validator('timelines')
     def check_timeline_segment_count(cls, timelines, values):
+        # Pydantic v1 style validation. For v2, use model_validator(mode='before') or similar
         if 'total_scenes_per_timeline' in values:
             expected_scenes = values['total_scenes_per_timeline']
             for i, timeline in enumerate(timelines):
@@ -188,7 +187,7 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
             "audio_config": {
                 "audio_encoding": "LINEAR16", # Required format for WAV output
                 "sample_rate_hertz": AUDIO_SAMPLING_RATE,
-                # "voice": voice if voice else "aura-asteria-en" # Optional: Specify voice if needed and available
+                # "voice": voice if voice else "aura-asteria-en" # Optional
             }
         }
 
@@ -206,18 +205,17 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
             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 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
+                     return None
 
         if not collected_audio:
             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.
+        # Write the collected audio bytes into a WAV file.
         with wave_file_writer(output_filename, rate=AUDIO_SAMPLING_RATE) as wf:
             wf.writeframes(bytes(collected_audio))
         logger.info(f"   ✅ [{task_id}] Audio saved: {os.path.basename(output_filename)} ({len(collected_audio)} bytes)")
@@ -228,7 +226,6 @@ 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
@@ -253,6 +250,7 @@ def generate_story_sequence_chrono(
         f"Clearly state the divergence reason for each timeline (except potentially the first)."
     )
 
+    # Updated prompt with stricter image_prompt length guidance
     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}".
@@ -264,7 +262,7 @@ def generate_story_sequence_chrono(
     4.  {divergence_instruction}
     5.  Maintain a consistent visual style across all scenes and timelines: **'Simple, friendly kids animation style with bright colors and rounded shapes'**, unless a `timeline_visual_modifier` subtly alters it.
     6.  Each scene's narration (`audio_text`) should be a single, concise sentence (approx. 5-10 seconds spoken length, max 30 words).
-    7.  Image prompts (`image_prompt`) should be descriptive (15-35 words), focusing on the non-human character(s), setting, action, and visual style. Explicitly mention the main character(s) for consistency.
+    7.  Image prompts (`image_prompt`) should be descriptive **and concise (target 15-35 words MAXIMUM)**, focusing only on the non-human character(s), setting, action, and essential visual style elements for *this specific scene*. Explicitly mention the main character(s) for consistency. **Do NOT repeat the general 'Simple, friendly kids animation style...' description in every image prompt unless it is essential for a specific visual change**; rely on the overall style instruction and the optional `timeline_visual_modifier`.
     8.  `character_description` should briefly describe recurring non-human characters mentioned *in the specific scene's image prompt* (name, key visual features). Keep consistent within a timeline.
 
     **Output Format:**
@@ -274,7 +272,7 @@ def generate_story_sequence_chrono(
     ```json
     {json.dumps(ChronoWeaveResponse.schema(), indent=2)}
     ```
-    """ # Using .schema() which is the Pydantic v1 way, adjust if using v2 (.model_json_schema())
+    """ # Using .schema() for Pydantic v1. Use .model_json_schema() for v2.
 
     try:
         # Use the standard client (GenerativeModel instance) for text generation
@@ -282,16 +280,12 @@ def generate_story_sequence_chrono(
             contents=prompt,
             generation_config=genai.types.GenerationConfig(
                 response_mime_type="application/json",
-                temperature=0.7 # Add some creativity
+                temperature=0.7
             )
         )
 
-        # Debugging: Log raw response
-        # logger.debug(f"Raw Gemini Response Text:\n{response.text}")
-
         # 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}")
@@ -302,14 +296,11 @@ def generate_story_sequence_chrono(
         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
+            # Use parse_obj for Pydantic v1, or 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!")
@@ -318,7 +309,7 @@ def generate_story_sequence_chrono(
             logger.error(f"JSON structure validation failed: {val_err}")
             logger.error(f"Received Data:\n{json.dumps(raw_data, indent=2)}")
             st.error(f"🚨 The generated story structure is invalid: {val_err}", icon="🧬")
-            st.json(raw_data) # Show the invalid structure
+            st.json(raw_data) # Show the invalid structure that failed validation
             return None
 
     except genai.types.generation_types.BlockedPromptException as bpe:
@@ -328,8 +319,6 @@ def generate_story_sequence_chrono(
     except Exception as e:
         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
 
 
@@ -340,76 +329,80 @@ def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1", task_id: str =
     """
     logger.info(f"🖼️ [{task_id}] Requesting image for: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
 
-    # Refined prompt incorporating negative constraints and style guidance
+    # Refined prompt - relies on the story generator to provide concise prompts now
+    # Still includes base style and negative constraints as reinforcement
     full_prompt = (
         f"Generate an image in a child-friendly, simple animation style with bright colors and rounded shapes. "
         f"Ensure absolutely NO humans or human-like figures are present. Focus on animals or objects. "
-        f"Aspect ratio should be {aspect_ratio}. " # Explicitly state aspect ratio in prompt too
-        f"Prompt: {prompt}"
+        f"Aspect ratio should be {aspect_ratio}. "
+        f"Scene Description: {prompt}" # Use the potentially shorter prompt from story gen
     )
 
     try:
-        # Use the standard client's generate_content method.
         response = client_standard.generate_content(
             full_prompt,
              generation_config=genai.types.GenerationConfig(
                  candidate_count=1,
-                 # Add other config like temperature if desired
              ),
-             # 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
-        # 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
+        safety_ratings = []
+        block_reason = None
+
+        # Check response structure - adjust based on actual API behavior
+        if hasattr(response, 'candidates') and response.candidates:
+            candidate = response.candidates[0]
+            if hasattr(candidate, 'content') and candidate.content and hasattr(candidate.content, 'parts') and candidate.content.parts:
+                part = candidate.content.parts[0]
+                if hasattr(part, 'inline_data') and part.inline_data and hasattr(part.inline_data, 'data'):
+                    image_bytes = part.inline_data.data
+            if hasattr(candidate, 'safety_ratings'):
+                safety_ratings = candidate.safety_ratings
+            # Finish reason might also indicate issues (e.g., SAFETY)
+            # if hasattr(candidate, 'finish_reason') and candidate.finish_reason != 'STOP': ...
+
+        # Check prompt feedback for blocking outside of candidates
+        if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
+            if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason != 'BLOCK_REASON_UNSPECIFIED':
+                 block_reason = response.prompt_feedback.block_reason.name # Get the name of the enum
+            if hasattr(response.prompt_feedback, 'safety_ratings'):
+                 # Combine prompt feedback ratings with candidate ratings if necessary
+                 safety_ratings.extend(response.prompt_feedback.safety_ratings)
+
 
         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="⚠️")
-
+                # Log safety flags if present
+                filtered_ratings = [f"{r.category.name}: {r.probability.name}" for r in safety_ratings if hasattr(r,'probability') and r.probability.name != 'NEGLIGIBLE']
+                if filtered_ratings:
+                    logger.warning(f"   ⚠️ [{task_id}] Image generated but flagged by safety filters: {', '.join(filtered_ratings)}.")
+                    st.warning(f"Image for scene {task_id} flagged by safety filters: {', '.join(filtered_ratings)}", icon="⚠️")
                 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 blocking or other issues
-            block_reason = None
-            prompt_feedback = getattr(response, 'prompt_feedback', None)
-            if prompt_feedback:
-                block_reason = getattr(prompt_feedback, 'block_reason', None)
-
+            # If no image bytes, determine why
             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="🚫")
             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="🖼️")
-                 # Log the full response for debugging
-                 # logger.debug(f"Full Imagen response object: {response}")
+                 # Check for safety flags even if no block reason explicitly given
+                 filtered_ratings = [f"{r.category.name}: {r.probability.name}" for r in safety_ratings if hasattr(r,'probability') and r.probability.name != 'NEGLIGIBLE']
+                 if filtered_ratings:
+                      logger.warning(f"   ⚠️ [{task_id}] Image generation failed, safety filters triggered: {', '.join(filtered_ratings)}. Prompt: '{prompt[:70]}...'")
+                      st.warning(f"Image generation failed for scene {task_id}, safety filters triggered: {', '.join(filtered_ratings)}", icon="⚠️")
+                 else:
+                      logger.warning(f"   ⚠️ [{task_id}] No image data received, unknown reason. Prompt: '{prompt[:70]}...'")
+                      st.warning(f"No image data received for scene {task_id}, reason unclear.", icon="🖼️")
+                      # logger.debug(f"Full Imagen response object: {response}")
             return None
 
     except genai.types.generation_types.BlockedPromptException as 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
@@ -426,7 +419,7 @@ st.sidebar.header("⚙️ Configuration")
 if GOOGLE_API_KEY:
     st.sidebar.success("Google API Key Loaded", icon="✅")
 else:
-    st.sidebar.error("Google API Key Missing!", icon="🚨") # Should not be reached if st.stop() works
+    st.sidebar.error("Google API Key Missing!", icon="🚨")
 
 # Story Parameters
 theme = st.sidebar.text_input("📖 Story Theme:", "A curious squirrel finds a mysterious, glowing acorn")
@@ -437,9 +430,6 @@ divergence_prompt = st.sidebar.text_input("↔️ Divergence Hint (Optional):",
 # 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
 
 generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY), use_container_width=True)
@@ -473,35 +463,33 @@ if generate_button:
             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.")
+            # Structure generated and validated successfully by the function
+            # st.success(...) is now inside generate_story_sequence_chrono on success
 
             # --- 2. Process Each Timeline ---
             overall_start_time = time.time()
             all_timelines_successful = True # Assume success initially
 
-            # Use st.status for collapsible progress updates
             with st.status("Generating assets and composing videos...", expanded=True) as status:
 
                 for timeline_index, timeline in enumerate(chrono_response.timelines):
                     timeline_id = timeline.timeline_id
                     divergence = timeline.divergence_reason
                     segments = timeline.segments
-                    timeline_label = f"Timeline {timeline_id}" # Consistent label
+                    timeline_label = f"Timeline {timeline_id}"
                     st.subheader(f"Processing {timeline_label}: {divergence}")
                     logger.info(f"--- Processing {timeline_label} (Index: {timeline_index}) ---")
-                    generation_errors[timeline_id] = [] # Initialize error list
+                    generation_errors[timeline_id] = []
 
-                    temp_image_files = {} # {scene_id: path}
-                    temp_audio_files = {} # {scene_id: path}
-                    video_clips = [] # List of moviepy clips
+                    temp_image_files = {}
+                    temp_audio_files = {}
+                    video_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
+                        task_id = f"T{timeline_id}_S{scene_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})** ---")
@@ -509,16 +497,18 @@ if generate_button:
 
                         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}")
 
                         # --- 2a. Image Generation ---
-                        generated_image: Optional[Image.Image] = None # Define before spinner
+                        generated_image: Optional[Image.Image] = None
                         with st.spinner(f"[{task_id}] Generating image... 🎨"):
-                            combined_prompt = f"{segment.image_prompt}. {segment.character_description}"
+                            # Combine prompt using the (hopefully shorter) prompt from the structure
+                            combined_prompt = segment.image_prompt # Use directly
+                            if segment.character_description: # Add character desc if present
+                                combined_prompt += f" Featuring: {segment.character_description}"
                             if segment.timeline_visual_modifier:
-                                combined_prompt += f" Visual style hint: {segment.timeline_visual_modifier}."
+                                combined_prompt += f" Style hint: {segment.timeline_visual_modifier}."
                             generated_image = generate_image_imagen(combined_prompt, aspect_ratio, task_id)
 
                         if generated_image:
@@ -536,7 +526,7 @@ 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.")
-                            continue # Skip audio/video for this scene
+                            continue
 
                         # --- 2b. Audio Generation ---
                         generated_audio_path: Optional[str] = None
@@ -544,7 +534,6 @@ if generate_button:
                             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()
                                     generated_audio_path = asyncio.run(
                                         generate_audio_live_async(segment.audio_text, audio_path_temp, audio_voice)
                                     )
@@ -570,7 +559,6 @@ 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 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])
@@ -578,16 +566,16 @@ if generate_button:
                                         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
+                                continue
 
                         # --- 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
+                            audio_clip_instance = None
+                            image_clip_instance = None
+                            composite_clip = None
                             try:
                                 if not os.path.exists(img_path): raise FileNotFoundError(f"Image file not found: {img_path}")
                                 if not os.path.exists(aud_path): raise FileNotFoundError(f"Audio file not found: {aud_path}")
@@ -597,21 +585,18 @@ if generate_button:
                                 image_clip_instance = ImageClip(np_image).set_duration(audio_clip_instance.duration)
 
                                 composite_clip = image_clip_instance.set_audio(audio_clip_instance)
-                                video_clips.append(composite_clip) # Add the clip to be concatenated later
+                                video_clips.append(composite_clip)
                                 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.")
-                                # 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)
 
@@ -619,23 +604,20 @@ if generate_button:
 
                     # --- 2d. Assemble Timeline Video ---
                     timeline_duration = time.time() - timeline_start_time
-                    # 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
+                        final_timeline_video = None
                         try:
-                            # Concatenate the collected clips
                             final_timeline_video = concatenate_videoclips(video_clips, method="compose")
                             final_timeline_video.write_videofile(
                                 output_filename,
                                 fps=VIDEO_FPS,
                                 codec=VIDEO_CODEC,
                                 audio_codec=AUDIO_CODEC,
-                                logger=None # Suppress moviepy console spam
+                                logger=None
                             )
                             final_video_paths[timeline_id] = output_filename
                             logger.info(f"   ✅ [{timeline_label}] Final video saved: {os.path.basename(output_filename)}")
@@ -647,11 +629,10 @@ if generate_button:
                             all_timelines_successful = False
                             generation_errors[timeline_id].append(f"Timeline {timeline_id}: Final video assembly failed.")
                         finally:
-                            # 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:
-                                    if clip: # Check if clip object exists
+                                    if clip:
                                       if clip.audio: clip.audio.close()
                                       clip.close()
                                 except Exception as e_close:
@@ -668,29 +649,27 @@ if generate_button:
                         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, so scene_success_count < len(segments)
+                    else: # Some scenes failed
                         error_count = len(segments) - scene_success_count
                         logger.warning(f"[{timeline_label}] Encountered errors in {error_count} scene(s). Skipping final video assembly.")
                         st.warning(f"{timeline_label} had errors in {error_count} scene(s). Final video not assembled.", icon="⚠️")
                         all_timelines_successful = False
 
-                    # Log accumulated errors for the timeline if any occurred
                     if generation_errors[timeline_id]:
                          logger.error(f"Summary of errors in {timeline_label}: {generation_errors[timeline_id]}")
 
                 # --- End of Timelines Loop ---
 
-                # Final status update
                 overall_duration = time.time() - overall_start_time
                 if all_timelines_successful and final_video_paths:
                     status_msg = f"ChronoWeave Generation Complete! ({len(final_video_paths)} videos in {overall_duration:.2f}s)"
                     status.update(label=status_msg, state="complete", expanded=False)
                     logger.info(status_msg)
-                elif final_video_paths: # Some videos made, but errors occurred
+                elif final_video_paths:
                     status_msg = f"ChronoWeave Partially Complete ({len(final_video_paths)} videos, some errors occurred). Total time: {overall_duration:.2f}s"
                     status.update(label=status_msg, state="warning", expanded=True)
                     logger.warning(status_msg)
-                else: # No videos made
+                else:
                     status_msg = f"ChronoWeave Generation Failed. No videos produced. Total time: {overall_duration:.2f}s"
                     status.update(label=status_msg, state="error", expanded=True)
                     logger.error(status_msg)
@@ -699,12 +678,11 @@ if generate_button:
             st.header("🎬 Generated Timelines")
             if final_video_paths:
                 sorted_timeline_ids = sorted(final_video_paths.keys())
-                # 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
+                    col = cols[idx % num_cols]
                     video_path = final_video_paths[timeline_id]
                     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"
@@ -722,9 +700,8 @@ if generate_button:
                                 data=video_bytes,
                                 file_name=f"chronoweave_timeline_{timeline_id}.mp4",
                                 mime="video/mp4",
-                                key=f"download_btn_{timeline_id}" # Unique key for download button
+                                key=f"download_btn_{timeline_id}"
                             )
-                            # 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]:
@@ -738,7 +715,6 @@ 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 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")
@@ -760,10 +736,9 @@ 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 likely already shown by generate_story_sequence_chrono
-            logger.error("Story generation failed, cannot proceed.")
+            # Error message already shown by generate_story_sequence_chrono or validation
+            logger.error("Story generation or validation failed, cannot proceed.")
         else:
-            # 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.")