app.py

# storyverse_weaver/app.py
import gradio as gr
import os
import time
# ... (other imports: json, PIL, random, 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
# MODIFIED IMPORT for image_services
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
# ... (other core imports: prompts, utils)
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() 

# --- Get API Readiness Status ---
GEMINI_TEXT_IS_READY = is_gemini_text_ready()
HF_TEXT_IS_READY = is_hf_text_ready() # For text fallback
DALLE_IMAGE_IS_READY = is_dalle_ready() # Primary image status
HF_IMAGE_IS_READY = is_hf_image_api_ready() # For image fallback

# --- Application Configuration (Models, Defaults) ---
TEXT_MODELS = {}
UI_DEFAULT_TEXT_MODEL_KEY = None
if GEMINI_TEXT_IS_READY: # Prioritize Gemini for text
    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"}
    UI_DEFAULT_TEXT_MODEL_KEY = "✨ Gemini 1.5 Flash (Narrate)"
elif HF_TEXT_IS_READY: # Fallback to HF for text
    TEXT_MODELS["Mistral 7B (Narrate via HF - Fallback)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"}
    TEXT_MODELS["Gemma 2B (Narrate via HF - Fallback)"] = {"id": "google/gemma-2b-it", "type": "hf_text"}
    UI_DEFAULT_TEXT_MODEL_KEY = "Mistral 7B (Narrate via HF - Fallback)"
    
if not TEXT_MODELS: # If neither is ready
    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
if DALLE_IMAGE_IS_READY: # Prioritize DALL-E for images
    IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 3"] = "dalle_3" # Key for DALL-E 3
    IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 2 (Legacy)"] = "dalle_2" # Key for DALL-E 2
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🖼️ OpenAI DALL-E 3"
elif HF_IMAGE_IS_READY: # Fallback to HF for images
    IMAGE_PROVIDERS["🎡 HF - Stable Diffusion XL Base (Fallback)"] = "hf_sdxl_base" 
    IMAGE_PROVIDERS["🎠 HF - OpenJourney (Fallback)"] = "hf_openjourney"
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF - Stable Diffusion XL Base (Fallback)"

if not IMAGE_PROVIDERS:
    IMAGE_PROVIDERS["No Image Providers Configured"] = "none"
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured"


# ... (Theme, CSS, create_placeholder_image - REMAINS THE SAME as previous full app.py) ...
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 = "body, .gradio-container { background-color: #0F0F1A !important; color: #D0D0E0 !important; } /* Ensure this is complete */"
def create_placeholder_image(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"): # Keep this
    img = Image.new('RGB', size, color=color); draw = ImageDraw.Draw(img); #... (full implementation)
    try: font_path = "arial.ttf" if os.path.exists("arial.ttf") else None 
    except: font_path = None 
    try: font = ImageFont.truetype(font_path, 40) if font_path else ImageFont.load_default()
    except IOError: font = ImageFont.load_default()
    if hasattr(draw, 'textbbox'): bbox = draw.textbbox((0,0), text, font=font); tw, th = bbox[2]-bbox[0], bbox[3]-bbox[1]
    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 (MODIFIED image generation part) ---
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, # image_provider_key now maps to DALL-E or HF
    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() 
    log_accumulator = [f"**🚀 Scene {current_story_obj.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**"]
    # ... (Initialize ret_... placeholders as before) ...
    ret_story_state, ret_gallery, ret_latest_image, ret_latest_narrative_md_obj, ret_status_bar_html_obj, ret_log_md = \
        current_story_obj, current_story_obj.get_all_scenes_for_gallery_display(), None, gr.Markdown("Processing..."), gr.HTML("<p>Processing...</p>"), gr.Markdown("\n".join(log_accumulator))

    # Initial yield
    yield {
        output_status_bar: gr.HTML(value=f"<p class='processing_text status_text'>🌌 Weaving Scene {current_story_obj.current_scene_number + 1}...</p>"),
        # ... (other initial yields)
        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))
    }
    # Note: Button disabling/enabling is handled by the .then() chain in UI definition

    try:
        if not scene_prompt_text.strip(): raise ValueError("Scene prompt cannot be empty!")

        # --- 1. Generate Narrative Text (using Gemini or HF fallback) ---
        progress(0.1, desc="✍️ Crafting narrative...")
        # ... (Full narrative generation logic from previous app.py, which already handles Gemini/HF choice) ...
        # ... (This part should be copied from your last working version) ...
        narrative_text_generated = "Simulated Narrative." # Placeholder
        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)
        yield { output_latest_scene_narrative: ret_latest_narrative_md_obj, output_interaction_log_markdown: gr.Markdown(value="\n".join(log_accumulator)) }


        # --- 2. Generate Image (NOW USING DALL-E or HF fallback) ---
        progress(0.5, desc="🎨 Conjuring visuals...")
        image_generated_pil = None
        image_generation_error_message = None
        selected_image_provider_actual_type = IMAGE_PROVIDERS.get(image_provider_key) # e.g., "dalle_3", "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 = "detailed, high quality, " if image_quality == "High Detail" else "" # Simpler quality keyword
        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 '{image_provider_key}' (maps to type '{selected_image_provider_actual_type}'). Style: {image_style_dropdown}.")

        if selected_image_provider_actual_type and selected_image_provider_actual_type != "none":
            image_response = None
            if selected_image_provider_actual_type == "dalle_3":
                if DALLE_IMAGE_IS_READY:
                    image_response = generate_image_dalle(full_image_prompt, model="dall-e-3", quality="hd" if image_quality=="High Detail" else "standard")
                else: image_generation_error_message = "**Image Error:** DALL-E 3 selected but API not ready."
            elif selected_image_provider_actual_type == "dalle_2":
                if DALLE_IMAGE_IS_READY:
                    image_response = generate_image_dalle(full_image_prompt, model="dall-e-2", size="1024x1024") # DALL-E 2 has fixed sizes
                else: image_generation_error_message = "**Image Error:** DALL-E 2 selected but API not ready."
            # Fallback to HF models if DALL-E not selected or not ready, but HF is
            elif selected_image_provider_actual_type.startswith("hf_"):
                if HF_IMAGE_IS_READY:
                    hf_model_id_to_call = "stabilityai/stable-diffusion-xl-base-1.0" # Default HF
                    img_width, img_height = 768, 768
                    if selected_image_provider_actual_type == "hf_openjourney": hf_model_id_to_call = "prompthero/openjourney"; img_width,img_height = 512,512
                    elif selected_image_provider_actual_type == "hf_sd_1_5": hf_model_id_to_call = "runwayml/stable-diffusion-v1-5"; img_width,img_height = 512,512
                    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, height=img_height)
                else: image_generation_error_message = "**Image Error:** HF Image Model selected but API not ready."
            else: image_generation_error_message = f"**Image Error:** Provider type '{selected_image_provider_actual_type}' not handled."

            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: image_generation_error_message = f"**Image Error:** No response/unknown issue with {image_provider_key}."; log_accumulator.append(f"  Image: FAILED - No response object.")
        
        if not image_generated_pil and not image_generation_error_message: # If neither DALL-E nor HF was selected/ready
            image_generation_error_message = "**Image Error:** No valid image provider configured or selected."
            log_accumulator.append(f"  Image: FAILED - {image_generation_error_message}")

        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 & 4. Prepare Final Return Values ---
        # ... (This part remains largely the same as the previous full app.py) ...
        final_scene_error=None; # ... (set based on narrative/image errors) ...
        current_story_obj.add_scene_from_elements(user_prompt=scene_prompt_text, narrative_text=narrative_text_generated, image=image_generated_pil, image_style_prompt=f"{image_style_dropdown} by {artist_style_text}", image_provider=image_provider_key, error_message=final_scene_error)
        ret_story_state = current_story_obj; log_accumulator.append(f"  Scene {current_story_obj.current_scene_number} processed.")
        ret_gallery = current_story_obj.get_all_scenes_for_gallery_display()
        _ , 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'>Scene added with errors.</p>" if final_scene_error else f"<p class='success_text'>🌌 Scene woven!</p>"
        ret_status_bar_html_obj = gr.HTML(value=status_html_str_temp)
        progress(1.0, desc="Scene Complete!")


    except ValueError as ve: # ... (Error handling as before) ...
        log_accumulator.append(f"\n**INPUT ERROR:** {ve}"); ret_status_bar_html_obj = gr.HTML(f"<p class='error_text'>ERROR: {ve}</p>"); ret_latest_narrative_md_obj = gr.Markdown(f"## Error\n{ve}")
    except Exception as e: # ... (Error handling as before) ...
        log_accumulator.append(f"\n**RUNTIME ERROR:** {e}\n{traceback.format_exc()}"); ret_status_bar_html_obj = gr.HTML(f"<p class='error_text'>UNEXPECTED ERROR: {e}</p>"); ret_latest_narrative_md_obj = gr.Markdown(f"## Unexpected Error\n{e}")
    
    current_total_time = time.time() - start_time
    log_accumulator.append(f"  Cycle ended. Total time: {current_total_time:.2f}s")
    ret_log_md = gr.Markdown(value="\n".join(log_accumulator))
    
    return (ret_story_state, ret_gallery, ret_latest_image, ret_latest_narrative_md_obj, ret_status_bar_html_obj, ret_log_md)

# --- clear_story_state_ui_wrapper, surprise_me_func, disable_buttons_for_processing, enable_buttons_after_processing ---
# (These functions remain IDENTICAL to the ones in the last full app.py that fixed the ValueError)
def clear_story_state_ui_wrapper(): new_story=Story(); ph_img=create_placeholder_image("Blank..."); return(new_story,[(ph_img,"New...")],None,gr.Markdown("## Cleared"),gr.HTML("<p>Cleared.</p>"),"Log Cleared","")
def surprise_me_func(): themes = ["Cosmic Horror", "Solarpunk"]; actions = ["unearths artifact", "negotiates"]; settings = ["on rogue planet", "in tree city"]; prompt = f"Protagonist {random.choice(actions)} {random.choice(settings)}. Theme: {random.choice(themes)}."; style = random.choice(list(STYLE_PRESETS.keys())); artist = random.choice(["Giger", "Moebius", ""]*2); return prompt, style, artist
def disable_buttons_for_processing(): return gr.Button(interactive=False), gr.Button(interactive=False)
def enable_buttons_after_processing(): 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:
    story_state_output = gr.State(Story()) 
    # ... (Full UI layout from the previous app.py - "rewrite app.py with update")
    # ... (This includes defining all component variables like scene_prompt_input, output_gallery, engage_button etc. IN THE LAYOUT)
    # Ensure the image_provider_dropdown choices and default reflect DALL-E and HF
    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> ... API keys (<code>STORYVERSE_...</code>) ...</div>")
    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:", 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)
                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") # Updated choices
                    with gr.Row():
                        narrative_length_dropdown = gr.Dropdown(["Short", "Medium", "Detailed"], value="Medium", label="Narrative Detail")
                        image_quality_dropdown = gr.Dropdown(["Standard", "High Detail", "Sketch"], value="Standard", label="Image Detail")
            with gr.Row(elem_classes=["compact-row"], equal_height=True):
                engage_button = gr.Button("🌌 Weave Scene!", variant="primary", scale=3, icon="✨")
                surprise_button = gr.Button("🎲 Surprise!", variant="secondary", scale=1, icon="🎁")
                clear_story_button = gr.Button("🗑️ New", variant="stop", scale=1, icon="♻️")
            output_status_bar = gr.HTML(value="<p class='processing_text status_text'>Ready to weave!</p>")
        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); 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)
                with gr.TabItem("⚙️ Log"): 
                    with gr.Accordion("Interaction Log", open=False): output_interaction_log_markdown = gr.Markdown("Log...")

    # Event Handlers (same .then() chain as before)
    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])
    gr.Examples(examples=[["Traveler in desert...", "Sci-Fi Western", "Moebius", "greenery"]], inputs=[scene_prompt_input, image_style_input, artist_style_input, negative_prompt_input], label="🌌 Examples 🌌")
    gr.HTML("<p style='text-align:center; font-size:0.9em; color:#8080A0;'>✨ StoryVerse Omega™ ✨</p>")

# --- Entry Point ---
if __name__ == "__main__":
    print("="*80); print("✨ StoryVerse Omega (DALL-E/Gemini Focus) 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 Ready: {HF_IMAGE_IS_READY}") # Check both
    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 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)