Update app.py

app.py

@@ -10,26 +10,26 @@ import tempfile
 st.set_page_config(page_title="Storyteller for Kids", layout="centered")
 st.title("🖼️ ➡️ 📖 Interactive Storyteller")
 
-# —––––––– Model loading + warm-up
+# —––––––– Load & warm models
 @st.cache_resource
 def load_pipelines():
-    # 1) Original BLIP-base for captions
+    # 1) BLIP-base for image captions
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
         device=0
     )
-    # 2) Instruction-tuned Flan-T5 small for stories
+    # 2) Flan-T5-Large for instruction following
     storyteller = pipeline(
         "text2text-generation",
-        model="google/flan-t5-small",
+        model="google/flan-t5-large",
         device=0
     )
 
-    # Warm up so first real request is faster
-    dummy = Image.new("RGB", (384, 384), color=(128, 128, 128))
+    # Warm up so first real call is faster
+    dummy = Image.new("RGB", (384, 384), color=(128,128,128))
     captioner(dummy)
-    storyteller("Tell me something", max_new_tokens=1)
+    storyteller("Hello", max_new_tokens=1)
 
     return captioner, storyteller
 
@@ -38,8 +38,9 @@ captioner, storyteller = load_pipelines()
 # —––––––– Main UI
 uploaded = st.file_uploader("Upload an image:", type=["jpg","jpeg","png"])
 if uploaded:
-    # 1) Load + downsize
-    image = Image.open(uploaded).convert("RGB").resize((384, 384), Image.LANCZOS)
+    # 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
@@ -47,32 +48,37 @@ if uploaded:
         cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # 3) Story generation
+    # 3) Story — stronger, clearer prompt
     prompt = (
-        f"Here is an image description: “{cap}”.\n"
-        "Write a playful, 80–100 word story for 3–10 year-olds\n\n"
+        f"Here’s an image description: “{cap}”.\n\n"
+        "Write a playful, 80–100 word story for 3–10 year-old children.\n"
+        "- Focus only on the panda and what it’s doing.\n"
+        "- Do not introduce any other characters (no kids, no parents).\n"
+        "- Be vivid: mention the panda’s feelings or the crunchy meat.\n\n"
         "Story:"
     )
-    with st.spinner("✍️ Generating story..."):
+    with st.spinner("✍️ Writing story..."):
         out = storyteller(
             prompt,
-            max_new_tokens=150,
+            max_new_tokens=130,
             do_sample=True,
             temperature=0.7,
             top_p=0.9,
-            repetition_penalty=1.2,
+            top_k=50,
+            repetition_penalty=1.3,
             no_repeat_ngram_size=3
         )
-        # strip off the prompt so you only get the story
-        story = out[0]["generated_text"].split("Story:")[-1].strip()
+        # strip prompt prefix — keep only the generated story
+        raw = out[0]["generated_text"]
+        story = raw.split("Story:")[-1].strip()
 
     st.markdown("**Story:**")
     st.write(story)
 
-    # 4) Text-to-Speech
+    # 4) Text-to-Speech (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)
         tmp.flush()
-    st.audio(tmp.name, format="audio/mp3")
+    st.audio(tmp.name, format="audio/mp3")