Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import pipeline
 from PIL import Image
 import numpy as np
 
-# Stage 1: Image to Text (Captioning)
+# Stage 1: Image to Caption
 @st.cache_resource
 def load_image_caption_model():
     return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
@@ -13,21 +13,21 @@ def generate_caption(image):
     result = caption_model(image)
     return result[0]['generated_text']
 
-# Stage 2: Text to Story (Children-friendly)
+# Stage 2: Caption to Story
 @st.cache_resource
 def load_story_generator():
-    return pipeline("text2text-generation", model="google/flan-t5-small", max_length=100)
+    return pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator", max_length=120)
 
 def text2story(description):
     story_gen = load_story_generator()
-    prompt = f"Generate a short children's story about {description}."
+    prompt = f"A short and fun children's story about {description}."
     story = story_gen(prompt)[0]['generated_text']
     return story
 
-# Stage 3: Story to Audio (Lightweight & Compatible)
+# Stage 3: Story to Speech
 @st.cache_resource
 def load_tts():
-    return pipeline("text-to-speech", model="coqui/tts-en-simply-tts")
+    return pipeline("text-to-speech", model="espnet/kan-bayashi_ljspeech_vits", framework="espnet")
 
 def story_to_audio(story_text):
     tts = load_tts()
@@ -40,7 +40,6 @@ def story_to_audio(story_text):
 def main():
     st.set_page_config(page_title="Kids Story Creator", layout="centered")
     st.title("🧒 Kids Story Creator 📖")
-    st.write("Upload a picture and let us create a short story with voice for children aged 3–10!")
 
     uploaded_image = st.file_uploader("Upload an image (jpg/jpeg/png):", type=["jpg", "jpeg", "png"])
 
@@ -48,19 +47,18 @@ def main():
         image = Image.open(uploaded_image).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        if st.button("Generate Story"):
-            with st.spinner("Step 1: Generating image description..."):
-                caption = generate_caption(image)
-                st.write(f"**Caption:** {caption}")
+        with st.spinner("Step 1: Generating description..."):
+            caption = generate_caption(image)
+        st.success(f"Caption: {caption}")
 
-            with st.spinner("Step 2: Creating children's story..."):
-                story = text2story(caption)
-                st.write("**Generated Story:**")
-                st.write(story)
+        with st.spinner("Step 2: Generating a short story..."):
+            story = text2story(caption)
+        st.success("Here's your story:")
+        st.write(story)
 
-            with st.spinner("Step 3: Generating audio..."):
-                audio, sample_rate = story_to_audio(story)
-                st.audio(audio, format="audio/wav", sample_rate=sample_rate)
+        with st.spinner("Step 3: Converting story to audio..."):
+            audio, sample_rate = story_to_audio(story)
+        st.audio(audio, format="audio/wav", sample_rate=sample_rate)
 
 if __name__ == "__main__":
     main()