Update app.py

app.py

@@ -2,123 +2,276 @@
 import streamlit as st
 from core.gemini_handler import GeminiHandler
 from core.visual_engine import VisualEngine
-from core.prompt_engineering import create_story_breakdown_prompt, create_image_prompt_from_scene
+from core.prompt_engineering import (
+    create_story_breakdown_prompt,
+    create_image_prompt_from_scene_data,
+    create_scene_regeneration_prompt,
+    create_visual_regeneration_prompt
+)
 import os
+import json # For debugging and display
 
 # --- Configuration & Initialization ---
-st.set_page_config(page_title="CineGen AI", layout="wide")
+st.set_page_config(page_title="CineGen AI Pro", layout="wide")
 
-# Load API Key from Streamlit secrets
-try:
-    GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
-except KeyError:
-    st.error("GEMINI_API_KEY not found in secrets. Please add it to your Streamlit Cloud secrets.")
-    st.stop()
+# --- Global State Variables --- (Using session state for persistence)
+if 'GEMINI_API_KEY' not in st.session_state:
+    try:
+        st.session_state.GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
+    except KeyError:
+        st.error("GEMINI_API_KEY not found in secrets. Please add it.")
+        st.stop()
 
-gemini = GeminiHandler(api_key=GEMINI_API_KEY)
-visualizer = VisualEngine(output_dir="temp_generated_media") # Use a temp dir for HF Spaces
+if 'gemini_handler' not in st.session_state:
+    st.session_state.gemini_handler = GeminiHandler(api_key=st.session_state.GEMINI_API_KEY)
+if 'visual_engine' not in st.session_state:
+    st.session_state.visual_engine = VisualEngine(output_dir="temp_cinegen_media")
 
-# --- UI Elements ---
-st.title("🎬 CineGen AI: Your Pocket AI Film Studio")
-st.markdown("Let's turn your ideas into cinematic visuals!")
-
-with st.sidebar:
-    st.header("🎨 Creative Controls")
-    user_idea = st.text_area("Enter your core idea/story prompt:", "A lone astronaut discovers a glowing alien artifact on Mars, a sense of wonder and slight dread.", height=100)
-    genre = st.selectbox("Select Genre:", ["Sci-Fi", "Fantasy", "Thriller", "Drama", "Comedy"], index=0)
-    mood = st.selectbox("Select Mood:", ["Suspenseful", "Mysterious", "Epic", "Uplifting", "Dark"], index=0)
-    generate_button = st.button("✨ Generate Cinematic Concept", type="primary")
-
-# --- Session State for Storing Results ---
-if 'story_scenes' not in st.session_state:
-    st.session_state.story_scenes = None
-if 'image_prompts' not in st.session_state:
-    st.session_state.image_prompts = []
-if 'generated_images_paths' not in st.session_state:
+# Story and generated content
+if 'story_scenes' not in st.session_state: # List of scene dicts
+    st.session_state.story_scenes = []
+if 'scene_image_prompts' not in st.session_state: # List of strings
+    st.session_state.scene_image_prompts = []
+if 'generated_images_paths' not in st.session_state: # List of file paths
     st.session_state.generated_images_paths = []
 if 'video_path' not in st.session_state:
     st.session_state.video_path = None
 
-# --- Main Logic ---
-if generate_button and user_idea:
-    st.session_state.story_scenes = None # Reset previous results
-    st.session_state.image_prompts = []
+# For Character Consistency (placeholders)
+if 'character_definitions' not in st.session_state:
+    st.session_state.character_definitions = {} # e.g., {"Eva": "Astronaut, red hair..."}
+
+# For Style Transfer (placeholders)
+if 'style_reference_description' not in st.session_state:
+    st.session_state.style_reference_description = None
+
+# --- Helper Functions ---
+def initialize_new_story():
+    st.session_state.story_scenes = []
+    st.session_state.scene_image_prompts = []
     st.session_state.generated_images_paths = []
     st.session_state.video_path = None
 
-    with st.spinner("Phase 1: Gemini is drafting the script and scene breakdown... 📜"):
-        story_prompt = create_story_breakdown_prompt(user_idea, genre, mood)
-        st.session_state.story_scenes = gemini.generate_story_breakdown(story_prompt)
-
-    if st.session_state.story_scenes:
-        st.toast("Script and scene breakdown complete!", icon="✅")
-        
-        with st.spinner("Phase 2: Gemini is crafting visual prompts for keyframes... 🖼️"):
-            for i, scene in enumerate(st.session_state.story_scenes):
-                scene_desc_for_image = f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')}. Setting: {scene.get('setting_description', 'N/A')}"
-                img_gen_prompt_text = create_image_prompt_from_scene(scene_desc_for_image, scene.get('visual_style_suggestion', 'cinematic'))
-                
-                # Get the actual image description from Gemini
-                # For this demo, we'll use the scene's key_action and setting as the "description" for the placeholder
-                # In a real app, Gemini would generate a detailed image prompt here.
-                # image_description_from_gemini = gemini.generate_image_prompt(img_gen_prompt_text)
-                # For simplicity, we use parts of the scene data directly for the placeholder image text
-                image_description_for_placeholder = f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')}\nStyle: {scene.get('visual_style_suggestion', 'cinematic')}"
-                
-                if image_description_for_placeholder: # was image_description_from_gemini
-                    st.session_state.image_prompts.append(image_description_for_placeholder)
-                    # Simulate image generation by creating a placeholder
-                    img_path = visualizer.create_placeholder_image(
-                        image_description_for_placeholder, # use the generated desc
-                        f"scene_{scene.get('scene_number', i+1)}_placeholder.png"
+def generate_visual_for_scene(scene_index, scene_data):
+    """Generates image prompt and placeholder image for a specific scene."""
+    img_gen_prompt_text = create_image_prompt_from_scene_data(
+        scene_data,
+        st.session_state.character_definitions,
+        st.session_state.style_reference_description
+    )
+    try:
+        actual_image_prompt = st.session_state.gemini_handler.generate_image_prompt(img_gen_prompt_text)
+        if actual_image_prompt:
+            if scene_index < len(st.session_state.scene_image_prompts):
+                st.session_state.scene_image_prompts[scene_index] = actual_image_prompt
+            else:
+                st.session_state.scene_image_prompts.append(actual_image_prompt)
+
+            img_path = st.session_state.visual_engine.create_placeholder_image(
+                f"Scene {scene_data.get('scene_number', scene_index + 1)}: {actual_image_prompt[:100]}...", # Show part of the actual prompt
+                f"scene_{scene_data.get('scene_number', scene_index + 1)}_placeholder.png"
+            )
+            if scene_index < len(st.session_state.generated_images_paths):
+                st.session_state.generated_images_paths[scene_index] = img_path
+            else:
+                st.session_state.generated_images_paths.append(img_path)
+            return True
+    except Exception as e:
+        st.error(f"Error generating visual for Scene {scene_data.get('scene_number', scene_index + 1)}: {e}")
+    return False
+
+# --- UI Sidebar ---
+with st.sidebar:
+    st.title("🎬 CineGen AI Pro")
+    st.markdown("### Creative Controls")
+    user_idea = st.text_area("Enter your core story idea:", "A detective in a cyberpunk city investigates a rogue AI.", height=100, key="user_idea_input")
+    genre = st.selectbox("Genre:", ["Sci-Fi", "Fantasy", "Noir", "Thriller", "Drama"], index=2, key="genre_select")
+    mood = st.selectbox("Mood:", ["Suspenseful", "Mysterious", "Gritty", "Hopeful"], index=2, key="mood_select")
+    num_scenes_val = st.slider("Number of Scenes:", 1, 10, 3, key="num_scenes_slider")
+
+    if st.button("✨ Generate Full Story Concept", type="primary", key="generate_full_story_btn"):
+        initialize_new_story()
+        with st.spinner("Phase 1: Gemini is drafting the script & scene breakdown... 📜"):
+            story_prompt_text = create_story_breakdown_prompt(user_idea, genre, mood, num_scenes_val)
+            try:
+                st.session_state.story_scenes = st.session_state.gemini_handler.generate_story_breakdown(story_prompt_text)
+                st.toast("Script breakdown complete!", icon="✅")
+
+                # Initialize placeholders for prompts and images
+                st.session_state.scene_image_prompts = [""] * len(st.session_state.story_scenes)
+                st.session_state.generated_images_paths = [""] * len(st.session_state.story_scenes)
+
+            except Exception as e:
+                st.error(f"Failed to generate story breakdown: {e}")
+                st.session_state.story_scenes = [] # Clear on failure
+
+        if st.session_state.story_scenes:
+            with st.spinner("Phase 2: Generating initial visual concepts... 🖼️"):
+                for i, scene_data in enumerate(st.session_state.story_scenes):
+                    generate_visual_for_scene(i, scene_data)
+                st.toast("Initial visual concepts generated!", icon="🖼️")
+
+    st.markdown("---")
+    st.markdown("### Advanced Options (Conceptual)")
+    # Character Consistency UI
+    st.subheader("Character Consistency")
+    char_name_input = st.text_input("Character Name (e.g., Eva)", key="char_name")
+    char_desc_input = st.text_area("Character Description (for visual consistency)", key="char_desc", height=80)
+    if st.button("Add/Update Character", key="add_char_btn"):
+        if char_name_input and char_desc_input:
+            st.session_state.character_definitions[char_name_input] = char_desc_input
+            st.success(f"Character '{char_name_input}' defined.")
+        else:
+            st.warning("Please provide both name and description.")
+    if st.session_state.character_definitions:
+        st.write("Defined Characters:")
+        st.json(st.session_state.character_definitions)
+
+    # Style Transfer UI
+    st.subheader("Style Transfer")
+    style_ref_text = st.text_area("Describe Style (e.g., 'Van Gogh inspired, swirling brushstrokes')", key="style_text_input", height=80)
+    if st.button("Apply Textual Style", key="apply_style_btn"):
+        st.session_state.style_reference_description = style_ref_text
+        st.success("Style reference applied. Re-generate visuals to see changes.")
+    # conceptual_style_image = st.file_uploader("Or Upload Style Reference Image (Future Feature)", type=['png', 'jpg'])
+
+
+# --- Main Content Area ---
+st.header("📝 Cinematic Storyboard")
+
+if not st.session_state.story_scenes:
+    st.info("Enter your idea in the sidebar and click 'Generate Full Story Concept' to begin.")
+else:
+    for i, scene_data in enumerate(st.session_state.story_scenes):
+        scene_num = scene_data.get('scene_number', i + 1)
+        st.subheader(f"Scene {scene_num}: {scene_data.get('emotional_beat', 'N/A')}")
+
+        col1, col2 = st.columns([2,3]) # Left for text, Right for image and edit
+
+        with col1: # Scene Details
+            with st.expander("Scene Details", expanded=True):
+                st.markdown(f"**Setting:** {scene_data.get('setting_description', 'N/A')}")
+                st.markdown(f"**Characters:** {', '.join(scene_data.get('characters_involved', []))}")
+                st.markdown(f"**Key Action:** {scene_data.get('key_action', 'N/A')}")
+                st.markdown(f"**Dialogue Snippet:** {scene_data.get('dialogue_snippet', 'N/A')}")
+                st.markdown(f"**Visual Style:** {scene_data.get('visual_style_suggestion', 'N/A')}")
+                st.markdown(f"**Camera:** {scene_data.get('camera_angle_suggestion', 'N/A')}")
+                if i < len(st.session_state.scene_image_prompts) and st.session_state.scene_image_prompts[i]:
+                     with st.popover("View Image Prompt"):
+                        st.code(st.session_state.scene_image_prompts[i])
+
+
+        with col2: # Image and Interactive Editing
+            if i < len(st.session_state.generated_images_paths) and st.session_state.generated_images_paths[i] and os.path.exists(st.session_state.generated_images_paths[i]):
+                st.image(st.session_state.generated_images_paths[i], caption=f"Visual Concept for Scene {scene_num}")
+            else:
+                st.info("Visual concept not yet generated or path is invalid.")
+
+            # Interactive Storyboarding UI
+            with st.popover(f"✏️ Edit Scene {scene_num} Script"):
+                st.markdown("#### Modify Scene Script Details")
+                feedback_script = st.text_area(
+                    "Describe changes to the scene (e.g., 'Make the action more intense', 'Change dialogue to be more hopeful')",
+                    key=f"script_feedback_{i}",
+                    height=100
+                )
+                if st.button(f"🔄 Regenerate Scene {scene_num} Script", key=f"regen_script_btn_{i}"):
+                    if feedback_script:
+                        with st.spinner(f"Gemini is rewriting Scene {scene_num}..."):
+                            regen_prompt = create_scene_regeneration_prompt(
+                                scene_data, feedback_script, st.session_state.story_scenes
+                            )
+                            try:
+                                updated_scene_data = st.session_state.gemini_handler.regenerate_scene_script_details(regen_prompt)
+                                st.session_state.story_scenes[i] = updated_scene_data # Update the specific scene
+                                st.toast(f"Scene {scene_num} script updated!", icon="✍️")
+                                # Automatically regenerate visuals for this updated scene
+                                generate_visual_for_scene(i, updated_scene_data)
+                                st.rerun() # Rerun to reflect changes immediately
+                            except Exception as e:
+                                st.error(f"Error regenerating scene script: {e}")
+                    else:
+                        st.warning("Please provide feedback for script regeneration.")
+
+            with st.popover(f"🎨 Edit Scene {scene_num} Visuals"):
+                st.markdown("#### Modify Visual Concept")
+                feedback_visual = st.text_area(
+                    "Describe visual changes (e.g., 'Make it nighttime', 'Add more fog', 'Use a wider angle')",
+                    key=f"visual_feedback_{i}",
+                    height=100
+                )
+                # Advanced Camera Control (Simple Example)
+                adv_camera = st.selectbox("Adv. Camera (Conceptual):",
+                                          ["Default", "Low Angle", "High Angle", "Dutch Angle", "Close-up"],
+                                          key=f"adv_cam_{i}")
+                if st.button(f"🔄 Regenerate Scene {scene_num} Visuals", key=f"regen_visual_btn_{i}"):
+                    if feedback_visual or adv_camera != "Default":
+                        full_visual_feedback = feedback_visual
+                        if adv_camera != "Default":
+                            full_visual_feedback += f" Emphasize camera: {adv_camera}."
+
+                        with st.spinner(f"Gemini is re-imagining Scene {scene_num} visuals..."):
+                            # We need the current image prompt for this scene to refine it
+                            current_image_prompt = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else ""
+                            if not current_image_prompt: # If no prompt, generate one first
+                                img_gen_prompt_text = create_image_prompt_from_scene_data(
+                                    scene_data,
+                                    st.session_state.character_definitions,
+                                    st.session_state.style_reference_description
+                                )
+                                current_image_prompt = st.session_state.gemini_handler.generate_image_prompt(img_gen_prompt_text)
+                                st.session_state.scene_image_prompts[i] = current_image_prompt
+
+
+                            regen_visual_prompt_text = create_visual_regeneration_prompt(
+                                current_image_prompt,
+                                full_visual_feedback,
+                                scene_data # Pass current scene data for context
+                            )
+                            try:
+                                new_image_prompt = st.session_state.gemini_handler.regenerate_image_prompt_from_feedback(regen_visual_prompt_text)
+                                st.session_state.scene_image_prompts[i] = new_image_prompt # Update the prompt
+
+                                # Now generate the new placeholder image with the new prompt
+                                img_path = st.session_state.visual_engine.create_placeholder_image(
+                                    f"Scene {scene_num}: {new_image_prompt[:100]}...",
+                                    f"scene_{scene_num}_placeholder_v{len(st.session_state.generated_images_paths)}.png" # Avoid overwriting, simple versioning
+                                )
+                                st.session_state.generated_images_paths[i] = img_path
+                                st.toast(f"Visuals for Scene {scene_num} updated!", icon="🎨")
+                                st.rerun()
+                            except Exception as e:
+                                st.error(f"Error regenerating visuals: {e}")
+                    else:
+                        st.warning("Please provide feedback or select an advanced camera angle for visual regeneration.")
+
+        st.markdown("---") # Separator between scenes
+
+    # Video Generation Button (after all scenes are displayed)
+    if st.session_state.story_scenes and st.session_state.generated_images_paths and \
+       all(p and os.path.exists(p) for p in st.session_state.generated_images_paths):
+        if st.button("🎬 Assemble Animatic Video", key="assemble_video_btn"):
+            with st.spinner("Assembling video... This might take a moment."):
+                valid_image_paths = [p for p in st.session_state.generated_images_paths if p and os.path.exists(p)]
+                if valid_image_paths:
+                    st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
+                        valid_image_paths,
+                        output_filename="cinegen_pro_output.mp4",
+                        duration_per_image=3
                     )
-                    st.session_state.generated_images_paths.append(img_path)
+                    st.toast("Animatic video assembled!", icon="🎞️")
+                    st.balloons()
                 else:
-                    st.warning(f"Could not generate image prompt/description for Scene {scene.get('scene_number', i+1)}.")
-        st.toast("Visual prompts and placeholder images generated!", icon="🖼️")
-
-        if st.session_state.generated_images_paths:
-            with st.spinner("Phase 3: Assembling the cinematic video... 🎞️"):
-                st.session_state.video_path = visualizer.create_video_from_images(
-                    st.session_state.generated_images_paths,
-                    output_filename="cinegen_output.mp4",
-                    duration_per_image=3 # Show each image for 3 seconds
-                )
-            st.toast("Video assembled!", icon="🎬")
-            st.balloons()
-        else:
-            st.error("No images were generated, cannot create video.")
+                    st.error("No valid images found to assemble video.")
     else:
-        st.error("Failed to generate story breakdown from Gemini. Please check logs or try a different prompt.")
-
-# --- Display Results ---
-if st.session_state.story_scenes:
-    st.header("📝 Generated Story Breakdown")
-    for i, scene in enumerate(st.session_state.story_scenes):
-        with st.expander(f"Scene {scene.get('scene_number', i+1)}: {scene.get('key_action', 'N/A')[:50]}...", expanded=i==0):
-            st.markdown(f"**Setting:** {scene.get('setting_description', 'N/A')}")
-            st.markdown(f"**Characters:** {', '.join(scene.get('characters_involved', []))}")
-            st.markdown(f"**Key Action:** {scene.get('key_action', 'N/A')}")
-            st.markdown(f"**Visual Style Suggestion:** {scene.get('visual_style_suggestion', 'N/A')}")
-            st.markdown(f"**Camera Angle Suggestion:** {scene.get('camera_angle_suggestion', 'N/A')}")
-            
-            if i < len(st.session_state.image_prompts):
-                 st.markdown(f"**Image Concept (Placeholder Text):**")
-                 st.code(st.session_state.image_prompts[i])
-            if i < len(st.session_state.generated_images_paths):
-                st.image(st.session_state.generated_images_paths[i], caption=f"Visual Concept for Scene {scene.get('scene_number', i+1)}")
-
-
-if st.session_state.video_path and os.path.exists(st.session_state.video_path):
-    st.header("🎬 Generated Cinematic Concept Video")
-    video_file = open(st.session_state.video_path, 'rb')
-    video_bytes = video_file.read()
-    st.video(video_bytes)
-    st.markdown(f"Download your video: [{os.path.basename(st.session_state.video_path)}]({st.session_state.video_path})") # This download might not work directly on Spaces due to pathing, but video will play.
-    # For download on Spaces, you might need to serve it differently or provide a direct link if Spaces keeps it.
-    # A simpler way for Spaces is to just let the user right-click the video player and save.
-else:
-    if generate_button and not st.session_state.video_path and st.session_state.generated_images_paths:
-        st.warning("Video generation might have failed or no images were available.")
-    elif not generate_button :
-        st.info("Enter your idea and click 'Generate Cinematic Concept' to begin!")
+        if st.session_state.story_scenes:
+            st.warning("Some visual concepts are missing. Please ensure all scenes have visuals before assembling video.")
+
+
+    if st.session_state.video_path and os.path.exists(st.session_state.video_path):
+        st.header("🎬 Generated Animatic Video")
+        video_file = open(st.session_state.video_path, 'rb')
+        video_bytes = video_file.read()
+        st.video(video_bytes)
+        # For download on Spaces, let user right-click. The path might be tricky for direct download link.
+        st.markdown(f"Video saved at: `{st.session_state.video_path}` (within the Space's file system)")