Update app.py

app.py

@@ -11,8 +11,20 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
-    story_text = pipe(text)[0]['generated_text']
+    pipe = pipeline("text-generation", 
+                    model="pranavpsv/genre-story-generator-v2",
+                    max_new_tokens=100,  
+                    truncation=True)      
+    story_text = pipe(text, 
+                      do_sample=True,     
+                      temperature=0.9,    
+                      top_k=50,           
+                      num_return_sequences=1)[0]['generated_text']
+    
+    # Make sure the story is complete
+    last_punctuation = max(story_text.rfind("."), story_text.rfind("!"), story_text.rfind("?"))
+    if last_punctuation != -1:
+        story_text = story_text[:last_punctuation+1]
     return story_text
 
 # text2audio
@@ -36,21 +48,21 @@ def main():
 
 
         #Stage 1: Image to Text
-        st.text('Processing img2text...')
+        st.text('🔍 Analyzing image...')
         scenario = img2text(uploaded_file.name)
         st.write(scenario)
 
         #Stage 2: Text to Story
-        st.text('Generating a story...')
+        st.text('📖 Generating story...')
         story = text2story(scenario)
         st.write(story)
 
         #Stage 3: Story to Audio data
-        st.text('Generating audio data...')
+        st.text('🎧 Converting to audio...')
         audio_data =text2audio(story)
 
         # Play button
-        if st.button("Play Audio"):
+        if st.button("Play Audio Story"):
             st.audio(audio_data['audio'], 
                      format="audio/wav", 
                      start_time=0,