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

def query_hf_image_generation(api_key, prompt):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {"inputs": prompt}

    response = requests.post(API_URL, headers=headers, json=data)

    if response.status_code != 200:
        return f"Error: Received HTTP {response.status_code} - {response.text}"

    try:
        result = response.json()
    except ValueError:
        return f"Error decoding JSON: Unexpected response format {response.text}"

    if 'error' in result:
        return f"Error: {result['error']}"

    if 'data' in result:
        try:
            base64_string = result['data'][0]
            base64_data = base64_string.split(",")[1] if "," in base64_string else base64_string
            image_data = base64.b64decode(base64_data)
            image = Image.open(BytesIO(image_data))
            return image
        except Exception as e:
            return f"Error processing image data: {e}"
    else:
        return "Error: Missing 'data' in the response."

iface = gr.Interface(
    fn=query_hf_image_generation,
    inputs=[
        gr.Textbox(label="Hugging Face API Key", placeholder="Enter your Hugging Face API Key here..."),
        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Stable Diffusion XL Image Generator",
    description="Enter your API Key and a prompt to generate an image using the Stable Diffusion XL model from Hugging Face."
)

iface.launch(share=True)