Update app.py

app.py

@@ -5,50 +5,50 @@ import random
 from PIL import Image
 import os
 
-# API Configuration for Hugging Face
+# API Configuration
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 headers = {"Authorization": f"Bearer {os.getenv('HF')}"}
 
+# Function to query the Hugging Face API
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code != 200:
-        st.error("Error calling the Hugging Face API.")
-        st.stop()
     return response.content
 
-def generate_image(prompt: str) -> Image.Image:
-    # Use a random seed so even identical prompts produce different images.
+def generate_image(prompt):
+    # Generate a random seed to ensure a new image each time
     random_seed = random.randint(0, 999999)
     payload = {
         "inputs": prompt,
-        "parameters": {"seed": random_seed}
+        "parameters": {
+            "seed": random_seed
+        }
     }
     image_bytes = query(payload)
     return Image.open(io.BytesIO(image_bytes))
 
-st.title("AI Image Generator 🎨🚀")
-st.write("Enter a prompt below and generate an AI-generated image using Hugging Face!")
-
-prompt = st.text_input("Enter your prompt:", "Astronaut riding a horse")
-
-if st.button("Generate Image"):
-    if prompt:
-        st.write("Generating image... Please wait ⏳")
-        image = generate_image(prompt)
-        st.image(image, caption="Generated Image", use_container_width=True)
-        
-        # Convert image to bytes for download
-        img_buffer = io.BytesIO()
-        image.save(img_buffer, format="PNG")
-        img_buffer.seek(0)
-        st.download_button(
-            label="Download Image 📥",
-            data=img_buffer,
-            file_name="generated_image.png",
-            mime="image/png"
-        )
-    else:
-        st.warning("Please enter a prompt before generating an image.")
-
-st.write("---")
-st.write("Powered by [Hugging Face](https://huggingface.co/) 🚀")
+# Default prompt (used since we're not using a query parameter anymore)
+default_prompt = "Astronaut riding a horse"
+
+# Use session_state to avoid regenerating the image on every rerun
+if "image" not in st.session_state:
+    st.session_state.image = generate_image(default_prompt)
+
+if st.button("Generate"):
+    st.session_state.image = generate_image(default_prompt)
+
+# Display the image using the updated parameter
+image = st.session_state.image
+st.image(image, caption="Result", use_container_width=True)
+
+# Prepare image for download
+img_buffer = io.BytesIO()
+image.save(img_buffer, format="PNG")
+img_buffer.seek(0)
+st.download_button(
+    label="Download",
+    data=img_buffer,
+    file_name="generated_image.png",
+    mime="image/png"
+)
+
+st.write("---")