Update app.py

app.py

@@ -1,54 +1,71 @@
 # app.py
 import streamlit as st
 from core.gemini_handler import GeminiHandler
-from core.visual_engine import VisualEngine # Your updated VisualEngine
+from core.visual_engine import VisualEngine # Your updated VisualEngine with DALL-E prep
 from core.prompt_engineering import (
     create_story_breakdown_prompt,
-    create_image_prompt_from_scene_data, # Generates the TEXT prompt for AI/placeholder
+    create_image_prompt_from_scene_data,
     create_scene_regeneration_prompt,
-    create_visual_regeneration_prompt   # Regenerates the TEXT prompt
+    create_visual_regeneration_prompt
 )
 import os
-# import json # Only if you need to st.json for debugging, not essential for core logic
+# import json # Only if you need to st.json for debugging
 
 # --- 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
 if 'GEMINI_API_KEY' not in st.session_state:
     try:
         st.session_state.GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
-    except KeyError: # For local dev if secrets aren't set yet
+    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: # Catch any other secrets-related errors
-        st.error(f"Error loading secrets: {e}")
+    except Exception as e: 
+        st.error(f"Error loading Gemini API Key from secrets: {e}")
         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")
+    # --- TRY TO SET OPENAI API KEY FOR VISUAL ENGINE ---
+    openai_key_from_secrets = None
+    try:
+        if "OPENAI_API_KEY" in st.secrets:
+            openai_key_from_secrets = st.secrets["OPENAI_API_KEY"]
+    except Exception as e:
+        print(f"Could not access st.secrets for OPENAI_API_KEY: {e}") # Log, don't stop app
+
+    if not openai_key_from_secrets and "OPENAI_API_KEY" in os.environ: # Fallback to env var
+        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)
+        # st.sidebar.info("OpenAI API Key loaded. DALL-E generation enabled.") # Optional UI feedback
+    else:
+        st.session_state.visual_engine.set_openai_api_key(None) 
+        st.sidebar.warning("OpenAI API Key for DALL-E not found. Visuals will be placeholders.")
+    # --- END OF OPENAI API KEY SETUP ---
 
-# Story and generated content
-if 'story_scenes' not in st.session_state: # List of scene dicts
+# 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: # List of textual image prompts
+if 'scene_image_prompts' not in st.session_state: 
     st.session_state.scene_image_prompts = []
-if 'generated_images_paths' not in st.session_state: # List of file paths for generated visuals
+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
-
-# For Character Consistency (placeholders)
 if 'character_definitions' not in st.session_state:
     st.session_state.character_definitions = {} 
-
-# For Style Transfer (placeholders)
 if 'style_reference_description' not in st.session_state:
     st.session_state.style_reference_description = None
 
@@ -58,40 +75,27 @@ def initialize_new_story():
     st.session_state.scene_image_prompts = []
     st.session_state.generated_images_paths = []
     st.session_state.video_path = None
-    # Optional: Clear old media files
-    # if os.path.exists(st.session_state.visual_engine.output_dir):
-    #     for f_name in os.listdir(st.session_state.visual_engine.output_dir):
-    #         try:
-    #             os.remove(os.path.join(st.session_state.visual_engine.output_dir, f_name))
-    #         except OSError: # File might be in use by video player, handle gracefully
-    #             print(f"Could not remove {f_name}, it might be in use.")
-
 
 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.
-    """
     scene_num_for_log = scene_data.get('scene_number', scene_index + 1)
-    st.info(f"Processing visual concept for Scene {scene_num_for_log} (v{version_count})...")
+    # This st.info will appear briefly in the main app area below the "Generate" button.
+    # Consider moving detailed progress to sidebar or using st.status.
+    # st.info(f"Processing visual concept for Scene {scene_num_for_log} (v{version_count})...") 
     
     textual_image_prompt = ""
     if is_regeneration and scene_index < len(st.session_state.scene_image_prompts):
         textual_image_prompt = st.session_state.scene_image_prompts[scene_index]
-        st.caption(f"Using (potentially refined) prompt for Scene {scene_num_for_log}.")
     else:
         textual_image_prompt = create_image_prompt_from_scene_data(
             scene_data,
             st.session_state.character_definitions,
             st.session_state.style_reference_description
         )
-        st.caption(f"Generated initial prompt for Scene {scene_num_for_log}.")
 
     if not textual_image_prompt:
         st.error(f"Failed to create/retrieve textual image prompt for Scene {scene_num_for_log}.")
         return False
 
-    # Update session state for the textual prompt if it's new or changed
     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("")
@@ -99,21 +103,21 @@ def generate_visual_for_scene_wrapper(scene_index, scene_data, is_regeneration=F
 
     image_filename = f"scene_{scene_num_for_log}_visual_v{version_count}.png"
     
+    # The actual generation (AI or placeholder) happens here
     generated_image_path = st.session_state.visual_engine.generate_image_visual(
         textual_image_prompt,
         image_filename
     )
 
-    # Ensure generated_images_paths is long enough
     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.success(f"Visual concept for Scene {scene_num_for_log} (v{version_count}) available: {os.path.basename(generated_image_path)}")
+        # st.success(f"Visual for Scene {scene_num_for_log} (v{version_count}) ready.") # Can be noisy
         st.session_state.generated_images_paths[scene_index] = generated_image_path
         return True
     else:
-        st.error(f"Visual generation failed for Scene {scene_num_for_log} (v{version_count}). Path: {generated_image_path}")
+        st.warning(f"Visual generation failed or path invalid for Scene {scene_num_for_log} (v{version_count}).")
         st.session_state.generated_images_paths[scene_index] = None
         return False
 
@@ -126,46 +130,52 @@ with st.sidebar:
     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")
 
-    if st.button("✨ Generate Full Story Concept", type="primary", key="generate_full_story_btn"):
+    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.")
         else:
-            with st.spinner("Phase 1: Gemini is drafting the script & scene breakdown... 📜"):
+            with st.status("Generating story...", expanded=True) as status_main:
+                st.write("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="✅")
+                    status_main.update(label="Script breakdown complete! ✅", 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:
-                    st.error(f"Failed to generate story breakdown: {e}")
+                    status_main.update(label=f"Failed to generate story breakdown: {e}", state="error")
                     st.session_state.story_scenes = []
+                    st.stop() # Stop further execution if script fails
 
-            if st.session_state.story_scenes:
-                with st.spinner("Phase 2: Generating initial visual concepts... 🖼️"):
+                if st.session_state.story_scenes:
+                    st.write("Phase 2: Generating initial visual concepts... 🖼️")
                     success_count = 0
                     for i_loop, scene_data_loop_var in enumerate(st.session_state.story_scenes):
+                        st.write(f"Generating visual for Scene {scene_data_loop_var.get('scene_number', i_loop + 1)}...")
                         if generate_visual_for_scene_wrapper(i_loop, 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="🖼️")
+                        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. Some failed.", state="warning", expanded=False)
                     else:
-                        st.warning(f"{success_count}/{len(st.session_state.story_scenes)} visual concepts generated. Some may have failed.")
+                        status_main.update(label="Visual concept 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") # Unique key
-        char_desc_input = st.text_area("Character Description (for visual consistency)", key="char_desc_adv_input", height=80) # Unique key
-        if st.button("Add/Update Character", key="add_char_adv_btn"): # Unique key
+        char_name_input = st.text_input("Character Name (e.g., Eva)", key="char_name_adv_input")
+        char_desc_input = st.text_area("Character Description (for visual consistency)", key="char_desc_adv_input", height=80)
+        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] = char_desc_input
-                st.success(f"Character '{char_name_input}' defined.")
+                st.session_state.character_definitions[char_name_input.strip()] = char_desc_input.strip()
+                st.success(f"Character '{char_name_input.strip()}' defined.")
             else:
                 st.warning("Please provide both name and description.")
         if st.session_state.character_definitions:
@@ -174,9 +184,9 @@ with st.sidebar:
                 st.markdown(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_adv_input", height=80) # Unique key
-        if st.button("Apply Textual Style", key="apply_style_adv_btn"): # Unique key
-            st.session_state.style_reference_description = style_ref_text
+        style_ref_text = st.text_area("Describe Style (e.g., 'Van Gogh inspired, swirling brushstrokes')", key="style_text_adv_input", height=80)
+        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.")
 
 # --- Main Content Area ---
@@ -187,14 +197,12 @@ if not st.session_state.story_scenes:
 else:
     for i, scene_data_display in enumerate(st.session_state.story_scenes):
         scene_num_display = scene_data_display.get('scene_number', i + 1)
-        # Create a more robust unique part for keys, sanitizing special characters
         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] # Take first 20 alphanumeric/space chars
+        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]) # Ratio for text vs image
+        col1, col2 = st.columns([2, 3])
 
         with col1: 
             with st.expander("Scene Details", expanded=True):
@@ -210,41 +218,42 @@ else:
                      with st.popover("View Image Prompt Text"): 
                         st.markdown(f"**Textual Prompt for Image Generation:**")
                         st.code(current_textual_prompt, language='text')
-
         with col2: 
             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):
                 st.image(current_image_path, caption=f"Visual Concept for Scene {scene_num_display}")
             else:
                 if st.session_state.story_scenes: 
-                    st.info("Visual concept for this scene is pending or failed.")
+                    st.caption("Visual for this scene is pending or failed to generate.")
 
-            # Removed 'key' argument from st.popover calls
             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}"): 
                     if feedback_script:
-                        with st.spinner(f"Gemini is rewriting Scene {scene_num_display}..."):
+                        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
                             )
                             try:
                                 updated_scene_data = st.session_state.gemini_handler.regenerate_scene_script_details(regen_prompt)
                                 st.session_state.story_scenes[i] = updated_scene_data 
-                                st.toast(f"Scene {scene_num_display} script updated!", icon="✍️")
+                                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 (ValueError, AttributeError, TypeError): pass
+                                    except : pass
                                 
-                                generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=True, version_count=version)
+                                if generate_visual_for_scene_wrapper(i, updated_scene_data, is_regeneration=True, version_count=version):
+                                    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)
                                 st.rerun()
                             except Exception as e:
-                                st.error(f"Error regenerating scene script: {e}")
+                                status_script_regen.update(label=f"Error regenerating scene script: {e}", state="error")
                     else:
                         st.warning("Please provide feedback for script regeneration.")
 
@@ -257,48 +266,50 @@ else:
                     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}"): 
                     if feedback_visual:
-                        with st.spinner(f"Refining image prompt for Scene {scene_num_display}..."):
+                        with st.status(f"Refining prompt & regenerating visual for Scene {scene_num_display}...", expanded=True) as status_visual_regen:
                             prompt_refinement_request = create_visual_regeneration_prompt(
                                 prompt_for_edit_visuals, feedback_visual, scene_data_display 
                             )
                             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 
-                                st.toast(f"Image prompt for Scene {scene_num_display} refined!", icon="💡")
+                                status_visual_regen.update(label="Image 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 (ValueError, AttributeError, TypeError): pass
+                                    except : pass
                                 
-                                generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=version)
+                                if generate_visual_for_scene_wrapper(i, scene_data_display, is_regeneration=True, version_count=version): # Pass current scene_data, as only prompt changed
+                                    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)
                                 st.rerun()
                             except Exception as e:
-                                st.error(f"Error refining image prompt or regenerating visual: {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.")
         st.markdown("---")
     
-    # Video Generation Button
     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 Animatic Video", key="assemble_video_btn", type="primary"):
-            with st.spinner("Assembling video... This might take a moment."):
+        if st.button("🎬 Assemble Animatic Video", key="assemble_video_btn", type="primary", use_container_width=True):
+            with st.status("Assembling video...", expanded=False) as status_video:
                 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_for_video,
                         output_filename="cinegen_pro_animatic.mp4",
-                        duration_per_image=3
+                        duration_per_image=3 # Adjust as needed
                     )
                     if st.session_state.video_path and os.path.exists(st.session_state.video_path):
-                        st.toast("Animatic video assembled!", icon="🎞️")
+                        status_video.update(label="Animatic video assembled! 🎉", state="complete")
                         st.balloons()
                     else:
-                        st.error("Video assembly failed. Check application logs for details from VisualEngine.")
+                        status_video.update(label="Video assembly failed. Check logs.", state="error")
                 else:
-                    st.error("No valid images found to assemble video. Please ensure visuals were generated successfully.")
+                    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.")
 
@@ -308,17 +319,15 @@ else:
             with open(st.session_state.video_path, 'rb') as video_file_obj:
                 video_bytes_content = video_file_obj.read()
             st.video(video_bytes_content)
-            # Download button
             with open(st.session_state.video_path, "rb") as fp_download_video:
                 st.download_button(
-                    label="Download Animatic Video",
-                    data=fp_download_video,
-                    file_name=os.path.basename(st.session_state.video_path), # e.g., "cinegen_pro_animatic.mp4"
-                    mime="video/mp4"
+                    label="Download Animatic Video", data=fp_download_video,
+                    file_name=os.path.basename(st.session_state.video_path), mime="video/mp4",
+                    use_container_width=True
                 )
         except Exception as e:
             st.error(f"Error displaying or preparing video for download: {e}")
 
 # --- Footer ---
 st.sidebar.markdown("---")
-st.sidebar.caption("CineGen AI Pro | Powered by Gemini & Streamlit")
+st.sidebar.caption("CineGen AI Pro | Powered by Gemini & Streamlit.")