Update app.py

app.py

@@ -4,22 +4,30 @@ import io
 import random
 from PIL import Image
 import os
+import json
 
 # 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)
-    return response.content
-
-# Check for 'text' parameter in URL
-query_params = st.query_params
-prompt_from_url = query_params.get('text')
+    
+    # Check for API errors
+    if response.status_code != 200:
+        try:
+            error = response.json()
+            return {"error": error.get("error", f"API Error: {response.status_code}")}
+        except:
+            return {"error": f"API Error: {response.status_code} - {response.text}"}
+    
+    # Check if response is actually an image
+    if 'image' not in response.headers.get('Content-Type', ''):
+        return {"error": "Unexpected non-image response from API"}
+    
+    return {"image": response.content}
 
 def generate_image(prompt):
-    # Generate a random seed to force a new image
     random_seed = random.randint(0, 999999)
     payload = {
         "inputs": prompt,
@@ -27,25 +35,37 @@ def generate_image(prompt):
             "seed": random_seed
         }
     }
-    image_bytes = query(payload)
-    return Image.open(io.BytesIO(image_bytes))
+    
+    result = query(payload)
+    
+    if "error" in result:
+        st.error(f"API Error: {result['error']}")
+        return None
+    
+    try:
+        return Image.open(io.BytesIO(result["image"]))
+    except Exception as e:
+        st.error(f"Failed to process image: {str(e)}")
+        return None
 
-if prompt_from_url:
-    # Generate and display image when 'text' parameter exists
-    prompt = prompt_from_url
-    image = generate_image(prompt)
-    st.image(image, caption="Generated Image", use_container_width=True)
+# Check for 'text' parameter in URL
+query_params = st.query_params
+prompt_from_url = query_params.get('text')
 
-    # Provide download link
-    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"
-    )
+if prompt_from_url:
+    image = generate_image(prompt_from_url)
+    if image:
+        st.image(image, caption="Generated Image", use_container_width=True)
+        
+        # Provide download link
+        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:
-    # Show instructions when no parameter is provided
     st.info("Add a 'text' parameter to the URL to generate an image. Example: ?text=astronaut+riding+a+horse")