Update app.py

app.py

@@ -1,71 +1,60 @@
-# storyverse_weaver/app.py
-import gradio as gr
+# storyverse_weaver_streamlit/app_st.py (example name)
+import streamlit as st
+from PIL import Image, ImageDraw, ImageFont
 import os
 import time
-import json 
-from PIL import Image, ImageDraw, ImageFont 
 import random
-import traceback 
 
-# --- Core Logic Imports ---
+# --- Assuming your core logic is in a sibling 'core' directory ---
+# You might need to adjust sys.path if running locally vs. deployed
+# import sys
+# sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # If core is one level up
+
 from core.llm_services import initialize_text_llms, is_gemini_text_ready, is_hf_text_ready, generate_text_gemini, generate_text_hf
-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 # CRITICAL: Ensure this is your updated Story class
+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 # Your existing Story and Scene classes
 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 Services ---
-initialize_text_llms()
-initialize_image_llms() 
+# --- Initialize Services ONCE ---
+# Use Streamlit's caching for resource-heavy initializations if they don't depend on session state
+@st.cache_resource # Caches the result across sessions/reruns if inputs don't change
+def load_ai_services():
+    print("--- Initializing AI Services (Streamlit Cache Resource) ---")
+    initialize_text_llms()
+    initialize_image_llms()
+    # Return status flags to be stored in session_state or used directly
+    return {
+        "gemini_text_ready": is_gemini_text_ready(),
+        "hf_text_ready": is_hf_text_ready(),
+        "dalle_image_ready": is_dalle_ready(),
+        "hf_image_ready": is_hf_image_api_ready()
+    }
 
-# --- Get API Readiness Status ---
-GEMINI_TEXT_IS_READY = is_gemini_text_ready()
-HF_TEXT_IS_READY = is_hf_text_ready()
-DALLE_IMAGE_IS_READY = is_dalle_ready() 
-HF_IMAGE_IS_READY = is_hf_image_api_ready()
+ai_services_status = load_ai_services()
 
 # --- Application Configuration (Models, Defaults) ---
-# (This section remains the same - ensure TEXT_MODELS, UI_DEFAULT_TEXT_MODEL_KEY, etc. are defined)
+# (Similar logic to your Gradio app.py for populating TEXT_MODELS, IMAGE_PROVIDERS etc.)
 TEXT_MODELS = {}
 UI_DEFAULT_TEXT_MODEL_KEY = None
-if GEMINI_TEXT_IS_READY:
-    TEXT_MODELS["✨ Gemini 1.5 Flash (Narrate)"] = {"id": "gemini-1.5-flash-latest", "type": "gemini"}
-if HF_TEXT_IS_READY:
-    TEXT_MODELS["Mistral 7B (Narrate via HF)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"}
-if TEXT_MODELS: 
-    UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0]
-    if GEMINI_TEXT_IS_READY and "✨ Gemini 1.5 Flash (Narrate)" in TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = "✨ Gemini 1.5 Flash (Narrate)"
-    elif HF_TEXT_IS_READY and "Mistral 7B (Narrate via HF)" in TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = "Mistral 7B (Narrate via HF)"
-else:
-    TEXT_MODELS["No Text Models Configured"] = {"id": "dummy_text_error", "type": "none"}
-    UI_DEFAULT_TEXT_MODEL_KEY = "No Text Models Configured"
+# ... (Populate based on ai_services_status["gemini_text_ready"], ai_services_status["hf_text_ready"]) ...
+if ai_services_status["gemini_text_ready"]: TEXT_MODELS["✨ Gemini 1.5 Flash (Narrate)"] = {"id": "gemini-1.5-flash-latest", "type": "gemini"} # etc.
+if ai_services_status["hf_text_ready"]: TEXT_MODELS["Mistral 7B (Narrate via HF)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"} # etc.
+if TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0] # Simplified default
 
-IMAGE_PROVIDERS = {} 
+IMAGE_PROVIDERS = {}
 UI_DEFAULT_IMAGE_PROVIDER_KEY = None
-if DALLE_IMAGE_IS_READY: 
-    IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 3"] = "dalle_3" 
-    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🖼️ OpenAI DALL-E 3"
-elif HF_IMAGE_IS_READY: 
-    IMAGE_PROVIDERS["🎡 HF - SDXL Base"] = "hf_sdxl_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"
-
-
-# --- Gradio UI Theme and CSS ---
-# (omega_theme and omega_css definitions remain THE SAME as the last full app.py version)
-omega_theme = gr.themes.Base(font=[gr.themes.GoogleFont("Lexend Deca")], primary_hue=gr.themes.colors.purple).set(body_background_fill="#0F0F1A", block_background_fill="#1A1A2E", slider_color="#A020F0")
-omega_css = """ /* ... Paste your full omega_css string here ... 
-body, .gradio-container .gr-button span { white-space: nowrap !important; overflow: hidden; text-overflow: ellipsis; display: inline-block; max-width: 90%; }
-.gradio-container .gr-button { display: flex; align-items: center; justify-content: center; }
-.gradio-container .gr-button svg { margin-right: 4px !important; }
-*/ """
-
-
-# --- Helper: Placeholder Image Creation ---
-def create_placeholder_image(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
+# ... (Populate based on ai_services_status["dalle_image_ready"], ai_services_status["hf_image_ready"]) ...
+if ai_services_status["dalle_image_ready"]: IMAGE_PROVIDERS["🖼️ DALL-E 3"] = "dalle_3" #etc.
+if ai_services_status["hf_image_ready"]: IMAGE_PROVIDERS["🎡 HF - SDXL Base"] = "hf_sdxl_base" #etc.
+if IMAGE_PROVIDERS: UI_DEFAULT_IMAGE_PROVIDER_KEY = list(IMAGE_PROVIDERS.keys())[0] # Simplified default
+
+
+# --- Helper: Placeholder Image (can be same as before) ---
+@st.cache_data # Cache placeholder images
+def create_placeholder_image_st(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
+    # ... (same PIL logic as before) ...
     img = Image.new('RGB', size, color=color); draw = ImageDraw.Draw(img)
     try: font_path = "arial.ttf" if os.path.exists("arial.ttf") else None 
     except: font_path = None 
@@ -75,276 +64,222 @@ def create_placeholder_image(text="Processing...", size=(512, 512), color="#2323
     else: tw, th = draw.textsize(text, font=font) 
     draw.text(((size[0]-tw)/2, (size[1]-th)/2), text, font=font, fill=text_color); return img
 
-# --- StoryVerse Weaver Orchestrator ---
-def add_scene_to_story_orchestrator(
-    current_story_obj: Story, scene_prompt_text: str, image_style_dropdown: str, artist_style_text: str,
-    negative_prompt_text: str, text_model_key: str, image_provider_key: str,
-    narrative_length: str, image_quality: str,
-    progress=gr.Progress(track_tqdm=True)
-):
-    start_time = time.time() 
-    if not current_story_obj: current_story_obj = Story() # Ensure story object exists
-    
-    log_accumulator = [f"**🚀 Scene {current_story_obj.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**"]
-    
-    # --- Initialize values for the final return tuple ---
-    # These correspond to the `outputs` list of `engage_button.click()`
-    # Order: story_state_output, output_gallery, output_latest_scene_image, 
-    #        output_latest_scene_narrative, output_status_bar, output_interaction_log_markdown
-    
-    # Get initial gallery state based on current story object
-    # This ensures that if we error out early, the gallery doesn't just disappear if it had items
-    current_gallery_items = current_story_obj.get_all_scenes_for_gallery_display()
-    if not current_gallery_items: # Handle initially empty story for gallery
-        placeholder_gallery_img = create_placeholder_image("Start Weaving!", size=(180,180), color="#1A1A2E")
-        current_gallery_items = [(placeholder_gallery_img, "Your StoryVerse awaits!")]
-
-    # These will be updated and form the basis of the final 'return'
-    ret_story_state = current_story_obj 
-    ret_gallery = current_gallery_items 
-    ret_latest_image = None
-    ret_latest_narrative_md_obj = gr.Markdown(value="## Processing...\nNarrative being woven...")
-    ret_status_bar_html_obj = gr.HTML(value="<p class='processing_text status_text'>Processing...</p>")
-    # ret_log_md is built up
-
-    # Initial UI update via yield (buttons disabled by .then() chain)
-    yield {
-        output_status_bar: gr.HTML(value=f"<p class='processing_text status_text'>🌌 Weaving Scene {current_story_obj.current_scene_number + 1}...</p>"),
-        output_latest_scene_image: gr.Image(value=create_placeholder_image("🎨 Conjuring visuals...")),
-        output_latest_scene_narrative: gr.Markdown(value=" Musing narrative..."),
-        output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator))
+# --- 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 (Streamlit Edition)!"]
+if 'latest_scene_image' not in st.session_state:
+    st.session_state.latest_scene_image = None
+if 'latest_scene_narrative' not in st.session_state:
+    st.session_state.latest_scene_narrative = "Describe your first scene to begin!"
+if 'processing_scene' not in st.session_state:
+    st.session_state.processing_scene = False
+
+# --- Page Configuration (Do this ONCE at the top) ---
+st.set_page_config(
+    page_title="✨ StoryVerse Weaver ✨",
+    page_icon="🌌",
+    layout="wide", # "wide" or "centered"
+    initial_sidebar_state="expanded" # "auto", "expanded", "collapsed"
+)
+
+# --- Custom CSS for Dark Theme "WOW" ---
+# (You'd inject this using st.markdown(..., unsafe_allow_html=True) or a separate CSS file)
+streamlit_omega_css = """
+<style>
+    /* Base dark theme */
+    body { color: #D0D0E0; background-color: #0F0F1A; }
+    .stApp { background-color: #0F0F1A; }
+    h1, h2, h3, h4, h5, h6 { color: #C080F0; }
+    .stTextInput > div > div > input, .stTextArea > div > div > textarea, .stSelectbox > div > div > select {
+        background-color: #2A2A4A; color: #E0E0FF; border: 1px solid #4A4A6A; border-radius: 8px;
+    }
+    .stButton > button {
+        background: linear-gradient(135deg, #7F00FF 0%, #E100FF 100%) !important;
+        color: white !important; border: none !important; border-radius: 8px !important;
+        padding: 0.5em 1em !important; font-weight: 600 !important;
+        box-shadow: 0 4px 8px rgba(0,0,0,0.15) !important;
     }
+    .stButton > button:hover { transform: scale(1.03) translateY(-1px); box-shadow: 0 8px 16px rgba(127,0,255,0.3) !important; }
+    /* Add more specific styles for sidebar, expanders, image display etc. */
+    .main .block-container { padding-top: 2rem; padding-bottom: 2rem; padding-left: 3rem; padding-right: 3rem; max-width: 1400px; margin: auto;}
+    .stImage > img { border-radius: 12px; box-shadow: 0 6px 15px rgba(0,0,0,0.25); max-height: 600px;}
+    .stExpander { background-color: #1A1A2E; 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;}
+</style>
+"""
+st.markdown(streamlit_omega_css, unsafe_allow_html=True)
 
-    try:
-        if not scene_prompt_text.strip():
-            raise ValueError("Scene prompt cannot be empty!")
-
-        # --- 1. Generate Narrative Text ---
-        progress(0.1, desc="✍️ Crafting narrative...")
-        narrative_text_generated = f"Narrative Error: Init failed." # Default
-        # ... (Full narrative generation logic from your previous working app.py)
-        # ... (This part should call generate_text_gemini or generate_text_hf and update narrative_text_generated)
-        text_model_info = TEXT_MODELS.get(text_model_key)
-        if text_model_info and text_model_info["type"] != "none":
-            system_p = get_narrative_system_prompt("default"); prev_narrative = current_story_obj.get_last_scene_narrative(); user_p = format_narrative_user_prompt(scene_prompt_text, prev_narrative)
-            log_accumulator.append(f"  Narrative: Using {text_model_key} ({text_model_info['id']}).")
-            text_response = None
-            if text_model_info["type"] == "gemini": text_response = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
-            elif text_model_info["type"] == "hf_text": text_response = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
-            if text_response and text_response.success: narrative_text_generated = basic_text_cleanup(text_response.text); log_accumulator.append(f"  Narrative: Success.")
-            elif text_response: narrative_text_generated = f"**Narrative Error ({text_model_key}):** {text_response.error}"; log_accumulator.append(f"  Narrative: FAILED - {text_response.error}")
-            else: log_accumulator.append(f"  Narrative: FAILED - No response from {text_model_key}.")
-        else: narrative_text_generated = "**Narrative Error:** Text model unavailable."; log_accumulator.append(f"  Narrative: FAILED - Model '{text_model_key}' unavailable.")
-
-        ret_latest_narrative_str_content = f"## Scene Idea: {scene_prompt_text}\n\n{narrative_text_generated}"
-        ret_latest_narrative_md_obj = gr.Markdown(value=ret_latest_narrative_str_content) # Prepare for final return
-        yield { output_latest_scene_narrative: ret_latest_narrative_md_obj,
-                output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)) }
-
-        # --- 2. Generate Image ---
-        progress(0.5, desc="🎨 Conjuring visuals...")
-        image_generated_pil = None
-        image_generation_error_message = None
-        # ... (Full image generation logic from your previous working app.py) ...
-        # ... (This part should call generate_image_dalle or generate_image_hf_model and update image_generated_pil)
-        selected_image_provider_key_from_ui = image_provider_key
-        selected_image_provider_type = IMAGE_PROVIDERS.get(selected_image_provider_key_from_ui)
-        image_content_prompt_for_gen = narrative_text_generated if narrative_text_generated and "Error" not in narrative_text_generated else scene_prompt_text
-        quality_keyword = "ultra detailed, " if image_quality == "High Detail" else ("concept sketch, " if image_quality == "Sketch Concept" else "")
-        full_image_prompt = format_image_generation_prompt(quality_keyword + image_content_prompt_for_gen[:350], image_style_dropdown, artist_style_text)
-        log_accumulator.append(f"  Image: Attempting with provider key '{selected_image_provider_key_from_ui}' (maps to type '{selected_image_provider_type}').")
-        if selected_image_provider_type and selected_image_provider_type != "none": # Actual call logic
-            image_response = None # ... (call DALL-E or HF based on selected_image_provider_type)
-            if selected_image_provider_type.startswith("dalle_"):
-                if DALLE_IMAGE_IS_READY: image_response = generate_image_dalle(full_image_prompt, model="dall-e-3" if selected_image_provider_type == "dalle_3" else "dall-e-2")
-                else: image_generation_error_message = "**Image Error:** DALL-E selected but not ready."
-            elif selected_image_provider_type.startswith("hf_"):
-                if HF_IMAGE_IS_READY: 
-                    hf_model_id = "stabilityai/stable-diffusion-xl-base-1.0" # Default
-                    if selected_image_provider_type == "hf_openjourney": hf_model_id = "prompthero/openjourney"
-                    elif selected_image_provider_type == "hf_sd_1_5": hf_model_id = "runwayml/stable-diffusion-v1-5"
-                    image_response = generate_image_hf_model(full_image_prompt, model_id=hf_model_id, negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS)
-                else: image_generation_error_message = "**Image Error:** HF Image selected but not ready."
-            # ... (process image_response)
-            if image_response and image_response.success: image_generated_pil = image_response.image; log_accumulator.append("  Image: Success.")
-            elif image_response: image_generation_error_message = f"**Image Error:** {image_response.error}"; log_accumulator.append(f"  Image: FAILED - {image_response.error}")
-            elif not image_generation_error_message: image_generation_error_message = "**Image Error:** No response/unknown issue."
-        else: image_generation_error_message = "**Image Error:** No valid image provider."
-
-        ret_latest_image = image_generated_pil if image_generated_pil else create_placeholder_image("Image Gen Failed", color="#401010")
-        yield { output_latest_scene_image: gr.Image(value=ret_latest_image),
-                output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)) }
-
-        # --- 3. Add Scene to Story Object ---
-        final_scene_error = None
-        if image_generation_error_message and "**Narrative Error**" in narrative_text_generated : final_scene_error = f"Narrative: {narrative_text_generated.split('**')[-1].strip()} \nImage: {image_generation_error_message.split('**')[-1].strip()}"
-        elif "**Narrative Error**" in narrative_text_generated: final_scene_error = narrative_text_generated
-        elif image_generation_error_message: final_scene_error = image_generation_error_message
-        
-        current_story_obj.add_scene_from_elements(
-            user_prompt=scene_prompt_text,
-            narrative_text=narrative_text_generated if "**Narrative Error**" not in narrative_text_generated else "(Narrative generation failed, see error log)",
-            image=image_generated_pil,
-            image_style_prompt=f"{image_style_dropdown}{f', by {artist_style_text}' if artist_style_text and artist_style_text.strip() else ''}",
-            image_provider=selected_image_provider_key_from_ui,
-            error_message=final_scene_error
+
+# --- Main App UI & Logic ---
+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 Omega help you weave captivating scenes with narrative and imagery. Ensure API keys (<code>STORYVERSE_...</code>) are correctly set in your environment/secrets!</div>", unsafe_allow_html=True)
+
+
+# --- Sidebar for Inputs & Configuration ---
+with st.sidebar:
+    st.header("🎨 Scene Weaver Panel")
+    
+    with st.form("scene_input_form"):
+        scene_prompt_text = st.text_area(
+            "Scene Vision (Description, Dialogue, Action):", 
+            height=200,
+            placeholder="e.g., Amidst swirling cosmic dust, Captain Eva pilots her damaged starfighter..."
         )
-        ret_story_state = current_story_obj 
-        log_accumulator.append(f"  Scene {current_story_obj.current_scene_number} processed and added to story object.")
         
-        # --- 4. Prepare Final Values for Return Tuple ---
-        gallery_tuples_final = current_story_obj.get_all_scenes_for_gallery_display()
-        processed_gallery_tuples = []
-        if not gallery_tuples_final: # Ensure gallery is not empty for Gradio if story just started
-            placeholder_gallery_img = create_placeholder_image("Your Story Begins!", size=(180,180), color="#1A1A2E")
-            processed_gallery_tuples = [(placeholder_gallery_img, "First scene pending or just added!")]
+        st.subheader("Visual Style")
+        col_style1, col_style2 = st.columns(2)
+        with col_style1:
+            image_style_dropdown = st.selectbox("Style Preset:", options=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), index=0)
+        with col_style2:
+            artist_style_text = st.text_input("Artistic Inspiration (Optional):", placeholder="e.g., Moebius")
+        
+        negative_prompt_text = st.text_area("Exclude from Image (Negative Prompt):", value=COMMON_NEGATIVE_PROMPTS, height=100)
+
+        with st.expander("⚙️ Advanced AI Configuration", expanded=False):
+            text_model_key = st.selectbox("Narrative AI Engine:", options=list(TEXT_MODELS.keys()), index=0 if UI_DEFAULT_TEXT_MODEL_KEY in TEXT_MODELS else (list(TEXT_MODELS.keys()).index(UI_DEFAULT_TEXT_MODEL_KEY) if UI_DEFAULT_TEXT_MODEL_KEY else 0) )
+            image_provider_key = st.selectbox("Visual AI Engine:", options=list(IMAGE_PROVIDERS.keys()), index=0 if UI_DEFAULT_IMAGE_PROVIDER_KEY in IMAGE_PROVIDERS else (list(IMAGE_PROVIDERS.keys()).index(UI_DEFAULT_IMAGE_PROVIDER_KEY) if UI_DEFAULT_IMAGE_PROVIDER_KEY else 0) )
+            narrative_length = st.selectbox("Narrative Detail:", options=["Short (1 paragraph)", "Medium (2-3 paragraphs)", "Detailed (4+ paragraphs)"], index=1)
+            image_quality = st.selectbox("Image Detail/Style:", options=["Standard", "High Detail", "Sketch Concept"], index=0)
+        
+        submit_scene_button = st.form_submit_button("🌌 Weave This Scene!", use_container_width=True, type="primary", disabled=st.session_state.processing_scene)
+
+    if st.button("🎲 Surprise Me!", use_container_width=True, disabled=st.session_state.processing_scene):
+        sur_prompt, sur_style, sur_artist = surprise_me_func() # Assuming this is defined as before
+        # Need to update the actual input widget values; Streamlit doesn't directly map outputs to inputs like Gradio's Examples
+        # This requires a more involved way to update widget states, or just display the suggestion.
+        # For simplicity, we'll just show what it would generate. A real app might use st.experimental_rerun or callbacks.
+        st.info(f"Surprise Idea: Prompt='{sur_prompt}', Style='{sur_style}', Artist='{sur_artist}'\n(Copy these into the fields above!)")
+
+
+    if st.button("🗑️ New Story", use_container_width=True, disabled=st.session_state.processing_scene):
+        st.session_state.story_object = Story()
+        st.session_state.current_log = ["Story Cleared. Ready for a new verse!"]
+        st.session_state.latest_scene_image = None
+        st.session_state.latest_scene_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene!"
+        st.experimental_rerun() # Rerun the script to refresh the UI
+
+    with st.expander("🔧 AI Services Status", expanded=False):
+        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:
-            for img_item, cap_text in gallery_tuples_final:
-                if img_item is None:
-                    gallery_placeholder = create_placeholder_image(f"S{cap_text.split(':')[0][1:]}\nError/NoImg", size=(180,180), color="#2A2A4A")
-                    processed_gallery_tuples.append((gallery_placeholder, cap_text))
-                else:
-                    processed_gallery_tuples.append((img_item, cap_text))
-        ret_gallery = processed_gallery_tuples
+            if text_llm_ok: st.success("Text Generation Service(s) Ready.")
+            else: st.warning("Text Generation Service(s) NOT Ready.")
+            if image_gen_ok: st.success("Image Generation Service(s) Ready.")
+            else: st.warning("Image Generation Service(s) NOT Ready.")
 
 
-        _ , latest_narr_for_display_final_str_temp = current_story_obj.get_latest_scene_details_for_display()
-        ret_latest_narrative_md_obj = gr.Markdown(value=latest_narr_for_display_final_str_temp)
-        
-        status_html_str_temp = f"<p class='error_text status_text'>Scene {current_story_obj.current_scene_number} added with errors.</p>" if final_scene_error else f"<p class='success_text status_text'>🌌 Scene {current_story_obj.current_scene_number} woven!</p>"
-        ret_status_bar_html_obj = gr.HTML(value=status_html_str_temp)
-        
-        progress(1.0, desc="Scene Complete!")
-
-    except ValueError as ve:
-        log_accumulator.append(f"\n**INPUT/CONFIG ERROR:** {ve}")
-        ret_status_bar_html_obj = gr.HTML(value=f"<p class='error_text status_text'>❌ CONFIGURATION ERROR: {ve}</p>")
-        ret_latest_narrative_md_obj = gr.Markdown(value=f"## Error\n{ve}")
-    except Exception as e:
-        log_accumulator.append(f"\n**UNEXPECTED RUNTIME ERROR:** {type(e).__name__} - {e}\n{traceback.format_exc()}")
-        ret_status_bar_html_obj = gr.HTML(value=f"<p class='error_text status_text'>❌ UNEXPECTED ERROR: {type(e).__name__}. Check logs.</p>")
-        ret_latest_narrative_md_obj = gr.Markdown(value=f"## Unexpected Error\n{type(e).__name__}: {e}\nSee log for details.")
+# --- Main Display Area ---
+st.markdown("---")
+st.markdown("### 🖼️ **Your Evolving StoryVerse**", unsafe_allow_html=True) # For potential custom class via CSS
+
+if st.session_state.processing_scene:
+    st.info("🌌 Weaving your scene... Please wait.")
+    # Could use st.spinner("Weaving your scene...")
+
+# Display Latest Scene
+if st.session_state.latest_scene_image or st.session_state.latest_scene_narrative != "Describe your first scene to begin!":
+    st.subheader("🌠 Latest Scene")
+    if st.session_state.latest_scene_image:
+        st.image(st.session_state.latest_scene_image, use_column_width=True, caption="Latest Generated Image")
+    st.markdown(st.session_state.latest_scene_narrative, unsafe_allow_html=True)
+    st.markdown("---")
+
+
+# Display Story Scroll (Gallery)
+if st.session_state.story_object and st.session_state.story_object.scenes:
+    st.subheader("📚 Story Scroll")
+    # Streamlit doesn't have a direct "Gallery" like Gradio. We display images in columns.
+    num_columns = 3
+    cols = st.columns(num_columns)
+    scenes_for_gallery = st.session_state.story_object.get_all_scenes_for_gallery_display() # Ensure this returns (PIL.Image or None, caption)
+
+    for i, (img, caption) in enumerate(scenes_for_gallery):
+        with cols[i % num_columns]:
+            if img:
+                st.image(img, caption=caption if caption else f"Scene {i+1}", use_column_width=True)
+            elif caption: # If no image but caption (e.g. error)
+                st.caption(caption) # Display caption as text
+else:
+    st.caption("Your story scroll is empty. Weave your first scene!")
+
+
+# Interaction Log
+with st.expander("⚙️ Interaction Log", expanded=False):
+    st.markdown("\n\n".join(st.session_state.current_log), unsafe_allow_html=True)
+
+
+# --- Logic for Form Submission ---
+if submit_scene_button and scene_prompt_text.strip(): # Check if form submitted and prompt is not empty
+    st.session_state.processing_scene = True
+    st.session_state.current_log.append(f"**🚀 New Scene Request - {time.strftime('%H:%M:%S')}**")
+    st.experimental_rerun() # Rerun to show "processing" state and disable button
+
+    # ---- This is where the main generation logic happens ----
+    # It's similar to add_scene_to_story_orchestrator but updates session_state
     
-    current_total_time = time.time() - start_time
-    log_accumulator.append(f"  Cycle ended at {time.strftime('%H:%M:%S')}. Total time: {current_total_time:.2f}s")
-    ret_log_md = gr.Markdown(value="\n".join(log_accumulator)) # Prepare final log content
+    # 1. Generate Narrative
+    current_narrative = f"Narrative Error: Init failed for '{scene_prompt_text[:30]}...'"
+    text_model_info = TEXT_MODELS.get(text_model_key)
+    if text_model_info and text_model_info["type"] != "none":
+        system_p = get_narrative_system_prompt("default")
+        prev_narrative = st.session_state.story_object.get_last_scene_narrative()
+        user_p = format_narrative_user_prompt(scene_prompt_text, prev_narrative)
+        st.session_state.current_log.append(f"  Narrative: Using {text_model_key} ({text_model_info['id']}).")
+        text_response = None
+        if text_model_info["type"] == "gemini" and ai_services_status["gemini_text_ready"]: text_response = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
+        elif text_model_info["type"] == "hf_text" and ai_services_status["hf_text_ready"]: text_response = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
+        
+        if text_response and text_response.success: current_narrative = basic_text_cleanup(text_response.text); st.session_state.current_log.append(f"  Narrative: Success.")
+        elif text_response: current_narrative = f"**Narrative Error ({text_model_key}):** {text_response.error}"; st.session_state.current_log.append(f"  Narrative: FAILED - {text_response.error}")
+        else: st.session_state.current_log.append(f"  Narrative: FAILED - No response from {text_model_key}.")
+    else: current_narrative = "**Narrative Error:** Text model unavailable."; st.session_state.current_log.append(f"  Narrative: FAILED - Model '{text_model_key}' unavailable.")
     
-    # Final return for the .click() handler's `outputs` list
-    return (
-        ret_story_state,
-        ret_gallery, # This is now processed_gallery_tuples
-        ret_latest_image, # This is the PIL image or placeholder
-        ret_latest_narrative_md_obj, # This is a gr.Markdown object
-        ret_status_bar_html_obj,     # This is a gr.HTML object
-        ret_log_md                   # This is a gr.Markdown object
-    )
+    st.session_state.latest_scene_narrative = f"## Scene Idea: {scene_prompt_text}\n\n{current_narrative}"
 
-def clear_story_state_ui_wrapper():
-    print("DEBUG: clear_story_state_ui_wrapper called")
-    new_story = Story()
-    placeholder_img = create_placeholder_image("Your StoryVerse is a blank canvas...", color="#1A1A2E", text_color="#A0A0C0")
-    cleared_gallery = [(placeholder_img, "Your StoryVerse is new and untold...")]
-    initial_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene idea..."
-    status_msg = "<p class='processing_text status_text'>📜 Story Cleared.</p>"
-    return (new_story, cleared_gallery, None, gr.Markdown(initial_narrative), gr.HTML(status_msg), "Log Cleared.", "")
-
-def surprise_me_func():
-    print("DEBUG: surprise_me_func called")
-    themes = ["Cosmic Horror", "Solarpunk Utopia", "Mythic Fantasy", "Noir Detective"]; actions = ["unearths an artifact", "negotiates"]; settings = ["on a rogue planet", "in a city in a tree"]; prompt = f"A protagonist {random.choice(actions)} {random.choice(settings)}. Theme: {random.choice(themes)}."; style = random.choice(list(STYLE_PRESETS.keys())); artist = random.choice(["H.R. Giger", "Moebius", ""]*2)
-    print(f"DEBUG: surprise_me_func returning: Prompt='{prompt}', Style='{style}', Artist='{artist}'")
-    return prompt, style, artist
-
-def disable_buttons_for_processing():
-    print("DEBUG: Disabling buttons")
-    return gr.Button(interactive=False), gr.Button(interactive=False)
-
-def enable_buttons_after_processing():
-    print("DEBUG: Enabling buttons")
-    return gr.Button(interactive=True), gr.Button(interactive=True)
-
-# --- Gradio UI Definition ---
-with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨") as story_weaver_demo:
-    # Define Python variables for UI components
-    story_state_output = gr.State(Story()) 
-    
-    gr.Markdown("<div align='center'><h1>✨ StoryVerse Omega ✨</h1>\n<h3>Craft Immersive Multimodal Worlds with AI</h3></div>")
-    gr.HTML("<div class='important-note'><strong>Welcome, Worldsmith!</strong> Describe your vision, choose your style, and let Omega help you weave captivating scenes with narrative and imagery. Ensure API keys (<code>STORYVERSE_...</code>) are correctly set in Space Secrets!</div>")
+    # 2. Generate Image
+    generated_image_pil = None
+    image_gen_error = None
+    selected_image_provider_actual_type = IMAGE_PROVIDERS.get(image_provider_key)
+    image_content_prompt = current_narrative if current_narrative and "Error" not in current_narrative else scene_prompt_text
+    quality_kw = "ultra detailed, " if image_quality == "High Detail" else ("concept sketch, " if image_quality == "Sketch Concept" else "")
+    full_img_prompt = format_image_generation_prompt(quality_kw + image_content_prompt[:350], image_style_dropdown, artist_style_text)
+    st.session_state.current_log.append(f"  Image: Attempting with {image_provider_key} (type '{selected_image_provider_actual_type}').")
 
-    with gr.Accordion("🔧 AI Services Status & Info", open=False):
-        status_text_list = []; text_llm_ok = (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY); image_gen_ok = (DALLE_IMAGE_IS_READY or HF_IMAGE_IS_READY)
-        if not text_llm_ok and not image_gen_ok: status_text_list.append("<p style='color:#FCA5A5;font-weight:bold;'>⚠️ CRITICAL: NO AI SERVICES CONFIGURED.</p>")
-        else:
-            if text_llm_ok: status_text_list.append("<p style='color:#A7F3D0;'>✅ Text Generation Ready.</p>")
-            else: status_text_list.append("<p style='color:#FCD34D;'>⚠️ Text Generation NOT Ready.</p>")
-            if image_gen_ok: status_text_list.append("<p style='color:#A7F3D0;'>✅ Image Generation Ready.</p>")
-            else: status_text_list.append("<p style='color:#FCD34D;'>⚠️ Image Generation NOT Ready.</p>")
-        gr.HTML("".join(status_text_list))
-
-    with gr.Row(equal_height=False, variant="panel"):
-        with gr.Column(scale=7, min_width=450): 
-            gr.Markdown("### 💡 **Craft Your Scene**", elem_classes="input-section-header")
-            with gr.Group(): scene_prompt_input = gr.Textbox(lines=7, label="Scene Vision (Description, Dialogue, Action):", placeholder="e.g., Amidst swirling cosmic dust...")
-            with gr.Row(elem_classes=["compact-row"]):
-                with gr.Column(scale=2): image_style_input = gr.Dropdown(choices=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), value="Default (Cinematic Realism)", label="Visual Style", allow_custom_value=True)
-                with gr.Column(scale=2): artist_style_input = gr.Textbox(label="Artistic Inspiration (Optional):", placeholder="e.g., Moebius...")
-            negative_prompt_input = gr.Textbox(lines=2, label="Exclude from Image:", value=COMMON_NEGATIVE_PROMPTS)
-            with gr.Accordion("⚙️ Advanced AI Configuration", open=False):
-                 with gr.Group():
-                    text_model_dropdown = gr.Dropdown(choices=list(TEXT_MODELS.keys()), value=UI_DEFAULT_TEXT_MODEL_KEY, label="Narrative AI Engine")
-                    image_provider_dropdown = gr.Dropdown(choices=list(IMAGE_PROVIDERS.keys()), value=UI_DEFAULT_IMAGE_PROVIDER_KEY, label="Visual AI Engine")
-                    with gr.Row():
-                        narrative_length_dropdown = gr.Dropdown(["Short (1 paragraph)", "Medium (2-3 paragraphs)", "Detailed (4+ paragraphs)"], value="Medium (2-3 paragraphs)", label="Narrative Detail")
-                        image_quality_dropdown = gr.Dropdown(["Standard", "High Detail", "Sketch Concept"], value="Standard", label="Image Detail/Style")
-            with gr.Row(elem_classes=["compact-row"], equal_height=True):
-                engage_button = gr.Button("🌌 Weave!", variant="primary", scale=3, icon="✨") # Shorter text
-                surprise_button = gr.Button("🎲 Surprise!", variant="secondary", scale=1, icon="🎁")
-                clear_story_button = gr.Button("🗑️ New", variant="stop", scale=1, icon="♻️") # Shorter text
-            output_status_bar = gr.HTML(value="<p class='processing_text status_text'>Ready to weave your first masterpiece!</p>")
+    if selected_image_provider_actual_type and selected_image_provider_actual_type != "none":
+        img_response = None
+        if selected_image_provider_actual_type.startswith("dalle_") and ai_services_status["dalle_image_ready"]:
+            dalle_model = "dall-e-3" if selected_image_provider_actual_type == "dalle_3" else "dall-e-2"
+            img_response = generate_image_dalle(full_img_prompt, model=dalle_model, quality="hd" if image_quality=="High Detail" else "standard")
+        elif selected_image_provider_actual_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_image_provider_actual_type == "hf_openjourney": hf_model_id="prompthero/openjourney";iw,ih=512,512
+            img_response = generate_image_hf_model(full_img_prompt, model_id=hf_model_id, negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS, width=iw, height=ih)
         
-        with gr.Column(scale=10, min_width=700): 
-            gr.Markdown("### 🖼️ **Your StoryVerse**", elem_classes="output-section-header")
-            with gr.Tabs():
-                with gr.TabItem("🌠 Latest Scene"): 
-                    output_latest_scene_image = gr.Image(label="Latest Image", type="pil", interactive=False, height=512, show_label=False, show_download_button=True, elem_classes=["panel_image"])
-                    output_latest_scene_narrative = gr.Markdown() 
-                with gr.TabItem("📚 Story Scroll"): 
-                    output_gallery = gr.Gallery(label="Story Scroll", show_label=False, columns=4, object_fit="cover", height=700, preview=True, allow_preview=True, elem_classes=["gallery_output"])
-                with gr.TabItem("⚙️ Log"): 
-                    with gr.Accordion("Interaction Log", open=False): 
-                        output_interaction_log_markdown = gr.Markdown("Log...")
-
-    # Event Handlers
-    engage_button.click(fn=disable_buttons_for_processing, outputs=[engage_button, surprise_button], queue=False)\
-        .then(fn=add_scene_to_story_orchestrator, 
-              inputs=[story_state_output, scene_prompt_input, image_style_input, artist_style_input, negative_prompt_input, text_model_dropdown, image_provider_dropdown, narrative_length_dropdown, image_quality_dropdown], 
-              outputs=[story_state_output, output_gallery, output_latest_scene_image, output_latest_scene_narrative, output_status_bar, output_interaction_log_markdown])\
-        .then(fn=enable_buttons_after_processing, outputs=[engage_button, surprise_button], queue=False)
-    
-    clear_story_button.click(fn=clear_story_state_ui_wrapper, 
-                             outputs=[story_state_output, output_gallery, output_latest_scene_image, output_latest_scene_narrative, output_status_bar, output_interaction_log_markdown, scene_prompt_input])
-    
-    surprise_button.click(fn=surprise_me_func, 
-                            outputs=[scene_prompt_input, image_style_input, artist_style_input])
+        if img_response and img_response.success: generated_image_pil = img_response.image; st.session_state.current_log.append(f"  Image: Success from {img_response.provider}.")
+        elif img_response: image_gen_error = f"**Image Error:** {img_response.error}"; st.session_state.current_log.append(f"  Image: FAILED - {img_response.error}")
+        else: image_gen_error = "**Image Error:** No response/unknown issue."; st.session_state.current_log.append(f"  Image: FAILED - No response object.")
+    else: image_gen_error = "**Image Error:** No valid provider."; st.session_state.current_log.append(f"  Image: FAILED - No provider configured.")
+
+    st.session_state.latest_scene_image = generated_image_pil if generated_image_pil else create_placeholder_image("Image Gen Failed", color="#401010")
+
+    # 3. Add to Story Object
+    scene_err = None
+    if image_gen_error and "**Narrative Error**" in current_narrative: scene_err = f"{current_narrative}\n{image_gen_error}"
+    elif "**Narrative Error**" in current_narrative: scene_err = current_narrative
+    elif image_gen_error: scene_err = image_gen_error
     
-    gr.Examples(
-        examples=[
-            ["A lone, weary traveler on a mechanical steed crosses a vast, crimson desert under twin suns. Dust devils dance in the distance.", "Sci-Fi Western", "Moebius", "greenery, water, modern city"],
-            ["Deep within an ancient, bioluminescent forest, a hidden civilization of sentient fungi perform a mystical ritual around a pulsating crystal.", "Psychedelic Fantasy", "Alex Grey", "technology, buildings, roads"],
-            ["A child sits on a crescent moon, fishing for stars in a swirling nebula. A friendly space whale swims nearby.", "Whimsical Cosmic", "James Jean", "realistic, dark, scary"],
-            ["A grand, baroque library where the books fly freely and whisper forgotten lore to those who listen closely.", "Magical Realism", "Remedios Varo", "minimalist, simple, technology"]
-        ],
-        inputs=[scene_prompt_input, image_style_input, artist_style_input, negative_prompt_input], 
-        label="🌌 Example Universes to Weave 🌌",
+    st.session_state.story_object.add_scene_from_elements(
+        user_prompt=scene_prompt_text, narrative_text=current_narrative, image=generated_image_pil,
+        image_style_prompt=f"{image_style_dropdown}{f', by {artist_style_text}' if artist_style_text else ''}",
+        image_provider=image_provider_key, error_message=scene_err
     )
-    gr.HTML("<div style='text-align:center; margin-top:30px; padding-bottom:20px;'><p style='font-size:0.9em; color:#8080A0;'>✨ StoryVerse Omega™ - Weaving Worlds with Words and Pixels ✨</p></div>")
-
-# --- Entry Point ---
-if __name__ == "__main__":
-    print("="*80); print("✨ StoryVerse Omega (Full App with Fixes) Launching... ✨")
-    print(f"  Gemini Text Ready: {GEMINI_TEXT_IS_READY}"); print(f"  HF Text Ready: {HF_TEXT_IS_READY}")
-    print(f"  DALL-E Image Ready: {DALLE_IMAGE_IS_READY}"); print(f"  HF Image API Ready: {HF_IMAGE_IS_READY}")
-    if not (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY) or not (DALLE_IMAGE_IS_READY or HF_IMAGE_IS_READY): 
-        print("  🔴 WARNING: Not all primary/fallback AI services configured.")
-    print(f"  Default Text Model: {UI_DEFAULT_TEXT_MODEL_KEY}"); print(f"  Default Image Provider: {UI_DEFAULT_IMAGE_PROVIDER_KEY}")
-    print("="*80)
-    story_weaver_demo.launch(debug=True, server_name="0.0.0.0", share=False)
+    st.session_state.current_log.append(f"  Scene {st.session_state.story_object.current_scene_number} processed.")
+    st.session_state.processing_scene = False
+    st.experimental_rerun() # Rerun to update the main display with new scene and re-enable button
+
+elif submit_scene_button and not scene_prompt_text.strip(): # If form submitted but prompt is empty
+    st.warning("Please enter a scene vision/prompt!")