Update app.py

app.py

@@ -1,59 +1,84 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
-import time
 
+# Cache model loading for performance
+@st.cache_resource
+def load_image_to_text():
+    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+
+@st.cache_resource
+def load_story_generator():
+    return pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
+
+@st.cache_resource
+def load_tts():
+    return pipeline("text-to-speech", model="facebook/mms-tts-eng")
+
+# Generation functions with error handling
 def generate_caption(image):
-    image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    caption = image_to_text(image)[0]["generated_text"]
-    return caption
+    image_to_text = load_image_to_text()
+    try:
+        caption = image_to_text(image)[0]["generated_text"]
+        return caption
+    except Exception as e:
+        st.error(f"Oops! Something went wrong while generating the caption: {e}")
+        return None
 
 def generate_story(caption):
-    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
-    story = pipe(caption)[0]['generated_text']
-    return story
+    story_generator = load_story_generator()
+    try:
+        prompt = f"Once upon a time, there was {caption}. "
+        story = story_generator(prompt, max_length=200, do_sample=True, temperature=0.7)[0]['generated_text']
+        return story
+    except Exception as e:
+        st.error(f"Oops! Something went wrong while generating the story: {e}")
+        return None
 
 def generate_audio(story):
-    pipe = pipeline("text-to-speech", model="facebook/mms-tts-eng")
-    audio = pipe(story)
-    return audio
+    tts = load_tts()
+    try:
+        audio = tts(story)
+        return audio
+    except Exception as e:
+        st.error(f"Oops! Something went wrong while generating the audio: {e}")
+        return None
 
 # Streamlit UI
-
-# Title of the Streamlit app
-st.title("Upload your image for an instant storytelling!")
-
-# Write a description
-st.write("This app is designed for 3-10 year-old kids by allowing them uploading image for fun storytelling entertainment.")
+st.markdown("<h1 style='text-align: center; color: blue;'>📸✨ Storyteller for Kids! ✨📸</h1>", unsafe_allow_html=True)
+st.markdown("<p style='text-align: center;'>Upload a fun picture and I’ll tell you a magical story about it!</p>", unsafe_allow_html=True)
 
 # File uploader for image
-uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])
 
 if uploaded_file is not None:
     # Display the uploaded image
     image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_container_width=True)
-    
-    # Generate Image Caption
-    image_caption = generate_caption(image)
-
-    # Display results
-    st.subheader("Image Caption:")
-    st.write(f"{image_caption}")
-
-    # Generate Story
-    story_telling = generate_story(image_caption)
-
-    # Display results
-    st.subheader("Story:")
-    st.write(f"{story_telling}")
-
-    # Generate Audio
-    audio = generate_audio(story_telling)
-
-    # Display an audio file with a spinner effect
-    if st.button("Play Audio"):
-        st.audio(audio['audio'], 
-                format="audio/wav", 
-                start_time=0, 
-                sample_rate = audio['sampling_rate'])
+    st.image(image, caption="Your Uploaded Image", use_column_width=True)
+
+    # Generate and display caption with spinner
+    with st.spinner("Creating a caption..."):
+        generated_caption = generate_caption(image)
+
+    if generated_caption:
+        # Allow user to edit the caption
+        caption_input = st.text_area("Caption:", value=generated_caption, height=100)
+        st.write("Feel free to change the caption to make your own story!")
+
+        # Generate story when button is clicked
+        if st.button("Generate Story"):
+            with st.spinner("Writing a magical story..."):
+                story = generate_story(caption_input)
+
+            if story:
+                st.subheader("Your Story:")
+                st.write(story)
+                st.download_button("Download Story", story, file_name="my_story.txt")
+
+                # Generate audio on demand
+                if st.button("Generate Audio"):
+                    with st.spinner("Turning your story into sound..."):
+                        audio = generate_audio(story)
+                    if audio:
+                        st.audio(audio['audio'], format="audio/wav", start_time=0, sample_rate=audio['sampling_rate'])
+                        st.download_button("Download Audio", audio['audio'], file_name="my_story.wav", mime="audio/wav")