import gradio as gr import requests import io from io import BytesIO import os from PIL import Image API_URL = os.environ.get("API_URL") 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: # 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): gr.Info("Image generation started") 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. 'Astronaut riding a horse'") 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()