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