Update app.py

app.py

@@ -8,9 +8,42 @@ API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
 api_key = os.environ.get('API_KEY')
 headers = {"Authorization": f"Bearer {api_key}"}
 
+# Define custom Exception class for better error handling
+class QueryError(Exception):
+    pass
+
 def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.content
+    try:
+      # Make sure we have valid JSON data before sending the request
+      assert type(payload) == dict
+      
+      # Send the POST request to the API URL
+      response = requests.post(API_URL, headers=headers, json=payload)
+    
+      # Check if the status code indicates success (HTTP Status Code 2xx)
+      if not str(response.status_code).startswith("2"):
+          raise QueryError(f"Query failed! Response status code was '{response.status_code}'")
+        
+      else:
+           # Return the raw bytes from the response object
+            return response.content
+            
+    except AssertionError:
+        print("Invalid Payload Error: Please provide a dictionary.")
+    except RequestException as e:
+        print("Request Failed: ", e)
+    except ConnectionError as ce:
+        print("Connection Error: Unable to connect to the API.", ce)
+    except Timeout as t:
+        print("Timeout Error: Request timed out while trying to reach the API.", t)
+    except TooManyRedirects as tmr:
+        print("Too Many Redirects Error: Exceeded maximum number of redirects.", tmr)
+    except HTTPError as he:
+        print("HTTP Error: Invalid HTTP response.", he)
+    except QueryError as qe:
+        print(qe)
+    except Exception as ex:
+        print("Unknown Error occurred: ", ex)
 
 def generate_image_from_prompt(prompt_text):
     image_bytes = query({"inputs": prompt_text})