Update app.py

app.py

@@ -104,35 +104,31 @@ except Exception as e:
 
 # --- Define Pydantic Schemas (Using V2 Syntax) ---
 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=250, # Keep increased limit
-                              description="Concise visual description (target 15-35 words). Focus on non-human characters, setting, action, style.")
-    audio_text: str = Field(..., min_length=5, max_length=150, description="Single sentence of narration/dialogue (max 30 words).")
-    # Increased max_length for character_description
-    character_description: str = Field(..., max_length=250, # <-- Increased from 100
-                                      description="Brief description of key non-human characters/objects in *this* scene's prompt (target under 20 words).")
-    timeline_visual_modifier: Optional[str] = Field(None, max_length=50, description="Optional subtle visual style hint.")
-
-    # Pydantic v2 style field validator
+    scene_id: int = Field(..., ge=0)
+    image_prompt: str = Field(..., min_length=10, max_length=250) # Keep increased limits
+    audio_text: str = Field(..., min_length=5, max_length=150)
+    character_description: str = Field(..., max_length=250) # Keep increased limits
+    timeline_visual_modifier: Optional[str] = Field(None, max_length=50)
+
     @field_validator('image_prompt')
     @classmethod
     def image_prompt_no_humans(cls, v: str) -> str:
         if any(word in v.lower() for word in ["person", "people", "human", "man", "woman", "boy", "girl", "child"]):
-            logger.warning(f"Image prompt '{v[:50]}...' may contain human descriptions. Relying on API-level controls & prompt instructions.")
+            logger.warning(f"Image prompt '{v[:50]}...' may contain human descriptions.")
         return v
 
 class Timeline(BaseModel):
-    timeline_id: int = Field(..., ge=0, description="Unique identifier for this timeline.")
+    timeline_id: int = Field(..., ge=0)
+    # Keep min_length=5 for divergence_reason, rely on improved prompt
     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.")
+    segments: List[StorySegment] = Field(..., min_items=1)
 
 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.")
+    core_theme: str = Field(..., min_length=5)
+    timelines: List[Timeline] = Field(..., min_items=1)
+    total_scenes_per_timeline: int = Field(..., gt=0)
 
-    # Pydantic v2 style model validator
-    @model_validator(mode='after') # Use 'after' to access validated fields
+    @model_validator(mode='after')
     def check_timeline_segment_count(self) -> 'ChronoWeaveResponse':
         expected_scenes = self.total_scenes_per_timeline
         for i, timeline in enumerate(self.timelines):
@@ -149,7 +145,7 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
     try:
         wf = wave.open(filename, "wb")
         wf.setnchannels(channels)
-        wf.setsampwidth(sample_width) # 2 bytes for 16-bit audio
+        wf.setsampwidth(sample_width)
         wf.setframerate(rate)
         yield wf
     except Exception as e:
@@ -157,10 +153,8 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
         raise
     finally:
         if wf:
-            try:
-                wf.close()
-            except Exception as e_close:
-                logger.error(f"Error closing wave file {filename}: {e_close}")
+            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]:
@@ -170,37 +164,24 @@ 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:
-        config = {
-            "response_modalities": ["AUDIO"],
-            "audio_config": {
-                "audio_encoding": "LINEAR16",
-                "sample_rate_hertz": AUDIO_SAMPLING_RATE,
-            }
-        }
-        directive_prompt = (
-            "Narrate the following sentence directly and engagingly. "
-            "Do not add any introductory or concluding remarks. "
-            "Speak only the sentence itself:\n\n"
-            f'"{api_text}"'
-        )
+        config = {"response_modalities": ["AUDIO"], "audio_config": {"audio_encoding": "LINEAR16", "sample_rate_hertz": AUDIO_SAMPLING_RATE}}
+        directive_prompt = f"Narrate the following sentence directly and engagingly. Do not add any introductory or concluding remarks. Speak only the sentence itself:\n\n\"{api_text}\""
 
         async with live_model.connect(config=config) as session:
             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)
+                if response.audio_chunk and response.audio_chunk.data: collected_audio.extend(response.audio_chunk.data)
                 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
 
         if not collected_audio:
-            logger.warning(f"⚠️ [{task_id}] No audio data received for: '{api_text[:60]}...'")
+            logger.warning(f"⚠️ [{task_id}] No audio data received.")
             st.warning(f"No audio data generated for scene {task_id}.", icon="🔊")
             return None
 
-        with wave_file_writer(output_filename, rate=AUDIO_SAMPLING_RATE) as wf:
-            wf.writeframes(bytes(collected_audio))
+        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
 
@@ -224,13 +205,13 @@ def generate_story_sequence_chrono(
     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}")
 
+    # Updated divergence instruction to guide the first timeline's reason
     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)."
+        f"Clearly state the divergence reason for each timeline. **For the first timeline (timeline_id 0), use a descriptive reason like 'Initial path' or 'Baseline scenario' that is at least 5 characters long.**"
     )
 
-    # Updated prompt with stricter guidance on description lengths
     prompt = f"""
     Act as an expert narrative designer specializing in short, visual, branching stories for children.
     Create a story based on the core theme: "{theme}".
@@ -240,10 +221,10 @@ def generate_story_sequence_chrono(
     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 **and concise (target 15-35 words MAXIMUM)**, focusing only on the non-human character(s), setting, action, and essential visual style elements for *this specific scene*. Explicitly mention the main character(s) for consistency. **Do NOT repeat the general style description** unless essential for a specific visual change.
-    8.  `character_description` should **very briefly** describe recurring non-human characters mentioned *in the specific scene's image prompt* (name, key visual features). **Keep descriptions extremely concise (e.g., "Nutsy: fluffy squirrel"). Aim for under 20 words total per scene.** Keep consistent within a timeline.
+    5.  Maintain a consistent visual style: **'Simple, friendly kids animation style with bright colors and rounded shapes'**, unless a `timeline_visual_modifier` subtly alters it.
+    6.  `audio_text` should be a single, concise sentence (max 30 words).
+    7.  `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*. **Do NOT repeat the general style description** unless essential.
+    8.  `character_description` should **very briefly** describe recurring non-human characters mentioned *in the scene's image prompt* (name, key features). **Keep descriptions extremely concise (target under 20 words total per scene).**
 
     **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.
@@ -252,7 +233,7 @@ def generate_story_sequence_chrono(
     ```json
     {json.dumps(ChronoWeaveResponse.model_json_schema(), indent=2)}
     ```
-    """ # Using .model_json_schema() for Pydantic v2.
+    """
 
     try:
         response = client_standard.generate_content(
@@ -266,32 +247,31 @@ def generate_story_sequence_chrono(
         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 story structure from AI: {json_err}", icon="📄")
+            logger.error(f"Failed to decode JSON: {json_err}")
+            logger.error(f"Response Text:\n{response.text}")
+            st.error(f"🚨 Failed to parse story structure: {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}")
+             logger.error(f"Error processing response text: {e}")
              st.error(f"🚨 Error processing AI response: {e}", icon="📄")
              return None
 
         try:
-            # Use model_validate for Pydantic v2
             validated_data = ChronoWeaveResponse.model_validate(raw_data)
             logger.info("✅ Story structure generated and validated successfully!")
             st.success("✅ Story structure generated and validated!")
             return validated_data
         except ValidationError as val_err:
-            logger.error(f"JSON structure validation failed: {val_err}")
+            logger.error(f"JSON 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.error(f"🚨 Generated story structure invalid: {val_err}", icon="🧬")
             st.json(raw_data)
             return None
 
     except genai.types.generation_types.BlockedPromptException as bpe:
          logger.error(f"Story generation prompt blocked: {bpe}")
-         st.error("🚨 Story generation prompt blocked (safety filters). Try rephrasing.", icon="🚫")
+         st.error("🚨 Story generation prompt blocked (safety filters).", icon="🚫")
          return None
     except Exception as e:
         logger.exception("Error during story sequence generation:")
@@ -301,113 +281,88 @@ def generate_story_sequence_chrono(
 
 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."""
-    logger.info(f"🖼️ [{task_id}] Requesting image for: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
+    logger.info(f"🖼️ [{task_id}] Requesting image: '{prompt[:70]}...' (Aspect: {aspect_ratio})")
 
     full_prompt = (
         f"Generate an image in a child-friendly, simple animation style with bright colors and rounded shapes. "
-        f"Ensure absolutely NO humans or human-like figures are present. Focus on animals or objects. "
-        f"Aspect ratio should be {aspect_ratio}. "
-        f"Scene Description: {prompt}"
+        f"Ensure absolutely NO humans or human-like figures. Focus on animals or objects. "
+        f"Aspect ratio {aspect_ratio}. Scene: {prompt}"
     )
 
     try:
         response = client_standard.generate_content(
-            full_prompt,
-             generation_config=genai.types.GenerationConfig(
-                 candidate_count=1,
-             ),
+            full_prompt, generation_config=genai.types.GenerationConfig(candidate_count=1)
         )
 
-        image_bytes = None
-        safety_ratings = []
-        block_reason = None
-        finish_reason = None # Store finish reason if available
+        image_bytes, safety_ratings, block_reason, finish_reason = None, [], None, None
 
-        # Consolidate response checking
         if hasattr(response, 'candidates') and response.candidates:
             candidate = response.candidates[0]
-            if hasattr(candidate, 'finish_reason'):
-                finish_reason = candidate.finish_reason.name if hasattr(candidate.finish_reason, 'name') else str(candidate.finish_reason)
+            if hasattr(candidate, 'finish_reason'): finish_reason = getattr(candidate.finish_reason, 'name', str(candidate.finish_reason))
             if hasattr(candidate, 'content') and candidate.content and hasattr(candidate.content, 'parts') and candidate.content.parts:
                 part = candidate.content.parts[0]
-                if hasattr(part, 'inline_data') and part.inline_data and hasattr(part.inline_data, 'data'):
-                    image_bytes = part.inline_data.data
-            if hasattr(candidate, 'safety_ratings'):
-                safety_ratings = candidate.safety_ratings
+                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
 
         if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
-            if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason.name != 'BLOCK_REASON_UNSPECIFIED':
-                 block_reason = response.prompt_feedback.block_reason.name
-            if hasattr(response.prompt_feedback, 'safety_ratings'):
-                 safety_ratings.extend(response.prompt_feedback.safety_ratings)
-
+            if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason.name != 'BLOCK_REASON_UNSPECIFIED': block_reason = response.prompt_feedback.block_reason.name
+            if hasattr(response.prompt_feedback, 'safety_ratings'): 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.")
+                logger.info(f"   ✅ [{task_id}] Image generated.")
                 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 flagged by safety filters: {', '.join(filtered_ratings)}.")
-                    st.warning(f"Image for scene {task_id} flagged: {', '.join(filtered_ratings)}", icon="⚠️")
+                    logger.warning(f"   ⚠️ [{task_id}] Image flagged: {', '.join(filtered_ratings)}.")
+                    st.warning(f"Image {task_id} flagged: {', '.join(filtered_ratings)}", icon="⚠️")
                 return image
             except Exception as img_err:
-                logger.error(f"   ❌ [{task_id}] Failed to decode generated image data: {img_err}")
-                st.warning(f"Failed to decode image data for scene {task_id}.", icon="🖼️")
+                logger.error(f"   ❌ [{task_id}] Failed to decode image data: {img_err}")
+                st.warning(f"Failed decode image data {task_id}.", icon="🖼️")
                 return None
         else:
-            # Enhanced logging for failure reasons
             fail_reason = "Unknown reason."
-            if block_reason:
-                fail_reason = f"Blocked (Reason: {block_reason})."
-            elif finish_reason and finish_reason not in ['STOP', 'FINISH_REASON_UNSPECIFIED']:
-                 fail_reason = f"Finished early (Reason: {finish_reason})."
+            if block_reason: fail_reason = f"Blocked (Reason: {block_reason})."
+            elif finish_reason and finish_reason not in ['STOP', 'FINISH_REASON_UNSPECIFIED']: fail_reason = f"Finished early (Reason: {finish_reason})."
             else:
-                 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:
-                     fail_reason = f"Safety filters triggered: {', '.join(filtered_ratings)}."
-
-            logger.warning(f"   ⚠️ [{task_id}] No image data received. Reason: {fail_reason} Prompt: '{prompt[:70]}...'")
-            st.warning(f"No image data received for scene {task_id}. Reason: {fail_reason}", icon="🖼️")
-            # Uncomment to log full response only on unknown failures
-            # if fail_reason == "Unknown reason.":
-            #      logger.debug(f"   ⚠️ [{task_id}] Full API response object: {response}")
+                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: fail_reason = f"Safety filters triggered: {', '.join(filtered_ratings)}."
+
+            # Add the full response logging here for persistent unknown failures
+            if fail_reason == "Unknown reason.":
+                 logger.warning(f"   ⚠️ [{task_id}] Full API response object: {response}")
+
+            logger.warning(f"   ⚠️ [{task_id}] No image data. Reason: {fail_reason} Prompt: '{prompt[:70]}...'")
+            st.warning(f"No image data for {task_id}. Reason: {fail_reason}", icon="🖼️")
             return None
 
     except genai.types.generation_types.BlockedPromptException as bpe:
          logger.error(f"   ❌ [{task_id}] Image generation blocked (exception): {bpe}")
-         st.error(f"Image generation blocked for scene {task_id} (exception).", icon="🚫")
+         st.error(f"Image generation blocked for {task_id} (exception).", icon="🚫")
          return None
     except Exception as e:
         logger.exception(f"   ❌ [{task_id}] Image generation failed unexpectedly: {e}")
-        st.error(f"Image generation failed for scene {task_id}: {e}", icon="🖼️")
+        st.error(f"Image generation failed for {task_id}: {e}", icon="🖼️")
         return None
 
-
 # --- Streamlit UI Elements ---
+# (Identical to previous version - No changes needed here)
 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="🚨")
-
+if GOOGLE_API_KEY: st.sidebar.success("Google API Key Loaded", icon="✅")
+else: st.sidebar.error("Google API Key Missing!", icon="🚨")
 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)
 num_timelines = st.sidebar.slider("🌿 Number of Timelines:", min_value=1, max_value=4, value=2)
 divergence_prompt = st.sidebar.text_input("↔️ Divergence Hint (Optional):", placeholder="e.g., What if a bird tried to steal it?")
-
 st.sidebar.subheader("🎨 Visual & Audio Settings")
 aspect_ratio = st.sidebar.selectbox("🖼️ Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0)
 audio_voice = None
-
 generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY), use_container_width=True)
-
 st.sidebar.markdown("---")
 st.sidebar.info("⏳ Generation can take several minutes.", icon="⏳")
 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 ---
 if generate_button:
     if not theme:
@@ -467,10 +422,8 @@ if generate_button:
                         generated_image: Optional[Image.Image] = None
                         with st.spinner(f"[{task_id}] Generating image... 🎨"):
                             combined_prompt = segment.image_prompt
-                            if segment.character_description:
-                                combined_prompt += f" Featuring: {segment.character_description}"
-                            if segment.timeline_visual_modifier:
-                                combined_prompt += f" Style hint: {segment.timeline_visual_modifier}."
+                            if segment.character_description: combined_prompt += f" Featuring: {segment.character_description}"
+                            if segment.timeline_visual_modifier: combined_prompt += f" Style hint: {segment.timeline_visual_modifier}."
                             generated_image = generate_image_imagen(combined_prompt, aspect_ratio, task_id)
 
                         if generated_image:
@@ -481,85 +434,73 @@ if generate_button:
                                 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: {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.")
+                                st.error(f"Failed save image {task_id}.", icon="💾")
+                                scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Img save fail.")
                         else:
-                            # Error/warning already logged by generate_image_imagen
-                            scene_has_error = True
-                            generation_errors[timeline_id].append(f"Scene {scene_id+1}: Image generation failed.")
-                            continue
+                            scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Img gen fail.")
+                            continue # Skip rest of scene processing
 
                         # --- 2b. Audio Generation ---
                         generated_audio_path: Optional[str] = None
-                        if not scene_has_error:
+                        if not scene_has_error: # Only proceed 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:
-                                    generated_audio_path = asyncio.run(
-                                        generate_audio_live_async(segment.audio_text, audio_path_temp, audio_voice)
-                                    )
+                                    generated_audio_path = asyncio.run(generate_audio_live_async(segment.audio_text, audio_path_temp, audio_voice))
                                 except RuntimeError as e:
-                                    logger.error(f"   ❌ [{task_id}] Asyncio runtime error: {e}")
-                                    st.error(f"Asyncio error during audio gen for {task_id}: {e}", icon="⚡")
-                                    scene_has_error = True
-                                    generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio async error.")
+                                    logger.error(f"   ❌ [{task_id}] Asyncio error: {e}")
+                                    st.error(f"Asyncio audio error {task_id}: {e}", icon="⚡")
+                                    scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Audio async err.")
                                 except Exception as e:
                                     logger.exception(f"   ❌ [{task_id}] Unexpected audio error: {e}")
-                                    st.error(f"Unexpected audio error for {task_id}: {e}", icon="💥")
-                                    scene_has_error = True
-                                    generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation error.")
+                                    st.error(f"Unexpected audio error {task_id}: {e}", icon="💥")
+                                    scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Audio gen err.")
 
                             if generated_audio_path:
                                 temp_audio_files[scene_id] = generated_audio_path
                                 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:
-                                # Error/warning already logged by generate_audio_live_async
-                                scene_has_error = True
-                                generation_errors[timeline_id].append(f"Scene {scene_id+1}: Audio generation failed.")
+                                    with open(generated_audio_path, 'rb') as ap: st.audio(ap.read(), format='audio/wav')
+                                except Exception as e: logger.warning(f"   ⚠️ [{task_id}] Audio preview error: {e}")
+                            else: # Audio generation failed
+                                scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Audio gen fail.")
+                                # Clean up corresponding image file
                                 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}] Failed to remove image after audio failure: {e}")
-                                continue
+                                    try: os.remove(temp_image_files[scene_id]); logger.info(f"   🗑️ [{task_id}] Removed image due to audio fail."); del temp_image_files[scene_id]
+                                    except OSError as e: logger.warning(f"   ⚠️ [{task_id}] Failed remove image after audio fail: {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}...")
-                            img_path = temp_image_files[scene_id]
-                            aud_path = temp_audio_files[scene_id]
+                            st.write(f"   🎬 Creating video clip S{scene_id+1}...")
+                            img_path, aud_path = temp_image_files[scene_id], temp_audio_files[scene_id]
                             audio_clip_instance, image_clip_instance, composite_clip = None, None, None
                             try:
-                                if not os.path.exists(img_path): raise FileNotFoundError(f"Image file missing: {img_path}")
-                                if not os.path.exists(aud_path): raise FileNotFoundError(f"Audio file missing: {aud_path}")
-
+                                if not os.path.exists(img_path): raise FileNotFoundError(f"Img missing: {img_path}")
+                                if not os.path.exists(aud_path): raise FileNotFoundError(f"Aud missing: {aud_path}")
                                 audio_clip_instance = AudioFileClip(aud_path)
                                 np_image = np.array(Image.open(img_path))
                                 image_clip_instance = ImageClip(np_image).set_duration(audio_clip_instance.duration)
                                 composite_clip = image_clip_instance.set_audio(audio_clip_instance)
                                 video_clips.append(composite_clip)
-                                logger.info(f"      ✅ [{task_id}] Video clip created (Duration: {audio_clip_instance.duration:.2f}s).")
-                                st.write(f"      ✅ Clip created (Duration: {audio_clip_instance.duration:.2f}s).")
+                                logger.info(f"      ✅ [{task_id}] Clip created (Dur: {audio_clip_instance.duration:.2f}s).")
+                                st.write(f"      ✅ Clip created (Dur: {audio_clip_instance.duration:.2f}s).")
                                 scene_success_count += 1
-
                             except Exception as e:
-                                logger.exception(f"      ❌ [{task_id}] Failed to create video clip: {e}")
-                                st.error(f"Failed to create video clip for {task_id}: {e}", icon="🎬")
-                                scene_has_error = True
-                                generation_errors[timeline_id].append(f"Scene {scene_id+1}: Video clip creation failed.")
+                                logger.exception(f"      ❌ [{task_id}] Failed clip creation: {e}")
+                                st.error(f"Failed create clip {task_id}: {e}", icon="🎬")
+                                scene_has_error = True; generation_errors[timeline_id].append(f"S{scene_id+1}: Clip creation fail.")
+                                # Cleanup resources from failed clip attempt
                                 if audio_clip_instance: audio_clip_instance.close()
                                 if image_clip_instance: image_clip_instance.close()
                                 try:
                                      if os.path.exists(img_path): os.remove(img_path)
                                      if os.path.exists(aud_path): os.remove(aud_path)
-                                except OSError as e_rem: logger.warning(f"   ⚠️ [{task_id}] Failed to remove files after clip error: {e_rem}")
+                                except OSError as e_rem: logger.warning(f"   ⚠️ [{task_id}] Failed remove files after clip error: {e_rem}")
 
                     # --- 2d. Assemble Timeline Video ---
                     timeline_duration = time.time() - timeline_start_time
                     if video_clips and scene_success_count == len(segments):
+                        # ... (Video assembly logic - same as before) ...
                         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)...")
@@ -573,11 +514,10 @@ if generate_button:
                             st.success(f"✅ Video for {timeline_label} completed in {timeline_duration:.2f}s.")
                         except Exception as e:
                             logger.exception(f"   ❌ [{timeline_label}] Failed to write final video: {e}")
-                            st.error(f"Failed to assemble video for {timeline_label}: {e}", icon="📼")
-                            all_timelines_successful = False
-                            generation_errors[timeline_id].append(f"Timeline {timeline_id}: Final video assembly failed.")
+                            st.error(f"Failed assemble video {timeline_label}: {e}", icon="📼")
+                            all_timelines_successful = False; generation_errors[timeline_id].append(f"Timeline {timeline_id}: Video assembly failed.")
                         finally:
-                            logger.debug(f"[{timeline_label}] Closing {len(video_clips)} source clips and final video...")
+                            logger.debug(f"[{timeline_label}] Closing clips...")
                             for i, clip in enumerate(video_clips):
                                 try:
                                     if clip:
@@ -589,88 +529,70 @@ if generate_button:
                                      if final_timeline_video.audio: final_timeline_video.audio.close()
                                      final_timeline_video.close()
                                 except Exception as e_close_final: logger.warning(f"   ⚠️ [{timeline_label}] Error closing final video object: {e_close_final}")
-
                     elif not video_clips:
-                        logger.warning(f"[{timeline_label}] No video clips generated. Skipping final assembly.")
-                        st.warning(f"No scenes processed for {timeline_label}. Video cannot be created.", icon="🚫")
-                        all_timelines_successful = False
+                         logger.warning(f"[{timeline_label}] No clips generated. Skipping assembly.")
+                         st.warning(f"No scenes processed for {timeline_label}. Cannot create video.", icon="🚫")
+                         all_timelines_successful = False
                     else: # Some scenes failed
                         error_count = len(segments) - scene_success_count
-                        logger.warning(f"[{timeline_label}] Errors in {error_count} scene(s). Skipping final video assembly.")
-                        st.warning(f"{timeline_label} had errors in {error_count} scene(s). Final video not assembled.", icon="⚠️")
+                        logger.warning(f"[{timeline_label}] Errors in {error_count} scene(s). Skipping assembly.")
+                        st.warning(f"{timeline_label} had errors in {error_count} scene(s). Video not assembled.", icon="⚠️")
                         all_timelines_successful = False
-
-                    if generation_errors[timeline_id]:
-                         logger.error(f"Summary of errors in {timeline_label}: {generation_errors[timeline_id]}")
+                    if generation_errors[timeline_id]: logger.error(f"Error summary {timeline_label}: {generation_errors[timeline_id]}")
 
                 # --- End of Timelines Loop ---
                 overall_duration = time.time() - overall_start_time
-                if all_timelines_successful and final_video_paths:
-                    status_msg = f"ChronoWeave Generation Complete! ({len(final_video_paths)} videos in {overall_duration:.2f}s)"
-                    status.update(label=status_msg, state="complete", expanded=False); logger.info(status_msg)
-                elif final_video_paths:
-                    status_msg = f"ChronoWeave Partially Complete ({len(final_video_paths)} videos, errors occurred). Time: {overall_duration:.2f}s"
-                    status.update(label=status_msg, state="warning", expanded=True); logger.warning(status_msg)
-                else:
-                    status_msg = f"ChronoWeave Generation Failed. No videos produced. Time: {overall_duration:.2f}s"
-                    status.update(label=status_msg, state="error", expanded=True); logger.error(status_msg)
+                # ... (Final status update logic - same as before) ...
+                if all_timelines_successful and final_video_paths: status_msg = f"ChronoWeave 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_msg = f"ChronoWeave Partially Complete ({len(final_video_paths)} videos, errors). Time: {overall_duration:.2f}s"; status.update(label=status_msg, state="warning", expanded=True); logger.warning(status_msg)
+                else: status_msg = f"ChronoWeave Failed. No videos. Time: {overall_duration:.2f}s"; status.update(label=status_msg, state="error", expanded=True); logger.error(status_msg)
 
             # --- 3. Display Results ---
             st.header("🎬 Generated Timelines")
             if final_video_paths:
+                # ... (Display logic - same as before) ...
                 sorted_timeline_ids = sorted(final_video_paths.keys())
                 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]
                     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"
-
+                    reason = timeline_data.divergence_reason if timeline_data else "Unknown"
                     with col:
-                        st.subheader(f"Timeline {timeline_id}")
-                        st.caption(f"Divergence: {reason}")
+                        st.subheader(f"Timeline {timeline_id}"); st.caption(f"Divergence: {reason}")
                         try:
-                            with open(video_path, 'rb') as video_file: video_bytes = video_file.read()
+                            with open(video_path, 'rb') as vf: video_bytes = vf.read()
                             st.video(video_bytes)
-                            logger.info(f"Displaying video for Timeline {timeline_id}")
-                            st.download_button(label=f"Download T{timeline_id} Video", data=video_bytes, file_name=f"chronoweave_timeline_{timeline_id}.mp4", mime="video/mp4", key=f"dl_{timeline_id}")
+                            logger.info(f"Displaying video T{timeline_id}")
+                            st.download_button(f"Download T{timeline_id}", video_bytes, f"timeline_{timeline_id}.mp4", "video/mp4", key=f"dl_{timeline_id}")
                             if generation_errors.get(timeline_id):
                                 with st.expander(f"⚠️ View {len(generation_errors[timeline_id])} Issues"):
-                                     for error_msg in generation_errors[timeline_id]: st.warning(f"- {error_msg}")
-                        except FileNotFoundError:
-                             logger.error(f"Video file not found for display: {video_path}")
-                             st.error(f"Error: Video file missing for T{timeline_id}.", icon="🚨")
-                        except Exception as e:
-                             logger.exception(f"Could not display video {video_path}: {e}")
-                             st.error(f"Error displaying video T{timeline_id}: {e}", icon="🚨")
+                                     for err in generation_errors[timeline_id]: st.warning(f"- {err}")
+                        except FileNotFoundError: logger.error(f"Video file missing: {video_path}"); st.error(f"Error: Video file missing T{timeline_id}.", icon="🚨")
+                        except Exception as e: logger.exception(f"Display error {video_path}: {e}"); st.error(f"Error display T{timeline_id}: {e}", icon="🚨")
             else:
                 st.warning("No final videos were successfully generated.")
+                # ... (Error summary display - same as before) ...
                 all_errors = [msg for err_list in generation_errors.values() for msg in err_list]
                 if all_errors:
                     st.subheader("Summary of Generation Issues")
                     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}")
+                            if errors: st.error(f"T{tid}:"); [st.error(f"  - {msg}") for msg in errors]
 
             # --- 4. Cleanup ---
-            st.info(f"Attempting to clean up temporary directory: {temp_dir}")
+            st.info(f"Attempting cleanup: {temp_dir}")
             try:
                 shutil.rmtree(temp_dir)
-                logger.info(f"✅ Temporary directory removed: {temp_dir}")
+                logger.info(f"✅ Temp dir removed: {temp_dir}")
                 st.success("✅ Temporary files cleaned up.")
             except Exception as e:
-                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="⚠️")
+                logger.error(f"⚠️ Failed remove temp dir {temp_dir}: {e}")
+                st.warning(f"Could not remove temp files: {temp_dir}.", icon="⚠️")
 
-        elif not chrono_response:
-            logger.error("Story generation or validation failed, cannot proceed.")
-        else:
-            st.error("An unexpected issue occurred after story generation.", icon="🛑")
-            logger.error("Chrono_response existed but was falsy in main logic.")
+        elif not chrono_response: logger.error("Story generation/validation failed.")
+        else: st.error("Unexpected issue post-generation.", icon="🛑"); logger.error("Chrono_response truthy but invalid state.")
 
 else:
-    st.info("Configure settings in the sidebar and click '✨ Generate ChronoWeave ✨' to start.")
+    st.info("Configure settings and click '✨ Generate ChronoWeave ✨' to start.")