Update app.py

app.py

@@ -1,30 +1,29 @@
-# app.py (Streamlit version for StoryVerse Weaver - FINAL FULL VERSION)
-import streamlit as st # FIRST STREAMLIT IMPORT
-from PIL import Image, ImageDraw, ImageFont 
+# app.py (Streamlit version for StoryVerse Weaver - Robust Error Handling)
+import streamlit as st
+from PIL import Image, ImageDraw, ImageFont
 import os
 import time
 import random
-import traceback 
+import traceback
 
 # --- Page Configuration (MUST BE THE VERY FIRST STREAMLIT COMMAND) ---
 st.set_page_config(
     page_title="✨ StoryVerse Weaver ✨",
-    page_icon="😻", 
-    layout="wide", 
-    initial_sidebar_state="expanded" 
+    page_icon="😻",
+    layout="wide",
+    initial_sidebar_state="expanded"
 )
-# --- END OF PAGE CONFIG ---
 
-# --- Core Logic Imports (NOW AFTER st.set_page_config) ---
+# --- Core Logic Imports ---
 from core.llm_services import initialize_text_llms, is_gemini_text_ready, is_hf_text_ready, generate_text_gemini, generate_text_hf, LLMTextResponse
 from core.image_services import initialize_image_llms, is_dalle_ready, is_hf_image_api_ready, generate_image_dalle, generate_image_hf_model, ImageGenResponse
-from core.story_engine import Story, Scene 
+from core.story_engine import Story, Scene
 from prompts.narrative_prompts import get_narrative_system_prompt, format_narrative_user_prompt
 from prompts.image_style_prompts import STYLE_PRESETS, COMMON_NEGATIVE_PROMPTS, format_image_generation_prompt
 from core.utils import basic_text_cleanup
 
 # --- Initialize AI Services (Cached) ---
-@st.cache_resource 
+@st.cache_resource
 def load_ai_services_config():
     print("--- Initializing AI Services (Streamlit Cache Resource) ---")
     initialize_text_llms()
@@ -43,10 +42,8 @@ TEXT_MODELS = {}
 UI_DEFAULT_TEXT_MODEL_KEY = None
 if AI_SERVICES_STATUS["gemini_text_ready"]:
     TEXT_MODELS["✨ Gemini 1.5 Flash (Narrate)"] = {"id": "gemini-1.5-flash-latest", "type": "gemini"}
-    TEXT_MODELS["Legacy Gemini 1.0 Pro (Narrate)"] = {"id": "gemini-1.0-pro-latest", "type": "gemini"}
-if AI_SERVICES_STATUS["hf_text_ready"]:
+if AI_SERVICES_STATUS["hf_text_ready"]: # Used if Gemini not ready or as an option
     TEXT_MODELS["Mistral 7B (Narrate via HF)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"}
-    TEXT_MODELS["Gemma 2B (Narrate via HF)"] = {"id": "google/gemma-2b-it", "type": "hf_text"}
 
 if TEXT_MODELS:
     if AI_SERVICES_STATUS["gemini_text_ready"] and "✨ Gemini 1.5 Flash (Narrate)" in TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = "✨ Gemini 1.5 Flash (Narrate)"
@@ -60,54 +57,17 @@ IMAGE_PROVIDERS = {}
 UI_DEFAULT_IMAGE_PROVIDER_KEY = None
 if AI_SERVICES_STATUS["dalle_image_ready"]: 
     IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 3"] = "dalle_3"
-    IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 2 (Legacy)"] = "dalle_2"
     UI_DEFAULT_IMAGE_PROVIDER_KEY = "🖼️ OpenAI DALL-E 3"
 elif AI_SERVICES_STATUS["hf_image_ready"]: 
-    IMAGE_PROVIDERS["🎡 HF - Stable Diffusion XL Base"] = "hf_sdxl_base" 
+    IMAGE_PROVIDERS["🎡 HF - SDXL Base"] = "hf_sdxl_base" 
     IMAGE_PROVIDERS["🎠 HF - OpenJourney"] = "hf_openjourney"
-    IMAGE_PROVIDERS["🌌 HF - Stable Diffusion v1.5"] = "hf_sd_1_5" 
-    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF - Stable Diffusion XL Base"
-    
+    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF - SDXL Base"
 if not IMAGE_PROVIDERS: 
     IMAGE_PROVIDERS["No Image Providers Configured"] = "none"
     UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured"
-elif not UI_DEFAULT_IMAGE_PROVIDER_KEY and IMAGE_PROVIDERS : 
-    UI_DEFAULT_IMAGE_PROVIDER_KEY = list(IMAGE_PROVIDERS.keys())[0]
 
-# --- Custom CSS for Dark Theme "WOW" ---
-streamlit_omega_css = """
-<style>
-    body { color: #D0D0E0; background-color: #0F0F1A; font-family: 'Lexend Deca', sans-serif;}
-    .stApp { background-color: #0F0F1A; } 
-    h1 { font-size: 2.8em !important; text-align: center; color: transparent !important; background: linear-gradient(135deg, #A020F0 0%, #E040FB 100%); -webkit-background-clip: text; background-clip: text; margin-bottom: 5px !important; letter-spacing: -1px;}
-    h3 { color: #C080F0 !important; text-align: center; font-weight: 400; margin-bottom: 25px !important;}
-    .main .block-container { padding-top: 2rem; padding-bottom: 2rem; padding-left: 2rem; padding-right: 2rem; max-width: 1400px; margin: auto; background-color: #1A1A2E; border-radius: 15px; box-shadow: 0 8px 24px rgba(0,0,0,0.15);}
-    [data-testid="stSidebar"] { background-color: #131325; border-right: 1px solid #2A2A4A; padding: 1rem;}
-    [data-testid="stSidebar"] .stMarkdown h3 { color: #D0D0FF !important; font-size: 1.5em; border-bottom: 2px solid #7F00FF; padding-bottom: 5px;}
-    .stTextInput > div > div > input, .stTextArea > div > div > textarea, .stSelectbox > div > div > div[data-baseweb="select"] > div, div[data-baseweb="input"] > input { background-color: #2A2A4A !important; color: #E0E0FF !important; border: 1px solid #4A4A6A !important; border-radius: 8px !important;}
-    .stButton > button {
-        background: linear-gradient(135deg, #7F00FF 0%, #E100FF 100%) !important;
-        color: white !important; border: none !important; border-radius: 8px !important;
-        padding: 0.7em 1.3em !important; font-weight: 600 !important;
-        box-shadow: 0 4px 8px rgba(0,0,0,0.15) !important;
-        transition: all 0.2s ease-in-out; width: 100%;
-    }
-    .stButton > button:hover { transform: scale(1.03) translateY(-1px); box-shadow: 0 8px 16px rgba(127,0,255,0.3) !important; }
-    .stButton > button:disabled { background: #4A4A6A !important; color: #8080A0 !important; cursor: not-allowed; }
-    button[kind="secondary"] { background: #4A4A6A !important; color: #E0E0FF !important;} /* For other buttons if not primary */
-    button[kind="secondary"]:hover { background: #5A5A7A !important; }
-    .stImage > img { border-radius: 12px; box-shadow: 0 6px 15px rgba(0,0,0,0.25); max-height: 550px; margin: auto; display: block;}
-    .stExpander { background-color: #161628; border: 1px solid #2A2A4A; border-radius: 12px; margin-bottom: 1em;}
-    .stExpander header { font-size: 1.1em; font-weight: 500; color: #D0D0FF;}
-    .important-note { background-color: rgba(127,0,255,0.1); border-left: 5px solid #7F00FF; padding: 15px; margin-bottom:20px; color: #E0E0FF; border-radius: 6px;}
-    .status-message { padding: 10px; border-radius: 6px; margin-top: 10px; text-align: center; font-weight: 500; }
-    .status-success { background-color: #104010; color: #B0FFB0; border: 1px solid #208020; }
-    .status-error { background-color: #401010; color: #FFB0B0; border: 1px solid #802020; }
-    .status-processing { background-color: #102040; color: #B0D0FF; border: 1px solid #204080; }
-    .gallery-col img { border-radius: 8px; box-shadow: 0 2px 6px rgba(0,0,0,0.15); margin-bottom: 5px;}
-    .gallery-col .stCaption { font-size: 0.85em; text-align: center; color: #A0A0C0;}
-</style>
-"""
+# --- Custom CSS ---
+streamlit_omega_css = """ /* ... Your full omega_css string here ... */ """
 st.markdown(streamlit_omega_css, unsafe_allow_html=True)
 
 # --- Helper: Placeholder Image Creation ---
@@ -123,71 +83,47 @@ def create_placeholder_image_st(text="Processing...", size=(512, 512), color="#2
     draw.text(((size[0]-tw)/2, (size[1]-th)/2), text, font=font, fill=text_color); return img
 
 # --- Initialize Session State ---
-if 'story_object' not in st.session_state: st.session_state.story_object = Story()
-if 'current_log' not in st.session_state: st.session_state.current_log = ["Welcome to StoryVerse Weaver!"]
-if 'latest_scene_image_pil' not in st.session_state: st.session_state.latest_scene_image_pil = None
-if 'latest_scene_narrative' not in st.session_state: st.session_state.latest_scene_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene idea in the panel on the left!"
-if 'status_message' not in st.session_state: st.session_state.status_message = {"text": "Ready to weave your first masterpiece!", "type": "processing"}
-if 'processing_scene' not in st.session_state: st.session_state.processing_scene = False
-if 'form_scene_prompt' not in st.session_state: st.session_state.form_scene_prompt = ""
-if 'form_image_style' not in st.session_state: st.session_state.form_image_style = "Default (Cinematic Realism)"
-if 'form_artist_style' not in st.session_state: st.session_state.form_artist_style = ""
+ss = st.session_state # Shortcut
+if 'story_object' not in ss: ss.story_object = Story()
+if 'current_log' not in ss: ss.current_log = ["Welcome to StoryVerse Weaver!"]
+if 'latest_scene_image_pil' not in ss: ss.latest_scene_image_pil = None
+if 'latest_scene_narrative' not in ss: ss.latest_scene_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene idea in the panel on the left!"
+if 'status_message' not in ss: ss.status_message = {"text": "Ready to weave your first masterpiece!", "type": "processing"}
+if 'processing_scene' not in ss: ss.processing_scene = False
+if 'form_scene_prompt' not in ss: ss.form_scene_prompt = ""
+if 'form_image_style' not in ss: ss.form_image_style = "Default (Cinematic Realism)"
+if 'form_artist_style' not in ss: ss.form_artist_style = ""
 
 # --- UI Definition ---
 st.markdown("<div align='center'><h1>✨ StoryVerse Weaver ✨</h1>\n<h3>Craft Immersive Multimodal Worlds with AI</h3></div>", unsafe_allow_html=True)
-st.markdown("<div class='important-note'><strong>Welcome, Worldsmith!</strong> Describe your vision, choose your style, and let Weaver help you craft captivating scenes with narrative and imagery. Ensure API keys (<code>STORYVERSE_...</code>) are correctly set in your environment/secrets!</div>", unsafe_allow_html=True)
+st.markdown("<div class='important-note'><strong>Welcome, Worldsmith!</strong> Describe your vision, choose your style, and let Weaver help you craft captivating scenes. Ensure API keys (<code>STORYVERSE_...</code>) are set in your environment/secrets!</div>", unsafe_allow_html=True)
 
 # --- Sidebar for Inputs & Configuration ---
 with st.sidebar:
     st.markdown("### 💡 **Craft Your Scene**")
-    with st.form("scene_input_form_key", clear_on_submit=False): # Keep values after submit for reference
-        scene_prompt_text_val = st.text_area("Scene Vision (Description, Dialogue, Action):", 
-                                             value=st.session_state.form_scene_prompt, height=150,
-                                             placeholder="e.g., A lone astronaut discovers a glowing alien artifact...")
+    with st.form("scene_input_form_key", clear_on_submit=False):
+        scene_prompt_text_val = st.text_area("Scene Vision:", value=ss.form_scene_prompt, height=150)
         st.markdown("#### 🎨 Visual Style")
         col_style1, col_style2 = st.columns(2)
-        with col_style1:
-            image_style_val = st.selectbox("Style Preset:", 
-                                           options=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), 
-                                           index=(["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))).index(st.session_state.form_image_style) if st.session_state.form_image_style in (["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))) else 0)
-        with col_style2:
-            artist_style_val = st.text_input("Artistic Inspiration (Optional):", 
-                                             value=st.session_state.form_artist_style, placeholder="e.g., Moebius")
-        
-        negative_prompt_val = st.text_area("Exclude from Image (Negative Prompt):", value=COMMON_NEGATIVE_PROMPTS, height=80)
-
+        with col_style1: image_style_val = st.selectbox("Style Preset:", options=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), index=(["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))).index(ss.form_image_style) if ss.form_image_style in (["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))) else 0)
+        with col_style2: artist_style_val = st.text_input("Artistic Inspiration:", value=ss.form_artist_style)
+        negative_prompt_val = st.text_area("Exclude from Image:", value=COMMON_NEGATIVE_PROMPTS, height=80)
         with st.expander("⚙️ Advanced AI Configuration", expanded=False):
-            text_model_val = st.selectbox("Narrative AI Engine:", options=list(TEXT_MODELS.keys()), 
-                                          index=list(TEXT_MODELS.keys()).index(UI_DEFAULT_TEXT_MODEL_KEY) if UI_DEFAULT_TEXT_MODEL_KEY in TEXT_MODELS else 0)
-            image_provider_val = st.selectbox("Visual AI Engine:", options=list(IMAGE_PROVIDERS.keys()), 
-                                              index=list(IMAGE_PROVIDERS.keys()).index(UI_DEFAULT_IMAGE_PROVIDER_KEY) if UI_DEFAULT_IMAGE_PROVIDER_KEY in IMAGE_PROVIDERS else 0)
-            narrative_length_val = st.selectbox("Narrative Detail:", ["Short (1 paragraph)", "Medium (2-3 paragraphs)", "Detailed (4+ paragraphs)"], index=1)
-            image_quality_val = st.selectbox("Image Detail:", ["Standard", "High Detail", "Sketch Concept"], index=0)
-        
-        submit_button_val = st.form_submit_button("🌌 Weave Scene!", use_container_width=True, type="primary", disabled=st.session_state.processing_scene)
+            text_model_val = st.selectbox("Narrative Engine:", options=list(TEXT_MODELS.keys()), index=list(TEXT_MODELS.keys()).index(UI_DEFAULT_TEXT_MODEL_KEY) if UI_DEFAULT_TEXT_MODEL_KEY in TEXT_MODELS else 0)
+            image_provider_val = st.selectbox("Visual Engine:", options=list(IMAGE_PROVIDERS.keys()), index=list(IMAGE_PROVIDERS.keys()).index(UI_DEFAULT_IMAGE_PROVIDER_KEY) if UI_DEFAULT_IMAGE_PROVIDER_KEY in IMAGE_PROVIDERS else 0)
+            narrative_length_val = st.selectbox("Narrative Detail:", ["Short", "Medium", "Detailed"], index=1)
+            image_quality_val = st.selectbox("Image Detail:", ["Standard", "High Detail", "Sketch"], index=0)
+        submit_button_val = st.form_submit_button("🌌 Weave Scene!", use_container_width=True, type="primary", disabled=ss.processing_scene)
 
-    # Buttons outside the form
-    col_btn1, col_btn2 = st.columns(2)
-    with col_btn1:
-        if st.button("🎲 Surprise Me!", use_container_width=True, disabled=st.session_state.processing_scene, key="sidebar_surprise_btn"):
-            sur_prompt, sur_style, sur_artist = surprise_me_func()
-            st.session_state.form_scene_prompt = sur_prompt
-            # To make selectbox update, its key should also be updated or use a callback if Streamlit version supports
-            st.session_state.form_image_style = sur_style 
-            st.session_state.form_artist_style = sur_artist
-            st.rerun()
-    with col_btn2:
-        if st.button("🗑️ New Story", use_container_width=True, disabled=st.session_state.processing_scene, key="sidebar_clear_btn", type="secondary"): # type="secondary" is conceptual
-            st.session_state.story_object = Story()
-            st.session_state.current_log = ["Story Cleared."]
-            st.session_state.latest_scene_image_pil = None
-            st.session_state.latest_scene_narrative = "## ✨ New Story Begins ✨"
-            st.session_state.status_message = {"text": "📜 Story Cleared.", "type": "processing"}
-            st.session_state.form_scene_prompt = ""
-            st.rerun() 
+    col_btn_s, col_btn_c = st.columns(2)
+    with col_btn_s:
+        if st.button("🎲 Surprise Me!", use_container_width=True, disabled=ss.processing_scene, key="sidebar_surprise_btn"):
+            sur_prompt, sur_style, sur_artist = surprise_me_func(); ss.form_scene_prompt = sur_prompt; ss.form_image_style = sur_style; ss.form_artist_style = sur_artist; st.rerun()
+    with col_btn_c:
+        if st.button("🗑️ New Story", use_container_width=True, disabled=ss.processing_scene, key="sidebar_clear_btn"):
+            ss.story_object = Story(); ss.current_log = ["Cleared."]; ss.latest_scene_image_pil = None; ss.latest_scene_narrative = "## New Story"; ss.status_message = {"text": "Cleared", "type": "processing"}; ss.form_scene_prompt = ""; st.rerun()
 
     with st.expander("🔧 AI Services Status", expanded=False):
-        # ... (API status HTML generation as before, using AI_SERVICES_STATUS) ...
         text_llm_ok, image_gen_ok = (AI_SERVICES_STATUS["gemini_text_ready"] or AI_SERVICES_STATUS["hf_text_ready"]), (AI_SERVICES_STATUS["dalle_image_ready"] or AI_SERVICES_STATUS["hf_image_ready"])
         if not text_llm_ok and not image_gen_ok: st.error("CRITICAL: NO AI SERVICES CONFIGURED.")
         else:
@@ -196,56 +132,54 @@ with st.sidebar:
             if image_gen_ok: st.success("Image Generation Ready.")
             else: st.warning("Image Generation NOT Ready.")
 
-
 # --- Main Display Area ---
-st.markdown("---") # Horizontal line
+st.markdown("---")
 st.markdown("### 🖼️ **Your Evolving StoryVerse**", unsafe_allow_html=True) 
 
-# Status Bar
-status_type = st.session_state.status_message.get("type", "processing")
-st.markdown(f"<p class='status-message status-{status_type}'>{st.session_state.status_message['text']}</p>", unsafe_allow_html=True)
+status_type = ss.status_message.get("type", "processing")
+st.markdown(f"<p class='status-message status-{status_type}'>{ss.status_message['text']}</p>", unsafe_allow_html=True)
 
-# Tabs for display
 tab_latest, tab_scroll, tab_log = st.tabs(["🌠 Latest Scene", "📚 Story Scroll", "⚙️ Interaction Log"])
 
 with tab_latest:
-    if st.session_state.processing_scene and st.session_state.latest_scene_image_pil is None:
-        st.image(create_placeholder_image_st("🎨 Conjuring visuals..."), use_column_width="always")
-    elif st.session_state.latest_scene_image_pil:
-        st.image(st.session_state.latest_scene_image_pil, use_column_width="always", caption="Latest Generated Image")
+    if ss.processing_scene and ss.latest_scene_image_pil is None :
+        st.image(create_placeholder_image_st("🎨 Conjuring visuals..."), use_container_width=True)
+    elif ss.latest_scene_image_pil:
+        st.image(ss.latest_scene_image_pil, use_container_width=True, caption="Latest Generated Image")
     else:
-        st.image(create_placeholder_image_st("Describe a scene to begin!", size=(512,300), color="#1A1A2E"), use_column_width="always")
-    st.markdown(st.session_state.latest_scene_narrative, unsafe_allow_html=True)
+        st.image(create_placeholder_image_st("Describe a scene to begin!", size=(512,300), color="#1A1A2E"), use_container_width=True)
+    st.markdown(ss.latest_scene_narrative, unsafe_allow_html=True)
 
 with tab_scroll:
-    if st.session_state.story_object and st.session_state.story_object.scenes:
-        num_columns_gallery = st.slider("Gallery Columns:", 1, 5, 3, key="gallery_cols_slider") 
-        gallery_cols = st.columns(num_columns_gallery)
-        scenes_for_gallery = st.session_state.story_object.get_all_scenes_for_gallery_display()
-        for i, (img, caption) in enumerate(scenes_for_gallery):
-            with gallery_cols[i % num_columns_gallery]:
-                if img: st.image(img, caption=caption if caption else f"Scene {i+1}", use_column_width="always", output_format="PNG")
-                elif caption: st.caption(f"Scene {i+1} (No Image)\n{caption}") 
-                else: st.caption(f"Scene {i+1} (No image or caption)")
+    if ss.story_object and ss.story_object.scenes:
+        st.subheader("Story Scroll")
+        num_cols_gallery = st.slider("Gallery Columns:", 1, 5, 3, key="gallery_cols_slider")
+        gallery_cols_list = st.columns(num_cols_gallery)
+        scenes_for_gallery_data = ss.story_object.get_all_scenes_for_gallery_display()
+        for i, (img, caption) in enumerate(scenes_for_gallery_data):
+            with gallery_cols_list[i % num_cols_gallery]:
+                display_img_gallery = img
+                if img is None: display_img_gallery = create_placeholder_image_st(f"S{i+1}\nNo Image", size=(180,180), color="#2A2A4A")
+                st.image(display_img_gallery, caption=caption if caption else f"Scene {i+1}", use_container_width=True, output_format="PNG")
     else:
         st.caption("Your story scroll is empty. Weave your first scene!")
 
 with tab_log:
-    log_display = "\n\n---\n\n".join(st.session_state.current_log[::-1][:50]) 
-    st.markdown(log_display, unsafe_allow_html=True)
+    log_display_text = "\n\n---\n\n".join(ss.current_log[::-1][:50]) 
+    st.markdown(log_display_text, unsafe_allow_html=True)
 
 # --- Logic for Form Submission ---
 if submit_button_val: 
     if not scene_prompt_text_val.strip():
-        st.session_state.status_message = {"text": "Scene prompt cannot be empty!", "type": "error"}
+        ss.status_message = {"text": "Scene prompt cannot be empty!", "type": "error"}
         st.rerun() 
     else:
-        st.session_state.processing_scene = True
-        st.session_state.status_message = {"text": f"🌌 Weaving Scene {st.session_state.story_object.current_scene_number + 1}...", "type": "processing"}
-        st.session_state.current_log.append(f"**🚀 Scene {st.session_state.story_object.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**")
-        st.rerun() # Rerun to show "processing" and disable button (will re-execute from top)
+        ss.processing_scene = True
+        ss.status_message = {"text": f"🌌 Weaving Scene {ss.story_object.current_scene_number + 1}...", "type": "processing"}
+        ss.current_log.append(f"**🚀 Scene {ss.story_object.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**")
+        st.rerun() # Rerun to show "processing" and disable button
 
-        # Store form values to use after the rerun, as widget values reset
+        # ---- Main Generation Logic ----
         _scene_prompt = scene_prompt_text_val
         _image_style = image_style_val
         _artist_style = artist_style_val
@@ -255,66 +189,79 @@ if submit_button_val:
         _narr_length = narrative_length_val
         _img_quality = image_quality_val
         
-        # ---- Main Generation Logic ----
         current_narrative_text = f"Narrative Error: Init failed."
         generated_image_pil = None
-        image_gen_error_msg = None
-        final_scene_error_msg = None # For consolidating errors
+        image_gen_error_msg = None # Specific to image generation part
+        final_scene_error_overall = None # Overall error for the scene
 
         # 1. Generate Narrative
-        with st.spinner("✍️ Crafting narrative... (This may take a moment)"):
-            text_model_info = TEXT_MODELS.get(_text_model)
-            if text_model_info and text_model_info["type"] != "none":
-                system_p = get_narrative_system_prompt("default"); prev_narr = st.session_state.story_object.get_last_scene_narrative(); user_p = format_narrative_user_prompt(_scene_prompt, prev_narr)
-                st.session_state.current_log.append(f"  Narrative: Using {_text_model} ({text_model_info['id']}). Length: {_narr_length}")
-                text_resp = None
-                if text_model_info["type"] == "gemini" and AI_SERVICES_STATUS["gemini_text_ready"]: text_resp = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
-                elif text_model_info["type"] == "hf_text" and AI_SERVICES_STATUS["hf_text_ready"]: text_resp = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
-                if text_resp and text_resp.success: current_narrative_text = basic_text_cleanup(text_resp.text); st.session_state.current_log.append("  Narrative: Success.")
-                elif text_resp: current_narrative_text = f"**Narrative Error ({_text_model}):** {text_resp.error}"; st.session_state.current_log.append(f"  Narrative: FAILED - {text_resp.error}")
-                else: st.session_state.current_log.append(f"  Narrative: FAILED - No response.")
-            else: current_narrative_text = "**Narrative Error:** Model unavailable."; st.session_state.current_log.append(f"  Narrative: FAILED - Model '{_text_model}' unavailable.")
-        st.session_state.latest_scene_narrative = f"## Scene Idea: {_scene_prompt}\n\n{current_narrative_text}"
+        try:
+            with st.spinner("✍️ Crafting narrative... (This may take a moment)"):
+                text_model_info = TEXT_MODELS.get(_text_model)
+                if text_model_info and text_model_info["type"] != "none":
+                    system_p = get_narrative_system_prompt("default"); prev_narr = ss.story_object.get_last_scene_narrative(); user_p = format_narrative_user_prompt(_scene_prompt, prev_narr)
+                    ss.current_log.append(f"  Narrative: Using {_text_model} ({text_model_info['id']}).")
+                    text_resp = None
+                    if text_model_info["type"] == "gemini" and AI_SERVICES_STATUS["gemini_text_ready"]: text_resp = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
+                    elif text_model_info["type"] == "hf_text" and AI_SERVICES_STATUS["hf_text_ready"]: text_resp = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
+                    if text_resp and text_resp.success: current_narrative_text = basic_text_cleanup(text_resp.text); ss.current_log.append("  Narrative: Success.")
+                    elif text_resp: current_narrative_text = f"**Narrative Error ({_text_model}):** {text_resp.error}"; ss.current_log.append(f"  Narrative: FAILED - {text_resp.error}")
+                    else: ss.current_log.append(f"  Narrative: FAILED - No response.")
+                else: current_narrative_text = "**Narrative Error:** Model unavailable."; ss.current_log.append(f"  Narrative: FAILED - Model '{_text_model}' unavailable.")
+            ss.latest_scene_narrative = f"## Scene Idea: {_scene_prompt}\n\n{current_narrative_text}"
+        except Exception as e_narr:
+            current_narrative_text = f"**Narrative Generation Exception:** {type(e_narr).__name__} - {str(e_narr)}"
+            ss.current_log.append(f"  CRITICAL NARRATIVE ERROR: {traceback.format_exc()}")
+            ss.latest_scene_narrative = f"## Scene Idea: {_scene_prompt}\n\n{current_narrative_text}"
+            final_scene_error_overall = current_narrative_text
 
         # 2. Generate Image
-        with st.spinner("🎨 Conjuring visuals... (This may take a moment)"):
-            selected_img_prov_type = IMAGE_PROVIDERS.get(_image_provider)
-            img_content_prompt = current_narrative_text if current_narrative_text and "Error" not in current_narrative_text else _scene_prompt
-            quality_kw = "ultra detailed, " if _img_quality == "High Detail" else ("concept sketch, " if _img_quality == "Sketch Concept" else "")
-            full_img_prompt_for_gen = format_image_generation_prompt(quality_kw + img_content_prompt[:350], _image_style, _artist_style)
-            st.session_state.current_log.append(f"  Image: Using {_image_provider} (type '{selected_img_prov_type}').")
-            if selected_img_prov_type and selected_img_prov_type != "none":
-                img_resp = None
-                if selected_img_prov_type.startswith("dalle_") and AI_SERVICES_STATUS["dalle_image_ready"]:
-                    dalle_model = "dall-e-3" if selected_img_prov_type == "dalle_3" else "dall-e-2"
-                    img_resp = generate_image_dalle(full_img_prompt_for_gen, model=dalle_model)
-                elif selected_img_prov_type.startswith("hf_") and AI_SERVICES_STATUS["hf_image_ready"]:
-                    hf_model_id = "stabilityai/stable-diffusion-xl-base-1.0" # Default
-                    if selected_img_prov_type == "hf_openjourney": hf_model_id="prompthero/openjourney"
-                    img_resp = generate_image_hf_model(full_img_prompt_for_gen, model_id=hf_model_id, negative_prompt=_negative_prompt)
-                if img_resp and img_resp.success: generated_image_pil = img_resp.image; st.session_state.current_log.append("  Image: Success.")
-                elif img_resp: image_gen_error_msg = f"**Image Error:** {img_resp.error}"; st.session_state.current_log.append(f"  Image: FAILED - {img_resp.error}")
-                else: image_gen_error_msg = "**Image Error:** No response."; st.session_state.current_log.append("  Image: FAILED - No response.")
-            else: image_gen_error_msg = "**Image Error:** No provider configured."; st.session_state.current_log.append(f"  Image: FAILED - No provider.")
-        st.session_state.latest_scene_image_pil = generated_image_pil if generated_image_pil else create_placeholder_image_st("Image Gen Failed", color="#401010")
+        try:
+            with st.spinner("🎨 Conjuring visuals... (This may take a moment)"):
+                selected_img_prov_type = IMAGE_PROVIDERS.get(_image_provider)
+                img_content_prompt = current_narrative_text if current_narrative_text and "Error" not in current_narrative_text else _scene_prompt
+                quality_kw = "ultra detailed, " if _img_quality == "High Detail" else ("concept sketch, " if _img_quality == "Sketch Concept" else "")
+                full_img_prompt_for_gen = format_image_generation_prompt(quality_kw + img_content_prompt[:350], _image_style, _artist_style)
+                ss.current_log.append(f"  Image: Using {_image_provider} (type '{selected_img_prov_type}').")
+                if selected_img_prov_type and selected_img_prov_type != "none":
+                    img_resp = None
+                    if selected_img_prov_type.startswith("dalle_") and AI_SERVICES_STATUS["dalle_image_ready"]:
+                        dalle_model = "dall-e-3" if selected_img_prov_type == "dalle_3" else "dall-e-2"
+                        img_resp = generate_image_dalle(full_img_prompt_for_gen, model=dalle_model)
+                    elif selected_img_prov_type.startswith("hf_") and AI_SERVICES_STATUS["hf_image_ready"]:
+                        hf_model_id = "stabilityai/stable-diffusion-xl-base-1.0"; iw,ih=768,768
+                        if selected_img_prov_type == "hf_openjourney": hf_model_id="prompthero/openjourney";iw,ih=512,512
+                        img_resp = generate_image_hf_model(full_img_prompt_for_gen, model_id=hf_model_id, negative_prompt=_negative_prompt, width=iw, height=ih)
+                    if img_resp and img_resp.success: generated_image_pil = img_resp.image; ss.current_log.append("  Image: Success.")
+                    elif img_resp: image_gen_error_msg = f"**Image Error:** {img_resp.error}"; ss.current_log.append(f"  Image: FAILED - {img_resp.error}")
+                    else: image_gen_error_msg = "**Image Error:** No response."; ss.current_log.append("  Image: FAILED - No response.")
+                else: image_gen_error_msg = "**Image Error:** No provider configured."; ss.current_log.append(f"  Image: FAILED - No provider.")
+            ss.latest_scene_image_pil = generated_image_pil if generated_image_pil else create_placeholder_image_st("Image Gen Failed", color="#401010")
+        except Exception as e_img:
+            image_gen_error_msg = f"**Image Generation Exception:** {type(e_img).__name__} - {str(e_img)}"
+            ss.current_log.append(f"  CRITICAL IMAGE ERROR: {traceback.format_exc()}")
+            ss.latest_scene_image_pil = create_placeholder_image_st("Image Gen Exception", color="#401010")
+            if not final_scene_error_overall: final_scene_error_overall = image_gen_error_msg
+            else: final_scene_error_overall += f"\n{image_gen_error_msg}"
 
         # 3. Add Scene
-        if image_gen_error_msg and "**Narrative Error**" in current_narrative_text: final_scene_error_msg = f"Narrative: {current_narrative_text.split('**')[-1].strip()} \nImage: {image_gen_error_msg.split('**')[-1].strip()}"
-        elif "**Narrative Error**" in current_narrative_text: final_scene_error_msg = current_narrative_text
-        elif image_gen_error_msg: final_scene_error_msg = image_gen_error_msg
+        if not final_scene_error_overall: # If no major error from narrative already
+            if image_gen_error_msg and "**Narrative Error**" in current_narrative_text: final_scene_error_overall = f"{current_narrative_text}\n{image_gen_error_msg}"
+            elif "**Narrative Error**" in current_narrative_text: final_scene_error_overall = current_narrative_text
+            elif image_gen_error_msg: final_scene_error_overall = image_gen_error_msg
         
-        st.session_state.story_object.add_scene_from_elements(
+        ss.story_object.add_scene_from_elements(
             user_prompt=_scene_prompt, narrative_text=current_narrative_text, image=generated_image_pil,
             image_style_prompt=f"{_image_style}{f', by {_artist_style}' if _artist_style else ''}",
-            image_provider=_image_provider, error_message=final_scene_error_msg
+            image_provider=_image_provider, error_message=final_scene_error_overall
         )
-        st.session_state.current_log.append(f"  Scene {st.session_state.story_object.current_scene_number} processed.")
+        ss.current_log.append(f"  Scene {ss.story_object.current_scene_number} processed.")
         
         # 4. Set final status message
-        if final_scene_error_msg: 
-            st.session_state.status_message = {"text": f"Scene {st.session_state.story_object.current_scene_number} added with errors.", "type": "error"}
+        if final_scene_error_overall: 
+            ss.status_message = {"text": f"Scene {ss.story_object.current_scene_number} generated with errors. Check log.", "type": "error"}
         else: 
-            st.session_state.status_message = {"text": f"🌌 Scene {st.session_state.story_object.current_scene_number} woven successfully!", "type": "success"}
+            ss.status_message = {"text": f"🌌 Scene {ss.story_object.current_scene_number} woven successfully!", "type": "success"}
         
-        st.session_state.processing_scene = False
-        st.rerun() # Final rerun to update UI with results and re-enable button
+        ss.processing_scene = False
+        st.rerun()