Update app.py

app.py

@@ -68,6 +68,43 @@ with st.sidebar:
 # Define functions
 #------------------------------------------------------------------------
 
+# SIMPLE CODE
+# def query(payload):
+#     response = requests.post(API_URL, headers=headers, json=payload)
+#     if response.status_code != 200:
+#         st.error(f"Error: {response.status_code} - {response.text}")
+#         return None
+#     return response.content
+
+# def generate_image(prompt):
+#     image_bytes = query({"inputs": prompt})
+#     if image_bytes:
+#         return Image.open(io.BytesIO(image_bytes))
+#     return None
+
+# def main():
+#     st.title("Stxtement | Image Generation")
+
+#     prompt = st.text_input("Enter a prompt for image generation:")
+    
+#     if st.button("Generate Image"):
+#         if prompt:
+#             image = generate_image(prompt)
+#             if image:
+#                 st.image(image, caption="Generated Image")
+#         else:
+#             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:
@@ -84,16 +121,42 @@ def generate_image(prompt):
 def main():
     st.title("Stxtement | Image Generation")
 
-    prompt = st.text_input("Enter a prompt for image generation:")
-    
+    # 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"] = ""
+
+    # 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"):
         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")
         else:
             st.warning("Please enter a prompt.")
 
+    # Show download button if an image is generated
+    if st.session_state["image"]:
+        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"
+        )
+
+    # 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
+
 #------------------------------------------------------------------------
 # Main Guard
 #------------------------------------------------------------------------