Update app.py

app.py

@@ -26,14 +26,14 @@ from typing import List, Optional, Literal, Dict, Any
 
 # 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
 
 # Type hints
 import typing_extensions as typing
 
 # Async support for Streamlit/Google API
 import nest_asyncio
-nest_asyncio.apply() # Apply patch for asyncio in environments like Streamlit/Jupyter
+nest_asyncio.apply()
 
 # --- Logging Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -48,19 +48,14 @@ Generate multiple, branching story timelines from a single theme using AI, compl
 """)
 
 # --- Constants ---
-# Text/JSON Model
-TEXT_MODEL_ID = "models/gemini-1.5-flash" # Or "gemini-1.5-pro"
-# Audio Model Config
-AUDIO_MODEL_ID = "models/gemini-1.5-flash" # Model used for audio tasks
+TEXT_MODEL_ID = "models/gemini-1.5-flash"
+AUDIO_MODEL_ID = "models/gemini-1.5-flash"
 AUDIO_SAMPLING_RATE = 24000
-# Image Model Config
-IMAGE_MODEL_ID = "imagen-3" # <<< NOTE: Likely needs Vertex AI SDK access
+IMAGE_MODEL_ID = "imagen-3" # <<< NOTE: Requires Vertex AI SDK access
 DEFAULT_ASPECT_RATIO = "1:1"
-# Video Config
 VIDEO_FPS = 24
 VIDEO_CODEC = "libx264"
 AUDIO_CODEC = "aac"
-# File Management
 TEMP_DIR_BASE = ".chrono_temp"
 
 # --- API Key Handling ---
@@ -70,10 +65,8 @@ try:
     logger.info("Google API Key loaded from Streamlit secrets.")
 except KeyError:
     GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
-    if GOOGLE_API_KEY:
-        logger.info("Google API Key loaded from environment variable.")
-    else:
-        st.error("🚨 **Google API Key Not Found!** Please configure it.", icon="🚨"); st.stop()
+    if GOOGLE_API_KEY: logger.info("Google API Key loaded from environment variable.")
+    else: st.error("🚨 **Google API Key Not Found!** Please configure it.", icon="🚨"); st.stop()
 
 # --- Initialize Google Clients ---
 try:
@@ -83,14 +76,11 @@ try:
     logger.info(f"Initialized text/JSON model handle: {TEXT_MODEL_ID}.")
     live_model = genai.GenerativeModel(AUDIO_MODEL_ID)
     logger.info(f"Initialized audio model handle: {AUDIO_MODEL_ID}.")
-    image_model_genai = genai.GenerativeModel(IMAGE_MODEL_ID) # Retained but likely needs Vertex SDK
+    image_model_genai = genai.GenerativeModel(IMAGE_MODEL_ID)
     logger.info(f"Initialized google-generativeai handle for image model: {IMAGE_MODEL_ID} (May require Vertex AI SDK).")
     # ---> TODO: Initialize Vertex AI client here if switching SDK <---
-
-except AttributeError as ae:
-     logger.exception("AttributeError during Client Init."); st.error(f"🚨 Init Error: {ae}. Update library?", icon="🚨"); st.stop()
-except Exception as e:
-    logger.exception("Failed to initialize Google Clients/Models."); st.error(f"🚨 Failed Init: {e}", icon="🚨"); st.stop()
+except AttributeError as ae: logger.exception("AttributeError during Client Init."); st.error(f"🚨 Init Error: {ae}. Update library?", icon="🚨"); st.stop()
+except Exception as e: logger.exception("Failed to initialize Google Clients/Models."); st.error(f"🚨 Failed Init: {e}", icon="🚨"); st.stop()
 
 # --- Define Pydantic Schemas (Using V2 Syntax) ---
 class StorySegment(BaseModel):
@@ -116,21 +106,33 @@ class ChronoWeaveResponse(BaseModel):
     def check_timeline_segment_count(self) -> 'ChronoWeaveResponse':
         expected = self.total_scenes_per_timeline
         for i, t in enumerate(self.timelines):
-            if len(t.segments) != expected: raise ValueError(f"Timeline {i} ID {t.timeline_id}: Expected {expected} segments, found {len(t.segments)}.")
+            if len(t.segments) != expected: raise ValueError(f"Timeline {i} ID {t.timeline_id}: Expected {expected}, found {len(t.segments)}.")
         return self
 
 # --- Helper Functions ---
+
+# CORRECTED wave_file_writer function with proper indentation
 @contextlib.contextmanager
 def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLING_RATE, sample_width: int = 2):
     """Context manager to safely write WAV files."""
     wf = None
     try:
-        wf = wave.open(filename, "wb"); wf.setnchannels(channels); wf.setsampwidth(sample_width); wf.setframerate(rate)
+        # Indented correctly
+        wf = wave.open(filename, "wb")
+        wf.setnchannels(channels)
+        wf.setsampwidth(sample_width)
+        wf.setframerate(rate)
         yield wf
-    except Exception as e: logger.error(f"Error wave file {filename}: {e}"); raise
+    except Exception as e:
+        logger.error(f"Error wave file {filename}: {e}")
+        raise
     finally:
-        if wf: try: wf.close()
-        except Exception as e_close: logger.error(f"Error closing wave file {filename}: {e_close}")
+        if wf:
+            # Indented correctly
+            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]:
@@ -138,16 +140,10 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
     collected_audio = bytearray(); task_id = os.path.basename(output_filename).split('.')[0]
     logger.info(f"🎙️ [{task_id}] Requesting audio: '{api_text[:60]}...'")
     try:
-        # CORRECTED config structure for audio generation <<<<<<-------
-        config = {
-            "response_modalities": ["AUDIO"],
-            # Removed 'audio_config' nesting
-            "audio_encoding": "LINEAR16",
-            "sample_rate_hertz": AUDIO_SAMPLING_RATE,
-            # Add other parameters like "voice" here directly if needed
-        }
+        # Corrected config structure
+        config = {"response_modalities": ["AUDIO"], "audio_encoding": "LINEAR16", "sample_rate_hertz": AUDIO_SAMPLING_RATE}
         directive_prompt = f"Narrate directly: \"{api_text}\""
-        async with live_model.connect(config=config) as session: # Pass corrected config
+        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)
@@ -157,11 +153,7 @@ async def generate_audio_live_async(api_text: str, output_filename: str, voice:
         logger.info(f"   ✅ [{task_id}] Audio saved: {os.path.basename(output_filename)} ({len(collected_audio)} bytes)")
         return output_filename
     except genai.types.generation_types.BlockedPromptException as bpe: logger.error(f"   ❌ [{task_id}] Audio blocked: {bpe}"); st.error(f"Audio blocked {task_id}.", icon="🔇"); return None
-    # Catch TypeError specifically for config issues
-    except TypeError as te:
-         logger.exception(f"   ❌ [{task_id}] Audio config TypeError: {te}")
-         st.error(f"Audio configuration error for {task_id} (TypeError): {te}. Check library version/config structure.", icon="⚙️")
-         return None
+    except TypeError as te: logger.exception(f"   ❌ [{task_id}] Audio config TypeError: {te}"); st.error(f"Audio config error {task_id} (TypeError): {te}. Check library/config.", icon="⚙️"); return None
     except Exception as e: logger.exception(f"   ❌ [{task_id}] Audio failed: {e}"); st.error(f"Audio failed {task_id}: {e}", icon="🔊"); return None
 
 
@@ -253,7 +245,7 @@ if generate_button:
 
                         # --- 2b. Audio Generation ---
                         generated_audio_path: Optional[str] = None
-                        if not scene_has_error: # Should not be reached currently due to image fail
+                        if not scene_has_error: # Should not be reached currently
                             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))
@@ -319,7 +311,7 @@ if generate_button:
                                 scene_errors = [err for err in generation_errors[timeline_id] if not err.startswith(f"T{timeline_id}:")]
                                 if scene_errors:
                                      with st.expander(f"⚠️ View {len(scene_errors)} Scene Issues"):
-                                         for err in scene_errors: st.warning(f"- {err}")
+                                         for err in scene_errors: st.warning(f"- {err}") # Use standard loop
                         except FileNotFoundError: logger.error(f"Video missing: {video_path}"); st.error(f"Error: Video missing T{timeline_id}.", icon="🚨")
                         except Exception as e: logger.exception(f"Display error {video_path}: {e}"); st.error(f"Display error T{timeline_id}: {e}", icon="🚨")
             else: # No videos generated