Update app.py

app.py

@@ -10,55 +10,58 @@ import tempfile
 st.set_page_config(page_title="Storyteller for Kids", layout="centered")
 st.title("🖼️ ➡️ 📖 Interactive Storyteller")
 
-# —––––––– Load and warm pipelines
+# —––––––– Load & warm pipelines
 @st.cache_resource
 def load_pipelines():
-    # BLIP-base for captions
+    # 1) BLIP-base for captions
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=-1  # GPU if available, else -1
+        device=0  # set to -1 if you only have CPU
     )
-    # Flan-T5-Large for stories
-    storyteller = pipeline(
-        "text2text-generation",
-        model="google/flan-t5-large",
+    # 2) DeepSeek-R1-Distill (Qwen-1.5B) for stories
+    ds_storyteller = pipeline(
+        "text-generation",
+        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+        trust_remote_code=True,
         device=0
     )
-    # Warm-up runs so user-facing calls are fast
+
+    # Warm-up both so the first real request is faster
     dummy = Image.new("RGB", (384, 384), color=(128, 128, 128))
     captioner(dummy)
-    storyteller("Warm up", max_new_tokens=1)
-    return captioner, storyteller
+    ds_storyteller("Warm up", max_new_tokens=1)
+
+    return captioner, ds_storyteller
 
-captioner, storyteller = load_pipelines()
+captioner, ds_storyteller = load_pipelines()
 
 # —––––––– Main UI
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
 if uploaded:
-    # 1) Preprocess image
+    # 1) Preprocess & display
     image = Image.open(uploaded).convert("RGB")
     image = image.resize((384, 384), Image.LANCZOS)
     st.image(image, caption="Your image", use_container_width=True)
 
-    # 2) Caption
-    with st.spinner("🔍 Generating caption…"):
+    # 2) Generate caption
+    with st.spinner("🔍 Generating caption..."):
         cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # 3) Build a dynamic prompt
+    # 3) Build prompt
     prompt = (
         f"Here is an image description: “{cap}”.\n"
         "Write an 80–100 word playful story for 3–10 year-old children that:\n"
-        "1) Describes the scene and subject from the description.\n"
-        "2) Explains what the subject is doing and how it feels.\n"
+        "1) Describes the scene and main subject.\n"
+        "2) Explains what it’s doing and how it feels.\n"
         "3) Concludes with a fun, imaginative ending.\n\n"
         "Story:"
     )
 
-    # 4) Generate the story
-    with st.spinner("✍️ Writing the story…"):
-        output = storyteller(
+    # 4) Generate story via DeepSeek
+    with st.spinner("✍️ Generating story with DeepSeek..."):
+        out = ds_storyteller(
             prompt,
             max_new_tokens=120,
             do_sample=True,
@@ -68,13 +71,13 @@ if uploaded:
             repetition_penalty=1.2,
             no_repeat_ngram_size=3
         )
-        story = output[0]["generated_text"].strip()
+        story = out[0]["generated_text"].strip()
 
     st.markdown("**Story:**")
     st.write(story)
 
-    # 5) Text-to-Speech
-    with st.spinner("🔊 Converting to speech…"):
+    # 5) Text-to-Speech via gTTS
+    with st.spinner("🔊 Converting to speech..."):
         tts = gTTS(text=story, lang="en")
         tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
         tts.write_to_fp(tmp)