import gradio as gr
import requests
import base64
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv
import os

# Load the .env file
load_dotenv()

# Get the API key from the environment variable
api_key = os.getenv("API_KEY")

# Define the endpoint URL
invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"

# Define the function that interacts with the API
def generate_image(prompt, seed=0, sampler="K_EULER_ANCESTRAL", steps=2):
    # Define the headers with the correct API key
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Accept": "application/json",
    }

    # Define the payload for the POST request
    payload = {
        "text_prompts": [{"text": prompt}],
        "seed": seed,
        "sampler": sampler,
        "steps": steps
    }

    # Make the POST request to the API
    response = requests.post(invoke_url, headers=headers, json=payload)

    # Raise an error if the request was unsuccessful
    response.raise_for_status()

    # Get the response body
    response_body = response.json()

    # Extract the base64 string from the response
    base64_str = response_body['artifacts'][0]['base64']

    # Decode the base64 string
    image_data = base64.b64decode(base64_str)

    # Open the image using PIL
    image = Image.open(BytesIO(image_data))

    return image

# Define the Gradio interface
inputs = [
    gr.Textbox(label="Prompt", placeholder="Enter your image description here"),
    gr.Number(label="Seed", value=0, precision=0),
    gr.Dropdown(label="Sampler", choices=["K_EULER_ANCESTRAL", "DDIM", "PLMS"], value="K_EULER_ANCESTRAL"),
    gr.Slider(label="Steps", minimum=1, maximum=100, value=2)
]

outputs = gr.Image(label="Generated Image")

gr.Interface(
    fn=generate_image,
    inputs=inputs,
    outputs=outputs,
    title="Stable Diffusion XL Image Generator",
    description="Generate images using Stable Diffusion XL by providing a text prompt."
).launch()