Tech-Meld/Automated_Stable_Diffusion_3_Comparison

1import torch2from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler3from huggingface_hub import hf_hub_download4from safetensors.torch import load_file5import gradio as gr6from tqdm.auto import tqdm7import psutil8 9base = "stabilityai/stable-diffusion-xl-base-1.0"10repo = "ByteDance/SDXL-Lightning"11ckpt = "sdxl_lightning_4step_unet.safetensors"12 13# Load model.14unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cpu")15unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))16pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float32).to("cpu")17 18# Ensure sampler uses "trailing" timesteps.19pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")20 21def generate_images(prompt, num_inference_steps, guidance_scale, batch_size):22    with tqdm(total=num_inference_steps, desc="Inference Progress") as pbar:23        images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_size=batch_size, progress_bar=pbar).images24    return images25 26# Define Gradio interface27def get_cpu_info():28    cpu_name = psutil.cpu_freq().brand29    memory_available = psutil.virtual_memory().available // 1024 // 1024  # in MB30    return f"CPU: {cpu_name}, Memory: {memory_available} MB"31 32cpu_info_text = gr.Textbox(label="CPU Information", value=get_cpu_info(), interactive=False)33 34iface = gr.Interface(35    fn=generate_images,36    inputs=[37        gr.Textbox(label="Prompt"),38        gr.Slider(label="Num Inference Steps", minimum=1, maximum=50, step=1, value=4),39        gr.Slider(label="Guidance Scale", minimum=0, maximum=20, step=0.1, value=0),40        gr.Slider(label="Batch Size", minimum=1, maximum=4, step=1, value=1),41    ],42    outputs=[43        gr.Gallery(label="Generated Images"),44        cpu_info_text45    ],46    title="SDXL Lightning 4-Step Inference (CPU)",47    description="Generate images with Stable Diffusion XL Lightning 4-Step model on CPU.",48)49 50iface.launch()