Update app.py

app.py

@@ -1,45 +1,53 @@
 # app.py
 import streamlit as st
 from core.gemini_handler import GeminiHandler
-from core.visual_engine import VisualEngine
+from core.visual_engine import VisualEngine # Your updated VisualEngine
 from core.prompt_engineering import (
     create_story_breakdown_prompt,
-    create_image_prompt_from_scene_data,
+    create_image_prompt_from_scene_data, # Generates the TEXT prompt for AI/placeholder
     create_scene_regeneration_prompt,
-    create_visual_regeneration_prompt
+    create_visual_regeneration_prompt   # Regenerates the TEXT prompt
 )
 import os
-import json # For debugging and display
+# import json # Only if you need to st.json for debugging, not essential for core logic
 
 # --- Configuration & Initialization ---
-st.set_page_config(page_title="CineGen AI Pro", layout="wide")
+st.set_page_config(page_title="CineGen AI Pro", layout="wide", initial_sidebar_state="expanded")
 
-# --- Global State Variables --- (Using session state for persistence)
+# --- 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.")
+    except KeyError: # For local dev if secrets aren't set yet
+        if "GEMINI_API_KEY" in os.environ:
+             st.session_state.GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
+        else:
+            st.error("GEMINI_API_KEY not found in secrets or environment variables. Please add it.")
+            st.stop()
+    except Exception as e: # Catch any other secrets-related errors
+        st.error(f"Error loading secrets: {e}")
         st.stop()
 
+
 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:
+    # Ensure VisualEngine is initialized (it will print font loading status)
     st.session_state.visual_engine = VisualEngine(output_dir="temp_cinegen_media")
 
 # 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
+if 'scene_image_prompts' not in st.session_state: # List of textual image prompts
     st.session_state.scene_image_prompts = []
-if 'generated_images_paths' not in st.session_state: # List of file paths
+if 'generated_images_paths' not in st.session_state: # List of file paths for generated visuals
     st.session_state.generated_images_paths = []
 if 'video_path' not in st.session_state:
     st.session_state.video_path = None
 
 # For Character Consistency (placeholders)
 if 'character_definitions' not in st.session_state:
-    st.session_state.character_definitions = {} # e.g., {"Eva": "Astronaut, red hair..."}
+    st.session_state.character_definitions = {} 
 
 # For Style Transfer (placeholders)
 if 'style_reference_description' not in st.session_state:
@@ -51,90 +59,137 @@ def initialize_new_story():
     st.session_state.scene_image_prompts = []
     st.session_state.generated_images_paths = []
     st.session_state.video_path = None
-
-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
+    # Clear old media files if any (optional, good for space management on persistent storage)
+    # For Hugging Face Spaces, temp files are usually cleared on restart/rebuild anyway.
+    # if os.path.exists(st.session_state.visual_engine.output_dir):
+    #     for f in os.listdir(st.session_state.visual_engine.output_dir):
+    #         os.remove(os.path.join(st.session_state.visual_engine.output_dir, f))
+
+
+def generate_visual_for_scene_wrapper(scene_index, scene_data, is_regeneration=False, version_count=1):
+    """
+    Wrapper to generate the textual image prompt and then the actual visual.
+    Updates session state for prompts and image paths.
+    """
+    st.info(f"Generating visual concept for Scene {scene_data.get('scene_number', scene_index + 1)}...")
+    
+    # 1. Get/Regenerate the detailed image prompt TEXT
+    if is_regeneration and scene_index < len(st.session_state.scene_image_prompts):
+        # For regeneration, we assume the prompt text is already updated in session state
+        # (e.g., by a call to create_visual_regeneration_prompt and Gemini)
+        textual_image_prompt = st.session_state.scene_image_prompts[scene_index]
+    else:
+        textual_image_prompt = create_image_prompt_from_scene_data(
+            scene_data,
+            st.session_state.character_definitions,
+            st.session_state.style_reference_description
+        )
+        # Optional: Call Gemini here to refine this initial prompt if desired
+        # textual_image_prompt = st.session_state.gemini_handler.generate_image_prompt(textual_image_prompt)
+
+    if not textual_image_prompt:
+        st.error(f"Failed to create textual image prompt for Scene {scene_data.get('scene_number', scene_index + 1)}.")
+        return False
+
+    # Update session state for the textual prompt
+    if scene_index < len(st.session_state.scene_image_prompts):
+        st.session_state.scene_image_prompts[scene_index] = textual_image_prompt
+    else:
+        # Ensure list is long enough before appending (shouldn't happen if initialized correctly)
+        while len(st.session_state.scene_image_prompts) <= scene_index:
+             st.session_state.scene_image_prompts.append("")
+        st.session_state.scene_image_prompts[scene_index] = textual_image_prompt
+
+
+    # 2. Generate the actual visual (AI or placeholder) using the textual prompt
+    image_filename = f"scene_{scene_data.get('scene_number', scene_index + 1)}_visual_v{version_count}.png"
+    
+    generated_image_path = st.session_state.visual_engine.generate_image_visual(
+        textual_image_prompt,
+        image_filename
     )
-    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
+
+    if generated_image_path and os.path.exists(generated_image_path):
+        st.success(f"Visual concept for Scene {scene_data.get('scene_number', scene_index + 1)} generated: {os.path.basename(generated_image_path)}")
+        if scene_index < len(st.session_state.generated_images_paths):
+            st.session_state.generated_images_paths[scene_index] = generated_image_path
+        else:
+            while len(st.session_state.generated_images_paths) <= scene_index:
+                st.session_state.generated_images_paths.append(None)
+            st.session_state.generated_images_paths[scene_index] = generated_image_path
+        return True
+    else:
+        st.error(f"Visual generation failed for Scene {scene_data.get('scene_number', scene_index + 1)} (path: {generated_image_path}).")
+        if scene_index < len(st.session_state.generated_images_paths):
+            st.session_state.generated_images_paths[scene_index] = None
+        else:
+            while len(st.session_state.generated_images_paths) <= scene_index:
+                st.session_state.generated_images_paths.append(None)
+            st.session_state.generated_images_paths[scene_index] = None
+        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")
+    user_idea = st.text_area("Enter your core story idea:", "A detective in a cyberpunk city investigates a rogue AI that believes it's the next step in evolution.", height=100, key="user_idea_input")
+    genre = st.selectbox("Genre:", ["Sci-Fi", "Fantasy", "Noir", "Thriller", "Drama", "Horror"], index=2, key="genre_select")
+    mood = st.selectbox("Mood:", ["Suspenseful", "Mysterious", "Gritty", "Epic", "Dark", "Hopeful"], index=2, key="mood_select")
+    num_scenes_val = st.slider("Number of Scenes:", 1, 5, 3, key="num_scenes_slider") # Max 5 for now
 
     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.")
+        if not user_idea.strip():
+            st.warning("Please enter a story idea.")
         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)
+            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 based on number of scenes
+                    num_actual_scenes = len(st.session_state.story_scenes)
+                    st.session_state.scene_image_prompts = [""] * num_actual_scenes
+                    st.session_state.generated_images_paths = [None] * num_actual_scenes
+
+                except Exception as e:
+                    st.error(f"Failed to generate story breakdown: {e}")
+                    st.session_state.story_scenes = []
+
+            if st.session_state.story_scenes:
+                with st.spinner("Phase 2: Generating initial visual concepts... 🖼️"):
+                    success_count = 0
+                    for i, scene_data_loop_var in enumerate(st.session_state.story_scenes):
+                        if generate_visual_for_scene_wrapper(i, scene_data_loop_var, version_count=1):
+                            success_count +=1
+                    if success_count == len(st.session_state.story_scenes):
+                        st.toast("Initial visual concepts generated!", icon="🖼️")
+                    else:
+                        st.warning(f"{success_count}/{len(st.session_state.story_scenes)} visual concepts generated. Some may have failed.")
 
-    # 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'])
 
+    st.markdown("---")
+    st.markdown("### Advanced Options (Conceptual)")
+    with st.expander("Character Consistency", expanded=False):
+        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:")
+            for char, desc in st.session_state.character_definitions.items():
+                st.caption(f"**{char}:** {desc}")
+
+    with st.expander("Style Transfer", expanded=False):
+        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 or full story to see changes.")
 
 # --- Main Content Area ---
 st.header("📝 Cinematic Storyboard")
@@ -142,136 +197,152 @@ 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')}")
+    for i, scene_data_display in enumerate(st.session_state.story_scenes):
+        scene_num_display = scene_data_display.get('scene_number', i + 1)
+        # Use a unique key part for expanders if scene numbers can repeat or change
+        expander_key_part = scene_data_display.get('key_action', f"scene{i}")[:20].replace(" ", "")
 
-        col1, col2 = st.columns([2,3]) # Left for text, Right for image and edit
+        st.subheader(f"Scene {scene_num_display}: {scene_data_display.get('emotional_beat', 'Untitled Scene')}")
+
+        col1, col2 = st.columns([2, 3])
 
         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')}")
+                st.markdown(f"**Setting:** {scene_data_display.get('setting_description', 'N/A')}")
+                st.markdown(f"**Characters:** {', '.join(scene_data_display.get('characters_involved', []))}")
+                st.markdown(f"**Key Action:** {scene_data_display.get('key_action', 'N/A')}")
+                st.markdown(f"**Dialogue Snippet:** `\"{scene_data_display.get('dialogue_snippet', '...')}\"`")
+                st.markdown(f"**Visual Style:** {scene_data_display.get('visual_style_suggestion', 'N/A')}")
+                st.markdown(f"**Camera:** {scene_data_display.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 st.popover("View Image Prompt Text"):
+                        st.markdown(f"**Textual Prompt for Image Generation:**")
+                        st.code(st.session_state.scene_image_prompts[i], language='text')
 
         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}")
+            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_display}")
             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 st.session_state.story_scenes: # Only show if story generation was attempted
+                    st.warning("Visual concept not generated or path is invalid for this scene.")
+
+            # Interactive Storyboarding UI using unique keys for each scene's widgets
+            popover_script_key = f"pop_script_{expander_key_part}_{i}"
+            popover_visual_key = f"pop_visual_{expander_key_part}_{i}"
+            
+            with st.popover(f"✏️ Edit Scene {scene_num_display} Script", key=popover_script_key):
+                feedback_script = st.text_area("Describe changes to script details:",
+                    key=f"script_feedback_{expander_key_part}_{i}", height=100)
+                if st.button(f"🔄 Regenerate Scene {scene_num_display} Script", key=f"regen_script_btn_{expander_key_part}_{i}"):
                     if feedback_script:
-                        with st.spinner(f"Gemini is rewriting Scene {scene_num}..."):
+                        with st.spinner(f"Gemini is rewriting Scene {scene_num_display}..."):
                             regen_prompt = create_scene_regeneration_prompt(
-                                scene_data, feedback_script, st.session_state.story_scenes
+                                scene_data_display, 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
+                                st.session_state.story_scenes[i] = updated_scene_data
+                                st.toast(f"Scene {scene_num_display} script updated!", icon="✍️")
+                                # Regenerate visuals for this updated scene with new script details
+                                # Increment version for the image filename
+                                current_version = 1
+                                if st.session_state.generated_images_paths[i]:
+                                    try: # Try to extract version from existing filename
+                                        base, ext = os.path.splitext(os.path.basename(st.session_state.generated_images_paths[i]))
+                                        if '_v' in base:
+                                            current_version = int(base.split('_v')[-1]) + 1
+                                    except ValueError: pass # Keep current_version as 1 if parsing fails
+                                
+                                generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=True, version_count=current_version)
+                                st.rerun()
                             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
+            with st.popover(f"🎨 Edit Scene {scene_num_display} Visuals", key=popover_visual_key):
+                current_prompt_for_edit = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else "No prompt yet."
+                st.caption("Current Image Prompt Text:")
+                st.code(current_prompt_for_edit, language='text')
+                
+                feedback_visual = st.text_area("Describe visual changes to apply to the prompt:",
+                    key=f"visual_feedback_{expander_key_part}_{i}", height=100)
+                if st.button(f"🔄 Regenerate Scene {scene_num_display} Visuals", key=f"regen_visual_btn_{expander_key_part}_{i}"):
+                    if feedback_visual:
+                        with st.spinner(f"Refining image prompt for Scene {scene_num_display}..."):
+                            # This creates a prompt FOR GEMINI to refine the image prompt
+                            prompt_refinement_request = create_visual_regeneration_prompt(
+                                current_prompt_for_edit,
+                                feedback_visual,
+                                scene_data_display 
                             )
                             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="🎨")
+                                # Gemini refines the textual image prompt
+                                refined_textual_image_prompt = st.session_state.gemini_handler.regenerate_image_prompt_from_feedback(prompt_refinement_request)
+                                st.session_state.scene_image_prompts[i] = refined_textual_image_prompt
+                                st.toast(f"Image prompt for Scene {scene_num_display} refined!", icon="💡")
+
+                                # Now generate the new visual using the refined prompt
+                                current_version = 1
+                                if st.session_state.generated_images_paths[i]:
+                                    try:
+                                        base, ext = os.path.splitext(os.path.basename(st.session_state.generated_images_paths[i]))
+                                        if '_v' in base:
+                                            current_version = int(base.split('_v')[-1]) + 1
+                                    except ValueError: pass
+                                
+                                generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=current_version)
                                 st.rerun()
                             except Exception as e:
-                                st.error(f"Error regenerating visuals: {e}")
+                                st.error(f"Error refining image prompt or regenerating visual: {e}")
                     else:
-                        st.warning("Please provide feedback or select an advanced camera angle for visual regeneration.")
-
-        st.markdown("---") # Separator between scenes
+                        st.warning("Please provide feedback for visual regeneration.")
+        st.markdown("---")
 
-    # 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"):
+    # Video Generation Button
+    if st.session_state.story_scenes and any(p for p in st.session_state.generated_images_paths): # Check if any image exists
+        if st.button("🎬 Assemble Animatic Video", key="assemble_video_btn", type="primary"):
             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:
+                valid_image_paths_for_video = [p for p in st.session_state.generated_images_paths if p and os.path.exists(p)]
+                if valid_image_paths_for_video:
                     st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
-                        valid_image_paths,
-                        output_filename="cinegen_pro_output.mp4",
+                        valid_image_paths_for_video,
+                        output_filename="cinegen_pro_animatic.mp4",
                         duration_per_image=3
                     )
-                    st.toast("Animatic video assembled!", icon="🎞️")
-                    st.balloons()
+                    if st.session_state.video_path and os.path.exists(st.session_state.video_path):
+                        st.toast("Animatic video assembled!", icon="🎞️")
+                        st.balloons()
+                    else:
+                        st.error("Video assembly failed. Check logs for details.")
                 else:
-                    st.error("No valid images found to assemble video.")
-    else:
-        if st.session_state.story_scenes:
-            st.warning("Some visual concepts are missing. Please ensure all scenes have visuals before assembling video.")
+                    st.error("No valid images found to assemble video. Please generate visuals for scenes first.")
+    elif st.session_state.story_scenes: # Story exists but no images yet
+         st.warning("Generate visuals for scenes before assembling the 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)")
+        try:
+            with open(st.session_state.video_path, 'rb') as video_file:
+                video_bytes = video_file.read()
+            st.video(video_bytes)
+            st.markdown(f"Video saved at: `{st.session_state.video_path}` (within the Space's file system)")
+            # Provide a download button for the video
+            with open(st.session_state.video_path, "rb") as fp_download:
+                st.download_button(
+                    label="Download Animatic Video",
+                    data=fp_download,
+                    file_name=os.path.basename(st.session_state.video_path),
+                    mime="video/mp4"
+                )
+        except Exception as e:
+            st.error(f"Error displaying or preparing video for download: {e}")
+
+
+# --- Footer or additional info ---
+st.sidebar.markdown("---")
+st.sidebar.info("CineGen AI by [Your Name/Company]. Powered by Gemini & Streamlit.")