Update app.py

app.py

@@ -13,62 +13,42 @@ import os
 # --- Configuration & Initialization ---
 st.set_page_config(page_title="CineGen AI Pro", layout="wide", initial_sidebar_state="expanded")
 
-# --- Global State Variables (Using session state for persistence) ---
-# Load Gemini API Key
+# --- Global State Variables ---
 if 'GEMINI_API_KEY' not in st.session_state:
-    try:
-        st.session_state.GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
-    except KeyError: 
-        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: 
-        st.error(f"Error loading Gemini API Key from secrets: {e}")
-        st.stop()
+    try: st.session_state.GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
+    except: 
+        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."); st.stop()
 
-# Initialize Gemini Handler
 if 'gemini_handler' not in st.session_state:
     st.session_state.gemini_handler = GeminiHandler(api_key=st.session_state.GEMINI_API_KEY)
 
-# Initialize Visual Engine and set OpenAI API Key
 if 'visual_engine' not in st.session_state:
     st.session_state.visual_engine = VisualEngine(output_dir="temp_cinegen_media")
-    openai_key_from_secrets = None
-    try:
-        if "OPENAI_API_KEY" in st.secrets:
-            openai_key_from_secrets = st.secrets["OPENAI_API_KEY"]
-    except AttributeError: 
-        print("st.secrets not available (likely local dev without secrets.toml). Checking environment variables.")
-    except Exception as e:
-        print(f"Could not access st.secrets for OPENAI_API_KEY: {e}")
+    openai_key = None
+    try: 
+        if "OPENAI_API_KEY" in st.secrets: openai_key = st.secrets["OPENAI_API_KEY"]
+    except AttributeError: print("st.secrets not available. Checking env for OpenAI key.")
+    except Exception as e: print(f"Error accessing st.secrets for OPENAI_API_KEY: {e}")
+    if not openai_key and "OPENAI_API_KEY" in os.environ: openai_key = os.environ["OPENAI_API_KEY"]
+    st.session_state.visual_engine.set_openai_api_key(openai_key)
+    if not openai_key and 'openai_warning_shown' not in st.session_state: 
+        st.sidebar.warning("OpenAI Key missing. DALL-E disabled, using placeholders.")
+        st.session_state.openai_warning_shown = True # Show only once per session
+    elif openai_key and 'openai_success_shown' not in st.session_state:
+        # st.sidebar.caption("OpenAI Key loaded. DALL-E ready.") # Can be a bit noisy
+        st.session_state.openai_success_shown = True
 
-    if not openai_key_from_secrets and "OPENAI_API_KEY" in os.environ: 
-        openai_key_from_secrets = os.environ["OPENAI_API_KEY"]
-    
-    if openai_key_from_secrets:
-        st.session_state.visual_engine.set_openai_api_key(openai_key_from_secrets)
-    else:
-        st.session_state.visual_engine.set_openai_api_key(None) 
-        st.sidebar.caption("OpenAI API Key for DALL-E not found. Visuals will be placeholders.")
 
-# Story and generated content state
-if 'story_scenes' not in st.session_state: 
-    st.session_state.story_scenes = []
-if 'scene_image_prompts' not in st.session_state: 
-    st.session_state.scene_image_prompts = []
-if 'generated_images_paths' not in st.session_state: 
-    st.session_state.generated_images_paths = []
-if 'video_path' not in st.session_state:
-    st.session_state.video_path = None
-if 'character_definitions' not in st.session_state:
-    st.session_state.character_definitions = {} 
-if 'style_reference_description' not in st.session_state:
-    st.session_state.style_reference_description = None
+if 'story_scenes' not in st.session_state: st.session_state.story_scenes = []
+if 'scene_image_prompts' not in st.session_state: st.session_state.scene_image_prompts = []
+if 'generated_images_paths' not in st.session_state: st.session_state.generated_images_paths = []
+if 'video_path' not in st.session_state: st.session_state.video_path = None
+if 'character_definitions' not in st.session_state: st.session_state.character_definitions = {} 
+if 'style_reference_description' not in st.session_state: st.session_state.style_reference_description = ""
 
 # --- Helper Functions ---
-def initialize_new_story():
+def initialize_new_story(): # Remains the same
     st.session_state.story_scenes = []
     st.session_state.scene_image_prompts = []
     st.session_state.generated_images_paths = []
@@ -77,26 +57,26 @@ def initialize_new_story():
 def generate_visual_for_scene_wrapper(scene_index, scene_data, is_regeneration=False, version_count=1):
     scene_num_for_log = scene_data.get('scene_number', scene_index + 1)
     textual_image_prompt = ""
+
     if is_regeneration and scene_index < len(st.session_state.scene_image_prompts) and st.session_state.scene_image_prompts[scene_index]:
-        textual_image_prompt = st.session_state.scene_image_prompts[scene_index]
-    else:
+        textual_image_prompt = st.session_state.scene_image_prompts[scene_index] # Use the already refined prompt
+    else: # Initial generation
         textual_image_prompt = create_image_prompt_from_scene_data(
             scene_data,
-            st.session_state.character_definitions,
-            st.session_state.style_reference_description
+            st.session_state.character_definitions, 
+            st.session_state.style_reference_description 
         )
-    if not textual_image_prompt:
-        return False
+    if not textual_image_prompt: return False
+
+    # Update session state with the textual prompt used (might be initial or refined)
     if scene_index >= len(st.session_state.scene_image_prompts):
-        while len(st.session_state.scene_image_prompts) <= scene_index:
-             st.session_state.scene_image_prompts.append("")
+        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
+
     image_filename = f"scene_{scene_num_for_log}_visual_v{version_count}.png"
-    generated_image_path = st.session_state.visual_engine.generate_image_visual(
-        textual_image_prompt, image_filename
-    )
-    while len(st.session_state.generated_images_paths) <= scene_index:
-        st.session_state.generated_images_paths.append(None)
+    generated_image_path = st.session_state.visual_engine.generate_image_visual(textual_image_prompt, image_filename)
+    
+    while len(st.session_state.generated_images_paths) <= scene_index: st.session_state.generated_images_paths.append(None)
     if generated_image_path and os.path.exists(generated_image_path):
         st.session_state.generated_images_paths[scene_index] = generated_image_path
         return True
@@ -109,211 +89,191 @@ with st.sidebar:
     st.title("🎬 CineGen AI Pro")
     st.markdown("### Creative Controls")
     user_idea = st.text_area("Enter your core story idea:", "A lone astronaut discovers a glowing alien artifact on Mars, a sense of wonder and slight dread.", height=100, key="user_idea_input")
-    genre = st.selectbox("Genre:", ["Sci-Fi", "Fantasy", "Noir", "Thriller", "Drama", "Horror", "Cyberpunk"], index=6, key="genre_select")
-    mood = st.selectbox("Mood:", ["Suspenseful", "Mysterious", "Gritty", "Epic", "Dark", "Hopeful", "Wonder"], index=6, key="mood_select")
-    num_scenes_val = st.slider("Number of Scenes:", 1, 5, 1, key="num_scenes_slider") # Default to 1 for faster testing
+    genre = st.selectbox("Genre:", ["Cyberpunk", "Sci-Fi", "Fantasy", "Noir", "Thriller", "Drama", "Horror"], index=0, key="genre_select")
+    mood = st.selectbox("Mood:", ["Wonder", "Suspenseful", "Mysterious", "Gritty", "Epic", "Dark", "Hopeful"], index=0, key="mood_select")
+    num_scenes_val = st.slider("Number of Scenes:", 1, 3, 1, key="num_scenes_slider") # Max 3 for DALL-E cost/time
 
     if st.button("✨ Generate Full Story Concept", type="primary", key="generate_full_story_btn", use_container_width=True):
         initialize_new_story()
-        if not user_idea.strip():
-            st.warning("Please enter a story idea.")
+        if not user_idea.strip(): st.warning("Please enter a story idea.")
         else:
-            with st.status("Generating story...", expanded=True) as status_main:
-                st.write("Phase 1: Gemini is drafting the script & scene breakdown... 📜")
+            with st.status("Generating story and visuals...", expanded=True) as status_main:
+                st.write("Phase 1: Drafting script with Gemini... 📜")
                 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)
-                    status_main.update(label="Script breakdown complete! ✅", state="running", expanded=True)
+                    status_main.update(label="Script complete! ✅ Generating visuals...", state="running")
                     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:
-                    status_main.update(label=f"Failed to generate story breakdown: {e}", state="error", expanded=True)
-                    st.session_state.story_scenes = [] 
-                    st.stop() 
+                    status_main.update(label=f"Script generation failed: {e}", state="error"); st.stop()
+                
                 if st.session_state.story_scenes:
-                    st.write("Phase 2: Generating initial visual concepts... 🖼️ (This may take time with DALL-E)")
+                    st.write("Phase 2: Creating visuals with DALL-E (this can take time)... 🖼️")
                     success_count = 0
-                    for i_loop, scene_data_loop_var in enumerate(st.session_state.story_scenes):
-                        scene_num_log = scene_data_loop_var.get('scene_number', i_loop + 1)
-                        st.write(f"Generating visual for Scene {scene_num_log}...")
-                        if generate_visual_for_scene_wrapper(i_loop, scene_data_loop_var, version_count=1):
-                            success_count +=1
-                            st.write(f"Visual for Scene {scene_num_log} done.")
-                        else:
-                            st.write(f"Visual for Scene {scene_num_log} failed.")
-                    if success_count == len(st.session_state.story_scenes):
-                        status_main.update(label="All concepts generated successfully! 🎉", state="complete", expanded=False)
-                    elif success_count > 0:
-                         status_main.update(label=f"{success_count}/{len(st.session_state.story_scenes)} visuals generated.", state="warning", expanded=False)
-                    else:
-                        status_main.update(label="Visual concept generation failed for all scenes.", state="error", expanded=False)
+                    for i, scene_data in enumerate(st.session_state.story_scenes):
+                        st.write(f"  Processing Scene {scene_data.get('scene_number', i + 1)}...")
+                        if generate_visual_for_scene_wrapper(i, scene_data, version_count=1): success_count +=1
+                    
+                    if success_count == len(st.session_state.story_scenes): status_main.update(label="All done! Story and visuals ready. 🎉", state="complete", expanded=False)
+                    elif success_count > 0: status_main.update(label=f"{success_count}/{len(st.session_state.story_scenes)} visuals OK.", state="warning", expanded=False)
+                    else: status_main.update(label="Visual generation failed for all scenes.", state="error", expanded=False)
 
     st.markdown("---")
     st.markdown("### Advanced Options")
     with st.expander("Character Consistency", expanded=False):
-        char_name_input = st.text_input("Character Name (e.g., Eva)", key="char_name_adv_input")
-        char_desc_input = st.text_area("Character Description (e.g., 'female, short red hair, green eyes, wearing a blue jumpsuit')", key="char_desc_adv_input", height=100)
+        char_name_input = st.text_input("Character Name (e.g., Jax)", key="char_name_adv_input")
+        char_desc_input = st.text_area("Character Description (e.g., 'male astronaut, rugged, dark short hair, blue eyes, wearing a worn white and orange spacesuit with a cracked visor helmet slung at his hip')", key="char_desc_adv_input", height=120)
         if st.button("Add/Update Character", key="add_char_adv_btn"):
             if char_name_input and char_desc_input:
-                st.session_state.character_definitions[char_name_input.strip().lower()] = char_desc_input.strip()
+                st.session_state.character_definitions[char_name_input.strip().lower()] = char_desc_input.strip() # Use lower for key
                 st.success(f"Character '{char_name_input.strip()}' defined.")
-            else:
-                st.warning("Please provide both name and description.")
+            else: st.warning("Provide both name and description.")
         if st.session_state.character_definitions:
             st.caption("Defined Characters:")
-            for char_key, desc_val in st.session_state.character_definitions.items():
-                st.markdown(f"**{char_key.title()}:** _{desc_val}_")
+            for char_key, desc_val in st.session_state.character_definitions.items(): st.markdown(f"**{char_key.title()}:** _{desc_val}_")
 
-    with st.expander("Style Transfer", expanded=False):
-        style_ref_text = st.text_area("Describe Visual Style (e.g., 'impressionistic oil painting, vibrant colors')", key="style_text_adv_input", height=100)
-        if st.button("Apply Textual Style", key="apply_style_adv_btn"):
-            st.session_state.style_reference_description = style_ref_text.strip()
-            st.success("Style reference applied. Re-generate visuals or full story to see changes.")
+    with st.expander("Style Controls", expanded=True): # Expanded by default
+        predefined_styles = {
+            "Default (Cinematic Photorealism)": "",
+            "Gritty Neo-Noir": "neo-noir aesthetic, gritty realism, deep shadows, rain-slicked streets, desaturated colors with vibrant neon highlights, anamorphic lens flares, film grain",
+            "Epic Fantasy Matte Painting": "epic fantasy matte painting style, vast landscapes, dramatic skies, painterly textures, rich colors, Lord of the Rings inspired, highly detailed",
+            "Impressionistic Dream": "impressionistic oil painting style, visible brushstrokes, soft focus, dreamlike atmosphere, pastel color palette, ethereal lighting",
+            "Vintage Sci-Fi (70s/80s Film)": "retro sci-fi aesthetic, analog film look, subtle film grain, practical effects feel, slightly desaturated colors, chrome and plastic textures, classic 70s/80s anamorphic lens flares, detailed spaceship interiors or alien landscapes",
+            "Modern Anime Key Visual": "high-detail modern anime key visual style, dynamic compositions, vibrant saturated colors, crisp cel-shaded look with intricate lighting and deep shadows, expressive character designs with detailed eyes"
+        }
+        selected_preset = st.selectbox("Base Style Preset:", options=list(predefined_styles.keys()), key="style_preset_select_adv")
+        custom_keywords = st.text_area("Add Custom Style Keywords:", key="custom_style_keywords_adv_input", height=80, placeholder="e.g., 'Dutch angle, fisheye lens, sepia tone'")
+        
+        current_style_desc = st.session_state.style_reference_description
+        if st.button("Apply & Set Styles", key="apply_styles_adv_btn"):
+            final_desc = predefined_styles[selected_preset]
+            if custom_keywords.strip():
+                final_desc = f"{final_desc}, {custom_keywords.strip()}" if final_desc else custom_keywords.strip()
+            st.session_state.style_reference_description = final_desc.strip()
+            current_style_desc = final_desc.strip() # Update for immediate display
+            if current_style_desc: st.success("Styles applied! Regenerate story or visuals.")
+            else: st.info("Default style (no specific additions).")
+        
+        if current_style_desc: st.caption(f"Active style prompt additions: \"{current_style_desc}\"")
+        else: st.caption("No specific style prompt additions active.")
 
-# --- Main Content Area ---
-st.header("📝 Cinematic Storyboard")
 
+# --- Main Content Area (largely same as previous full version, with minor tweaks) ---
+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.")
+    st.info("Enter your idea in the sidebar and click '✨ Generate Full Story Concept' to begin.")
 else:
     for i, scene_data_display in enumerate(st.session_state.story_scenes):
+        # ... (scene_num_display, unique_key_part, subheader, columns - same) ...
         scene_num_display = scene_data_display.get('scene_number', i + 1)
         action_summary = scene_data_display.get('key_action', f"scene{i}")
         key_part_raw = ''.join(e for e in action_summary if e.isalnum() or e.isspace())
         unique_key_part = key_part_raw.replace(" ", "_")[:20]
-
         st.subheader(f"Scene {scene_num_display}: {scene_data_display.get('emotional_beat', 'Untitled Scene')}")
         col1, col2 = st.columns([2, 3])
 
-        with col1: 
+        with col1: # Scene Details
             with st.expander("Scene Details", expanded=True):
                 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')}")
-                
+                st.markdown(f"**Visual Style Suggestion (from Gemini):** {scene_data_display.get('visual_style_suggestion', 'N/A')}")
+                st.markdown(f"**Camera Suggestion (from Gemini):** {scene_data_display.get('camera_angle_suggestion', 'N/A')}")
                 current_textual_prompt = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else None
                 if current_textual_prompt:
-                     with st.popover("View Image Prompt Text"): 
-                        st.markdown(f"**Textual Prompt for Image Generation:**")
-                        st.code(current_textual_prompt, language='text')
-        with col2: 
+                     with st.popover("View Full Image Prompt"): 
+                        st.markdown(f"**Full Textual Prompt Sent to DALL-E:**"); st.code(current_textual_prompt, language='text')
+        
+        with col2: # Image and Edit Buttons
             current_image_path = st.session_state.generated_images_paths[i] if i < len(st.session_state.generated_images_paths) else None
             if current_image_path and os.path.exists(current_image_path):
-                # MODIFICATION HERE: Removed use_column_width=True
                 st.image(current_image_path, caption=f"Visual Concept for Scene {scene_num_display}") 
             else:
-                if st.session_state.story_scenes: 
-                    st.caption("Visual for this scene is pending or failed to generate.")
+                if st.session_state.story_scenes: st.caption("Visual for this scene is pending or failed.")
 
+            # Edit Popovers (logic remains largely the same, calls generate_visual_for_scene_wrapper)
             with st.popover(f"✏️ Edit Scene {scene_num_display} Script"): 
-                feedback_script = st.text_area("Describe changes to script details:",
-                    key=f"script_feedback_{unique_key_part}_{i}", height=100) 
-                if st.button(f"🔄 Regenerate Scene {scene_num_display} Script", key=f"regen_script_btn_{unique_key_part}_{i}"): 
+                feedback_script = st.text_area("Describe script changes:", key=f"script_feedback_{unique_key_part}_{i}", height=100) 
+                if st.button(f"🔄 Update Scene {scene_num_display} Script", key=f"regen_script_btn_{unique_key_part}_{i}"): 
                     if feedback_script:
-                        with st.status(f"Rewriting Scene {scene_num_display} script...", expanded=True) as status_script_regen:
-                            regen_prompt = create_scene_regeneration_prompt(
-                                scene_data_display, feedback_script, st.session_state.story_scenes
-                            )
+                        with st.status(f"Updating Scene {scene_num_display}...", expanded=True) as status_script_regen:
+                            regen_prompt = create_scene_regeneration_prompt(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 
                                 status_script_regen.update(label="Script updated! Regenerating visual...", state="running")
                                 version = 1
                                 if current_image_path:
-                                    try: 
-                                        base, _ = os.path.splitext(os.path.basename(current_image_path))
-                                        if '_v' in base: version = int(base.split('_v')[-1]) + 1
-                                    except : pass
-                                if generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=True, version_count=version):
+                                    try: base, _ = os.path.splitext(os.path.basename(current_image_path)); version = int(base.split('_v')[-1]) + 1 if '_v' in base else 2
+                                    except: version = 2 # Default to v2 if parsing fails
+                                if generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=False, version_count=version): # is_regeneration=False to use new scene_data for prompt
                                     status_script_regen.update(label="Scene script & visual updated! 🎉", state="complete", expanded=False)
-                                else:
-                                    status_script_regen.update(label="Script updated, but visual failed.", state="warning", expanded=False)
+                                else: status_script_regen.update(label="Script updated, visual failed.", state="warning", expanded=False)
                                 st.rerun()
-                            except Exception as e:
-                                status_script_regen.update(label=f"Error regenerating scene script: {e}", state="error")
-                    else:
-                        st.warning("Please provide feedback for script regeneration.")
+                            except Exception as e: status_script_regen.update(label=f"Error: {e}", state="error")
+                    else: st.warning("Please provide feedback.")
 
             with st.popover(f"🎨 Edit Scene {scene_num_display} Visuals"): 
-                prompt_for_edit_visuals = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else "No prompt text available."
-                st.caption("Current Image Prompt Text:")
-                st.code(prompt_for_edit_visuals, language='text')
-                feedback_visual = st.text_area("Describe visual changes to apply to the prompt:",
-                    key=f"visual_feedback_{unique_key_part}_{i}", height=100) 
-                if st.button(f"🔄 Regenerate Scene {scene_num_display} Visuals", key=f"regen_visual_btn_{unique_key_part}_{i}"): 
+                prompt_for_edit_visuals = st.session_state.scene_image_prompts[i] if i < len(st.session_state.scene_image_prompts) else "No prompt."
+                st.caption("Current Full Image Prompt:"); st.code(prompt_for_edit_visuals, language='text')
+                feedback_visual = st.text_area("Describe visual changes for the prompt:", key=f"visual_feedback_{unique_key_part}_{i}", height=100) 
+                if st.button(f"🔄 Update Scene {scene_num_display} Visuals", key=f"regen_visual_btn_{unique_key_part}_{i}"): 
                     if feedback_visual:
-                        with st.status(f"Refining prompt & regenerating visual for Scene {scene_num_display}...", expanded=True) as status_visual_regen:
+                        with st.status(f"Refining prompt & regenerating visual...", expanded=True) as status_visual_regen:
                             prompt_refinement_request = create_visual_regeneration_prompt(
-                                prompt_for_edit_visuals, feedback_visual, scene_data_display 
+                                prompt_for_edit_visuals, feedback_visual, scene_data_display, 
+                                st.session_state.character_definitions, st.session_state.style_reference_description
                             )
                             try:
                                 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 
-                                status_visual_regen.update(label="Image prompt refined! Regenerating visual...", state="running")
+                                st.session_state.scene_image_prompts[i] = refined_textual_image_prompt # Update the prompt in session state
+                                status_visual_regen.update(label="Prompt refined! Regenerating visual...", state="running")
                                 version = 1
                                 if current_image_path:
-                                    try:
-                                        base, _ = os.path.splitext(os.path.basename(current_image_path))
-                                        if '_v' in base: version = int(base.split('_v')[-1]) + 1
-                                    except : pass
-                                if generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=version):
-                                    status_visual_regen.update(label="Visual for Scene updated! 🎉", state="complete", expanded=False)
-                                else:
-                                     status_visual_regen.update(label="Prompt refined, but visual failed.", state="warning", expanded=False)
+                                    try: base, _ = os.path.splitext(os.path.basename(current_image_path)); version = int(base.split('_v')[-1]) + 1 if '_v' in base else 2
+                                    except: version = 2
+                                if generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=version): # is_regeneration=True to use the refined prompt
+                                    status_visual_regen.update(label="Visual updated! 🎉", state="complete", expanded=False)
+                                else: status_visual_regen.update(label="Prompt refined, visual failed.", state="warning", expanded=False)
                                 st.rerun()
-                            except Exception as e:
-                                status_visual_regen.update(label=f"Error refining prompt or regenerating visual: {e}", state="error")
-                    else:
-                        st.warning("Please provide feedback for visual regeneration.")
+                            except Exception as e: status_visual_regen.update(label=f"Error: {e}", state="error")
+                    else: st.warning("Please provide feedback.")
         st.markdown("---")
     
+    # Video Generation Button and Display (remains same as previous full app.py)
     if st.session_state.story_scenes and any(p for p in st.session_state.generated_images_paths if p is not None):
         if st.button("🎬 Assemble Enhanced Animatic", key="assemble_enhanced_video_btn", type="primary", use_container_width=True):
-            with st.status("Assembling enhanced animatic video...", expanded=False) as status_video:
+            # ... (video assembly logic - same as before) ...
+            with st.status("Assembling enhanced animatic video...", expanded=True) as status_video: 
                 image_data_for_video = []
                 for idx, scene_info in enumerate(st.session_state.story_scenes):
                     img_path = st.session_state.generated_images_paths[idx] if idx < len(st.session_state.generated_images_paths) else None
                     if img_path and os.path.exists(img_path):
                         image_data_for_video.append({
-                            'path': img_path,
-                            'scene_num': scene_info.get('scene_number', idx + 1),
+                            'path': img_path, 'scene_num': scene_info.get('scene_number', idx + 1),
                             'key_action': scene_info.get('key_action', '') 
-                        })
+                        }); st.write(f"Adding Scene {scene_info.get('scene_number', idx + 1)} to video.")
                 if image_data_for_video:
-                    st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
-                        image_data_for_video, 
-                        output_filename="cinegen_pro_animatic_enhanced.mp4",
-                        duration_per_image=4, 
-                        fps=24 
-                    )
+                    st.write("Calling video engine..."); st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
+                        image_data_for_video, output_filename="cinegen_pro_animatic_enhanced.mp4", duration_per_image=4, fps=24 )
                     if st.session_state.video_path and os.path.exists(st.session_state.video_path):
-                        status_video.update(label="Enhanced animatic assembled! 🎉", state="complete")
-                        st.balloons()
-                    else:
-                        status_video.update(label="Enhanced video assembly failed. Check logs.", state="error")
-                else:
-                    status_video.update(label="No valid images to assemble video.", state="error")
-    elif st.session_state.story_scenes:
-         st.info("Generate visuals for scenes before assembling the video.")
-
+                        status_video.update(label="Enhanced animatic assembled! 🎉", state="complete", expanded=False); st.balloons()
+                    else: status_video.update(label="Video assembly failed. Check logs.", state="error", expanded=False)
+                else: status_video.update(label="No valid images to assemble video.", state="error", expanded=False)
+    elif st.session_state.story_scenes: st.info("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")
         try:
-            with open(st.session_state.video_path, 'rb') as video_file_obj:
-                video_bytes_content = video_file_obj.read()
+            with open(st.session_state.video_path, 'rb') as video_file_obj: video_bytes_content = video_file_obj.read()
             st.video(video_bytes_content, format="video/mp4") 
             with open(st.session_state.video_path, "rb") as fp_download_video:
-                st.download_button(
-                    label="Download Enhanced Animatic", data=fp_download_video,
+                st.download_button(label="Download Enhanced Animatic", data=fp_download_video,
                     file_name=os.path.basename(st.session_state.video_path), mime="video/mp4",
-                    use_container_width=True, key="download_video_btn" 
-                )
-        except Exception as e:
-            st.error(f"Error displaying or preparing video for download: {e}")
+                    use_container_width=True, key="download_video_btn_enhanced" )
+        except Exception as e: st.error(f"Error displaying video: {e}")
 
 # --- Footer ---
 st.sidebar.markdown("---")