Update app.py

app.py

@@ -13,7 +13,7 @@ def img2text(image):
 # Function to generate a story based on the caption
 def text2story(text):
     # Load the text generation model
-    story_model = pipeline("text-generation", model="gpt2")
+    story_model = pipeline("text-generation", model="distilbert/distilgpt2")
     # Generate a story based on the input text
     story_text = story_model(f"Once upon a time, {text}.", max_length=100, num_return_sequences=1)
     return story_text[0]["generated_text"]  # Return the generated story
@@ -52,4 +52,7 @@ if uploaded_file is not None:
 
     # Play button for the audio
     if st.button("Play Audio"):  # Create a button to play the audio
-        st.audio(audio_data['audio'], format="audio/wav", start_time=0, sample_rate=audio_data['sampling_rate'])  # Play the audio
+        st.audio(audio_data['audio'], 
+                 format="audio/wav", 
+                 start_time=0, 
+                 sample_rate=audio_data['sampling_rate'])  # Play the audio