Update app.py

app.py

@@ -1,7 +1,6 @@
 # import part
 import streamlit as st
 from transformers import pipeline
-from gtts import gTTS
 import os
 
 # function part
@@ -20,10 +19,11 @@ def img2text(url):
 # text2story
 def text2story(text):
     try:
-        story_generator = pipeline("text-generation", model="gpt2")
+        # Use DistilGPT2 for faster text generation
+        story_generator = pipeline("text-generation", model="distilgpt2")
         # Add a fun and happy prompt to guide the story generation
         prompt = f"One sunny day, {text}. "
-        story = story_generator(prompt, max_length=150, num_return_sequences=1)[0]["generated_text"]
+        story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]
         # Make the story more fun by adding a happy ending
         happy_story = story + " And everyone had a big smile on their faces at the end of the day! 😄🌈"
         return happy_story
@@ -34,17 +34,21 @@ def text2story(text):
 # text2audio
 def text2audio(story_text):
     try:
-        tts = gTTS(text=story_text, lang='en')
-        audio_file = "story_audio.mp3"
-        tts.save(audio_file)
+        # Use a fast TTS model from Hugging Face (e.g., Facebook's FastSpeech2)
+        tts_pipeline = pipeline("text-to-speech", model="facebook/fastspeech2-en-ljspeech")
+        audio_output = tts_pipeline(story_text)
+        audio_file = "story_audio.wav"
+        # Save the audio file
+        with open(audio_file, "wb") as f:
+            f.write(audio_output["audio"])
         return audio_file
     except Exception as e:
         st.error(f"Oops! Something went wrong while turning your story into audio. Please try again! 😊")
         return None
 
 # main part
-st.set_page_config(page_title="Fun Story Maker", page_icon="😊")
-st.header("😄 Fun Story Maker: Turn Your Picture into a Happy Story! 😄")
+st.set_page_config(page_title="Story Maker", page_icon="😊")
+st.header("😄 Story Maker: Turn Your Picture into a Happy Story! 😄")
 uploaded_file = st.file_uploader("📷 Choose a picture to create a fun story...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
@@ -73,6 +77,6 @@ if uploaded_file is not None:
             if audio_file:
                 # Play button
                 if st.button("🎧 Listen to Your Story!"):
-                    st.audio(audio_file, format="audio/mp3")
+                    st.audio(audio_file, format="audio/wav")
                 # Clean up the audio file after playing
                 os.remove(audio_file)