Update app.py

app.py

@@ -35,7 +35,7 @@ def img2text(url):
 # Function to generate a kid-friendly superhero story from the text caption
 def text2story(text):
     """
-    Generates a kid-friendly superhero story from the text caption using the pranavpsv/gpt2-genre-story-generator model.
+    Generates a superhero story from the text caption using the pranavpsv/gpt2-genre-story-generator model.
     
     Args:
         text (str): Text caption generated from the image.
@@ -49,7 +49,7 @@ def text2story(text):
         
         # Generate the story with the superhero genre
         prompt = f"<BOS> <superhero> {text}"  # Add genre tags to the prompt
-        story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
+        story = story_generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']
         
         # Remove <BOS> and <superhero> tags from the generated story
         story = story.replace("<BOS>", "").replace("<superhero>", "").strip()
@@ -61,6 +61,14 @@ def text2story(text):
         # Ensure the story is within 100 words by truncating if necessary
         story = " ".join(story.split()[:100])
         
+        # If the story is too short, regenerate it
+        if len(story.split()) < 50:  # Minimum 50 words
+            story = story_generator(prompt, max_length=200, num_return_sequences=1)[0]['generated_text']
+            story = story.replace("<BOS>", "").replace("<superhero>", "").strip()
+            if text in story:
+                story = story.replace(text, "").strip()
+            story = " ".join(story.split()[:100])
+        
         return story
     except Exception as e:
         st.error(f"Error generating story: {e}")  # Display error message if something goes wrong
@@ -109,6 +117,14 @@ def main():
     # Upload image
     uploaded_file = st.file_uploader("📷 **Upload your picture here!**", type=["jpg", "jpeg", "png"])
 
+    # Initialize session state variables
+    if 'scenario' not in st.session_state:
+        st.session_state.scenario = None
+    if 'story' not in st.session_state:
+        st.session_state.story = None
+    if 'audio_file' not in st.session_state:
+        st.session_state.audio_file = None
+
     if uploaded_file is not None:
         # Save the uploaded file to disk
         image_bytes = uploaded_file.getvalue()
@@ -120,22 +136,26 @@ def main():
 
         # Stage 1: Image to Text
         with st.spinner('✨ Turning your picture into words...'):
-            scenario = img2text(uploaded_file.name)  # Generate text caption from the image
-            if scenario:
-                st.write("**What we see:**", scenario)  # Display the generated caption
+            if st.session_state.scenario is None or uploaded_file.name != st.session_state.get('uploaded_file_name', None):
+                st.session_state.scenario = img2text(uploaded_file.name)  # Generate text caption from the image
+                st.session_state.uploaded_file_name = uploaded_file.name  # Store the uploaded file name
+
+            if st.session_state.scenario:
+                st.write("**What we see:**", st.session_state.scenario)  # Display the generated caption
 
                 # Stage 2: Text to Story
                 with st.spinner('📖 Creating a fun superhero story for you...'):
-                    story = text2story(scenario)  # Generate a superhero story from the caption
-                    if story:
-                        st.write("**Your superhero story:**", story)  # Display the generated story
+                    if st.session_state.story is None or uploaded_file.name != st.session_state.uploaded_file_name:
+                        st.session_state.story = text2story(st.session_state.scenario)  # Generate a superhero story from the caption
+
+                    if st.session_state.story:
+                        st.write("**Your superhero story:**", st.session_state.story)  # Display the generated story
 
                         # Stage 3: Story to Audio
                         with st.spinner('🎧 Turning your story into audio...'):
-                            # Generate audio file if it doesn't already exist in the session state
-                            if 'audio_file' not in st.session_state:
-                                st.session_state.audio_file = text2audio(story)
-                            
+                            if st.session_state.audio_file is None or uploaded_file.name != st.session_state.uploaded_file_name:
+                                st.session_state.audio_file = text2audio(st.session_state.story)  # Generate audio file
+
                             # Play button for the generated audio
                             if st.button("🎵 **Play Audio**"):
                                 if os.path.exists(st.session_state.audio_file):
@@ -143,11 +163,9 @@ def main():
                                 else:
                                     st.error("Audio file not found. Please try again.")  # Display error if audio file is missing
 
-        # Clean up temporary files (uploaded image and generated audio)
+        # Clean up temporary files (uploaded image)
         if os.path.exists(uploaded_file.name):
             os.remove(uploaded_file.name)  # Delete the uploaded image file
-        if 'audio_file' in st.session_state and os.path.exists(st.session_state.audio_file):
-            os.remove(st.session_state.audio_file)  # Delete the generated audio file
 
 # Run the application
 if __name__ == "__main__":