Update app.py

app.py

@@ -27,20 +27,20 @@ from pydantic import BaseModel, Field, ValidationError, field_validator, model_v
 # Video and audio processing
 from moviepy.editor import ImageClip, AudioFileClip, concatenate_videoclips
 
-# Google generative API and async patch
+# Google Generative AI library and async patch
 import google.generativeai as genai
 import nest_asyncio
-nest_asyncio.apply()  # Make asyncio work in Streamlit/Jupyter
+nest_asyncio.apply()  # Ensure asyncio works correctly in Streamlit/Jupyter
 
 # --- Logging Setup ---
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 # --- Constants & Configurations ---
-TEXT_MODEL_ID = "models/gemini-1.5-flash"  # Alternatively "gemini-1.5-pro"
-AUDIO_MODEL_ID = "models/gemini-1.5-flash"   # Audio generation uses the text model for now
+TEXT_MODEL_ID = "models/gemini-1.5-flash"   # Alternatively "gemini-1.5-pro"
+AUDIO_MODEL_ID = "models/gemini-1.5-flash"    # Synchronous generation for audio now
 AUDIO_SAMPLING_RATE = 24000
-IMAGE_MODEL_ID = "imagen-3"  # NOTE: Requires Vertex AI SDK update for production
+IMAGE_MODEL_ID = "imagen-3"  # NOTE: Requires Vertex AI SDK integration in the future
 DEFAULT_ASPECT_RATIO = "1:1"
 VIDEO_FPS = 24
 VIDEO_CODEC = "libx264"
@@ -60,7 +60,7 @@ class StorySegment(BaseModel):
     @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-related descriptors.")
+            logger.warning(f"Image prompt '{v[:50]}...' may include human-related descriptions.")
         return v
 
 
@@ -78,16 +78,18 @@ class ChronoWeaveResponse(BaseModel):
     @model_validator(mode="after")
     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, got {len(t.segments)}.")
+        for i, timeline in enumerate(self.timelines):
+            if len(timeline.segments) != expected:
+                raise ValueError(f"Timeline {i} (ID: {timeline.timeline_id}): Expected {expected} segments, got {len(timeline.segments)}.")
         return self
 
 
 # --- Helper Functions ---
 @contextlib.contextmanager
 def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLING_RATE, sample_width: int = 2):
-    """Safely writes a WAV file using a context manager."""
+    """
+    Safely writes a WAV file using a context manager.
+    """
     wf = None
     try:
         wf = wave.open(filename, "wb")
@@ -109,8 +111,8 @@ def wave_file_writer(filename: str, channels: int = 1, rate: int = AUDIO_SAMPLIN
 # --- ChronoWeave Generator Class ---
 class ChronoWeaveGenerator:
     """
-    Encapsulates the logic for generating branching narratives, processing assets (audio, image)
-    and assembling final videos.
+    Encapsulates the logic for generating branching narratives,
+    processing audio, images, and assembling video outputs.
     """
 
     def __init__(self, api_key: str):
@@ -134,7 +136,7 @@ class ChronoWeaveGenerator:
         """
         Generates a story structure as JSON using the text model and validates it via Pydantic.
         """
-        st.info(f"Generating {num_timelines} timeline(s) each with {num_scenes} scene(s) for theme: '{theme}'")
+        st.info(f"Generating {num_timelines} timeline(s) with {num_scenes} scene(s) for theme: '{theme}'")
         logger.info(f"Story generation request: Theme='{theme}', Timelines={num_timelines}, Scenes={num_scenes}")
 
         divergence_instruction = (
@@ -150,7 +152,7 @@ Instructions:
 4. {divergence_instruction}
 5. Style: **'Simple, friendly kids animation, bright colors, rounded shapes'** unless modified by `timeline_visual_modifier`.
 6. `audio_text`: One concise sentence (max 30 words).
-7. `image_prompt`: Descriptive prompt (15–35 words) that emphasizes scene elements. **Avoid repeating general style.**
+7. `image_prompt`: Descriptive prompt (15–35 words) emphasizing scene elements. **Avoid repeating general style.**
 8. `character_description`: Very brief (name and features; < 20 words).
 
 Output only a valid JSON object conforming exactly to this schema:
@@ -185,61 +187,55 @@ JSON Schema: ```json
 
     async def generate_audio(self, text: str, output_filename: str, voice: Optional[str] = None) -> Optional[str]:
         """
-        Asynchronously generates audio using the Gemini Live API.
+        Asynchronously generates audio by wrapping the synchronous generate_content call.
+        The call is executed using asyncio.to_thread to avoid blocking.
         """
         task_id = os.path.basename(output_filename).split(".")[0]
-        collected_audio = bytearray()
         logger.info(f"🎙️ [{task_id}] Generating audio for text: '{text[:60]}...'")
+
         try:
-            config = {
-                "response_modalities": ["AUDIO"],
-                "audio_config": {"audio_encoding": "LINEAR16", "sample_rate_hertz": AUDIO_SAMPLING_RATE},
-            }
-            directive = f"Narrate directly: \"{text}\""
-            async with self.client_audio.connect(config=config) as session:
-                await session.send_request([directive])
-                async for response in session.stream_content():
-                    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}] Audio error: {response.error}")
-                        st.error(f"Audio stream error {task_id}: {response.error}", icon="🔊")
-                        return None
-
-            if not collected_audio:
-                logger.warning(f"⚠️ [{task_id}] No audio data received.")
-                st.warning(f"No audio data for {task_id}.", icon="🔊")
+            # Define a synchronous function for audio generation.
+            def sync_generate_audio():
+                prompt = f"Narrate directly: \"{text}\""
+                response = self.client_audio.generate_content(
+                    contents=prompt,
+                    generation_config=genai.types.GenerationConfig(
+                        response_mime_type="application/octet-stream",
+                        temperature=0.7,
+                        audio_config={"audio_encoding": "LINEAR16", "sample_rate_hertz": AUDIO_SAMPLING_RATE}
+                    )
+                )
+                return response
+
+            # Execute the synchronous call in a separate thread.
+            response = await asyncio.to_thread(sync_generate_audio)
+
+            # Process the response. Adjust as necessary based on the API’s actual response structure.
+            if not response or not hasattr(response, "audio_chunk") or not response.audio_chunk.data:
+                logger.error(f"❌ [{task_id}] No audio data returned.")
+                st.error(f"Audio generation failed for {task_id}: No audio data.", icon="🔊")
                 return None
 
+            audio_data = response.audio_chunk.data
             with wave_file_writer(output_filename) as wf:
-                wf.writeframes(bytes(collected_audio))
-            logger.info(f"✅ [{task_id}] Audio saved: {os.path.basename(output_filename)} ({len(collected_audio)} bytes)")
+                wf.writeframes(audio_data)
+            logger.info(f"✅ [{task_id}] Audio saved: {os.path.basename(output_filename)} ({len(audio_data)} 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 for {task_id}.", icon="🔇")
         except Exception as e:
-            logger.exception(f"❌ [{task_id}] Audio generation failed: {e}")
+            logger.exception(f"❌ [{task_id}] Audio generation error: {e}")
             st.error(f"Audio generation failed for {task_id}: {e}", icon="🔊")
-        return None
+            return None
 
     async def generate_image_async(self, prompt: str, aspect_ratio: str, task_id: str) -> Optional[Image.Image]:
         """
-        Wraps the synchronous image generation function in a thread pool to allow asynchronous invocation.
-        Currently, this function is a stub pending Vertex AI SDK integration.
+        Placeholder for image generation.
+        Currently logs an error and returns None. Update this function once integrating Vertex AI SDK.
         """
-        loop = asyncio.get_event_loop()
         logger.info(f"🖼️ [{task_id}] Requesting image for prompt: '{prompt[:70]}...' (Aspect Ratio: {aspect_ratio})")
-        # Placeholder: the real implementation would call a Vertex AI SDK function.
-        def gen_image():
-            logger.error(f"❌ [{task_id}] Image generation not implemented. Update required for Vertex AI.")
-            return None
-
-        image_result = await loop.run_in_executor(None, gen_image)
-        if image_result is None:
-            st.error(f"Image generation for {task_id} skipped: Requires Vertex AI SDK implementation.", icon="🖼️")
-        return image_result
+        logger.error(f"❌ [{task_id}] Image generation not implemented. Update required for Vertex AI.")
+        st.error(f"Image generation for {task_id} skipped: Requires Vertex AI SDK implementation.", icon="🖼️")
+        return None
 
     async def process_scene(
         self,
@@ -250,7 +246,8 @@ JSON Schema: ```json
         audio_voice: Optional[str] = None,
     ) -> Tuple[Optional[str], Optional[str], Optional[Any], List[str]]:
         """
-        Processes a single scene: generates image and audio concurrently, creates a video clip if both succeed.
+        Processes a single scene: concurrently generates image and audio,
+        and then creates a video clip if both outputs are available.
         Returns a tuple of (image_path, audio_path, video_clip, [error messages]).
         """
         errors: List[str] = []
@@ -260,16 +257,18 @@ JSON Schema: ```json
         video_clip = None
 
         # Launch image and audio generation concurrently.
-        image_future = asyncio.create_task(self.generate_image_async(
-            prompt=f"{segment.image_prompt} Featuring: {segment.character_description} {'Style hint: ' + segment.timeline_visual_modifier if segment.timeline_visual_modifier else ''}",
-            aspect_ratio=aspect_ratio,
-            task_id=task_id,
-        ))
+        image_future = asyncio.create_task(
+            self.generate_image_async(
+                prompt=f"{segment.image_prompt} Featuring: {segment.character_description} " +
+                       (f"Style hint: {segment.timeline_visual_modifier}" if segment.timeline_visual_modifier else ""),
+                aspect_ratio=aspect_ratio,
+                task_id=task_id,
+            )
+        )
         audio_future = asyncio.create_task(self.generate_audio(segment.audio_text, audio_path, audio_voice))
 
         image_result, audio_result = await asyncio.gather(image_future, audio_future)
 
-        # Handle image result (if available, save and preview)
         if image_result:
             try:
                 image_result.save(image_path)
@@ -280,7 +279,6 @@ JSON Schema: ```json
         else:
             errors.append(f"Scene {segment.scene_id + 1}: Image generation failed.")
 
-        # Handle audio result and preview
         if audio_result:
             try:
                 with open(audio_result, "rb") as ap:
@@ -290,7 +288,6 @@ JSON Schema: ```json
         else:
             errors.append(f"Scene {segment.scene_id + 1}: Audio generation failed.")
 
-        # Create video clip if both image and audio exist.
         if not errors and os.path.exists(image_path) and os.path.exists(audio_path):
             try:
                 audio_clip = AudioFileClip(audio_path)
@@ -302,7 +299,6 @@ JSON Schema: ```json
                 logger.exception(f"❌ [{task_id}] Failed to create video clip: {e}")
                 errors.append(f"Scene {segment.scene_id + 1}: Video clip creation failed.")
             finally:
-                # Cleanup moviepy instances.
                 try:
                     if 'audio_clip' in locals():
                         audio_clip.close()
@@ -311,10 +307,12 @@ JSON Schema: ```json
                 except Exception:
                     pass
 
-        return (image_path if os.path.exists(image_path) else None,
-                audio_path if os.path.exists(audio_path) else None,
-                video_clip,
-                errors)
+        return (
+            image_path if os.path.exists(image_path) else None,
+            audio_path if os.path.exists(audio_path) else None,
+            video_clip,
+            errors,
+        )
 
     async def process_timeline(
         self,
@@ -324,9 +322,9 @@ JSON Schema: ```json
         audio_voice: Optional[str] = None,
     ) -> Tuple[Optional[str], List[str]]:
         """
-        Processes an entire timeline by concurrently processing all scenes,
-        and then assembling a final video if all scenes succeed.
-        Returns the final video path and a list of error messages.
+        Processes an entire timeline by concurrently processing all its scenes,
+        then assembling a final video if every scene produced a valid clip.
+        Returns a tuple of (final video path, list of error messages).
         """
         timeline_id = timeline.timeline_id
         scene_tasks = [
@@ -342,20 +340,14 @@ JSON Schema: ```json
             if clip is not None:
                 video_clips.append(clip)
 
-        # Assemble the timeline video only if every scene produced a valid clip.
         if video_clips and len(video_clips) == len(timeline.segments):
             output_filename = os.path.join(temp_dir, f"timeline_{timeline_id}_final.mp4")
             try:
                 final_video = concatenate_videoclips(video_clips, method="compose")
                 final_video.write_videofile(
-                    output_filename,
-                    fps=VIDEO_FPS,
-                    codec=VIDEO_CODEC,
-                    audio_codec=AUDIO_CODEC,
-                    logger=None
+                    output_filename, fps=VIDEO_FPS, codec=VIDEO_CODEC, audio_codec=AUDIO_CODEC, logger=None
                 )
                 logger.info(f"✅ Timeline {timeline_id} video saved: {output_filename}")
-                # Cleanup the clips.
                 for clip in video_clips:
                     clip.close()
                 final_video.close()
@@ -370,7 +362,7 @@ JSON Schema: ```json
 
 # --- Streamlit UI and Main Process ---
 def main():
-    # --- API Key Retrieval ---
+    # API Key Retrieval
     GOOGLE_API_KEY: Optional[str] = None
     try:
         GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
@@ -383,7 +375,6 @@ def main():
             st.error("🚨 **Google API Key Not Found!** Please configure it.", icon="🚨")
             st.stop()
 
-    # --- UI Configuration ---
     st.set_page_config(page_title="ChronoWeave", layout="wide", initial_sidebar_state="expanded")
     st.title("🌀 ChronoWeave: Advanced Branching Narrative Generator")
     st.markdown("""
@@ -415,7 +406,6 @@ def main():
             st.error("Please enter a story theme.", icon="👈")
             return
 
-        # Create a unique temporary directory for this run
         run_id = str(uuid.uuid4()).split('-')[0]
         temp_dir = os.path.join(TEMP_DIR_BASE, f"run_{run_id}")
         try:
@@ -425,7 +415,7 @@ def main():
             st.error(f"🚨 Failed to create temporary directory {temp_dir}: {e}", icon="📂")
             st.stop()
 
-        # Instantiate the ChronoWeave generator
+        # Instantiate ChronoWeaveGenerator and generate story structure.
         generator = ChronoWeaveGenerator(GOOGLE_API_KEY)
         chrono_response = None
         with st.spinner("Generating narrative structure... 🤔"):
@@ -440,30 +430,29 @@ def main():
         generation_errors: Dict[int, List[str]] = {}
 
         async def process_all_timelines():
-            timeline_tasks = {}
-            for timeline in chrono_response.timelines:
-                timeline_tasks[timeline.timeline_id] = asyncio.create_task(
+            timeline_tasks = {
+                timeline.timeline_id: asyncio.create_task(
                     generator.process_timeline(timeline, temp_dir, aspect_ratio, audio_voice)
                 )
-            return await asyncio.gather(*timeline_tasks.values(), return_exceptions=False)
+                for timeline in chrono_response.timelines
+            }
+            results = await asyncio.gather(*timeline_tasks.values(), return_exceptions=False)
+            return results
 
         with st.spinner("Processing scenes and assembling videos..."):
             timeline_results = asyncio.run(process_all_timelines())
 
-        # Collect results per timeline.
         for timeline, (video_path, errors) in zip(chrono_response.timelines, timeline_results):
             generation_errors[timeline.timeline_id] = errors
             if video_path:
                 final_video_paths[timeline.timeline_id] = video_path
 
         overall_duration = time.time() - overall_start_time
-        # Display status messages
         if final_video_paths:
             st.success(f"Complete! ({len(final_video_paths)} video(s) created in {overall_duration:.2f}s)")
         else:
             st.error(f"Failed. No final videos generated in {overall_duration:.2f}s")
 
-        # --- Display Final Videos ---
         st.header("🎬 Generated Timelines")
         if final_video_paths:
             sorted_ids = sorted(final_video_paths.keys())
@@ -506,7 +495,6 @@ def main():
                         for msg in errs:
                             st.error(f"  - {msg}")
 
-        # --- Cleanup ---
         st.info(f"Cleaning up temporary files: {temp_dir}")
         try:
             shutil.rmtree(temp_dir)
@@ -518,5 +506,6 @@ def main():
     else:
         st.info("Configure settings and click '✨ Generate ChronoWeave ✨' to start.")
 
+
 if __name__ == "__main__":
     main()