Update app.py

app.py

@@ -15,17 +15,27 @@ def query_hf_image_generation(api_key, prompt):
     }
 
     response = requests.post(API_URL, headers=headers, json=data)
+    if response.status_code != 200:
+        return "Error: Unexpected status code {} with message {}".format(response.status_code, response.text)
+
     result = response.json()
     
+    # Debug output to help figure out what result looks like
+    print("DEBUG:", result)
+
     # Check if the API response contains an error.
     if 'error' in result:
         return "Error: " + result['error'], None
 
     # Assuming the API returns an image in base64 format.
-    image_data = result['data'][0]  # Adjust this according to your specific API response structure
-    image_bytes = base64.b64decode(image_data.split(",")[1])
-    image = Image.open(BytesIO(image_bytes))
-    return image
+    if 'data' in result:
+        base64_image = result['data'][0]  # This might need to be adjusted based on API return structure
+        base64_data = base64_image.split(",")[1] if ',' in base64_image else base64_image
+        image_bytes = base64.b64decode(base64_data)
+        image = Image.open(BytesIO(image_bytes))
+        return image
+    else:
+        return "Error: 'data' not found in response."
 
 iface = gr.Interface(
     fn=query_hf_image_generation,