Update app.py

app.py

@@ -9,6 +9,8 @@ def img2text(url):
     try:
         image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
         text = image_to_text_model(url)[0]["generated_text"]
+        # Remove unwanted words like "illustration"
+        text = text.replace("illustration", "").strip()
         # Make the caption more fun and happy
         fun_caption = f"🌟 Wow! This picture shows {text.lower()}. Let’s turn it into a fun story! 🌟"
         return fun_caption
@@ -19,11 +21,13 @@ def img2text(url):
 # text2story
 def text2story(text):
     try:
-        # Use DistilGPT2 for faster text generation
-        story_generator = pipeline("text-generation", model="distilgpt2")
+        # Use a better model for text generation (e.g., GPT-Neo)
+        story_generator = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")
         # Add a fun and happy prompt to guide the story generation
         prompt = f"One sunny day, {text}. "
         story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]
+        # Remove any unwanted text (e.g., usernames, special characters)
+        story = " ".join([word for word in story.split() if not word.startswith("@") and not word.startswith("http")])
         # 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,8 +38,8 @@ def text2story(text):
 # text2audio
 def text2audio(story_text):
     try:
-        # Use a fast TTS model from Hugging Face (e.g., Facebook's FastSpeech2)
-        tts_pipeline = pipeline("text-to-speech", model="facebook/fastspeech2-en-ljspeech")
+        # Use a reliable TTS model (e.g., ESPnet's VITS model)
+        tts_pipeline = pipeline("text-to-speech", model="espnet/kan-bayashi_ljspeech_vits")
         audio_output = tts_pipeline(story_text)
         audio_file = "story_audio.wav"
         # Save the audio file