Update app.py

app.py

@@ -118,39 +118,48 @@ def main():
     if "prompt" not in st.session_state:
         st.session_state["prompt"] = ""
 
-    # Input field for the prompt
-    prompt = st.text_input("Enter a prompt for image generation:", value=st.session_state["prompt"])
-
-    if st.button("Generate Image"):
+    # Input field for the prompt, tied to session state
+    prompt = st.text_input(
+        "Enter a prompt for image generation:",
+        value=st.session_state["prompt"],
+        key="prompt_input"
+    )
+
+    # Generate Image button
+    if st.button("Generate Image", key="generate_button"):
         if prompt:
+            # Save prompt to session state
+            st.session_state["prompt"] = prompt
             # Add a spinner while generating the image
             with st.spinner('Generating image...'):
                 image = generate_image(prompt)
                 if image:
                     st.session_state["image"] = image  # Store generated image in session state
-                    st.image(image, caption="Generated Image")
-                    st.session_state["prompt"] = prompt  # Store the prompt in session state
         else:
             st.warning("Please enter a prompt.")
 
-    # Show download and reset buttons only if an image is generated
+    # Show the image if it exists in the session state
     if st.session_state["image"]:
+        st.image(st.session_state["image"], caption="Generated Image")
+
         # Download button
         image_bytes = io.BytesIO()
         st.session_state["image"].save(image_bytes, format='PNG')
-        
+
         st.download_button(
             label="Download Image",
             data=image_bytes.getvalue(),
             file_name="generated_image.png",
-            mime="image/png"
+            mime="image/png",
+            key="download_button"
         )
-        
+
         # Reset button
-        if st.button("Reset"):
+        if st.button("Reset", key="reset_button"):
             # Clear session state variables
             st.session_state["image"] = None
             st.session_state["prompt"] = ""
+            st.session_state["prompt_input"] = ""
 
             # Clear UI by resetting query params
             st.query_params = {}