Update app.py

app.py

@@ -25,24 +25,25 @@ def img2text(url):
     
     return text.strip()
 
-# Function to generate a kid-friendly story from the text caption
+# Function to generate a kid-friendly superhero story from the text caption
 def text2story(text):
     """
-    Generates a kid-friendly story from the text caption using the pranavpsv/gpt2-genre-story-generator model.
+    Generates a kid-friendly superhero story from the text caption using the pranavpsv/gpt2-genre-story-generator model.
     
     Args:
         text (str): Text caption generated from the image.
     
     Returns:
-        str: Generated story suitable for kids aged 3-10, within 100 words.
+        str: Generated superhero story suitable for kids aged 3-10, within 100 words.
     """
-
     # Load the text generation model
     story_generator = pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator")
     
-    # Generate the story
-    story = story_generator(f"<BOS> <superhero> {text}"))
+    # Generate the story with the superhero genre
+    prompt = f"<BOS> <superhero> {text}"
+    story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
     
+    return story
 
 # Function to convert text to audio using gTTS
 def text2audio(story_text):
@@ -65,15 +66,15 @@ def text2audio(story_text):
 # Main application
 st.set_page_config(page_title="Picture Stories 🎨📖", page_icon="🦄")
 st.title("Picture Stories 🎨📖")
-st.markdown("### Turn your pictures into fun stories and listen to them! 🎉")
+st.markdown("### Turn your pictures into fun superhero stories and listen to them! 🎉")
 
 # Instructions for kids
 st.markdown("""
 **How to use this app:**
 1. **Upload a picture** of something fun, like your favorite toy, a park, or your pet.
-2. Wait for the app to **create a story** from your picture.
+2. Wait for the app to **create a superhero story** from your picture.
 3. **Listen to the story** by clicking the "Play Audio" button.
-4. Enjoy your fun story! 🎧
+4. Enjoy your fun superhero story! 🎧
 """)
 
 # Upload image
@@ -94,9 +95,9 @@ if uploaded_file is not None:
     st.write("**What we see:**", scenario)
 
     # Stage 2: Text to Story
-    st.text('📖 Creating a fun story for you...')
+    st.text('📖 Creating a fun superhero story for you...')
     story = text2story(scenario)
-    st.write("**Your story:**", story)
+    st.write("**Your superhero story:**", story)
 
     # Stage 3: Story to Audio
     st.text('🎧 Turning your story into audio...')
@@ -106,5 +107,6 @@ if uploaded_file is not None:
     if st.button("🎵 **Play Audio**"):
         st.audio(audio_file, format="audio/mp3")
 
-    # Clean up the generated audio file
-    os.remove(audio_file)
+    # Clean up the generated audio file and uploaded image
+    os.remove(audio_file)
+    os.remove(uploaded_file.name)