Update app.py

app.py

@@ -1,39 +1,69 @@
+import gradio as gr
 import requests
+import io
 from io import BytesIO
-from PIL import Image
-import gradio as gr
 import os
+from PIL import Image
 
 API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
-api_key = os.environ.get('API_KEY', 'YOUR_API_KEY_HERE')
+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):
     try:
-        assert isinstance(payload, dict)
-        response = requests.post(API_URL, headers=headers, json=payload)
-        
-        if not str(response.status_code).startswith("2"):
-            raise QueryError(f"Query failed! Response status code was '{response.status_code}'")
-        
-        return response.content
+      # 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 requests.exceptions.RequestException as e:
+    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_images_from_prompt(prompt_text, num_images):
-    images = []
-    for _ in range(num_images):
-        image_bytes = query({"inputs": prompt_text})
-        img = Image.open(BytesIO(image_bytes))
-        images.append(img)
-    return images
+def generate_image_from_prompt(prompt_text):
+    image_bytes = query({"inputs": prompt_text})
+    img = BytesIO(image_bytes)   # Convert to BytesIO stream
+    pil_img = Image.open(img)    # Open the image using PIL library
+    return pil_img               # Return the converted PIL image
+
+title = "KVIImager 2.0 Demo 🎨"
+description = "This app uses Hugging Face AI model to generate an image based on the provided text prompt 🖼."
+
+input_prompt = gr.Textbox(label="Enter Prompt 📝", placeholder="E.g. 'A peaceful garden with a small cottage'")
+output_generated_image = gr.Image(label="Generated Image")
+
+iface = gr.Interface(
+    fn=generate_image_from_prompt,
+    inputs=input_prompt, 
+    outputs=output_generated_image, 
+    title=title,
+    description=description,
+    theme="soft"
+)
+iface.launch()