diffusers/benchmark

1import gradio as gr2import torch3from diffusers import AutoPipelineForText2Image4import time5 6USE_TORCH_COMPILE = True7 8dtype = torch.float169device = torch.device("cuda:0")10 11pipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", variant="fp16", torch_dtype=dtype)12pipeline.vae.register_to_config(force_upcast=False)13pipeline.to(device)14 15 16if USE_TORCH_COMPILE:17    pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)18    19def generate(num_images_per_prompt: int = 1):20    print("Start...")21    print("Torch version", torch.__version__)22    print("Torch CUDA version", torch.version.cuda)23 24    for _ in range(3):25        prompt = 77 * "a"26        num_inference_steps = 2027        start_time = time.time()28        pipeline(prompt, num_images_per_prompt=num_images_per_prompt, num_inference_steps=num_inference_steps).images29        end_time = time.time()30    31        print(f"For {num_inference_steps} steps", end_time - start_time)32        print("Avg per step", (end_time - start_time) / num_inference_steps)33 34 35with gr.Blocks(css="style.css") as demo:36    batch_size = gr.Slider(37        label="Batch size",38        minimum=0,39        maximum=16,40        step=1,41        value=1,42    )43    btn = gr.Button("Benchmark!").style(44        margin=False,45        rounded=(False, True, True, False),46        full_width=False,47    )48 49    btn.click(fn=generate, inputs=[batch_size])50 51demo.launch()52