Update app.py

app.py

@@ -26,6 +26,7 @@ STABILITY_API_IS_READY = STABILITY_API_CONFIGURED
 OPENAI_DALLE_IS_READY = OPENAI_DALLE_CONFIGURED
 
 # --- Application Configuration (Models, Defaults) ---
+# (This section remains the same - ensure TEXT_MODELS, UI_DEFAULT_TEXT_MODEL_KEY, etc. are defined)
 TEXT_MODELS = {}
 UI_DEFAULT_TEXT_MODEL_KEY = None
 if GEMINI_TEXT_IS_READY:
@@ -56,51 +57,16 @@ else:
     IMAGE_PROVIDERS["No Image Providers Configured"] = "none"
     UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured"
 
+
 # --- Gradio UI Theme and CSS ---
-omega_theme = gr.themes.Base(
-    font=[gr.themes.GoogleFont("Lexend Deca"), "ui-sans-serif", "system-ui", "sans-serif"],
-    primary_hue=gr.themes.colors.purple, secondary_hue=gr.themes.colors.pink, neutral_hue=gr.themes.colors.slate
-).set(
-    body_background_fill="#0F0F1A", block_background_fill="#1A1A2E", block_border_width="1px",
-    block_border_color="#2A2A4A", block_label_background_fill="#2A2A4A", input_background_fill="#2A2A4A",
-    input_border_color="#4A4A6A", button_primary_background_fill="linear-gradient(135deg, #7F00FF 0%, #E100FF 100%)",
-    button_primary_text_color="white", button_secondary_background_fill="#4A4A6A",
-    button_secondary_text_color="#E0E0FF", slider_color="#A020F0"
-)
-omega_css = """
-body, .gradio-container { background-color: #0F0F1A !important; color: #D0D0E0 !important; }
-.gradio-container { max-width: 1400px !important; margin: auto !important; border-radius: 20px; box-shadow: 0 10px 30px rgba(0,0,0,0.2); padding: 25px !important; border: 1px solid #2A2A4A;}
-.gr-panel, .gr-box, .gr-accordion { background-color: #1A1A2E !important; border: 1px solid #2A2A4A !important; border-radius: 12px !important; box-shadow: 0 4px 15px rgba(0,0,0,0.1);}
-.gr-markdown h1 { font-size: 2.8em !important; text-align: center; color: transparent; background: linear-gradient(135deg, #A020F0 0%, #E040FB 100%); -webkit-background-clip: text; background-clip: text; margin-bottom: 5px !important; letter-spacing: -1px;}
-.gr-markdown h3 { color: #C080F0 !important; text-align: center; font-weight: 400; margin-bottom: 25px !important;}
-.input-section-header { font-size: 1.6em; font-weight: 600; color: #D0D0FF; margin-top: 15px; margin-bottom: 8px; border-bottom: 2px solid #7F00FF; padding-bottom: 5px;}
-.output-section-header { font-size: 1.8em; font-weight: 600; color: #D0D0FF; margin-top: 15px; margin-bottom: 12px;}
-.gr-input input, .gr-input textarea, .gr-dropdown select, .gr-textbox textarea { background-color: #2A2A4A !important; color: #E0E0FF !important; border: 1px solid #4A4A6A !important; border-radius: 8px !important; padding: 10px !important;}
-.gr-button { border-radius: 8px !important; font-weight: 500 !important; transition: all 0.2s ease-in-out !important;}
-.gr-button-primary:hover { transform: scale(1.03) translateY(-1px) !important; box-shadow: 0 8px 16px rgba(127,0,255,0.3) !important; }
-.panel_image { border-radius: 12px !important; overflow: hidden; box-shadow: 0 6px 15px rgba(0,0,0,0.25) !important; background-color: #23233A;}
-.panel_image img { max-height: 600px !important; }
-.gallery_output { background-color: transparent !important; border: none !important; }
-.gallery_output .thumbnail-item { border-radius: 8px !important; box-shadow: 0 3px 8px rgba(0,0,0,0.2) !important; margin: 6px !important; transition: transform 0.2s ease; height: 180px !important; width: 180px !important;}
-.gallery_output .thumbnail-item:hover { transform: scale(1.05); }
-.status_text { font-weight: 500; padding: 12px 18px; text-align: center; border-radius: 8px; margin-top:12px; border: 1px solid transparent; font-size: 1.05em;}
-.error_text { background-color: #401010 !important; color: #FFB0B0 !important; border-color: #802020 !important; }
-.success_text { background-color: #104010 !important; color: #B0FFB0 !important; border-color: #208020 !important;}
-.processing_text { background-color: #102040 !important; color: #B0D0FF !important; border-color: #204080 !important;}
-.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;}
-.gr-tabitem { background-color: #1A1A2E !important; border-radius: 0 0 12px 12px !important; padding: 15px !important;}
-.gr-tab-button.selected { background-color: #2A2A4A !important; color: white !important; border-bottom: 3px solid #A020F0 !important; border-radius: 8px 8px 0 0 !important; font-weight: 600 !important;}
-.gr-tab-button { color: #A0A0C0 !important; border-radius: 8px 8px 0 0 !important;}
-.gr-accordion > .gr-block { border-top: 1px solid #2A2A4A !important; }
-.gr-markdown code { background-color: #2A2A4A !important; color: #C0C0E0 !important; padding: 0.2em 0.5em; border-radius: 4px; }
-.gr-markdown pre { background-color: #23233A !important; padding: 1em !important; border-radius: 6px !important; border: 1px solid #2A2A4A !important;}
-.gr-markdown pre > code { padding: 0 !important; background-color: transparent !important; }
-#surprise_button { background: linear-gradient(135deg, #ff7e5f 0%, #feb47b 100%) !important; font-weight:600 !important;}
-#surprise_button:hover { transform: scale(1.03) translateY(-1px) !important; box-shadow: 0 8px 16px rgba(255,126,95,0.3) !important; }
-"""
+# (omega_theme and omega_css definitions remain the same as the last full version)
+omega_theme = gr.themes.Base(font=[gr.themes.GoogleFont("Lexend Deca"), "ui-sans-serif"], primary_hue=gr.themes.colors.purple, secondary_hue=gr.themes.colors.pink, neutral_hue=gr.themes.colors.slate).set(body_background_fill="#0F0F1A", block_background_fill="#1A1A2E", block_border_width="1px", block_border_color="#2A2A4A", block_label_background_fill="#2A2A4A", input_background_fill="#2A2A4A", input_border_color="#4A4A6A", button_primary_background_fill="linear-gradient(135deg, #7F00FF 0%, #E100FF 100%)", button_primary_text_color="white", button_secondary_background_fill="#4A4A6A", button_secondary_text_color="#E0E0FF", slider_color="#A020F0")
+omega_css = """ /* ... Paste your full omega_css string here ... */ """
+
 
 # --- Helper: Placeholder Image Creation ---
 def create_placeholder_image(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
+    # ... (Full implementation from previous version) ...
     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 
@@ -117,33 +83,25 @@ def add_scene_to_story_orchestrator(
     narrative_length: str, image_quality: str,
     progress=gr.Progress(track_tqdm=True)
 ):
-    # Initialize start_time here
     start_time = time.time() 
-    
     if not current_story_obj: current_story_obj = Story() 
     
     log_accumulator = [f"**🚀 Scene {current_story_obj.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**"]
     
-    # Define placeholders for the final return tuple values
-    # This ensures they exist even if errors occur early.
-    # Order must match engage_button.click outputs list
-    ret_story_state = current_story_obj
-    ret_gallery = current_story_obj.get_all_scenes_for_gallery_display()
-    ret_latest_image = None
-    ret_latest_narrative = gr.Markdown(value="## Error\nProcessing did not complete.")
-    ret_status_bar = gr.HTML(value="<p class='error_text status_text'>An error occurred during processing.</p>")
-    ret_log = gr.Markdown(value="\n".join(log_accumulator))
-
-
-    # Initial UI update using direct component updates in yield by component variable name
-    yield {
+    # --- Define dictionary for UI updates ---
+    # Keys are the Python component variables
+    updates_dict = {
         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..."), visible=True),
         output_latest_scene_narrative: gr.Markdown(value=" Musing narrative...", visible=True),
-        engage_button: gr.Button(interactive=False), # Disable button
+        engage_button: gr.Button(interactive=False),
         surprise_button: gr.Button(interactive=False),
-        output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator))
+        output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)),
+        # Initialize other output components that are in the .click() outputs list
+        story_state_output: current_story_obj, # Pass current state
+        output_gallery: current_story_obj.get_all_scenes_for_gallery_display() # Initial gallery
     }
+    yield updates_dict
 
     try:
         if not scene_prompt_text.strip():
@@ -152,11 +110,10 @@ def add_scene_to_story_orchestrator(
         # --- 1. Generate Narrative Text ---
         progress(0.1, desc="✍️ Crafting narrative...")
         narrative_text_generated = f"Narrative Error: Init failed."
+        # ... (Full narrative generation logic, updating narrative_text_generated and log_accumulator) ...
         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)
+            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']}). Length: {narrative_length}")
             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)
@@ -166,20 +123,20 @@ def add_scene_to_story_orchestrator(
             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 = gr.Markdown(value=f"## Scene Idea: {scene_prompt_text}\n\n{narrative_text_generated}")
-        yield { output_latest_scene_narrative: ret_latest_narrative,
-                output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)) }
+        updates_dict[output_latest_scene_narrative] = gr.Markdown(value=f"## Scene Idea: {scene_prompt_text}\n\n{narrative_text_generated}")
+        updates_dict[output_interaction_log_markdown] = gr.Markdown(value="\n".join(log_accumulator))
+        yield updates_dict
 
         # --- 2. Generate Image ---
         progress(0.5, desc="🎨 Conjuring visuals...")
         image_generated_pil = None
         image_generation_error_message = None
+        # ... (Full image generation logic, updating image_generated_pil, image_generation_error_message, log_accumulator) ...
         selected_image_provider_type = IMAGE_PROVIDERS.get(image_provider_key)
         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, intricate, masterpiece, " if image_quality == "High Detail" else ("concept sketch, line art, " if image_quality == "Sketch Concept" else "")
+        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: Using {image_provider_key}. Style: {image_style_dropdown}. Artist: {artist_style_text or 'N/A'}.")
-
         if selected_image_provider_type and selected_image_provider_type != "none":
             image_response = None
             if selected_image_provider_type == "stability_ai": image_response = generate_image_stabilityai(full_image_prompt, negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS, steps=40 if image_quality=="High Detail" else 25)
@@ -189,12 +146,13 @@ def add_scene_to_story_orchestrator(
             else: image_generation_error_message = f"**Image Error:** No response from {image_provider_key}."; log_accumulator.append(f"  Image: FAILED - No response.")
         else: image_generation_error_message = "**Image Error:** Image provider unavailable."; log_accumulator.append(f"  Image: FAILED - Provider '{image_provider_key}' unavailable.")
 
-        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)) }
+        updates_dict[output_latest_scene_image] = gr.Image(value=image_generated_pil if image_generated_pil else create_placeholder_image("Image Gen Failed", color="#401010"))
+        updates_dict[output_interaction_log_markdown] = gr.Markdown(value="\n".join(log_accumulator))
+        yield updates_dict
 
         # --- 3. Add Scene to Story Object ---
         final_scene_error = None
+        # ... (Logic to set final_scene_error) ...
         if image_generation_error_message and "**Narrative Error**" in narrative_text_generated : final_scene_error = f"{narrative_text_generated}\n{image_generation_error_message}"
         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
@@ -207,93 +165,103 @@ def add_scene_to_story_orchestrator(
             image_provider=image_provider_key if selected_image_provider_type != "none" else "N/A",
             error_message=final_scene_error
         )
-        ret_story_state = current_story_obj 
-        log_accumulator.append(f"  Scene {current_story_obj.current_scene_number} processed and added.")
+        updates_dict[story_state_output] = current_story_obj 
+        log_accumulator.append(f"  Scene {current_story_obj.current_scene_number} processed.")
         
-        # --- 4. Prepare Final Values for Return Tuple ---
-        ret_gallery = current_story_obj.get_all_scenes_for_gallery_display()
+        # --- 4. Prepare Final Values for UI Update via Yield ---
+        updates_dict[output_gallery] = current_story_obj.get_all_scenes_for_gallery_display()
         _ , latest_narr_for_display_final_str = current_story_obj.get_latest_scene_details_for_display()
-        ret_latest_narrative = gr.Markdown(value=latest_narr_for_display_final_str)
+        updates_dict[output_latest_scene_narrative] = gr.Markdown(value=latest_narr_for_display_final_str)
         
         status_html_str = 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 = gr.HTML(value=status_html_str)
+        updates_dict[output_status_bar] = gr.HTML(value=status_html_str)
         
         progress(1.0, desc="Scene Complete!")
 
     except ValueError as ve:
         log_accumulator.append(f"\n**INPUT/CONFIG ERROR:** {ve}")
-        ret_status_bar = gr.HTML(value=f"<p class='error_text status_text'>❌ CONFIGURATION ERROR: {ve}</p>")
-        ret_latest_narrative = gr.Markdown(value=f"## Error\n{ve}")
+        updates_dict[output_status_bar] = gr.HTML(value=f"<p class='error_text status_text'>❌ CONFIGURATION ERROR: {ve}</p>")
+        updates_dict[output_latest_scene_narrative] = 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 = gr.HTML(value=f"<p class='error_text status_text'>❌ UNEXPECTED ERROR: {type(e).__name__}. Check logs.</p>")
-        ret_latest_narrative = gr.Markdown(value=f"## Unexpected Error\n{type(e).__name__}: {e}\nSee log for details.")
+        updates_dict[output_status_bar] = gr.HTML(value=f"<p class='error_text status_text'>❌ UNEXPECTED ERROR: {type(e).__name__}. Check logs.</p>")
+        updates_dict[output_latest_scene_narrative] = gr.Markdown(value=f"## Unexpected Error\n{type(e).__name__}: {e}\nSee log for details.")
     finally:
         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 = gr.Markdown(value="\n".join(log_accumulator))
+        updates_dict[output_interaction_log_markdown] = gr.Markdown(value="\n".join(log_accumulator))
+        updates_dict[engage_button] = gr.Button(interactive=True)
+        updates_dict[surprise_button] = gr.Button(interactive=True)
         
-        # This yield is only for button states, does not affect the main 'return'
-        yield { 
-            engage_button: gr.Button(interactive=True), 
-            surprise_button: gr.Button(interactive=True),
-            output_interaction_log_markdown: ret_log # Update log one last time via yield
-        }
+        # Final yield containing all updates for components in the outputs list and others
+        yield updates_dict
     
-    # Final return for the .click() handler's `outputs` list
-    return ret_story_state, ret_gallery, ret_latest_image, ret_latest_narrative, ret_status_bar, ret_log
+    # No explicit return is needed if all updates are handled by the final yield.
+    # Gradio will use the last yielded dictionary for the components in the `outputs` list.
+    return # Or return None
 
 
 def clear_story_state_ui_wrapper():
-    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 in the panel to the left and let the AI help you weave your world!"
-    status_msg = "<p class='processing_text status_text'>📜 Story Cleared. A fresh canvas awaits your imagination!</p>"
-    return (new_story, cleared_gallery, None, gr.Markdown(value=initial_narrative), gr.HTML(value=status_msg), "Log Cleared. Ready for a new adventure!", "")
+    # ... (Same clear_story_state_ui_wrapper, ensure it returns a TUPLE matching outputs)
+    new_story = Story(); placeholder_img = create_placeholder_image("Blank canvas...", color="#1A1A2E", text_color="#A0A0C0")
+    cleared_gallery = [(placeholder_img, "StoryVerse is new...")]
+    initial_narrative = "## ✨ New Story ✨\nDescribe your first scene!"
+    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():
-    themes = ["Cosmic Horror", "Solarpunk Utopia", "Mythic Fantasy", "Noir Detective", "Silent Film Comedy", "Deep Sea Exploration", "Prehistoric Survival"]
-    actions = ["unearths an artifact of immense power", "negotiates with an interdimensional being", "solves an ancient riddle", "embarks on a perilous journey", "attends a secret festival", "witnesses a celestial event", "finds a hidden sanctuary"]
-    settings = ["on a rogue planet drifting through an empty void", "in a city built within a colossal, living tree", "within a library containing all possible books", "on a generation ship nearing its ancient destination", "in a dreamlike landscape where physics is suggestive", "at the bottom of a Mariana Trench-like abyss", "in a lush jungle teeming with dinosaurs"]
-    prompt = f"A protagonist {random.choice(actions)} {random.choice(settings)}. The overall theme is {random.choice(themes)}."
-    style = random.choice(list(STYLE_PRESETS.keys()))
-    artist = random.choice(["H.R. Giger", "Moebius", "Eyvind Earle", " Remedios Varo", "Alphonse Mucha", ""]*2)
-    return prompt, style, artist
+    # ... (Same surprise_me_func as before)
+    themes = ["Cosmic Horror", "Solarpunk Utopia"]; actions = ["unearths an artifact", "negotiates with a being"]; 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); return prompt, style, artist
+
 
 # --- Gradio UI Definition ---
-with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨ - AI Story & World Weaver") as story_weaver_demo:
-    # Define output components that will be updated by the orchestrator
+with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨") as story_weaver_demo:
+    # --- Define all UI components that will be inputs or outputs first ---
+    # This helps Gradio correctly map them when using `yield` with component variables as keys.
     story_state_output = gr.State(Story()) 
+    scene_prompt_input = gr.Textbox()      
+    image_style_input = gr.Dropdown()      
+    artist_style_input = gr.Textbox()      
+    negative_prompt_input = gr.Textbox()   
+    text_model_dropdown = gr.Dropdown()    
+    image_provider_dropdown = gr.Dropdown()
+    narrative_length_dropdown = gr.Dropdown() 
+    image_quality_dropdown = gr.Dropdown()    
+    output_gallery = gr.Gallery()
+    output_latest_scene_image = gr.Image()
+    output_latest_scene_narrative = gr.Markdown()
+    output_status_bar = gr.HTML()
+    output_interaction_log_markdown = gr.Markdown()
+    engage_button = gr.Button()
+    surprise_button = gr.Button()
+    clear_story_button = gr.Button() # Defined here, though not in outputs of add_scene
 
+    # --- Layout UI ---
     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>")
+    gr.HTML("<div class='important-note'><strong>Welcome, Worldsmith!</strong> ... Ensure API keys ...</div>")
 
     with gr.Accordion("🔧 AI Services Status & Info", open=False):
-        status_text_list = []
-        text_llm_ok, image_gen_ok = (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY), (STABILITY_API_IS_READY or OPENAI_DALLE_IS_READY)
+        # ... (API status HTML generation as before)
+        status_text_list = []; text_llm_ok, image_gen_ok = (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY), (STABILITY_API_IS_READY or OPENAI_DALLE_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 Service(s) Ready.</p>")
-            else: status_text_list.append("<p style='color:#FCD34D;'>⚠️ Text Generation Service(s) NOT Ready.</p>")
-            if image_gen_ok: status_text_list.append("<p style='color:#A7F3D0;'>✅ Image Generation Service(s) Ready.</p>")
-            else: status_text_list.append("<p style='color:#FCD34D;'>⚠️ Image Generation Service(s) NOT Ready.</p>")
+            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, Captain Eva pilots her damaged starfighter towards a colossal, ringed gas giant. Alarms blare. 'Just a little further,' she mutters, gripping the controls.")
-            
+            with gr.Group(): # Redefine components here to bind to the variables
+                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 Preset")
                 with gr.Column(scale=2):
-                    artist_style_input = gr.Textbox(label="Artistic Inspiration (Optional):", placeholder="e.g., Moebius, Zdzisław Beksiński")
-            
-            negative_prompt_input = gr.Textbox(lines=2, label="Exclude from Image (Negative Prompt):", placeholder="Default exclusions applied. Add more if needed.", value=COMMON_NEGATIVE_PROMPTS)
-
+                    artist_style_input = gr.Textbox(label="Artistic Inspiration (Optional):", placeholder="e.g., Moebius...")
+            negative_prompt_input = gr.Textbox(lines=2, label="Exclude from Image (Negative Prompt):", 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")
@@ -301,24 +269,20 @@ with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨
                     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):
+            with gr.Row(elem_classes=["compact-row"], equal_height=True): # Redefine buttons
                 engage_button = gr.Button("🌌 Weave This Scene!", variant="primary", scale=3, icon="✨")
                 surprise_button = gr.Button("🎲 Surprise Me!", variant="secondary", scale=1, icon="🎁")
                 clear_story_button = gr.Button("🗑️ New Story", variant="stop", scale=1, icon="♻️")
-            
             output_status_bar = gr.HTML(value="<p class='processing_text status_text'>Ready to weave your first masterpiece!</p>")
 
         with gr.Column(scale=10, min_width=700): 
             gr.Markdown("### 🖼️ **Your Evolving StoryVerse**", elem_classes="output-section-header")
             with gr.Tabs():
-                with gr.TabItem("🌠 Latest Scene", id="latest_scene_tab"):
+                with gr.TabItem("🌠 Latest Scene", id="latest_scene_tab"): # Redefine components
                     output_latest_scene_image = gr.Image(label="Latest Scene Image", type="pil", interactive=False, show_download_button=True, height=512, show_label=False, elem_classes=["panel_image"])
                     output_latest_scene_narrative = gr.Markdown() 
-                
                 with gr.TabItem("📚 Story Scroll", id="story_scroll_tab"):
                     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("⚙️ Interaction Log", id="log_tab"):
                     with gr.Accordion(label="Developer Interaction Log", open=False): 
                         output_interaction_log_markdown = gr.Markdown("Log will appear here...")
@@ -331,13 +295,18 @@ with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨
             text_model_dropdown, image_provider_dropdown,
             narrative_length_dropdown, image_quality_dropdown
         ],
-        outputs=[ # These components are updated by the FINAL return tuple of the orchestrator
+        outputs=[ # This list should match the components updated by the *final return* or *last yield*
             story_state_output, 
             output_gallery, 
             output_latest_scene_image, 
             output_latest_scene_narrative, 
             output_status_bar, 
             output_interaction_log_markdown
+            # The buttons `engage_button`, `surprise_button` are updated via yield dicts,
+            # so they don't strictly need to be in this `outputs` list if the last yield
+            # doesn't include them AND the function has a specific `return` for these outputs.
+            # However, for robustness, if a yield dict is the last thing the generator does,
+            # make sure its keys cover these.
         ]
     )
     clear_story_button.click(
@@ -355,31 +324,20 @@ with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨
         inputs=[],
         outputs=[scene_prompt_input, image_style_input, artist_style_input]
     )
-
-    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], 
+    gr.Examples( # ... (Examples as before)
+        examples=[["A lone traveler...", "Sci-Fi Western", "Moebius", "greenery"]],
+        inputs=[scene_prompt_input, image_style_input, artist_style_input, negative_prompt_input],
         label="🌌 Example Universes to Weave 🌌",
     )
-    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>")
+    gr.HTML("<div style='text-align:center; margin-top:30px; padding-bottom:20px;'><p style='font-size:0.9em; color:#8080A0;'>✨ StoryVerse Omega™</p></div>")
 
 # --- Entry Point ---
 if __name__ == "__main__":
+    # ... (Startup print messages as before)
+    print("="*80); print("✨ StoryVerse Omega™ Launching... ✨")
+    print(f"  Gemini Text Ready: {GEMINI_TEXT_IS_READY}"); print(f"  HF Text Ready: {HF_TEXT_IS_READY}")
+    print(f"  Stability AI Ready: {STABILITY_API_IS_READY}"); print(f"  DALL-E Ready: {OPENAI_DALLE_IS_READY}")
+    if not (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY) or not (STABILITY_API_IS_READY or OPENAI_DALLE_IS_READY): print("  🔴 WARNING: Not all services configured.")
+    print(f"  Default Text Model: {UI_DEFAULT_TEXT_MODEL_KEY}"); print(f"  Default Image Provider: {UI_DEFAULT_IMAGE_PROVIDER_KEY}")
     print("="*80)
-    print("✨ StoryVerse Omega™ - AI Story & World Weaver - Launching... ✨")
-    print(f"  Text LLM Ready (Gemini): {GEMINI_TEXT_IS_READY}")
-    print(f"  Text LLM Ready (HF): {HF_TEXT_IS_READY}")
-    print(f"  Image Provider Ready (Stability AI): {STABILITY_API_IS_READY}")
-    print(f"  Image Provider Ready (DALL-E): {OPENAI_DALLE_IS_READY}")
-    if not (GEMINI_TEXT_IS_READY or HF_TEXT_IS_READY) or not (STABILITY_API_IS_READY or OPENAI_DALLE_IS_READY):
-        print("  🔴 WARNING: Not all required AI services are configured correctly.")
-    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)