Update app.py

app.py

@@ -13,20 +13,20 @@ st.title("🖼️ ➡️ 📖 Interactive Storyteller")
 # —––––––– Model loading + warm-up
 @st.cache_resource
 def load_pipelines():
-    # 1) Original BLIP-base captioner
+    # 1) Original BLIP-base for captions
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=0  # set to -1 if CPU-only
+        device=0  # change to -1 if you only have CPU
     )
-    # 2) Lightweight GPT-Neo for stories
+    # 2) Small GPT-Neo for quick stories
     storyteller = pipeline(
         "text-generation",
         model="EleutherAI/gpt-neo-125M",
         device=0
     )
 
-    # Warm-up so first real request is fast
+    # Warm up both so the first real call is faster
     dummy = Image.new("RGB", (384, 384), color=(128, 128, 128))
     captioner(dummy)
     storyteller("Hello", max_new_tokens=1)
@@ -35,10 +35,10 @@ def load_pipelines():
 
 captioner, storyteller = load_pipelines()
 
-# —––––––– Image upload & processing
+# —––––––– Main UI
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
 if uploaded:
-    # 1) Load + downsize for faster vision encoding
+    # 1) Load + resize for faster encoding
     image = Image.open(uploaded).convert("RGB")
     image = image.resize((384, 384), Image.LANCZOS)
     st.image(image, caption="Your image", use_container_width=True)
@@ -48,26 +48,31 @@ if uploaded:
         cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # 3) Story generation (greedy for speed)
+    # 3) Story generation (sampling + repetition control)
     prompt = (
-        f"Write an 80–100 word playful story for 3–10 year-olds "
-        f"based on this description:\n\n“{cap}”\n\nStory:"
+        f"Write an 80–100 word fun story for 3–10 year-old children "
+        f"based on this description:\n\n“{cap}”\n\nStory: "
     )
     with st.spinner("✍️ Generating story..."):
         out = storyteller(
             prompt,
-            max_new_tokens=120,
-            do_sample=False
+            max_new_tokens=120,        # room for ~100 words
+            do_sample=True,            # enable sampling
+            temperature=0.8,           # creativity
+            top_p=0.9,                 # nucleus sampling
+            top_k=50,                  # limit to top 50 tokens
+            repetition_penalty=1.2,    # discourage exact repeats
+            no_repeat_ngram_size=3     # prevent 3-gram repeats
         )
-        story = out[0]["generated_text"].strip()
+        # strip off the prompt so only the story remains
+        story = out[0]["generated_text"][len(prompt):].strip()
     st.markdown("**Story:**")
     st.write(story)
 
-    # 4) Text-to-Speech via gTTS (network-based)
+    # 4) 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)
         tmp.flush()
     st.audio(tmp.name, format="audio/mp3")
-