Update app.py

app.py

@@ -17,58 +17,58 @@ def load_pipelines():
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=0  # GPU if available, else set -1 for CPU
+        device=-1  # or -1 if you only have CPU
     )
-    # 2) Flan-T5-Large for coherent, instruction-driven stories
+    # 2) Flan-T5-Large for instruction-driven stories
     storyteller = pipeline(
         "text2text-generation",
         model="google/flan-t5-large",
         device=0
     )
 
-    # Warm-up so first real inference is faster
+    # Warm-up to avoid first-call lag
     dummy = Image.new("RGB", (384, 384), color=(128, 128, 128))
     captioner(dummy)
-    storyteller("Hello", max_new_tokens=1, return_full_text=False)
+    storyteller("Warm up", max_new_tokens=1)
 
     return captioner, storyteller
 
 captioner, storyteller = load_pipelines()
 
 # —––––––– Main UI
-uploaded = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
+uploaded = st.file_uploader("Upload an image:", type=["jpg","jpeg","png"])
 if uploaded:
     # 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) Generate caption
+    # 2) Caption
     with st.spinner("🔍 Generating caption..."):
         cap = captioner(image)[0]["generated_text"].strip()
     st.markdown(f"**Caption:** {cap}")
 
-    # 3) Generate story with stronger prompt & no carry-over
+    # 3) Story
     prompt = (
-        f"Here is an image description: “{cap}”.\n"
+        f"Image description: “{cap}”.\n"
         "Write an 80–100 word playful story for 3–10 year-old children that:\n"
-        "1) Sets the scene with the panda and its surroundings.\n"
-        "2) Describes what the panda is doing and how it feels.\n"
-        "3) Wraps up with a fun conclusion.\n\n"
+        "1) Sets the scene around the panda.\n"
+        "2) Describes what it’s doing and how it feels.\n"
+        "3) Ends with a fun conclusion.\n\n"
         "Story:"
     )
     with st.spinner("✍️ Generating story..."):
         out = storyteller(
             prompt,
-            max_new_tokens=120,
+            max_new_tokens=130,
             do_sample=True,
             temperature=0.7,
             top_p=0.9,
             top_k=50,
-            repetition_penalty=1.2,
-            no_repeat_ngram_size=3,
-            return_full_text=False    # only the story text, not the prompt
+            repetition_penalty=1.3,
+            no_repeat_ngram_size=3
         )
+        # text2text pipeline returns only the generated part
         story = out[0]["generated_text"].strip()
 
     st.markdown("**Story:**")
@@ -81,4 +81,3 @@ if uploaded:
         tts.write_to_fp(tmp)
         tmp.flush()
     st.audio(tmp.name, format="audio/mp3")
-