Update app.py

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
+import io
 
-# Stage 1: Image to Caption
+# Stage 1: Image to Text (Captioning)
 @st.cache_resource
 def load_image_caption_model():
     return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
@@ -12,21 +13,21 @@ def generate_caption(image):
     result = caption_model(image)
     return result[0]['generated_text']
 
-# Stage 2: Caption to Story
+# Stage 2: Text to Story (Children-friendly)
 @st.cache_resource
 def load_story_generator():
-    return pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator", max_length=120)
+    return pipeline("text2text-generation", model="google/flan-t5-base", max_length=100)
 
 def text2story(description):
     story_gen = load_story_generator()
-    prompt = f"A short and fun children's story about {description}."
+    prompt = f"Generate a short and imaginative children's story about: {description}"
     story = story_gen(prompt)[0]['generated_text']
     return story
 
-# Stage 3: Story to Speech
+# Stage 3: Story to Speech (Lightweight & Compatible)
 @st.cache_resource
 def load_tts():
-    return pipeline("text-to-speech", model="facebook/fastspeech2-en-ljspeech")
+    return pipeline("text-to-speech", model="suno/bark-small")
 
 def story_to_audio(story_text):
     tts = load_tts()
@@ -46,16 +47,16 @@ def main():
         image = Image.open(uploaded_image).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        with st.spinner("Step 1: Generating description..."):
+        with st.spinner("Generating description..."):
             caption = generate_caption(image)
         st.success(f"Caption: {caption}")
 
-        with st.spinner("Step 2: Generating a short story..."):
+        with st.spinner("Generating story from caption..."):
             story = text2story(caption)
         st.success("Here's your story:")
         st.write(story)
 
-        with st.spinner("Step 3: Converting story to audio..."):
+        with st.spinner("Converting story to audio..."):
             audio, sample_rate = story_to_audio(story)
         st.audio(audio, format="audio/wav", sample_rate=sample_rate)