Update app.py

app.py

@@ -3,65 +3,76 @@
 import streamlit as st
 from PIL import Image
 from transformers import pipeline
-from gtts import gTTS
+import pyttsx3
 import tempfile
 
 # —––––––– Page config
 st.set_page_config(page_title="Storyteller for Kids", layout="centered")
 st.title("🖼️ ➡️ 📖 Interactive Storyteller")
 
-# —––––––– Cache model loading
+# —––––––– Model loading + warm-up
 @st.cache_resource
 def load_pipelines():
-    # 1) Image-to-text (captioning)
+    # 1) Smaller BLIP for captions
     captioner = pipeline(
         "image-to-text",
-        model="Salesforce/blip-image-captioning-base"
+        model="Salesforce/blip-image-captioning-small",
+        device=0
     )
-    # 2) Story generation with GPT-Neo 2.7B
+    # 2) Small GPT-Neo for stories
     storyteller = pipeline(
         "text-generation",
-        model="EleutherAI/gpt-neo-2.7B",
-        device=-1  # set to -1 if you only have CPU
+        model="EleutherAI/gpt-neo-125M",
+        device=0
     )
+
+    # Warm up both models once
+    dummy = Image.new("RGB", (384,384), color=(128,128,128))
+    captioner(dummy)
+    storyteller("Hello", max_new_tokens=1)
+
     return captioner, storyteller
 
+# —––––––– TTS engine init (offline)
+@st.cache_resource
+def init_tts_engine():
+    engine = pyttsx3.init()
+    engine.setProperty('rate', 150)
+    return engine
+
 captioner, storyteller = load_pipelines()
+tts_engine = init_tts_engine()
 
-# —––––––– Image upload
+# —––––––– Image upload & processing
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
 if uploaded:
+    # 1) Load + downsize
     image = Image.open(uploaded).convert("RGB")
-    st.image(image, caption="Your image", use_column_width=True)
+    image = image.resize((384, 384), Image.LANCZOS)
+    st.image(image, caption="Your image", use_container_width=True)
 
-    # —––––––– 1. Caption
-    with st.spinner("🔍 Looking at the image..."):
-        cap_outputs = captioner(image)
-        cap = cap_outputs[0].get("generated_text", "").strip()
+    # 2) Caption
+    with st.spinner("🔍 Generating caption..."):
+        cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # —––––––– 2. Story generation
+    # 3) Story
     prompt = (
-        "Write a playful, 80–100 word story for 3–10 year-old children "
-        f"based on this description:\n\n“{cap}”\n\nStory:"
+        f"Write a fun, 80–100 word story for kids based on:\n\n“{cap}”\n\nStory:"
     )
-    with st.spinner("✍️ Writing a story..."):
+    with st.spinner("✍️ Generating story..."):
         out = storyteller(
             prompt,
-            max_new_tokens=120,    # allow space for ~100 words
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.8,
-            num_return_sequences=1
+            max_new_tokens=120,
+            do_sample=False,     # greedy = fastest
         )
         story = out[0]["generated_text"].strip()
     st.markdown("**Story:**")
     st.write(story)
 
-    # —––––––– 3. Text-to-Speech
+    # 4) TTS
     with st.spinner("🔊 Converting to speech..."):
-        tts = gTTS(story, lang="en")
-        tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
-        tts.write_to_fp(tmp)
-        tmp.flush()
-    st.audio(tmp.name, format="audio/mp3")
+        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+        tts_engine.save_to_file(story, tmp.name)
+        tts_engine.runAndWait()
+    st.audio(tmp.name)