Update app.py

app.py

@@ -9,7 +9,8 @@ import traceback
 
 # --- 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
-from core.image_services import initialize_image_llms, STABILITY_API_CONFIGURED, OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, generate_image_stabilityai, generate_image_dalle, generate_image_hf_model, ImageGenResponse
+from core.image_services import initialize_image_llms, is_hf_image_api_ready, generate_image_hf_model, ImageGenResponse 
+# Removed STABILITY_API_CONFIGURED, OPENAI_DALLE_CONFIGURED, generate_image_stabilityai, generate_image_dalle
 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
@@ -17,21 +18,18 @@ from core.utils import basic_text_cleanup
 
 # --- Initialize Services ---
 initialize_text_llms()
-initialize_image_llms()
+initialize_image_llms() # This will initialize based on HF_TOKEN for images
 
 # --- Get API Readiness Status ---
 GEMINI_TEXT_IS_READY = is_gemini_text_ready()
-HF_TEXT_IS_READY = is_hf_text_ready() # For text models
-STABILITY_API_IS_READY = STABILITY_API_CONFIGURED
-OPENAI_DALLE_IS_READY = OPENAI_DALLE_CONFIGURED
-HF_IMAGE_IS_READY = HF_IMAGE_API_CONFIGURED # For image models via HF API
+HF_TEXT_IS_READY = is_hf_text_ready()
+HF_IMAGE_IS_READY = is_hf_image_api_ready() # Primary image status
 
 # --- Application Configuration (Models, Defaults) ---
 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"}
-    TEXT_MODELS["Legacy Gemini 1.0 Pro (Narrate)"] = {"id": "gemini-1.0-pro-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"}
     TEXT_MODELS["Gemma 2B (Narrate via HF)"] = {"id": "google/gemma-2b-it", "type": "hf_text"}
@@ -39,29 +37,21 @@ if HF_TEXT_IS_READY:
 if TEXT_MODELS:
     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)"
-    elif TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0]
+    else: UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0]
 else:
     TEXT_MODELS["No Text Models Configured"] = {"id": "dummy_text_error", "type": "none"}
     UI_DEFAULT_TEXT_MODEL_KEY = "No Text Models Configured"
 
-
 IMAGE_PROVIDERS = {}
 UI_DEFAULT_IMAGE_PROVIDER_KEY = None
-# Order of preference for default: Stability, then HF, then DALL-E (simulated)
-if STABILITY_API_IS_READY: IMAGE_PROVIDERS["🎨 Stability AI (SDXL)"] = "stability_ai"
 if HF_IMAGE_IS_READY: 
-    IMAGE_PROVIDERS["🎡 HF Model (SDXL Base)"] = "hf_image_sdxl_base" # Key for SDXL via HF
-    IMAGE_PROVIDERS["🎠 HF Model (OpenJourney)"] = "hf_image_openjourney" # Key for OpenJourney via HF
-if OPENAI_DALLE_IS_READY: IMAGE_PROVIDERS["🖼️ DALL-E 3 (Sim.)"] = "dalle"
-
-if IMAGE_PROVIDERS:
-    if STABILITY_API_IS_READY and "🎨 Stability AI (SDXL)" in IMAGE_PROVIDERS: UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎨 Stability AI (SDXL)"
-    elif HF_IMAGE_IS_READY and "🎡 HF Model (SDXL Base)" in IMAGE_PROVIDERS: UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF Model (SDXL Base)"
-    elif OPENAI_DALLE_IS_READY and "🖼️ DALL-E 3 (Sim.)" in IMAGE_PROVIDERS: UI_DEFAULT_IMAGE_PROVIDER_KEY = "🖼️ DALL-E 3 (Sim.)"
-    else: UI_DEFAULT_IMAGE_PROVIDER_KEY = list(IMAGE_PROVIDERS.keys())[0]
+    IMAGE_PROVIDERS["🎡 HF - Stable Diffusion XL Base"] = "hf_sdxl_base" 
+    IMAGE_PROVIDERS["🎠 HF - OpenJourney (Midjourney-like)"] = "hf_openjourney"
+    IMAGE_PROVIDERS["🌌 HF - Stable Diffusion v1.5"] = "hf_sd_1_5" # Another option
+    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF - Stable Diffusion XL Base"
 else:
-    IMAGE_PROVIDERS["No Image Providers Configured"] = "none"
-    UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured"
+    IMAGE_PROVIDERS["No Image Providers Configured (HF Token needed)"] = "none"
+    UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured (HF Token needed)"
 
 
 # --- Gradio UI Theme and CSS ---
@@ -85,7 +75,7 @@ body, .gradio-container { background-color: #0F0F1A !important; color: #D0D0E0 !
 .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 { padding-top: 10px !important; padding-bottom: 10px !important; } /* Ensure enough padding for icons/text */
+.gr-button-primary { padding-top: 10px !important; padding-bottom: 10px !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; }
@@ -110,7 +100,6 @@ body, .gradio-container { background-color: #0F0F1A !important; color: #D0D0E0 !
 
 # --- Helper: Placeholder Image Creation ---
 def create_placeholder_image(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
-    # ... (Full implementation 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 
@@ -120,7 +109,6 @@ 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,
@@ -133,20 +121,18 @@ def add_scene_to_story_orchestrator(
     
     log_accumulator = [f"**🚀 Scene {current_story_obj.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**"]
     
-    # Placeholders for the final return tuple
     ret_story_state = current_story_obj
     ret_gallery = current_story_obj.get_all_scenes_for_gallery_display()
-    ret_latest_image = create_placeholder_image("🎨 Conjuring visuals...") # Initial placeholder for image component
-    ret_latest_narrative_md_obj = gr.Markdown(value=" Musing narrative...")
-    ret_status_bar_html_obj = gr.HTML(value=f"<p class='processing_text status_text'>🌌 Weaving Scene {current_story_obj.current_scene_number + 1}...</p>")
-    # Log is built up
+    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 will be built up
 
-    # Initial yield for UI updates handled by the .then() chain for buttons
-    # This first yield is for immediate feedback on other components
+    # Initial yield for UI updates (buttons disabled by .then() chain)
     yield {
-        output_status_bar: ret_status_bar_html_obj,
-        output_latest_scene_image: gr.Image(value=ret_latest_image), # Show placeholder
-        output_latest_scene_narrative: ret_latest_narrative_md_obj,
+        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))
     }
 
@@ -166,21 +152,13 @@ def add_scene_to_story_orchestrator(
             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. (Snippet: {narrative_text_generated[:50]}...)")
-            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 object from {text_model_key}.")
-        else:
-            narrative_text_generated = "**Narrative Error:** Selected text model not available or misconfigured."
-            log_accumulator.append(f"  Narrative: FAILED - Model '{text_model_key}' not available.")
+            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 object from {text_model_key}.")
+        else: narrative_text_generated = "**Narrative Error:** Selected text model not available or misconfigured."; log_accumulator.append(f"  Narrative: FAILED - Model '{text_model_key}' not available.")
         
         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)
+        ret_latest_narrative_md_obj = gr.Markdown(value=ret_latest_narrative_str_content) # Update the component object
         yield { output_latest_scene_narrative: ret_latest_narrative_md_obj,
                 output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)) }
 
@@ -188,31 +166,51 @@ def add_scene_to_story_orchestrator(
         progress(0.5, desc="🎨 Conjuring visuals...")
         image_generated_pil = None
         image_generation_error_message = None
-        selected_image_provider_type = IMAGE_PROVIDERS.get(image_provider_key)
+        selected_image_provider_key_from_ui = image_provider_key # The key from the dropdown
+        selected_image_provider_type = IMAGE_PROVIDERS.get(selected_image_provider_key_from_ui) # The value (e.g., "hf_sdxl_base")
+
         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 "")
         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'}.")
+        log_accumulator.append(f"  Image: Using provider key '{selected_image_provider_key_from_ui}' (maps to type '{selected_image_provider_type}'). Style: {image_style_dropdown}.")
 
         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)
-            elif selected_image_provider_type == "dalle": image_response = generate_image_dalle(full_image_prompt, quality="hd" if image_quality=="High Detail" else "standard")
-            elif selected_image_provider_type == "hf_image_sdxl_base": image_response = generate_image_hf_model(full_image_prompt, model_id="stabilityai/stable-diffusion-xl-base-1.0", negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS)
-            elif selected_image_provider_type == "hf_image_openjourney": image_response = generate_image_hf_model(full_image_prompt, model_id="prompthero/openjourney", negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS, width=512, height=512)
-            
-            if image_response and image_response.success:
-                image_generated_pil = image_response.image
-                log_accumulator.append(f"  Image: Success from {image_response.provider}.")
-            elif image_response:
-                image_generation_error_message = f"**Image Error ({image_response.provider}):** {image_response.error}"
-                log_accumulator.append(f"  Image: FAILED - {image_response.error}")
+            if not HF_IMAGE_IS_READY and selected_image_provider_type.startswith("hf_"): # Check if HF image specifically is ready
+                image_generation_error_message = "**Image Error:** Hugging Face Image API not configured (check STORYVERSE_HF_TOKEN)."
+                log_accumulator.append(f"  Image: FAILED - {image_generation_error_message}")
             else:
-                image_generation_error_message = f"**Image Error:** No response object from {image_provider_key} service."
-                log_accumulator.append(f"  Image: FAILED - No response object from {image_provider_key}.")
+                image_response = None
+                hf_model_id_to_call = None
+                if selected_image_provider_type == "hf_sdxl_base": hf_model_id_to_call = "stabilityai/stable-diffusion-xl-base-1.0"
+                elif selected_image_provider_type == "hf_openjourney": hf_model_id_to_call = "prompthero/openjourney"; img_width, img_height = 512, 512
+                elif selected_image_provider_type == "hf_sd_1_5": hf_model_id_to_call = "runwayml/stable-diffusion-v1-5"; img_width, img_height = 512, 512
+                # Add other specific HF model mappings here
+                
+                if hf_model_id_to_call: # If it's an HF model type
+                    image_response = generate_image_hf_model(
+                        full_image_prompt, model_id=hf_model_id_to_call,
+                        negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS,
+                        width=img_width if 'img_width' in locals() else 768, 
+                        height=img_height if 'img_height' in locals() else 768
+                    )
+                # Add elif for other providers like StabilityAI or DALL-E if you re-enable them
+                # elif selected_image_provider_type == "stability_ai": ...
+                else:
+                    image_generation_error_message = f"**Image Error:** Provider type '{selected_image_provider_type}' not handled."
+                    log_accumulator.append(f"  Image: FAILED - {image_generation_error_message}")
+
+                if image_response and image_response.success:
+                    image_generated_pil = image_response.image
+                    log_accumulator.append(f"  Image: Success from {image_response.provider} (Model: {image_response.model_id_used}).")
+                elif image_response:
+                    image_generation_error_message = f"**Image Error ({image_response.provider} - {image_response.model_id_used}):** {image_response.error}"
+                    log_accumulator.append(f"  Image: FAILED - {image_response.error}")
+                elif not image_generation_error_message: # If no response and no specific error yet
+                    image_generation_error_message = f"**Image Error:** No response from image service for {selected_image_provider_key_from_ui}."
+                    log_accumulator.append(f"  Image: FAILED - No response object.")
         else:
-            image_generation_error_message = "**Image Error:** Selected image provider not available or misconfigured."
-            log_accumulator.append(f"  Image: FAILED - Provider '{image_provider_key}' unavailable.")
+            image_generation_error_message = "**Image Error:** No image provider selected or configured."
+            log_accumulator.append(f"  Image: FAILED - Provider key '{selected_image_provider_key_from_ui}' not found or type is 'none'.")
 
         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),
@@ -229,7 +227,7 @@ def add_scene_to_story_orchestrator(
             narrative_text=narrative_text_generated if "**Narrative Error**" not in narrative_text_generated else "(Narrative gen failed)",
             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=image_provider_key if selected_image_provider_type != "none" else "N/A",
+            image_provider=selected_image_provider_key_from_ui, # Use the UI key for display
             error_message=final_scene_error
         )
         ret_story_state = current_story_obj 
@@ -253,15 +251,20 @@ def add_scene_to_story_orchestrator(
         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.")
-    # `finally` block is removed from here; button re-enabling is handled by .then()
-
+    
+    # Final log update happens as part of the main return now
     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))
     
+    # This is the FINAL return. It must be a tuple matching the `outputs` list of engage_button.click()
     return (
-        ret_story_state, ret_gallery, ret_latest_image, 
-        ret_latest_narrative_md_obj, ret_status_bar_html_obj, ret_log_md
+        ret_story_state,
+        ret_gallery,
+        ret_latest_image, 
+        ret_latest_narrative_md_obj, 
+        ret_status_bar_html_obj,     
+        ret_log_md                   
     )
 
 def clear_story_state_ui_wrapper():
@@ -283,34 +286,13 @@ def enable_buttons_after_processing():
 
 # --- 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 that will be inputs or outputs
     story_state_output = gr.State(Story()) 
-    # Input Column Components
-    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 Column Components
-    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()
-    # Button components (updated via chained .then() events)
-    engage_button = gr.Button()
-    surprise_button = gr.Button()
-    clear_story_button = gr.Button()
 
-    # --- Layout the 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>")
 
     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 or HF_IMAGE_IS_READY) # Added HF_IMAGE_IS_READY
+        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 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 Service(s) Ready.</p>")
@@ -326,7 +308,7 @@ with gr.Blocks(theme=omega_theme, css=omega_css, title="✨ StoryVerse Omega ✨
                 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", allow_custom_value=True) # allow_custom_value for examples
+                    image_style_input = gr.Dropdown(choices=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), value="Default (Cinematic Realism)", label="Visual Style Preset", 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 (Negative Prompt):", value=COMMON_NEGATIVE_PROMPTS)