import gradio as gr
import requests
import io
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}"}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def generate_image_from_prompt(prompt_text):
    image_bytes = query({"inputs": prompt_text})
    generated_image = Image.open(io.BytesIO(image_bytes))
    return generated_image

title = "KVIImager 2.0 Demo 🎨"
description = "This app uses Hugging Face AI model to generate an image based on the provided text prompt 🖼️."
examples = [
    ["A peaceful garden with a small cottage 🏡"],
    ["A colorful abstract painting with geometric shapes 🎨"],
    ["A serene beach at sunset 🌅"]
]

css_styles = {
    'body': {
        'background-color': '#f4f4f4',
        'font-family': 'Arial, sans-serif'
    },
    'title': {
        'color': 'navy',
        'font-size': '36px',
        'text-align': 'center',
        'margin-bottom': '20px'
    },
    'textbox': {
        'border': '2px solid #008CBA',
        'border-radius': '5px',
        'padding': '10px',
        'margin-bottom': '20px',
        'width': '300px',
        'font-size': '16px'
    },
    'output_image': {
        'box-shadow': '2px 2px 5px #888888'
    }
}

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")

gr.Interface(
    generate_image_from_prompt, 
    inputs=input_prompt, 
    outputs=output_generated_image, 
    title=title,
    description=description,
    examples=examples
).launch(inline=True, css=css_styles)