Update app.py

app.py

@@ -1,38 +1,64 @@
 # Import necessary libraries
-import streamlit as st  # Streamlit for web application
-from transformers import pipeline  # Hugging Face transformer pipeline
-from PIL import Image  # Python Imaging Library for image handling
+import streamlit as st  # Streamlit for creating the web application
+from transformers import pipeline  # Pipeline for using Hugging Face models
+from PIL import Image  # PIL for image processing
 
-# Set the title of the Streamlit app
-st.set_page_config(page_title="Storytelling Friend",
-                   page_icon="haha")
- # Title of the application
+# Function to load models
+def load_models():
+    # Load the image to text model
+    caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")  # Load pre-trained image to text model
+    # Load the text generation model
+    story_model = pipeline("text-generation", model="gpt2")  # Load pre-trained text generation model
+    # Load the text-to-speech model
+    tts_model = pipeline("text-to-speech", model="facebook/tts_transformer-es-css10")  # Load a TTS model
+    return caption_model, story_model, tts_model  # Return all three models
 
-# Create a file uploader for the image
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])  # User uploads an image
+# Function to generate story from caption
+def generate_story(caption):
+    # Generate a story based on the caption
+    story = story_model(caption, max_length=100, num_return_sequences=1)[0]['generated_text']  # Generate the story
+    return story  # Return the generated story
 
-# Load the image captioning model
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")  # Load pre-trained model
+# Function to convert text to audio
+def text_to_audio(text, tts_model):
+    audio = tts_model(text)  # Generate audio from text using the TTS model
+    return audio  # Return the audio object
 
-# Load the text generation model
-story_model = pipeline("text-generation", model="gpt2")  # Load a text generation model
+# Function to process the uploaded image and generate a story
+def process_image(image, caption_model):
+    # Generate a caption from the uploaded image
+    caption = caption_model(image)[0]['caption']  # Get the caption from the model
+    # Generate a story from the caption
+    story = generate_story(caption)  # Call the story generation function
+    return caption, story  # Return both caption and story
 
-# Function to generate a story from the caption
-def generate_story(caption):  # Define a function to generate a story
-    story_input = f"Once upon a time, {caption}"  # Create a story prompt
-    story = story_model(story_input, max_length=150, num_return_sequences=1)[0]['generated_text']  # Generate the story
-    return story  # Return the generated story
+# Main part
+def main():
+    st.set_page_config(page_title="Storytelling Friend", page_icon="🦦") # Title of the application
+    st.write("Upload an image to generate a story!")  # Instructions for the user
 
-# Process the uploaded image and generate story
-if uploaded_file is not None:  # Check if a file is uploaded
-    image = Image.open(uploaded_file)  # Open the uploaded image
-    st.image(image, caption="Uploaded Image", use_column_width=True)  # Display the uploaded image
+    # Upload image section
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])  # File uploader for images
 
-    caption = caption_model(image)[0]['caption']  # Generate caption from the image
-    st.subheader("Image Caption:")  # Subtitle for the caption
-    st.write(caption)  # Display the caption
+    # Load models once
+    caption_model, story_model, tts_model = load_models()  # Load models
 
-    # Generate story based on the caption
-    story = generate_story(caption)  # Call the story generation function
-    st.subheader("Generated Story:")  # Subtitle for the generated story
-    st.write(story)  # Display the generated story
+    if uploaded_file is not None:
+        # Open and read the uploaded image
+        image = Image.open(uploaded_file)  # Open the uploaded image file
+        st.image(image, caption="Uploaded Image", use_column_width=True)  # Display the uploaded image
+
+        # Process the image and generate story
+        caption, story = process_image(image, caption_model)  # Get caption and story
+        st.subheader("Generated Caption:")  # Subheader for caption
+        st.write(caption)  # Display the caption
+        st.subheader("Generated Story:")  # Subheader for story
+        st.write(story)  # Display the generated story
+
+        # Convert story to audio and play it
+        audio = text_to_audio(story, tts_model)  # Convert story to audio
+        st.audio(audio, format='audio/wav')  # Play the audio
+
+# Run the app
+if __name__ == "__main__":
+    main()  # Call the main function to run the app