Update app.py

app.py

@@ -13,22 +13,22 @@ st.title("📖✨ Turn Images into Children's Stories")
 # —––––––– Load Pipelines (cached) —–––––––
 @st.cache_resource(show_spinner=False)
 def load_pipelines():
-    # 1) Image-captioning pipeline (BLIP)
+    # 1) Image captioning pipeline
     captioner = pipeline(
         task="image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=-1  # CPU; set to 0+ for GPU
+        device=-1
     )
-    # 2) Story-generation pipeline (T5-base Story)
+    
+    # 2) Story generation pipeline using verified model
     storyteller = pipeline(
         task="text2text-generation",
-        model="mrm8488/t5-base-finetuned-story-generation",
-        tokenizer="mrm8488/t5-base-finetuned-story-generation",
+        model="laxya007/story-generator-t5-small",
+        tokenizer="t5-small",
         device=-1,
+        max_length=200,
         temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.2,
-        max_new_tokens=150
+        do_sample=True
     )
     return captioner, storyteller
 
@@ -43,41 +43,43 @@ if uploaded:
     # Generate caption
     with st.spinner("🔍 Generating caption..."):
         cap = captioner(img)
-        caption = cap[0].get("generated_text", "").strip() if isinstance(cap, list) else ""
+        caption = cap[0].get("generated_text", "").strip()
     if not caption:
         st.error("😢 Couldn't understand this image. Try another one!")
         st.stop()
     st.success(f"**Caption:** {caption}")
 
-    # Build prompt and generate story
+    # Generate story
     prompt = f"generate story: {caption}"
     with st.spinner("📝 Writing story..."):
         start = time.time()
-        out = storyteller(prompt)
+        story = storyteller(prompt)[0]['generated_text']
         gen_time = time.time() - start
         st.text(f"⏱ Generated in {gen_time:.1f}s")
-    story = out[0].get("generated_text", "").strip()
 
-    # Enforce ≤100 words
+    # Format story output
+    story = story.replace("<pad>", "").replace("</s>", "").strip()
+    if story.startswith("generate story:"):
+        story = story[15:].strip()
+    
+    # Word limit enforcement
     words = story.split()
-    if len(words) > 100:
-        story = " ".join(words[:100]) + ("" if story.endswith('.') else ".")
+    story = " ".join(words[:100]) if len(words) > 100 else story
 
     # Display story
     st.subheader("📚 Your Magical Story")
     st.write(story)
 
-    # Convert to audio
+    # Audio conversion
     with st.spinner("🔊 Converting to audio..."):
         try:
             tts = gTTS(text=story, lang="en", slow=False)
-            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
-            tts.save(tmp.name)
-            st.audio(tmp.name, format="audio/mp3")
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
+                tts.save(tmp.name)
+                st.audio(tmp.name, format="audio/mp3")
         except Exception as e:
-            st.warning(f"⚠️ TTS failed: {e}")
+            st.warning(f"⚠️ Audio conversion failed: {str(e)}")
 
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
 
-