Update app.py

app.py

@@ -1,13 +1,13 @@
 # app.py
 import streamlit as st
 from core.gemini_handler import GeminiHandler
-from core.visual_engine import VisualEngine 
+from core.visual_engine import VisualEngine
 from core.prompt_engineering import (
-    create_cinematic_treatment_prompt, 
-    construct_dalle_prompt,           
-    create_narration_script_prompt_enhanced, 
-    create_scene_regeneration_prompt,   
-    create_visual_regeneration_prompt   
+    create_cinematic_treatment_prompt,
+    construct_dalle_prompt,
+    create_narration_script_prompt_enhanced,
+    create_scene_regeneration_prompt,
+    create_visual_regeneration_prompt
 )
 import os
 import logging
@@ -17,6 +17,20 @@ st.set_page_config(page_title="CineGen AI Ultra+", layout="wide", initial_sideba
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# <<< ADDED/MODIFIED START >>>
+# --- Global Definitions for New Features ---
+SHOT_TYPES_OPTIONS = [
+    "Director's Choice", "Establishing Shot", "Long Shot", "Full Shot",
+    "Medium Long Shot (Cowboy)", "Medium Shot", "Medium Close-up",
+    "Close-up", "Extreme Close-up", "Point of View (POV)",
+    "Over the Shoulder", "Tracking Shot", "Dolly Zoom", "Crane Shot",
+    "Aerial Shot", "Static Shot", "Dutch Angle", "Whip Pan"
+]
+DEFAULT_SCENE_DURATION_SECS = 5  # Default duration in seconds for each scene
+DEFAULT_SHOT_TYPE = "Director's Choice"
+# <<< ADDED/MODIFIED END >>>
+
+
 # --- Global State Variables & API Key Setup ---
 def load_api_key(key_name_streamlit, key_name_env, service_name):
     key = None; secrets_available = hasattr(st, 'secrets')
@@ -38,15 +52,15 @@ if 'services_initialized' not in st.session_state:
     st.session_state.ELEVENLABS_API_KEY = load_api_key("ELEVENLABS_API_KEY", "ELEVENLABS_API_KEY", "ElevenLabs")
     st.session_state.PEXELS_API_KEY = load_api_key("PEXELS_API_KEY", "PEXELS_API_KEY", "Pexels")
     st.session_state.ELEVENLABS_VOICE_ID_CONFIG = load_api_key("ELEVENLABS_VOICE_ID", "ELEVENLABS_VOICE_ID", "ElevenLabs Voice ID")
-    
-    if not st.session_state.GEMINI_API_KEY: 
+
+    if not st.session_state.GEMINI_API_KEY:
         st.error("CRITICAL: Gemini API Key is essential and missing!"); logger.critical("Gemini API Key missing. Halting."); st.stop()
-    
+
     try:
         st.session_state.gemini_handler = GeminiHandler(api_key=st.session_state.GEMINI_API_KEY)
         logger.info("GeminiHandler initialized successfully.")
     except Exception as e: st.error(f"Failed to init GeminiHandler: {e}"); logger.critical(f"GeminiHandler init failed: {e}", exc_info=True); st.stop()
-    
+
     try:
         default_voice_id = "Rachel" # A common fallback if no secret is set
         configured_voice_id = st.session_state.ELEVENLABS_VOICE_ID_CONFIG or default_voice_id
@@ -55,7 +69,6 @@ if 'services_initialized' not in st.session_state:
             default_elevenlabs_voice_id=configured_voice_id # Pass it to __init__
         )
         st.session_state.visual_engine.set_openai_api_key(st.session_state.OPENAI_API_KEY)
-        # set_elevenlabs_api_key will use the voice_id passed to __init__ or overridden by its own param
         st.session_state.visual_engine.set_elevenlabs_api_key(st.session_state.ELEVENLABS_API_KEY, voice_id_from_secret=st.session_state.ELEVENLABS_VOICE_ID_CONFIG)
         st.session_state.visual_engine.set_pexels_api_key(st.session_state.PEXELS_API_KEY)
         logger.info("VisualEngine initialized and API keys set (or attempted).")
@@ -72,19 +85,19 @@ for key, default_val in [
 ]:
     if key not in st.session_state: st.session_state[key] = default_val
 
-def initialize_new_project(): # Same
+def initialize_new_project():
     st.session_state.story_treatment_scenes, st.session_state.scene_dalle_prompts, st.session_state.generated_visual_paths = [], [], []
     st.session_state.video_path, st.session_state.overall_narration_audio_path, st.session_state.narration_script_display = None, None, ""
     logger.info("New project initialized.")
 
-def generate_visual_for_scene_core(scene_index, scene_data, version=1): # Same
+def generate_visual_for_scene_core(scene_index, scene_data, version=1):
     dalle_prompt = construct_dalle_prompt(scene_data, st.session_state.character_definitions, st.session_state.global_style_additions)
     if not dalle_prompt: logger.error(f"DALL-E prompt construction failed for scene {scene_data.get('scene_number', scene_index+1)}"); return False
     while len(st.session_state.scene_dalle_prompts) <= scene_index: st.session_state.scene_dalle_prompts.append("")
     while len(st.session_state.generated_visual_paths) <= scene_index: st.session_state.generated_visual_paths.append(None)
     st.session_state.scene_dalle_prompts[scene_index] = dalle_prompt
     filename = f"scene_{scene_data.get('scene_number', scene_index+1)}_visual_v{version}.png"
-    img_path = st.session_state.visual_engine.generate_image_visual(dalle_prompt, scene_data, filename) 
+    img_path = st.session_state.visual_engine.generate_image_visual(dalle_prompt, scene_data, filename)
     if img_path and os.path.exists(img_path):
         st.session_state.generated_visual_paths[scene_index] = img_path; logger.info(f"Visual generated for Scene {scene_data.get('scene_number', scene_index+1)}: {os.path.basename(img_path)}"); return True
     else:
@@ -97,7 +110,7 @@ with st.sidebar:
     user_idea = st.text_area("Core Story Idea / Theme:", "A lone wanderer searches for a mythical oasis in a vast, post-apocalyptic desert, haunted by mirages and mechanical scavengers.", height=120, key="user_idea_main_v5")
     genre = st.selectbox("Primary Genre:", ["Cyberpunk", "Sci-Fi", "Fantasy", "Noir", "Thriller", "Western", "Post-Apocalyptic", "Historical Drama", "Surreal"], index=6, key="genre_main_v5")
     mood = st.selectbox("Overall Mood:", ["Hopeful yet Desperate", "Mysterious & Eerie", "Gritty & Tense", "Epic & Awe-Inspiring", "Melancholy & Reflective", "Whimsical & Lighthearted"], index=0, key="mood_main_v5")
-    num_scenes = st.slider("Number of Key Scenes:", 1, 3, 1, key="num_scenes_main_v5")
+    num_scenes = st.slider("Number of Key Scenes:", 1, 10, 3, key="num_scenes_main_v5") # MODIFIED: Increased max scenes
     creative_guidance_options = {"Standard Director": "standard", "Artistic Visionary": "more_artistic", "Experimental Storyteller": "experimental_narrative"}
     selected_creative_guidance_key = st.selectbox("AI Creative Director Style:", options=list(creative_guidance_options.keys()), key="creative_guidance_select_v5")
     actual_creative_guidance = creative_guidance_options[selected_creative_guidance_key]
@@ -109,10 +122,23 @@ with st.sidebar:
             with st.status("AI Director is envisioning your masterpiece...", expanded=True) as status:
                 try:
                     status.write("Phase 1: Gemini crafting cinematic treatment... 📜"); logger.info("Phase 1: Cinematic Treatment Gen.")
+                    # Note: Consider updating create_cinematic_treatment_prompt to also ask Gemini
+                    # for 'suggested_shot_type' and 'estimated_duration_secs' for each scene.
                     treatment_prompt = create_cinematic_treatment_prompt(user_idea, genre, mood, num_scenes, actual_creative_guidance)
-                    treatment_result_json = st.session_state.gemini_handler.generate_story_breakdown(treatment_prompt)
-                    if not isinstance(treatment_result_json, list) or not treatment_result_json: raise ValueError("Gemini returned invalid scene list.")
-                    st.session_state.story_treatment_scenes = treatment_result_json; num_gen_scenes = len(st.session_state.story_treatment_scenes)
+                    treatment_result_json_raw = st.session_state.gemini_handler.generate_story_breakdown(treatment_prompt)
+                    if not isinstance(treatment_result_json_raw, list) or not treatment_result_json_raw: raise ValueError("Gemini returned invalid scene list.")
+
+                    # <<< ADDED/MODIFIED START >>>
+                    # Process raw scenes and add default shot type and duration
+                    processed_scenes = []
+                    for scene_data_from_gemini in treatment_result_json_raw:
+                        scene_data_from_gemini['shot_type'] = scene_data_from_gemini.get('suggested_shot_type', DEFAULT_SHOT_TYPE) # Use Gemini's suggestion if available
+                        scene_data_from_gemini['scene_duration_secs'] = scene_data_from_gemini.get('estimated_duration_secs', DEFAULT_SCENE_DURATION_SECS) # Use Gemini's suggestion if available
+                        processed_scenes.append(scene_data_from_gemini)
+                    st.session_state.story_treatment_scenes = processed_scenes
+                    # <<< ADDED/MODIFIED END >>>
+
+                    num_gen_scenes = len(st.session_state.story_treatment_scenes)
                     st.session_state.scene_dalle_prompts = [""]*num_gen_scenes; st.session_state.generated_visual_paths = [None]*num_gen_scenes
                     logger.info(f"Phase 1 complete. {num_gen_scenes} scenes."); status.update(label="Treatment complete! ✅ Generating visuals...", state="running")
 
@@ -122,37 +148,37 @@ with st.sidebar:
                         sc_num_log = sc_data.get('scene_number', i+1)
                         status.write(f"  Visual for Scene {sc_num_log}..."); logger.info(f"  Processing visual for Scene {sc_num_log}.")
                         if generate_visual_for_scene_core(i, sc_data, version=1): visual_successes += 1
-                    
+
                     current_status_label_ph2 = "Visuals ready! "
-                    next_step_state = "running" # Assume success unless visuals totally fail
-                    if visual_successes == 0 and num_gen_scenes > 0: 
-                        logger.error("Visual gen failed all scenes."); current_status_label_ph2 = "Visual gen FAILED for all scenes."; next_step_state="error"; 
-                        status.update(label=current_status_label_ph2, state=next_step_state, expanded=True); st.stop() # Stop if no visuals
-                    elif visual_successes < num_gen_scenes: 
+                    next_step_state = "running"
+                    if visual_successes == 0 and num_gen_scenes > 0:
+                        logger.error("Visual gen failed all scenes."); current_status_label_ph2 = "Visual gen FAILED for all scenes."; next_step_state="error";
+                        status.update(label=current_status_label_ph2, state=next_step_state, expanded=True); st.stop()
+                    elif visual_successes < num_gen_scenes:
                         logger.warning(f"Visuals partial ({visual_successes}/{num_gen_scenes})."); current_status_label_ph2 = f"Visuals partially generated ({visual_successes}/{num_gen_scenes}). "
-                    status.update(label=f"{current_status_label_ph2}Generating narration script...", state=next_step_state) # Propagate error state if needed
-                    if next_step_state == "error": st.stop() # Should have already stopped
+                    status.update(label=f"{current_status_label_ph2}Generating narration script...", state=next_step_state)
+                    if next_step_state == "error": st.stop()
 
                     status.write("Phase 3: Generating narration script..."); logger.info("Phase 3: Narration Script Gen.")
                     voice_style_for_prompt = st.session_state.get("selected_voice_style_for_generation", "cinematic_trailer")
                     narr_prompt = create_narration_script_prompt_enhanced(st.session_state.story_treatment_scenes, mood, genre, voice_style_for_prompt)
-                    st.session_state.narration_script_display = st.session_state.gemini_handler.generate_image_prompt(narr_prompt)
+                    st.session_state.narration_script_display = st.session_state.gemini_handler.generate_image_prompt(narr_prompt) # Assuming this is correct for narration text
                     logger.info("Narration script generated."); status.update(label="Narration script ready! Synthesizing voice...", state="running")
-                    
+
                     status.write("Phase 4: Synthesizing voice (ElevenLabs)... 🔊"); logger.info("Phase 4: Voice Synthesis.")
                     st.session_state.overall_narration_audio_path = st.session_state.visual_engine.generate_narration_audio(st.session_state.narration_script_display)
-                    
+
                     final_label = "All components ready! Storyboard below. 🚀"
-                    final_state_val = "complete" 
+                    final_state_val = "complete"
                     if not st.session_state.overall_narration_audio_path:
                         final_label = f"{current_status_label_ph2}Storyboard ready (Voiceover skipped or failed)."
                         logger.warning("Voiceover was skipped or failed.")
                     else: logger.info("Voiceover generated successfully.")
-                    status.update(label=final_label, state=final_state_val, expanded=False) # Always complete if this far
+                    status.update(label=final_label, state=final_state_val, expanded=False)
 
-                except ValueError as ve: logger.error(f"ValueError: {ve}", exc_info=True); status.update(label=f"Input or Gemini response error: {ve}", state="error", expanded=True); 
+                except ValueError as ve: logger.error(f"ValueError: {ve}", exc_info=True); status.update(label=f"Input or Gemini response error: {ve}", state="error", expanded=True);
                 except Exception as e: logger.error(f"Unhandled Exception: {e}", exc_info=True); status.update(label=f"An unexpected error occurred: {e}", state="error", expanded=True);
-    
+
     st.markdown("---"); st.markdown("### Fine-Tuning Options")
     with st.expander("Define Characters", expanded=False):
         char_name = st.text_input("Character Name", key="char_name_adv_ultra_v5"); char_desc = st.text_area("Visual Description", key="char_desc_adv_ultra_v5", height=100, placeholder="e.g., Jax: rugged male astronaut...")
@@ -160,14 +186,14 @@ with st.sidebar:
             if char_name and char_desc: st.session_state.character_definitions[char_name.strip().lower()] = char_desc.strip(); st.success(f"Char '{char_name.strip()}' saved.")
             else: st.warning("Name and description needed.")
         if st.session_state.character_definitions: st.caption("Current Characters:"); [st.markdown(f"**{k.title()}:** _{v}_") for k,v in st.session_state.character_definitions.items()]
-    
+
     with st.expander("Global Style Overrides", expanded=False):
         presets = { "Default (Director's Choice)": "", "Hyper-Realistic Gritty Noir": "hyper-realistic gritty neo-noir, extreme detail, deep dynamic shadows, complex reflections on wet surfaces, cinematic film grain, desaturated palette with isolated vibrant neon accents (e.g. red, cyan), anamorphic lens distortion, atmospheric haze.", "Surreal Dreamscape Fantasy": "surreal dreamscape, epic fantasy elements, painterly with photorealistic details, impossible architecture, bioluminescent flora, otherworldly color palette (e.g., magenta skies, turquoise rivers), style of Roger Dean meets Zdzisław Beksiński.", "Vintage Analog Sci-Fi": "70s/80s analog sci-fi film aesthetic, tangible practical effects look, subtle light leaks, lens flares, warm filmic tones mixed with cool blues, detailed retro-futuristic technology with chunky buttons and CRT screens."}
         sel_preset = st.selectbox("Base Style Preset:", options=list(presets.keys()), key="style_preset_adv_ultra_v5")
         custom_kw = st.text_area("Additional Custom Style Keywords:", key="custom_style_adv_ultra_v5", height=80, placeholder="e.g., 'Dutch angle'")
         cur_style = st.session_state.global_style_additions
         if st.button("Apply Global Styles", key="apply_styles_adv_ultra_v5"):
-            final_s = presets[sel_preset]; 
+            final_s = presets[sel_preset];
             if custom_kw.strip(): final_s = f"{final_s}, {custom_kw.strip()}" if final_s else custom_kw.strip()
             st.session_state.global_style_additions = final_s.strip(); cur_style = final_s.strip()
             if cur_style: st.success("Global styles applied!")
@@ -175,26 +201,25 @@ with st.sidebar:
         if cur_style: st.caption(f"Active global styles: \"{cur_style}\"")
 
     with st.expander("Voice & Narration Style", expanded=False):
-        # User can override the voice ID from secrets if they want to experiment
-        default_voice_from_engine = "Rachel" # Fallback if engine not init or no secret
+        default_voice_from_engine = "Rachel"
         if hasattr(st.session_state, 'visual_engine') and st.session_state.visual_engine:
             default_voice_from_engine = st.session_state.visual_engine.elevenlabs_voice_id
 
         user_voice_id_override = st.text_input(
-            "ElevenLabs Voice ID (optional override):", 
-            value=default_voice_from_engine, 
+            "ElevenLabs Voice ID (optional override):",
+            value=default_voice_from_engine,
             key="el_voice_id_override_v5",
             help=f"Defaulting to '{default_voice_from_engine}' from secrets/config. Enter a specific Voice ID from your ElevenLabs account to override."
         )
         prompt_v_styles = {"Cinematic Trailer": "cinematic_trailer", "Neutral Documentary": "documentary_neutral", "Character Introspection": "introspective_character"}
         sel_prompt_v_style_key = st.selectbox("Narration Script Style:", list(prompt_v_styles.keys()), key="narr_style_sel_v5", index=0)
-        
+
         if st.button("Set Narrator Voice & Style", key="set_voice_btn_ultra_v5"):
             final_voice_id_to_use = user_voice_id_override.strip()
-            if not final_voice_id_to_use: # If user cleared the input, revert to the one from secrets/default
+            if not final_voice_id_to_use:
                 final_voice_id_to_use = st.session_state.get("ELEVENLABS_VOICE_ID_CONFIG", "Rachel")
-            
-            if hasattr(st.session_state, 'visual_engine'): 
+
+            if hasattr(st.session_state, 'visual_engine'):
                 st.session_state.visual_engine.elevenlabs_voice_id = final_voice_id_to_use
             st.session_state.selected_voice_style_for_generation = prompt_v_styles[sel_prompt_v_style_key]
             st.success(f"Narrator Voice ID set to: {final_voice_id_to_use}. Script Style: {sel_prompt_v_style_key}")
@@ -209,69 +234,124 @@ if not st.session_state.story_treatment_scenes: st.info("Use the sidebar to gene
 else:
     for i_main, scene_content_display in enumerate(st.session_state.story_treatment_scenes):
         scene_n = scene_content_display.get('scene_number', i_main + 1); scene_t = scene_content_display.get('scene_title', 'Untitled')
-        key_base = f"s{scene_n}_{''.join(filter(str.isalnum, scene_t[:10]))}_v5" 
+        key_base = f"s{scene_n}_{''.join(filter(str.isalnum, scene_t[:10]))}_v5_{i_main}" # Ensure unique key base with index
         if "director_note" in scene_content_display and scene_content_display['director_note']: st.info(f"🎬 Director Note S{scene_n}: {scene_content_display['director_note']}")
         st.subheader(f"SCENE {scene_n}: {scene_t.upper()}"); col_d, col_v = st.columns([0.45, 0.55])
         with col_d:
-            with st.expander("📝 Scene Treatment", expanded=True):
+            with st.expander("📝 Scene Treatment & Controls", expanded=True): # MODIFIED Expander Title
                 st.markdown(f"**Beat:** {scene_content_display.get('emotional_beat', 'N/A')}"); st.markdown(f"**Setting:** {scene_content_display.get('setting_description', 'N/A')}"); st.markdown(f"**Chars:** {', '.join(scene_content_display.get('characters_involved', ['N/A']))}"); st.markdown(f"**Focus Moment:** _{scene_content_display.get('character_focus_moment', 'N/A')}_"); st.markdown(f"**Plot Beat:** {scene_content_display.get('key_plot_beat', 'N/A')}"); st.markdown(f"**Dialogue Hook:** `\"{scene_content_display.get('suggested_dialogue_hook', '...')}\"`"); st.markdown("---"); st.markdown(f"**Dir. Visual Style:** _{scene_content_display.get('PROACTIVE_visual_style_감독', 'N/A')}_"); st.markdown(f"**Dir. Camera:** _{scene_content_display.get('PROACTIVE_camera_work_감독', 'N/A')}_"); st.markdown(f"**Dir. Sound:** _{scene_content_display.get('PROACTIVE_sound_design_감독', 'N/A')}_")
+
+                # <<< ADDED/MODIFIED START >>>
+                st.markdown("---")
+                st.markdown("##### Shot & Pacing Controls")
+
+                # Shot Type Selection
+                current_shot_type = st.session_state.story_treatment_scenes[i_main].get('shot_type', DEFAULT_SHOT_TYPE)
+                try:
+                    shot_type_index = SHOT_TYPES_OPTIONS.index(current_shot_type)
+                except ValueError:
+                    shot_type_index = SHOT_TYPES_OPTIONS.index(DEFAULT_SHOT_TYPE) # Fallback if value is somehow invalid
+
+                new_shot_type = st.selectbox(
+                    "Dominant Shot Type:",
+                    options=SHOT_TYPES_OPTIONS,
+                    index=shot_type_index,
+                    key=f"shot_type_widget_{key_base}",
+                    help="Suggests the primary camera shot for this scene. Influences visual generation style if DALL-E prompt is ever made aware of it."
+                )
+                if new_shot_type != st.session_state.story_treatment_scenes[i_main]['shot_type']:
+                    st.session_state.story_treatment_scenes[i_main]['shot_type'] = new_shot_type
+                    # No st.rerun() needed unless other UI elements immediately depend on this change
+
+                # Scene Duration Control
+                current_duration = st.session_state.story_treatment_scenes[i_main].get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
+                new_duration = st.number_input(
+                    "Scene Duration (seconds):",
+                    min_value=1,
+                    max_value=300, # Max 5 minutes per scene image
+                    value=current_duration,
+                    step=1,
+                    key=f"duration_widget_{key_base}",
+                    help="Approximate duration this scene's visual will be shown in the final animatic."
+                )
+                if new_duration != st.session_state.story_treatment_scenes[i_main]['scene_duration_secs']:
+                    st.session_state.story_treatment_scenes[i_main]['scene_duration_secs'] = new_duration
+                    # No st.rerun() needed unless other UI elements immediately depend on this change
+                st.markdown("---")
+                # <<< ADDED/MODIFIED END >>>
+
                 cur_d_prompt = st.session_state.scene_dalle_prompts[i_main] if i_main < len(st.session_state.scene_dalle_prompts) else None
                 if cur_d_prompt:
                      with st.popover("👁️ DALL-E Prompt"): st.markdown(f"**DALL-E Prompt:**"); st.code(cur_d_prompt, language='text')
                 pexels_q = scene_content_display.get('pexels_search_query_감독', None)
                 if pexels_q: st.caption(f"Pexels Fallback Query: `{pexels_q}`")
+
         with col_v:
             cur_img_p = st.session_state.generated_visual_paths[i_main] if i_main < len(st.session_state.generated_visual_paths) else None
-            if cur_img_p and os.path.exists(cur_img_p): st.image(cur_img_p, caption=f"Scene {scene_n}: {scene_t}") 
+            if cur_img_p and os.path.exists(cur_img_p): st.image(cur_img_p, caption=f"Scene {scene_n}: {scene_t}")
             else:
                 if st.session_state.story_treatment_scenes: st.caption("Visual pending/failed.")
-            
-            with st.popover(f"✏️ Edit Scene {scene_n} Treatment"): 
-                fb_script = st.text_area("Changes to treatment:", key=f"treat_fb_{key_base}", height=150) 
-                if st.button(f"🔄 Update Scene {scene_n} Treatment", key=f"regen_treat_btn_{key_base}"): 
+
+            with st.popover(f"✏️ Edit Scene {scene_n} Treatment"):
+                fb_script = st.text_area("Changes to treatment:", key=f"treat_fb_{key_base}", height=150)
+                if st.button(f"🔄 Update Scene {scene_n} Treatment", key=f"regen_treat_btn_{key_base}"):
                     if fb_script:
                         with st.status(f"Updating Scene {scene_n}...", expanded=True) as s_treat_regen:
                             prompt_text = create_scene_regeneration_prompt(scene_content_display, fb_script, st.session_state.story_treatment_scenes)
                             try:
-                                updated_sc_data = st.session_state.gemini_handler.regenerate_scene_script_details(prompt_text)
-                                st.session_state.story_treatment_scenes[i_main] = updated_sc_data 
+                                updated_sc_data_raw = st.session_state.gemini_handler.regenerate_scene_script_details(prompt_text) # This is just the script part
+
+                                # <<< ADDED/MODIFIED START >>>
+                                # Preserve user-set shot type and duration
+                                original_shot_type = st.session_state.story_treatment_scenes[i_main].get('shot_type', DEFAULT_SHOT_TYPE)
+                                original_duration = st.session_state.story_treatment_scenes[i_main].get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
+
+                                # Merge Gemini's script updates with existing scene data, preserving our custom fields
+                                updated_sc_data = {**st.session_state.story_treatment_scenes[i_main], **updated_sc_data_raw}
+                                updated_sc_data['shot_type'] = updated_sc_data_raw.get('suggested_shot_type', original_shot_type) # If Gemini suggests it, use it, else keep old
+                                updated_sc_data['scene_duration_secs'] = updated_sc_data_raw.get('estimated_duration_secs', original_duration) # Same for duration
+
+                                st.session_state.story_treatment_scenes[i_main] = updated_sc_data
+                                # <<< ADDED/MODIFIED END >>>
+
                                 s_treat_regen.update(label="Treatment updated! Regenerating visual...", state="running")
                                 v_num = 1
-                                if cur_img_p and os.path.exists(cur_img_p): 
+                                if cur_img_p and os.path.exists(cur_img_p):
                                     try: b,_=os.path.splitext(os.path.basename(cur_img_p)); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
-                                    except (ValueError, IndexError, TypeError): v_num = 2 
-                                else: v_num = 1 
+                                    except (ValueError, IndexError, TypeError): v_num = 2
+                                else: v_num = 1
                                 if generate_visual_for_scene_core(i_main, updated_sc_data, version=v_num): s_treat_regen.update(label="Treatment & Visual Updated! 🎉", state="complete", expanded=False)
-                                else: s_treat_regen.update(label="Treatment updated, visual failed.", state="complete", expanded=False) 
+                                else: s_treat_regen.update(label="Treatment updated, visual failed.", state="complete", expanded=False)
                                 st.rerun()
                             except Exception as e_regen: s_treat_regen.update(label=f"Error: {e_regen}", state="error"); logger.error(f"Scene treatment regen error: {e_regen}", exc_info=True)
                     else: st.warning("Please provide feedback.")
 
-            with st.popover(f"🎨 Edit Scene {scene_n} Visual Prompt"): 
+            with st.popover(f"🎨 Edit Scene {scene_n} Visual Prompt"):
                 d_prompt_edit = st.session_state.scene_dalle_prompts[i_main] if i_main < len(st.session_state.scene_dalle_prompts) else "No DALL-E prompt."
                 st.caption("Current DALL-E Prompt:"); st.code(d_prompt_edit, language='text')
-                fb_visual = st.text_area("Changes for DALL-E prompt:", key=f"visual_fb_{key_base}", height=150) 
-                if st.button(f"🔄 Update Scene {scene_n} Visual", key=f"regen_visual_btn_{key_base}"): 
+                fb_visual = st.text_area("Changes for DALL-E prompt:", key=f"visual_fb_{key_base}", height=150)
+                if st.button(f"🔄 Update Scene {scene_n} Visual", key=f"regen_visual_btn_{key_base}"):
                     if fb_visual:
                         with st.status(f"Refining prompt & visual for Scene {scene_n}...", expanded=True) as s_visual_regen:
-                            ref_req_prompt = create_visual_regeneration_prompt(d_prompt_edit, fb_visual, scene_content_display, 
+                            ref_req_prompt = create_visual_regeneration_prompt(d_prompt_edit, fb_visual, scene_content_display,
                                 st.session_state.character_definitions, st.session_state.global_style_additions)
                             try:
                                 refined_d_prompt = st.session_state.gemini_handler.generate_image_prompt(ref_req_prompt)
-                                st.session_state.scene_dalle_prompts[i_main] = refined_d_prompt 
+                                st.session_state.scene_dalle_prompts[i_main] = refined_d_prompt
                                 s_visual_regen.update(label="DALL-E prompt refined! Regenerating visual...", state="running")
                                 v_num = 1
-                                if cur_img_p and os.path.exists(cur_img_p): 
+                                if cur_img_p and os.path.exists(cur_img_p):
                                     try: b,_=os.path.splitext(os.path.basename(cur_img_p)); v_num = int(b.split('_v')[-1])+1 if '_v' in b else 2
                                     except (ValueError, IndexError, TypeError): v_num=2
                                 else: v_num = 1
-                                if generate_visual_for_scene_core(i_main, scene_content_display, version=v_num): s_visual_regen.update(label="Visual Updated! 🎉", state="complete", expanded=False)
+                                # Pass the current scene_content_display, as DALL-E prompt is separate from treatment text
+                                if generate_visual_for_scene_core(i_main, st.session_state.story_treatment_scenes[i_main], version=v_num): s_visual_regen.update(label="Visual Updated! 🎉", state="complete", expanded=False)
                                 else: s_visual_regen.update(label="Prompt refined, visual failed.", state="complete", expanded=False)
                                 st.rerun()
                             except Exception as e_regen_vis: s_visual_regen.update(label=f"Error: {e_regen_vis}", state="error"); logger.error(f"Visual prompt regen error: {e_regen_vis}", exc_info=True)
                     else: st.warning("Please provide feedback.")
         st.markdown("---")
-    
+
     if st.session_state.story_treatment_scenes and any(p for p in st.session_state.generated_visual_paths if p is not None):
         if st.button("🎬 Assemble Narrated Cinematic Animatic", key="assemble_ultra_video_btn_v5", type="primary", use_container_width=True):
             with st.status("Assembling Ultra Animatic...", expanded=True) as status_vid:
@@ -279,21 +359,42 @@ else:
                 for i_v, sc_c in enumerate(st.session_state.story_treatment_scenes):
                     img_p_v = st.session_state.generated_visual_paths[i_v] if i_v < len(st.session_state.generated_visual_paths) else None
                     if img_p_v and os.path.exists(img_p_v):
-                        img_data_vid.append({'path':img_p_v, 'scene_num':sc_c.get('scene_number',i_v+1), 'key_action':sc_c.get('key_plot_beat','')}); status_vid.write(f"Adding Scene {sc_c.get('scene_number', i_v + 1)}.")
+                        # <<< ADDED/MODIFIED START >>>
+                        scene_duration = sc_c.get('scene_duration_secs', DEFAULT_SCENE_DURATION_SECS)
+                        img_data_vid.append({
+                            'path': img_p_v,
+                            'scene_num': sc_c.get('scene_number', i_v + 1),
+                            'key_action': sc_c.get('key_plot_beat', ''),
+                            'duration': scene_duration # Use per-scene duration
+                        })
+                        status_vid.write(f"Adding Scene {sc_c.get('scene_number', i_v + 1)} (Duration: {scene_duration}s).")
+                        # <<< ADDED/MODIFIED END >>>
+                    else:
+                        logger.warning(f"Skipping Scene {sc_c.get('scene_number', i_v + 1)} for video assembly: No valid visual path.")
+
+
                 if img_data_vid:
-                    status_vid.write("Calling video engine..."); st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
-                        img_data_vid, overall_narration_path=st.session_state.overall_narration_audio_path,
-                        output_filename="cinegen_ultra_animatic.mp4", duration_per_image=5, fps=24)
+                    status_vid.write("Calling video engine...");
+                    # <<< ADDED/MODIFIED START >>>
+                    # IMPORTANT: Assumes visual_engine.create_video_from_images is updated to accept
+                    # images_data (list of dicts with 'path' and 'duration') and no longer needs duration_per_image.
+                    st.session_state.video_path = st.session_state.visual_engine.create_video_from_images(
+                        images_data=img_data_vid, # Pass the list of dicts
+                        overall_narration_path=st.session_state.overall_narration_audio_path,
+                        output_filename="cinegen_ultra_animatic.mp4",
+                        fps=24 # Keep FPS or make it configurable
+                    )
+                    # <<< ADDED/MODIFIED END >>>
                     if st.session_state.video_path and os.path.exists(st.session_state.video_path): status_vid.update(label="Ultra animatic assembled! 🎉", state="complete", expanded=False); st.balloons()
                     else: status_vid.update(label="Video assembly failed. Check logs.", state="error", expanded=False); logger.error("Video assembly returned None or file does not exist.")
-                else: status_vid.update(label="No valid images for video.", state="error", expanded=False); logger.warning("No valid images found for video assembly.")
+                else: status_vid.update(label="No valid images with duration for video.", state="error", expanded=False); logger.warning("No valid images found for video assembly.")
     elif st.session_state.story_treatment_scenes: st.info("Generate visuals before assembling video.")
 
     if st.session_state.video_path and os.path.exists(st.session_state.video_path):
-        st.header("🎬 Generated Cinematic Animatic"); 
+        st.header("🎬 Generated Cinematic Animatic");
         try:
             with open(st.session_state.video_path, 'rb') as vf_obj: video_bytes = vf_obj.read()
-            st.video(video_bytes, format="video/mp4") 
+            st.video(video_bytes, format="video/mp4")
             with open(st.session_state.video_path, "rb") as fp_dl:
                 st.download_button(label="Download Ultra Animatic", data=fp_dl, file_name=os.path.basename(st.session_state.video_path), mime="video/mp4", use_container_width=True, key="download_ultra_video_btn_v5" )
         except Exception as e: st.error(f"Error displaying video: {e}"); logger.error(f"Error displaying video: {e}", exc_info=True)