def sdxl_benchmark()

in benchmark/muse_perf.py [0:0]

• 24 lines of code
• 4 McCabe index (conditional complexity)

def sdxl_benchmark(batch_size, timesteps, use_xformers, gpu_type):
    model = "stabilityai/stable-diffusion-xl-base-1.0"
    device = "cuda"
    dtype = torch.float16

    pipe = StableDiffusionXLPipeline.from_pretrained(model, torch_dtype=dtype)
    pipe = pipe.to(device)

    if use_xformers:
        pipe.enable_xformers_memory_efficient_attention()

    if gpu_type == "4090" and batch_size == 8:
        output_type = "latent"
    else:
        output_type = "pil"

    def benchmark_fn():
        pipe(prompt, num_inference_steps=timesteps, num_images_per_prompt=batch_size, output_type=output_type)

    pipe(prompt, num_inference_steps=2, num_images_per_prompt=batch_size, output_type=output_type)

    def fn():
        return Timer(
            stmt="benchmark_fn()",
            globals={"benchmark_fn": benchmark_fn},
            num_threads=num_threads,
            label=f"batch_size: {batch_size}, dtype: {dtype}, timesteps {timesteps}, use_xformers: {use_xformers}",
            description=model,
        ).blocked_autorange(min_run_time=1)

    return measure_max_memory_allocated(fn)