Update app.py

app.py

@@ -143,10 +143,6 @@ def main():
 
     with col2:
         if uploaded_file is not None:
-            # Reset session state variables when a new image is uploaded
-            st.session_state.story = ""
-            st.session_state.audio_file_path = ""
-
             # Display uploaded image
             st.markdown("## 🖼️ Your Image")
             bytes_data = uploaded_file.read()
@@ -154,53 +150,42 @@ def main():
                 file.write(bytes_data)
             st.image(uploaded_file, use_column_width=True)
 
-            # Generate image caption as soon as image is uploaded
-            with st.spinner("📝 Generating image caption..."):
-                try:
-                    scenario = img2txt("uploaded_image.jpg")
-                    st.session_state.caption = scenario  # Store caption in session state
-                    st.success("Image caption generated.")
-                except Exception as e:
-                    st.error(f"An error occurred while generating image caption: {str(e)}")
-                    st.warning("Please try again or contact support if the problem persists.")
-
-            # Display image caption
-            if st.session_state.caption:
-                st.markdown("## 📜 Image Caption")
-                st.write(st.session_state.caption)
-
-            # Generate Story button
-            if st.session_state.caption:
-                if st.button("🎨 Generate Story"):
-                    with st.spinner("📝 Generating story..."):
-                        try:
-                            # Create story prompt using st.session_state.caption
-                            prompt = f"""Based on the image description: '{st.session_state.caption}', 
-                            create a {preferences['genre']} story set in {preferences['setting']} 
-                            in {preferences['continent']}. The story should have a {preferences['tone']} 
-                            tone and explore the theme of {preferences['theme']}. The main conflict 
-                            should be {preferences['conflict']}. The story should have a {preferences['twist']} 
-                            and end with a {preferences['ending']} ending."""
-                            
-                            # Generate story
-                            story = txt2story(prompt, top_k=5, top_p=0.8, temperature=1.5)
-                            st.session_state.story = story  # Store story in session state
-                            st.success("Story generated.")
-
-                            # Convert to audio
-                            with st.spinner("🔊 Generating audio story..."):
-                                txt2speech(story)
-                                st.session_state.audio_file_path = "audio_story.mp3"  # Store audio path in session state
-                                st.success("Audio story generated.")
-
-                        except Exception as e:
-                            st.error(f"An error occurred: {str(e)}")
-                            st.warning("Please try again or contact support if the problem persists.")
+            # Process image and generate story
+            if st.button("🎨 Generate Story"):
+                with st.spinner("🤖 AI is working its magic..."):
+                    try:
+                        # Get image description
+                        scenario = img2txt("uploaded_image.jpg")
+                        st.session_state.caption = scenario  # Store caption in session state
+                        
+                        # Create story prompt
+                        prompt = f"""Based on the image description: '{scenario}', 
+                        create a {preferences['genre']} story set in {preferences['setting']} 
+                        in {preferences['continent']}. The story should have a {preferences['tone']} 
+                        tone and explore the theme of {preferences['theme']}. The main conflict 
+                        should be {preferences['conflict']}. The story should have a {preferences['twist']} 
+                        and end with a {preferences['ending']} ending."""
+                        
+                        # Generate story
+                        story = txt2story(prompt, top_k=5, top_p=0.8, temperature=1.5)
+                        st.session_state.story = story  # Store story in session state
+                        
+                        # Convert to audio
+                        txt2speech(story)
+                        st.session_state.audio_file_path = "audio_story.mp3"  # Store audio path in session state
+
+                    except Exception as e:
+                        st.error(f"An error occurred: {str(e)}")
+                        st.warning("Please try again or contact support if the problem persists.")
 
         # Display results if story exists in session state
         if st.session_state.story:
             st.markdown("---")
             
+            # Image caption
+            with st.expander("📜 Image Caption", expanded=True):
+                st.write(st.session_state.caption)
+            
             # Story text
             with st.expander("📖 Generated Story", expanded=True):
                 st.write(st.session_state.story)
@@ -211,7 +196,7 @@ def main():
 
             # Email section
             st.markdown("---")
-            st.markdown("## 📧 Get Story via Email")
+            st.markdown("## 📧 Get Story via Email - RUN LOCALLY")
             email = st.text_input(
                 "Enter your email address:",
                 help="We'll send you the story text and audio file"