Update app.py

app.py

@@ -14,26 +14,36 @@ def query_hf_image_generation(api_key, prompt):
         "inputs": prompt
     }
 
+    # Make the request
     response = requests.post(API_URL, headers=headers, json=data)
+
+    # Check if the response was successful
     if response.status_code != 200:
-        return "Error: Unexpected status code {} with message {}".format(response.status_code, response.text)
+        return f"Error: Received status code {response.status_code} with message: {response.text}"
+
+    # Try parsing JSON response
+    try:
+        result = response.json()
+    except ValueError as e:
+        return f"Error decoding JSON: {e}"
 
-    result = response.json()
-    
-    # Debug output to help figure out what result looks like
+    # Debug output to diagnose the structure of the returned 'result'
     print("DEBUG:", result)
 
-    # Check if the API response contains an error.
+    # Check if the API response contains an error message.
     if 'error' in result:
-        return "Error: " + result['error'], None
+        return f"Error: {result['error']}"
 
     # Assuming the API returns an image in base64 format.
     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
+        try:
+            base64_image = result['data'][0]
+            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
+        except Exception as e:
+            return f"Error processing image data: {e}"
     else:
         return "Error: 'data' not found in response."