Spaces:
Running
Running
| # Copyright 2025 Google LLC. Based on work by Yousif Ahmed. | |
| # Concept: ChronoWeave - Branching Narrative Generation | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 | |
| import streamlit as st | |
| import google.generativeai as genai | |
| import os | |
| import json | |
| import numpy as np | |
| from io import BytesIO | |
| import time | |
| import wave | |
| import contextlib | |
| import asyncio | |
| import uuid # For unique filenames | |
| import shutil # For cleaning up temp dirs | |
| # Image handling | |
| from PIL import Image | |
| # Video and audio processing | |
| from moviepy.editor import ImageClip, AudioFileClip, CompositeVideoClip, concatenate_videoclips | |
| # 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 | |
| # --- Configuration --- | |
| st.set_page_config(page_title="ChronoWeave", layout="wide") | |
| st.title("🌀 ChronoWeave: Branching Narrative Generator") | |
| st.markdown(""" | |
| Generate multiple, branching story timelines from a single theme using AI. | |
| Based on the work of Yousif Ahmed. Copyright 2025 Google LLC. | |
| """) | |
| # --- Constants --- | |
| MODEL = "models/gemini-1.5-flash" # Or other suitable text model supporting JSON | |
| # Using v1alpha for the Live API for audio output. | |
| AUDIO_MODEL_VERSION = 'v1alpha' # Must be alpha for audio modality | |
| IMAGE_MODEL_ID = "imagen-3" # Or your preferred Imagen model "imagen-3.0-generate-002" | |
| # --- API Key Handling --- | |
| try: | |
| # Preferred way to handle secrets in Streamlit sharing/HF Spaces | |
| GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"] | |
| os.environ['GOOGLE_API_KEY'] = GOOGLE_API_KEY | |
| except KeyError: | |
| st.error("🚨 Google API Key not found! Please add it as a Secret named 'GOOGLE_API_KEY' in your Hugging Face Space settings.", icon="🚨") | |
| st.stop() # Halt execution if no key | |
| # --- Initialize Google Client --- | |
| try: | |
| # Initialize the client with the API key | |
| genai.configure(api_key=GOOGLE_API_KEY) | |
| # Create separate clients or configure one for different API versions if needed | |
| # Client for Text/Imagen (standard API) | |
| client_standard = genai.GenerativeModel(MODEL) | |
| # Client for Live Audio (v1alpha) - requires different client init | |
| client_live = genai.Client( | |
| client_options={'api_endpoint': f'{AUDIO_MODEL_VERSION}.generativelanguage.googleapis.com'} | |
| ) | |
| # Note: As of recent updates, genai.configure might handle this better, | |
| # but separating clients or explicitly setting endpoints can be more robust. | |
| # Adjust based on the library version and observed behavior. | |
| except Exception as e: | |
| st.error(f"🚨 Failed to initialize Google AI Client: {e}", icon="🚨") | |
| st.stop() | |
| # --- Define Structured Output Schemas --- | |
| class StorySegment(typing.TypedDict): | |
| scene_id: int | |
| image_prompt: str | |
| audio_text: str | |
| character_description: str | |
| timeline_visual_modifier: typing.Optional[str] | |
| class Timeline(typing.TypedDict): | |
| timeline_id: int | |
| divergence_reason: str | |
| segments: list[StorySegment] | |
| class ChronoWeaveResponse(typing.TypedDict): | |
| core_theme: str | |
| timelines: list[Timeline] | |
| total_scenes_per_timeline: int | |
| # --- Helper Functions --- | |
| def wave_file(filename, channels=1, rate=24000, sample_width=2): | |
| """Context manager to write WAV files.""" | |
| with wave.open(filename, "wb") as wf: | |
| wf.setnchannels(channels) | |
| wf.setsampwidth(sample_width) | |
| wf.setframerate(rate) | |
| yield wf | |
| async def generate_audio_live_async(api_text, output_filename): | |
| """Generates audio using Gemini Live API (async version).""" | |
| collected_audio = bytearray() | |
| st.write(f"🎙️ Generating audio for: '{api_text[:50]}...'") # Log start | |
| try: | |
| # Use the 'client_live' specifically configured for v1alpha | |
| live_model = client_live.get_model(f"models/gemini-1.5-flash") # Specify model within the live client context | |
| config = { | |
| "response_modalities": ["AUDIO"] | |
| } | |
| # Connect to the Live API using the live client. | |
| async with live_model.connect(config=config) as session: | |
| await session.send_request([api_text]) # Simpler send for single prompt | |
| async for response in session.stream_content(): | |
| if response.audio_chunk: | |
| collected_audio.extend(response.audio_chunk.data) | |
| if not collected_audio: | |
| st.warning(f"⚠️ No audio data received for: '{api_text[:50]}...'") | |
| return None # Indicate failure | |
| audio_bytes = bytes(collected_audio) | |
| # Write the collected audio bytes into a WAV file. | |
| with wave_file(output_filename) as wf: | |
| wf.writeframes(audio_bytes) | |
| st.write(f" ✅ Audio saved: {os.path.basename(output_filename)}") | |
| return output_filename | |
| except Exception as e: | |
| st.error(f" ❌ Audio generation failed for '{api_text[:50]}...': {e}", icon="🚨") | |
| return None | |
| def generate_story_sequence_chrono(theme: str, num_scenes: int, num_timelines: int, divergence_prompt: str = "") -> ChronoWeaveResponse | None: | |
| """Generates branching story sequences using Gemini structured output.""" | |
| st.write(f"📚 Generating {num_timelines} timeline(s) for theme: '{theme}'...") | |
| divergence_instruction = f"Introduce divergence between timelines. {divergence_prompt}" if divergence_prompt else "Introduce natural points of divergence between timelines after the first scene or two." | |
| prompt = f''' | |
| As an expert narrative designer, create a branching story based on the theme: "{theme}". | |
| Generate exactly {num_timelines} distinct timelines, each containing exactly {num_scenes} scenes. | |
| Each scene should be approximately 5-10 seconds long when narrated. | |
| {divergence_instruction} Clearly state the reason for divergence for each timeline after the first. | |
| For each scene in each timeline, provide: | |
| - scene_id: An integer starting from 0 for the scene number within its timeline. | |
| - image_prompt: A concise (15-25 words) description for an image generation model. Focus on visual details, characters (animals/objects only, NO PEOPLE), background, and action. Maintain a consistent 'kids animation style' (e.g., simple, rounded shapes, bright colors) across all scenes and timelines unless specified by a timeline_visual_modifier. | |
| - audio_text: A single, engaging sentence of narration or dialogue for the scene (max 25 words). | |
| - character_description: Brief description of recurring characters (names, key features) mentioned in *this specific scene's image prompt*. Keep consistent within a timeline. (Max 30 words). | |
| - timeline_visual_modifier: (Optional, string or null) A *brief* hint if this timeline should have a slightly different visual feel from this scene onwards (e.g., "slightly darker lighting", "more cluttered background", "character looks worried"). Keep it subtle. Use null if no specific modifier. | |
| Constraint: Ensure the output strictly adheres to the following JSON schema. Do not include preamble or explanations outside the JSON structure. Respond ONLY with the JSON object. | |
| JSON Schema: | |
| {{ | |
| "type": "object", | |
| "properties": {{ | |
| "core_theme": {{"type": "string"}}, | |
| "timelines": {{ | |
| "type": "array", | |
| "items": {{ | |
| "type": "object", | |
| "properties": {{ | |
| "timeline_id": {{"type": "integer"}}, | |
| "divergence_reason": {{"type": "string"}}, | |
| "segments": {{ | |
| "type": "array", | |
| "items": {{ | |
| "type": "object", | |
| "properties": {{ | |
| "scene_id": {{"type": "integer"}}, | |
| "image_prompt": {{"type": "string"}}, | |
| "audio_text": {{"type": "string"}}, | |
| "character_description": {{"type": "string"}}, | |
| "timeline_visual_modifier": {{"type": ["string", "null"]}} | |
| }}, | |
| "required": ["scene_id", "image_prompt", "audio_text", "character_description", "timeline_visual_modifier"] | |
| }} | |
| }} | |
| }}, | |
| "required": ["timeline_id", "divergence_reason", "segments"] | |
| }} | |
| }}, | |
| "total_scenes_per_timeline": {{"type": "integer"}} | |
| }}, | |
| "required": ["core_theme", "timelines", "total_scenes_per_timeline"] | |
| }} | |
| ''' | |
| try: | |
| response = client_standard.generate_content( | |
| contents=prompt, | |
| generation_config=genai.types.GenerationConfig( | |
| response_mime_type="application/json", | |
| # Optional: Add temperature, etc. if needed | |
| ) | |
| # The schema can also be passed via generation_config in some versions/models | |
| # config={ | |
| # 'response_mime_type': 'application/json', | |
| # 'response_schema': ChronoWeaveResponse # Pass the TypedDict directly | |
| # } | |
| ) | |
| # Debugging: Print raw response text | |
| # st.text_area("Raw Gemini Response:", response.text, height=200) | |
| story_data = json.loads(response.text) # Assuming response.text contains the JSON string | |
| st.success("✅ Story structure generated successfully!") | |
| # Basic validation (can be more thorough) | |
| if 'timelines' in story_data and isinstance(story_data['timelines'], list): | |
| # Further validation could check segment structure, etc. | |
| return story_data # Return the parsed dictionary | |
| else: | |
| st.error("🚨 Generated story data is missing the 'timelines' list.", icon="🚨") | |
| return None | |
| except json.JSONDecodeError as e: | |
| st.error(f"🚨 Failed to decode JSON response from Gemini: {e}", icon="🚨") | |
| st.text_area("Problematic Response Text:", response.text if 'response' in locals() else "No response object.", height=150) | |
| return None | |
| except Exception as e: | |
| st.error(f"🚨 Error generating story sequence: {e}", icon="🚨") | |
| # Log the prompt potentially? Be careful with sensitive data if applicable. | |
| # st.text_area("Failed Prompt:", prompt, height=200) | |
| return None | |
| def generate_image_imagen(prompt: str, aspect_ratio: str = "1:1") -> Image.Image | None: | |
| """Generates an image using Imagen.""" | |
| st.write(f"🖼️ Generating image for: '{prompt[:60]}...'") | |
| try: | |
| # Use the standard client's dedicated image generation method | |
| response = client_standard.generate_content( | |
| f"Generate an image with the following prompt, ensuring a child-friendly animation style and NO human figures: {prompt}", | |
| generation_config=genai.types.GenerationConfig( | |
| candidate_count=1, # Generate one image | |
| # Imagen specific parameters are often passed differently or rely on model defaults | |
| # Check documentation for precise Imagen control via the unified API | |
| ), | |
| # If the model/API version requires specific image parameters: | |
| # tools=[genai.ImageParams(model=IMAGE_MODEL_ID, number_of_images=1, aspect_ratio=aspect_ratio, person_generation="DONT_ALLOW")] | |
| ) | |
| # Accessing image data might vary slightly depending on API response structure | |
| # This assumes response.parts contains the image data if successful | |
| if response.parts and response.parts[0].inline_data: | |
| image_bytes = response.parts[0].inline_data.data | |
| image = Image.open(BytesIO(image_bytes)) | |
| st.write(" ✅ Image generated.") | |
| return image | |
| else: | |
| # Check for safety blocks or other reasons for failure | |
| if response.prompt_feedback.block_reason: | |
| st.warning(f" ⚠️ Image generation blocked for prompt '{prompt[:60]}...'. Reason: {response.prompt_feedback.block_reason}", icon="⚠️") | |
| else: | |
| st.warning(f" ⚠️ No image data received for prompt '{prompt[:60]}...'.", icon="⚠️") | |
| # Debugging: st.write(response) | |
| return None | |
| except Exception as e: | |
| st.error(f" ❌ Image generation failed for '{prompt[:60]}...': {e}", icon="🚨") | |
| return None | |
| # --- Streamlit UI Elements --- | |
| st.sidebar.header("Configuration") | |
| # API Key display/check (already handled above, but sidebar is a good place) | |
| 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 shiny object") | |
| 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 it started raining?") | |
| aspect_ratio = st.sidebar.selectbox("Image Aspect Ratio:", ["1:1", "16:9", "9:16"], index=0) | |
| generate_button = st.sidebar.button("✨ Generate ChronoWeave ✨", type="primary", disabled=(not GOOGLE_API_KEY)) | |
| st.sidebar.markdown("---") | |
| st.sidebar.info("Note: Generation can take several minutes depending on settings.") | |
| # --- Main Logic --- | |
| if generate_button: | |
| if not theme: | |
| st.error("Please enter a story theme.", icon="👈") | |
| else: | |
| # Create a unique temporary directory for this run | |
| run_id = str(uuid.uuid4()) | |
| temp_dir = os.path.join(".", f"chrono_temp_{run_id}") # Create in current dir | |
| os.makedirs(temp_dir, exist_ok=True) | |
| st.write(f"Working directory: {temp_dir}") | |
| final_video_paths = {} # To store {timeline_id: video_path} | |
| with st.spinner("Generating narrative structure..."): | |
| chrono_data = generate_story_sequence_chrono(theme, num_scenes, num_timelines, divergence_prompt) | |
| if chrono_data and 'timelines' in chrono_data: | |
| st.success(f"Found {len(chrono_data['timelines'])} timelines. Processing each...") | |
| all_timelines_successful = True # Flag to track if all timelines worked | |
| # Use st.status for detailed progress | |
| with st.status("Generating assets and composing videos...", expanded=True) as status: | |
| for timeline in chrono_data['timelines']: | |
| timeline_id = timeline['timeline_id'] | |
| divergence = timeline['divergence_reason'] | |
| segments = timeline['segments'] | |
| st.subheader(f"Timeline {timeline_id}: {divergence}") | |
| temp_image_files = [] | |
| temp_audio_files = [] | |
| video_clips = [] | |
| timeline_successful = True # Flag for this specific timeline | |
| for i, segment in enumerate(segments): | |
| status.update(label=f"Processing Timeline {timeline_id}, Scene {i+1}/{num_scenes}...") | |
| scene_id = segment['scene_id'] | |
| image_prompt = segment['image_prompt'] | |
| audio_text = segment['audio_text'] | |
| char_desc = segment['character_description'] | |
| vis_mod = segment['timeline_visual_modifier'] | |
| st.write(f"--- Scene {i+1} (T{timeline_id}) ---") | |
| st.write(f"* **Image Prompt:** {image_prompt}" + (f" (Modifier: {vis_mod})" if vis_mod else "")) | |
| st.write(f"* **Audio Text:** {audio_text}") | |
| # st.write(f"* Character Desc: {char_desc}") # Can be verbose | |
| # --- Image Generation --- | |
| combined_prompt = f"{image_prompt} {char_desc}" | |
| if vis_mod: | |
| combined_prompt += f" Style hint: {vis_mod}" | |
| generated_image = generate_image_imagen(combined_prompt, aspect_ratio) | |
| if generated_image: | |
| image_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_image.png") | |
| generated_image.save(image_path) | |
| temp_image_files.append(image_path) | |
| st.image(generated_image, width=200) # Show thumbnail | |
| else: | |
| st.warning(f"Skipping scene {i+1} in timeline {timeline_id} due to image generation failure.") | |
| timeline_successful = False | |
| continue # Skip to next segment if image fails | |
| # --- Audio Generation --- | |
| # Add negative prompt to prevent conversational filler | |
| audio_negative_prompt = "Narrate the following sentence directly, with expression, without any introduction or closing remarks like 'Okay' or 'Here is the narration'. Just read the sentence:" | |
| full_audio_prompt = f"{audio_negative_prompt}\n{audio_text}" | |
| audio_path = os.path.join(temp_dir, f"t{timeline_id}_s{i}_audio.wav") | |
| # Run the async audio generation function | |
| try: | |
| generated_audio_path = asyncio.run(generate_audio_live_async(full_audio_prompt, audio_path)) | |
| except Exception as e: | |
| st.error(f"Asyncio error during audio gen: {e}") | |
| generated_audio_path = None | |
| if generated_audio_path: | |
| temp_audio_files.append(generated_audio_path) | |
| # st.audio(generated_audio_path) # Optional: Preview audio | |
| else: | |
| st.warning(f"Skipping video clip for scene {i+1} in timeline {timeline_id} due to audio generation failure.") | |
| # Clean up the image file for this failed scene segment | |
| if os.path.exists(image_path): | |
| os.remove(image_path) | |
| temp_image_files.remove(image_path) | |
| timeline_successful = False | |
| continue # Skip making video clip if audio fails | |
| # --- Create Video Clip --- | |
| try: | |
| st.write(" 🎬 Creating video clip...") | |
| audio_clip = AudioFileClip(generated_audio_path) | |
| # Ensure PIL Image is used if needed, or numpy array directly | |
| np_image = np.array(Image.open(image_path)) | |
| # Create ImageClip, ensure duration matches audio | |
| image_clip = ImageClip(np_image).set_duration(audio_clip.duration) | |
| # Handle potential size mismatch if needed (resize image_clip or set size explicitly) | |
| # image_clip = image_clip.resize(width=...) | |
| composite_clip = image_clip.set_audio(audio_clip) # Simpler composition | |
| video_clips.append(composite_clip) | |
| st.write(" ✅ Clip created.") | |
| except Exception as e: | |
| st.error(f" ❌ Failed to create video clip for scene {i+1} (T{timeline_id}): {e}", icon="🚨") | |
| timeline_successful = False | |
| # Don't break the whole timeline, just skip this clip maybe? Or mark timeline as failed. | |
| # --- Assemble Timeline Video --- | |
| if video_clips and timeline_successful: # Only assemble if clips were made and no major errors | |
| status.update(label=f"Composing final video for Timeline {timeline_id}...") | |
| st.write(f"🎞️ Assembling final video for Timeline {timeline_id}...") | |
| try: | |
| final_timeline_video = concatenate_videoclips(video_clips, method="compose") | |
| output_filename = os.path.join(temp_dir, f"timeline_{timeline_id}_final_video.mp4") | |
| # Use 'libx264' for broader compatibility, specify audio codec | |
| final_timeline_video.write_videofile(output_filename, fps=24, codec='libx264', audio_codec='aac') | |
| final_video_paths[timeline_id] = output_filename | |
| st.success(f" ✅ Video for Timeline {timeline_id} saved: {os.path.basename(output_filename)}") | |
| # Close clips to release resources | |
| for clip in video_clips: | |
| if hasattr(clip, 'close'): clip.close() | |
| if hasattr(clip, 'audio') and hasattr(clip.audio, 'close'): clip.audio.close() | |
| if hasattr(final_timeline_video, 'close'): final_timeline_video.close() | |
| except Exception as e: | |
| st.error(f" ❌ Failed to write final video for Timeline {timeline_id}: {e}", icon="🚨") | |
| all_timelines_successful = False | |
| elif not video_clips: | |
| st.warning(f"No video clips were successfully generated for Timeline {timeline_id}. Skipping final video assembly.") | |
| all_timelines_successful = False | |
| else: | |
| st.warning(f"Timeline {timeline_id} encountered errors. Skipping final video assembly.") | |
| all_timelines_successful = False | |
| # Intermediate cleanup for the timeline (optional, helps manage files) | |
| # for file in temp_audio_files: | |
| # if os.path.exists(file): os.remove(file) | |
| # for file in temp_image_files: | |
| # if os.path.exists(file): os.remove(file) | |
| # Final status update | |
| if all_timelines_successful and final_video_paths: | |
| status.update(label="ChronoWeave Generation Complete!", state="complete", expanded=False) | |
| elif final_video_paths: | |
| status.update(label="ChronoWeave Generation Partially Complete (some errors occurred).", state="warning", expanded=False) | |
| else: | |
| status.update(label="ChronoWeave Generation Failed.", state="error", expanded=False) | |
| # --- Display Results --- | |
| st.header("Generated Timelines") | |
| if final_video_paths: | |
| sorted_timeline_ids = sorted(final_video_paths.keys()) | |
| for timeline_id in sorted_timeline_ids: | |
| video_path = final_video_paths[timeline_id] | |
| # Find matching timeline divergence reason | |
| reason = "Unknown" | |
| for t in chrono_data.get('timelines', []): | |
| if t.get('timeline_id') == timeline_id: | |
| reason = t.get('divergence_reason', 'N/A') | |
| break | |
| st.subheader(f"Timeline {timeline_id}: {reason}") | |
| try: | |
| video_file = open(video_path, 'rb') | |
| video_bytes = video_file.read() | |
| st.video(video_bytes) | |
| video_file.close() | |
| except FileNotFoundError: | |
| st.error(f"Could not find video file: {video_path}", icon="🚨") | |
| except Exception as e: | |
| st.error(f"Could not display video {video_path}: {e}", icon="🚨") | |
| else: | |
| st.warning("No final videos were successfully generated.") | |
| # --- Cleanup --- | |
| st.write("Cleaning up temporary files...") | |
| try: | |
| shutil.rmtree(temp_dir) | |
| st.write(" ✅ Temporary files removed.") | |
| except Exception as e: | |
| st.warning(f" ⚠️ Could not remove temporary directory {temp_dir}: {e}", icon="⚠️") | |
| elif not chrono_data: | |
| st.error("Story generation failed. Cannot proceed.", icon="🛑") | |
| else: | |
| # This case might happen if chrono_data is returned but is malformed (e.g., no 'timelines' key) | |
| st.error("Story data seems malformed. Cannot proceed.", icon="🛑") | |
| # st.json(chrono_data) # Display the problematic data | |
| else: | |
| st.info("Configure settings in the sidebar and click 'Generate ChronoWeave'") |