Update app.py

app.py

@@ -109,37 +109,39 @@ def generate_image(prompt):
         return Image.open(io.BytesIO(image_bytes))
     return None
 
-def main():
-    st.title("Stxtement | Image Generation")
+def generate_image_callback():
+    prompt = st.session_state.get("prompt_input", "")
+    if prompt:
+        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
+    else:
+        st.warning("Please enter a prompt.")
+
+def reset_callback():
+    # Clear session state variables
+    st.session_state["prompt"] = ""
+    st.session_state["prompt_input"] = ""
+    st.session_state["image"] = None
 
-    # Initialize session state variables for prompt and image
-    if "image" not in st.session_state:
-        st.session_state["image"] = None
-    if "prompt" not in st.session_state:
-        st.session_state["prompt"] = ""
+def main():
+    st.title("MiImage Creator | Image Generation")
 
     # Input field for the prompt, tied to session state
-    prompt = st.text_input(
+    st.text_input(
         "Enter a prompt for image generation:",
-        value=st.session_state["prompt"],
+        value=st.session_state.get("prompt_input", ""),
         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
-        else:
-            st.warning("Please enter a prompt.")
+    # Generate Image button with callback
+    st.button("Generate Image", on_click=generate_image_callback)
 
     # Show the image if it exists in the session state
-    if st.session_state["image"]:
+    if st.session_state.get("image"):
         st.image(st.session_state["image"], caption="Generated Image")
 
         # Download button
@@ -150,19 +152,15 @@ def main():
             label="Download Image",
             data=image_bytes.getvalue(),
             file_name="generated_image.png",
-            mime="image/png",
-            key="download_button"
+            mime="image/png"
         )
 
         # Reset button
-        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 = {}
+        st.button("Reset", on_click=reset_callback)
+
+    # If no image but prompt exists, show reset button to clear prompt
+    elif st.session_state.get("prompt_input"):
+        st.button("Reset", on_click=reset_callback)
 
 #------------------------------------------------------------------------
 # Main Guard