Update app.py

app.py

@@ -13,40 +13,40 @@ st.title("🖼️ ➡️ 📖 Interactive Storyteller")
 # —––––––– Model loading + warm-up
 @st.cache_resource
 def load_pipelines():
-    # 1) Smaller BLIP for captions
+    # 1) Keep the original BLIP-base for captions
     captioner = pipeline(
         "image-to-text",
-        model="Salesforce/blip-image-captioning-small",
-        device=0
+        model="Salesforce/blip-image-captioning-base",
+        device=0  # if you have GPU; use -1 for CPU-only
     )
-    # 2) Small GPT-Neo for stories
+    # 2) Switch to a lightweight story model
     storyteller = pipeline(
         "text-generation",
         model="EleutherAI/gpt-neo-125M",
         device=0
     )
 
-    # Warm up both models once
-    dummy = Image.new("RGB", (384,384), color=(128,128,128))
+    # Warm up with a dummy run so first real call is fast
+    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)
+# —––––––– Initialize local TTS (offline)
 @st.cache_resource
 def init_tts_engine():
     engine = pyttsx3.init()
-    engine.setProperty('rate', 150)
+    engine.setProperty("rate", 150)   # words per minute
     return engine
 
 captioner, storyteller = load_pipelines()
 tts_engine = init_tts_engine()
 
-# —––––––– Image upload & processing
+# —––––––– Main UI
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
 if uploaded:
-    # 1) Load + downsize
+    # 1) Resize image to reduce BLIP load
     image = Image.open(uploaded).convert("RGB")
     image = image.resize((384, 384), Image.LANCZOS)
     st.image(image, caption="Your image", use_container_width=True)
@@ -56,21 +56,21 @@ if uploaded:
         cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # 3) Story
+    # 3) Story (greedy = fastest)
     prompt = (
-        f"Write a fun, 80–100 word story for kids based on:\n\n“{cap}”\n\nStory:"
+        f"Tell an 80–100 word fun story for 3–10 year-olds based on:\n\n“{cap}”\n\nStory:"
     )
     with st.spinner("✍️ Generating story..."):
         out = storyteller(
             prompt,
             max_new_tokens=120,
-            do_sample=False,     # greedy = fastest
+            do_sample=False
         )
         story = out[0]["generated_text"].strip()
     st.markdown("**Story:**")
     st.write(story)
 
-    # 4) TTS
+    # 4) TTS (local, no network)
     with st.spinner("🔊 Converting to speech..."):
         tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
         tts_engine.save_to_file(story, tmp.name)