Update app.py

app.py

@@ -96,15 +96,6 @@ with st.sidebar:
 #             st.warning("Please enter a prompt.")
 
 # COMPREHENSIVE CODE
-import streamlit as st
-import requests
-from PIL import Image
-import io
-
-#------------------------------------------------------------------------
-# Define functions
-#------------------------------------------------------------------------
-
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
     if response.status_code != 200:
@@ -128,34 +119,39 @@ def main():
         st.session_state["prompt"] = ""
 
     # Input field for the prompt
-    prompt = st.text_input("Enter a prompt for image generation:", value=st.session_state["prompt"])
+    prompt = st.text_input("Enter a prompt for image generation:", value="")
 
     if st.button("Generate Image"):
         if prompt:
             image = generate_image(prompt)
             if image:
                 st.session_state["image"] = image  # Store generated image in session state
-                st.session_state["prompt"] = prompt  # Store the prompt 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 button if an image is generated
+    # Show download and reset buttons only if an image is generated
     if st.session_state["image"]:
         image_bytes = io.BytesIO()
         st.session_state["image"].save(image_bytes, format='PNG')
+        
+        # Download button
         st.download_button(
             label="Download Image",
             data=image_bytes.getvalue(),
             file_name="generated_image.png",
             mime="image/png"
         )
-
-    # Add reset button to clear the prompt and image
-    if st.button("Reset"):
-        st.session_state["image"] = None
-        st.session_state["prompt"] = ""
-        st.experimental_rerun()  # Reset the app by rerunning it
+        
+        # Reset button
+        if st.button("Reset"):
+            # Clear session state variables
+            st.session_state["image"] = None
+            st.session_state["prompt"] = ""
+            
+            # Clear UI by updating query params (this will force rerun of the app)
+            st.experimental_set_query_params()
 
 #------------------------------------------------------------------------
 # Main Guard