import gradio as gr
import torch
import spaces
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

device = "cuda" if torch.cuda.is_available() else "cpu"
base = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ByteDance/SDXL-Lightning"
opts = {
    "1 Step"  : ("sdxl_lightning_1step_unet_x0.safetensors", 1),
    "2 Steps" : ("sdxl_lightning_2step_unet.safetensors", 2),
    "4 Steps" : ("sdxl_lightning_4step_unet.safetensors", 4),
    "8 Steps" : ("sdxl_lightning_8step_unet.safetensors", 8),
}

step_loaded = 4
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)

@spaces.GPU(enable_queue=True)
def generate_image(prompt, option):
    ckpt, step = opts[option]
    if step_loaded != step:
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
        step_loaded = step
    return pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]

with gr.Blocks() as demo:
    gr.HTML("<h1><center>SDXL-Lightning</center></h1>")
    gr.HTML("<p><center>Lightning-fast text-to-image generation.</center></p>")
    gr.HTML("<p><center><a href='https://huggingface.co/ByteDance/SDXL-Lightning'>https://huggingface.co/ByteDance/SDXL-Lightning</a></center></p>")
    
    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(
                label="Text prompt",
                scale=8
            )
            option = gr.Dropdown(
                label="Inference steps",
                choices=["1 Step", "2 Steps", "4 Steps", "8 Steps"],
                value="4 Steps",
                interactive=True
            )
            submit = gr.Button(
                scale=1,
                variant="primary"
            )
    
    img = gr.Image(label="SDXL-Lightening Generated Image")

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, option],
        outputs=img,
    )
    submit.click(
        fn=generate_image,
        inputs=[prompt, option],
        outputs=img,
    )
    
demo.queue().launch()