import gradio as gr
import requests
import io
from io import BytesIO
import os
from PIL import Image

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):
    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})
    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
)
iface.launch()