Update app.py

app.py

@@ -1,12 +1,11 @@
-import gradio as gr
 import requests
-import io
 from io import BytesIO
-import os
 from PIL import Image
+import gradio as gr
+import os
 
 API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
-api_key = os.environ.get('API_KEY')
+api_key = os.environ.get('API_KEY', 'YOUR_API_KEY_HERE')
 headers = {"Authorization": f"Bearer {api_key}"}
 
 class QueryError(Exception):
@@ -14,7 +13,7 @@ class QueryError(Exception):
 
 def query(payload):
     try:
-        assert type(payload) == dict
+        assert isinstance(payload, dict)
         response = requests.post(API_URL, headers=headers, json=payload)
         
         if not str(response.status_code).startswith("2"):
@@ -24,16 +23,8 @@ def query(payload):
         
     except AssertionError:
         print("Invalid Payload Error: Please provide a dictionary.")
-    except RequestException as e:
+    except requests.exceptions.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:
@@ -43,25 +34,6 @@ def generate_images_from_prompt(prompt_text, num_images):
     images = []
     for _ in range(num_images):
         image_bytes = query({"inputs": prompt_text})
-        img = BytesIO(image_bytes)
-        pil_img = Image.open(img)
-        images.append(pil_img)
-    return images
-
-title = "KVIImager 2.0 Demo 🎨"
-description = "This app uses Hugging Face AI model to generate images based on the provided text prompt 🖼."
-
-input_prompt = gr.Textbox(label="Enter Prompt 📝", placeholder="E.g. 'A peaceful garden with a small cottage'")
-input_slider = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Images per Prompt")
-
-output_generated_images = gr.Image(label="Generated Image", type="pil")
-
-iface = gr.Interface(
-    fn=generate_images_from_prompt,
-    inputs=[input_prompt, input_slider],
-    outputs=output_generated_images,
-    title=title,
-    description=description,
-    theme="soft"
-)
-iface.launch()
+        img = Image.open(BytesIO(image_bytes))
+        images.append(img)
+    return images