Update app.py

app.py

@@ -6,14 +6,6 @@ from transformers import pipeline
 from gtts import gTTS
 import tempfile
 
-# —––––––– Requirements —–––––––
-# streamlit>=1.20
-# pillow>=9.0
-# torch>=2.0.0
-# transformers>=4.30
-# sentencepiece>=0.1.97
-# gTTS>=2.3.1
-
 # —––––––– Page Setup —–––––––
 st.set_page_config(page_title="Magic Story Generator", layout="centered")
 st.title("📖✨ Turn Images into Children's Stories")
@@ -27,16 +19,15 @@ def load_pipelines():
         model="Salesforce/blip-image-captioning-base",
         device=-1  # CPU; set to 0+ for GPU
     )
-    # 2) Story-generation pipeline (StoryTeller)
+    # 2) Story-generation pipeline (T5-base Story)
     storyteller = pipeline(
         task="text2text-generation",
-        model="pranavpsv/StoryTeller",
-        tokenizer="pranavpsv/StoryTeller",
-        device=-1,                # CPU; set to 0+ for GPU
+        model="mrm8488/t5-base-finetuned-story-generation",
+        tokenizer="mrm8488/t5-base-finetuned-story-generation",
+        device=-1,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.2,
-        no_repeat_ngram_size=3,
         max_new_tokens=150
     )
     return captioner, storyteller
@@ -46,7 +37,6 @@ captioner, storyteller = load_pipelines()
 # —––––––– Main App —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
 if uploaded:
-    # Load and display the image
     img = Image.open(uploaded).convert("RGB")
     st.image(img, use_column_width=True)
 
@@ -60,10 +50,7 @@ if uploaded:
     st.success(f"**Caption:** {caption}")
 
     # Build prompt and generate story
-    prompt = (
-        f"Image description: {caption}\n"
-        "Write a coherent, 50-100 word children’s story that flows naturally."
-    )
+    prompt = f"generate story: {caption}"
     with st.spinner("📝 Writing story..."):
         start = time.time()
         out = storyteller(prompt)
@@ -93,3 +80,4 @@ if uploaded:
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
 
+