Update app.py

app.py

@@ -1,15 +1,37 @@
 import requests
+import io
+from PIL import Image
+import json
 
 API_URL = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
 headers = {"Authorization": "Bearer HF_API_KEY"}
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
-    return response.content
+    if response.status_code == 200:
+        content_type = response.headers.get("Content-Type", "")
+        if "image" in content_type:
+            return response.content  # Returns raw image data
+        else:
+            # Log or print the unexpected content (likely JSON error message)
+            print("Unexpected response format:", response.text)
+            return None
+    else:
+        print("API Error:", response.status_code, response.text)
+        return None
+
+# Request image data
 image_bytes = query({
     "inputs": "Astronaut riding a horse",
 })
-# You can access the image with PIL.Image for example
-import io
-from PIL import Image
-image = Image.open(io.BytesIO(image_bytes))
+
+# Check if image_bytes contains valid data before proceeding
+if image_bytes:
+    try:
+        # Open image if it's valid
+        image = Image.open(io.BytesIO(image_bytes))
+        image.show()  # Display the image (optional)
+    except Exception as e:
+        print("Error opening image:", e)
+else:
+    print("Failed to retrieve image.")