Update app.py

app.py

@@ -18,10 +18,11 @@ def load_pipelines():
         "image-to-text",
         model="Salesforce/blip-image-captioning-base"
     )
-    # 2) Story generation with Flan-T5
+    # 2) Story generation with GPT-Neo 2.7B
     storyteller = pipeline(
-        "text2text-generation",
-        model="google/flan-t5-base"
+        "text-generation",
+        model="EleutherAI/gpt-neo-2.7B",
+        device=-1  # set to -1 if you only have CPU
     )
     return captioner, storyteller
 
@@ -36,25 +37,24 @@ if uploaded:
     # —––––––– 1. Caption
     with st.spinner("🔍 Looking at the image..."):
         cap_outputs = captioner(image)
-        # BLIP returns a list of dicts with key "generated_text"
         cap = cap_outputs[0].get("generated_text", "").strip()
     st.markdown(f"**Caption:** {cap}")
 
     # —––––––– 2. Story generation
     prompt = (
-        "Write a 80–100 words story for kids"
+        "Write a playful, 80–100 word story for 3–10 year-old children "
         f"based on this description:\n\n“{cap}”\n\nStory:"
     )
     with st.spinner("✍️ Writing a story..."):
         out = storyteller(
-        prompt,
-        max_new_tokens=120, 
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.8,
-        num_return_sequences=1
-)
-    story = out[0]["generated_text"].strip()
+            prompt,
+            max_new_tokens=120,    # allow space for ~100 words
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.8,
+            num_return_sequences=1
+        )
+        story = out[0]["generated_text"].strip()
     st.markdown("**Story:**")
     st.write(story)