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

base = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ByteDance/SDXL-Lightning"
ckpt = "sdxl_lightning_4step_unet.safetensors"

# Load model.
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cpu")
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float32, variant="fp32").to("cpu")

# Ensure sampler uses "trailing" timesteps.
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")

def generate_images(prompt, num_inference_steps, guidance_scale, batch_size):
    images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_size=batch_size).images
    return images

# Define Gradio interface
iface = gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(label="Num Inference Steps", minimum=1, maximum=50, step=1, value=4),
        gr.Slider(label="Guidance Scale", minimum=0, maximum=20, step=0.1, value=0),
        gr.Slider(label="Batch Size", minimum=1, maximum=8, step=1, value=2),
    ],
    outputs=gr.Gallery(label="Generated Images"),
    title="SDXL Lightning 4-Step Inference (CPU)",
    description="Generate images with Stable Diffusion XL Lightning 4-Step model on CPU.",
)

iface.launch()