Update app.py

app.py

@@ -4,58 +4,32 @@ from transformers import pipeline
 from gtts import gTTS
 import os
 
-# Function to convert image to text using Hugging Face's BLIP model
+# Function to convert image to text
 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 generate a story from the text using T5-small model
+# Function to convert text to a kid-friendly story
 def text2story(text):
-    """
-    Expands the text caption into a full story using the T5-small model.
-    
-    Args:
-        text (str): Text caption generated from the image.
-    
-    Returns:
-        str: Generated story suitable for kids aged 3-10.
-    """
-    # Simplified and rephrased prompt to guide the model
-    prompt = f"Write a short, happy story for kids about {text}. The story should be cheerful, imaginative, and under 95 words."
+    # Initialize the text generation pipeline
+    text_generator = pipeline("text-generation", model="distilbert/distilgpt2")
     
-    text_generator = pipeline("text2text-generation", model="t5-small")
-    story = text_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
+    # Create 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}. The story should be suitable for children and should not contain any magical elements. Keep the story under 95 words."
     
-    # Remove any repetition of the prompt in the output
-    if prompt in story:
-        story = story.replace(prompt, "").strip()
+    # Generate the story
+    story = text_generator(prompt, max_length=95, num_return_sequences=1)[0]['generated_text']
     
     return story
 
-# Function to convert text to audio using gTTS
+# Function to convert text to audio
 def text2audio(story_text):
-    """
-    Converts the generated story text into audio using gTTS.
-    
-    Args:
-        story_text (str): The generated story text.
-    
-    Returns:
-        str: Path to the generated audio file.
-    """
+    # Convert the story text to audio using gTTS
     tts = gTTS(text=story_text, lang='en')
     audio_file = "story_audio.mp3"
     tts.save(audio_file)
+    
     return audio_file
 
 # Main application
@@ -76,10 +50,10 @@ if uploaded_file is not None:
     # Stage 1: Image to Text
     st.text('Processing image to text...')
     scenario = img2text(uploaded_file.name)
-    st.write("**Caption:**", scenario)
+    st.write("**Scenario:**", scenario)
 
     # Stage 2: Text to Story
-    st.text('Generating a fun story...')
+    st.text('Generating a fun story for kids...')
     story = text2story(scenario)
     st.write("**Story:**", story)