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
from tqdm.auto import tqdm
import psutil

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).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):
    with tqdm(total=num_inference_steps, desc="Inference Progress") as pbar:
        images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_size=batch_size, progress_bar=pbar).images
    return images

# Define Gradio interface
def get_cpu_info():
    cpu_name = psutil.cpu_freq().brand
    memory_available = psutil.virtual_memory().available // 1024 // 1024  # in MB
    return f"CPU: {cpu_name}, Memory: {memory_available} MB"

cpu_info_text = gr.Textbox(label="CPU Information", value=get_cpu_info(), interactive=False)

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=4, step=1, value=1),
    ],
    outputs=[
        gr.Gallery(label="Generated Images"),
        cpu_info_text
    ],
    title="SDXL Lightning 4-Step Inference (CPU)",
    description="Generate images with Stable Diffusion XL Lightning 4-Step model on CPU.",
)

iface.launch()