Update app.py

app.py

@@ -15,15 +15,17 @@ def img2text(url):
         text = text.replace(word, "")
     return text.strip()
 
-# Function to generate a story from text using GPT-2
+# Function to generate a story from text using GPT-Neo
 def text2story(text):
+    # Use the GPT-Neo model for storytelling
+    text_generator = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")
+    
     # Add a strong prompt to guide the model
     prompt = f"Write a short, happy, and fun story for kids aged 3-10 based on the following description: {text}. " \
              "The story should be cheerful, imaginative, and suitable for young children. " \
              "Avoid any scary or sad elements. Keep the story under 95 words and make sure it has a clear beginning, middle, and end."
     
-    # Use the GPT-2 model to generate the story
-    text_generator = pipeline("text-generation", model="gpt2")
+    # Generate the story
     story = text_generator(prompt, max_length=95, num_return_sequences=1)[0]["generated_text"]
     
     # Clean up the output to remove the prompt (if necessary)
@@ -42,6 +44,12 @@ st.set_page_config(page_title="Image to Story", page_icon="📖")
 st.header("📖 Image to Story")
 st.markdown("### Turn your image into a fun story!")
 
+# Initialize session state
+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
+
 uploaded_file = st.file_uploader("Select an Image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
@@ -58,20 +66,22 @@ if uploaded_file is not None:
     st.write("**What I see:**", scenario)
 
     # Stage 2: Text to Story
-    st.text('📝 Creating a story...')
-    story = text2story(scenario)
-    st.write("**Your Story:**", story)
+    if st.session_state.story is None or st.button("🔄 Generate New Story"):
+        st.text('📝 Creating a story...')
+        st.session_state.story = text2story(scenario)
+        st.write("**Your Story:**", st.session_state.story)
 
-    # Stage 3: Story to Audio
-    st.text('🎙️ Turning your story into audio...')
-    audio_file = text2audio(story)
+        # Stage 3: Story to Audio
+        st.text('🎙️ Turning your story into audio...')
+        st.session_state.audio_file = text2audio(st.session_state.story)
 
     # Play button for audio
-    if st.button("🎧 Listen to the Story"):
-        st.audio(audio_file, format="audio/mp3")
+    if st.session_state.audio_file and st.button("🎧 Listen to the Story"):
+        st.audio(st.session_state.audio_file, format="audio/mp3")
 
     # Clean up the generated audio file
-    os.remove(audio_file)
+    if st.session_state.audio_file and os.path.exists(st.session_state.audio_file):
+        os.remove(st.session_state.audio_file)
 
 # Add some fun prompts for kids
 st.markdown("### 🎨 Tips for a Great Story!")