Update app.py

app.py

@@ -1,5 +1,7 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from diffusers import StableDiffusionPipeline
+import torch
 import json
 import textwrap
 
@@ -12,29 +14,59 @@ st.set_page_config(
     layout="wide"
 )
 
-st.title("🎬 AI Story → Movie Scene Generator (Phase 1)")
+st.title("🎬 AI Story → Movie Scene Generator")
 st.write(
     """
-    Paste a short story, and this app will break it into **cinematic scenes** with:
-    - Scene title  
-    - Setting & mood  
-    - Characters  
-    - Short summary  
-    - A detailed **visual prompt** (for future image generation)  
+    Paste a short story, and this app will:
+    1. Break it into **cinematic scenes** (title, setting, characters, mood, summary).
+    2. Generate a **visual prompt** for each scene.
+    3. Turn prompts into **AI images** in either:
+       - 🧪 Anime-style visuals  
+       - 🎥 Realistic cinematic visuals  
     """
 )
 
 # =========================
-# 2. LOAD LLM (FLAN-T5) - CACHED
+# 2. SIDEBAR: VISUAL STYLE
+# =========================
+st.sidebar.header("Visual Style Settings")
+style = st.sidebar.selectbox(
+    "Choose visual style for images:",
+    ["Anime", "Cinematic Realistic"]
+)
+
+
+def build_styled_prompt(base_prompt: str, style: str) -> str:
+    """
+    Take the base visual prompt from the scene and inject style instructions.
+    """
+    base_prompt = base_prompt.strip()
+    if style == "Anime":
+        return (
+            base_prompt +
+            ", anime style, detailed 2D illustration, clean line art, vibrant colors, "
+            "studio anime, keyframe, sharp focus, highly detailed, dramatic lighting"
+        )
+    else:  # Cinematic Realistic
+        return (
+            base_prompt +
+            ", ultra realistic, cinematic lighting, 35mm film, depth of field, 4k, "
+            "high detail, dramatic shadows, film still, volumetric light, highly detailed"
+        )
+
+
+# =========================
+# 3. LOAD LLM (FLAN-T5) - CACHED
 # =========================
 
 @st.cache_resource
 def load_scene_model():
-    model_name = "google/flan-t5-base"  # good starting point, can upgrade later
+    model_name = "google/flan-t5-base"  # good starting point; can upgrade to -large later
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     return tokenizer, model
 
+
 tokenizer, scene_model = load_scene_model()
 
 
@@ -55,7 +87,7 @@ def generate_text(prompt: str, max_new_tokens: int = 256) -> str:
 
 
 # =========================
-# 3. STORY → CHUNKS → SCENES LOGIC
+# 4. STORY → CHUNKS → SCENES LOGIC
 # =========================
 
 def split_story_into_chunks(story_text: str, max_chars_per_chunk: int = 600):
@@ -138,7 +170,52 @@ def story_to_scenes(story_text: str):
 
 
 # =========================
-# 4. STREAMLIT UI
+# 5. LOAD STABLE DIFFUSION PIPELINE (IMAGE MODEL)
+# =========================
+
+@st.cache_resource
+def load_image_model():
+    """
+    Load Stable Diffusion pipeline for image generation.
+    Uses CPU on Spaces by default; will use GPU if available.
+    """
+    model_id = "runwayml/stable-diffusion-v1-5"
+
+    if torch.cuda.is_available():
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=dtype,
+        safety_checker=None  # can be customized if needed
+    )
+
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")
+    else:
+        pipe = pipe.to("cpu")
+
+    return pipe
+
+
+def generate_scene_image(prompt: str):
+    """
+    Generate a single image from a text prompt using Stable Diffusion.
+    """
+    pipe = load_image_model()
+    # You can tweak num_inference_steps and guidance_scale for quality/speed tradeoff
+    image = pipe(
+        prompt,
+        num_inference_steps=25,
+        guidance_scale=7.5
+    ).images[0]
+    return image
+
+
+# =========================
+# 6. STREAMLIT UI
 # =========================
 
 st.subheader("📝 Paste Your Story")
@@ -159,32 +236,59 @@ story_text = st.text_area(
     height=260
 )
 
-if st.button("🎬 Generate Scenes"):
+generate_clicked = st.button("🎬 Generate Scenes")
+
+if "scenes" not in st.session_state:
+    st.session_state["scenes"] = None
+
+if generate_clicked:
     if not story_text.strip():
         st.error("Please paste a story first.")
     else:
         with st.spinner("Breaking story into scenes..."):
             scenes = story_to_scenes(story_text)
-
+        st.session_state["scenes"] = scenes
         st.success(f"Generated {len(scenes)} scene(s).")
 
-        st.markdown("---")
-        st.subheader("📚 Generated Scenes")
+scenes = st.session_state.get("scenes", None)
+
+if scenes:
+    st.markdown("---")
+    st.subheader("📚 Generated Scenes & Visuals")
+
+    for scene in scenes:
+        scene_id = scene.get("scene_id", "?")
+        title = scene.get("title", f"Scene {scene_id}")
+        setting = scene.get("setting", "")
+        mood = scene.get("mood", "")
+        characters = scene.get("characters", [])
+        summary = scene.get("summary", "")
+        base_prompt = scene.get("visual_prompt", "")
+
+        styled_prompt = build_styled_prompt(base_prompt, style)
+
+        with st.expander(f"Scene {scene_id}: {title}", expanded=True):
+            st.markdown(f"**Setting:** {setting}")
+            st.markdown(f"**Mood:** {mood}")
+            st.markdown(f"**Characters:** {', '.join(characters) or 'N/A'}")
+
+            st.markdown("**Summary:**")
+            st.write(summary)
 
-        for scene in scenes:
-            scene_id = scene.get("scene_id", "?")
-            title = scene.get("title", f"Scene {scene_id}")
-            with st.expander(f"Scene {scene_id}: {title}", expanded=True):
-                st.markdown(f"**Setting:** {scene.get('setting', '')}")
-                st.markdown(f"**Mood:** {scene.get('mood', '')}")
-                st.markdown(f"**Characters:** {', '.join(scene.get('characters', [])) or 'N/A'}")
+            st.markdown("**Base Visual Prompt:**")
+            st.code(textwrap.fill(base_prompt, width=90), language="text")
 
-                st.markdown("**Summary:**")
-                st.write(scene.get("summary", ""))
+            st.markdown(f"**Styled Prompt for {style} Image:**")
+            st.code(textwrap.fill(styled_prompt, width=90), language="text")
 
-                st.markdown("**Visual Prompt (for future image generation):**")
-                st.code(textwrap.fill(scene.get("visual_prompt", ""), width=90), language="text")
+            img_btn = st.button(
+                f"🖼 Generate {style} Image for Scene {scene_id}",
+                key=f"img_btn_{scene_id}"
+            )
+            if img_btn:
+                with st.spinner("Generating image... This may take some time."):
+                    img = generate_scene_image(styled_prompt)
+                st.image(img, caption=f"Scene {scene_id} – {title} ({style})", use_column_width=True)
 
-        st.info("✅ Phase 1 complete: story → structured scenes.\n\nNext phases will turn these visual prompts into images and then a video.")
 else:
     st.info("Paste a story and click **Generate Scenes** to begin.")