Update app.py

app.py

@@ -4,25 +4,55 @@ from transformers import pipeline
 from gtts import gTTS
 import os
 
-# Function to convert image to text
+# Function to convert image to text using Hugging Face's BLIP model
 def img2text(url):
+    """
+    Converts an image to text using the Salesforce/blip-image-captioning-base model.
+    
+    Args:
+        url (str): Path to the image file.
+    
+    Returns:
+        str: Generated text caption from the image.
+    """
     image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     text = image_to_text_model(url)[0]["generated_text"]
     return text
 
-# Function to convert text to a story
+# Function to generate a kid-friendly story from the text caption
 def text2story(text):
-    # Initialize the text generation pipeline
-    text_generator = pipeline("text-generation", model="distilbert/distilgpt2")
+    """
+    Generates a kid-friendly story from the text caption using the aspis/gpt2-genre-story-generation model.
+    
+    Args:
+        text (str): Text caption generated from the image.
+    
+    Returns:
+        str: Generated story suitable for kids aged 3-10.
+    """
+    # Add a prompt to ensure the story is kid-friendly and happy
+    prompt = f"Write a happy and fun story for kids aged 3-10 based on the following scenario: {text}."
     
-    # Generate the story directly from the input text
-    story = text_generator(text, max_length=95, num_return_sequences=1)[0]['generated_text']
+    # Load the text generation model
+    story_generator = pipeline("text-generation", model="aspis/gpt2-genre-story-generation")
+    
+    # Generate the story
+    story = story_generator(prompt, max_length=95, num_return_sequences=1)[0]["generated_text"]
     
     return story
 
-# Function to convert text to audio
+# Function to convert text to audio using gTTS
 def text2audio(story_text):
-    # Convert the story text to audio using gTTS
+    """
+    Converts the generated story text to audio using gTTS.
+    
+    Args:
+        story_text (str): The generated story text.
+    
+    Returns:
+        str: Path to the generated audio file.
+    """
+    # Convert text to speech
     tts = gTTS(text=story_text, lang='en')
     audio_file = "story_audio.mp3"
     tts.save(audio_file)
@@ -32,7 +62,9 @@ def text2audio(story_text):
 # Main application
 st.set_page_config(page_title="Your Image to Audio Story",
                    page_icon="🦜")
-st.header("Turn Your Image into a Fun Audio Story!")
+st.header("Turn Your Image into a Fun Audio Story! 🎉")
+
+# Upload image
 uploaded_file = st.file_uploader("Select an Image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
@@ -47,7 +79,7 @@ if uploaded_file is not None:
     # Stage 1: Image to Text
     st.text('Processing image to text...')
     scenario = img2text(uploaded_file.name)
-    st.write("**Scenario:**", scenario)
+    st.write("**Caption:**", scenario)
 
     # Stage 2: Text to Story
     st.text('Generating a fun story for kids...')