Update app.py

app.py

@@ -1,27 +1,40 @@
 # Import necessary libraries
 import streamlit as st
 from transformers import pipeline
-import pyttsx3
+from gtts import gTTS
 import os
 
 # Function to convert image to text
 def img2text(url):
     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
+    
+    # Remove unwanted words like "illustration"
+    unwanted_words = ["illustration", "painting", "drawing", "artwork"]
+    for word in unwanted_words:
+        text = text.replace(word, "")
+    return text.strip()
 
 # Function to generate a story from text using GPT-2
 def text2story(text):
+    # Add a strong prompt to guide the model
+    prompt = f"Write a short, happy, and fun story for kids aged 3-10 based on the following description: {text}. " \
+             "The story should be cheerful, imaginative, and suitable for young children. " \
+             "Avoid any scary or sad elements. Keep the story under 95 words and make sure it has a clear beginning, middle, and end."
+    
+    # Use the GPT-2 model to generate the story
     text_generator = pipeline("text-generation", model="gpt2")
-    story = text_generator(text, max_length=95, num_return_sequences=1)[0]["generated_text"]
+    story = text_generator(prompt, max_length=95, num_return_sequences=1)[0]["generated_text"]
+    
+    # Clean up the output to remove the prompt (if necessary)
+    story = story.replace(prompt, "").strip()
     return story
 
-# Function to convert text to audio using pyttsx3
+# Function to convert text to audio using gTTS
 def text2audio(story_text):
-    engine = pyttsx3.init()
-    audio_file = "story_audio.wav"
-    engine.save_to_file(story_text, audio_file)
-    engine.runAndWait()
+    audio_file = "story_audio.mp3"
+    tts = gTTS(story_text, lang="en")
+    tts.save(audio_file)
     return audio_file
 
 # Main application
@@ -37,7 +50,7 @@ if uploaded_file is not None:
     with open(uploaded_file.name, "wb") as file:
         file.write(bytes_data)
 
-    st.image(uploaded_file, caption="Your Uploaded Image", use_column_width=True)
+    st.image(uploaded_file, caption="Your Uploaded Image", use_container_width=True)
 
     # Stage 1: Image to Text
     st.text('🖼️ Processing image...')
@@ -55,7 +68,7 @@ if uploaded_file is not None:
 
     # Play button for audio
     if st.button("🎧 Listen to the Story"):
-        st.audio(audio_file, format="audio/wav")
+        st.audio(audio_file, format="audio/mp3")
 
     # Clean up the generated audio file
     os.remove(audio_file)