import csv from argparse import ArgumentParser import torch from diffusers import ( AutoencoderKL, AutoPipelineForText2Image, LatentConsistencyModelPipeline, LCMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline, UNet2DConditionModel, ) from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS from torch.utils.benchmark import Compare, Timer from transformers import ( AutoTokenizer, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer, ) from muse import MaskGiTUViT, PipelineMuse, VQGANModel torch.manual_seed(0) torch.set_grad_enabled(False) torch.set_float32_matmul_precision("high") num_threads = torch.get_num_threads() prompt = "A high tech solarpunk utopia in the Amazon rainforest" do_sd15 = True do_sdxl = True do_ssd_1b = True do_sdxl_turbo = True do_sd_turbo = True do_muse = True do_wurst = True do_lcm = True def main(): args = ArgumentParser() args.add_argument("--device", choices=["4090", "a100"], required=True) args = args.parse_args() csv_data = [] for batch_size in [1, 8]: for timesteps in [12, 20]: for use_xformers in [False, True]: if do_sd15: out, mem_bytes = sd_benchmark(batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "stable_diffusion_1_5", out.median * 1000, args.device, timesteps, mem_bytes, 512, use_xformers, None, ] ) if do_sdxl: out, mem_bytes = sdxl_benchmark( batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers, gpu_type=args.device ) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "sdxl", out.median * 1000, args.device, timesteps, mem_bytes, 1024, use_xformers, None, ] ) if do_ssd_1b: out, mem_bytes = ssd_1b_benchmark( batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers, gpu_type=args.device ) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "ssd_1b", out.median * 1000, args.device, timesteps, mem_bytes, 1024, use_xformers, None, ] ) if do_muse: for resolution in [256, 512]: for use_fused_residual_norm in [False, True]: out, mem_bytes = muse_benchmark( resolution=resolution, batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers, use_fused_residual_norm=use_fused_residual_norm, ) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "muse", out.median * 1000, args.device, timesteps, mem_bytes, resolution, use_xformers, use_fused_residual_norm, ] ) if do_sdxl_turbo: for use_xformers in [True, False]: timesteps = 1 out, mem_bytes = sdxl_turbo_benchmark( batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers ) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "sdxl_turbo", out.median * 1000, args.device, timesteps, mem_bytes, 1024, use_xformers, None, ] ) if do_sd_turbo: for use_xformers in [True, False]: timesteps = 1 out, mem_bytes = sd_turbo_benchmark( batch_size=batch_size, timesteps=timesteps, use_xformers=use_xformers ) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "sd_turbo", out.median * 1000, args.device, timesteps, mem_bytes, 512, use_xformers, None, ] ) if do_wurst: for use_xformers in [False, True]: out, mem_bytes = wurst_benchmark(batch_size, use_xformers) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "wurst", out.median * 1000, args.device, "default", mem_bytes, 1024, use_xformers, None, ] ) if do_lcm: for timesteps in [4, 8]: for use_xformers in [False, True]: out, mem_bytes = lcm_benchmark(batch_size, timesteps, use_xformers) Compare([out]).print() print("*******") csv_data.append( [ batch_size, "lcm", out.median * 1000, args.device, timesteps, mem_bytes, 1024, use_xformers, None, ] ) with open("benchmark/artifacts/all.csv", "a", newline="") as csvfile: writer = csv.writer(csvfile) writer.writerows(csv_data) def muse_benchmark(resolution, batch_size, timesteps, use_xformers, use_fused_residual_norm): model = "williamberman/muse_research_run_benchmarking_512_output" device = "cuda" dtype = torch.float16 tokenizer = AutoTokenizer.from_pretrained(model, subfolder="text_encoder") text_encoder = CLIPTextModelWithProjection.from_pretrained(model, subfolder="text_encoder") text_encoder.to(device=device, dtype=dtype) vae = VQGANModel.from_pretrained(model, subfolder="vae") vae.to(device=device, dtype=dtype) transformer = MaskGiTUViT( use_fused_mlp=False, use_fused_residual_norm=use_fused_residual_norm, force_down_up_sample=resolution == 512, ) transformer = transformer.to(device=device, dtype=dtype) transformer.eval() if use_xformers: transformer.enable_xformers_memory_efficient_attention() pipe = PipelineMuse( tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, ) pipe.device = device pipe.dtype = dtype seq_len = (resolution // 16) ** 2 def benchmark_fn(): pipe(prompt, num_images_per_prompt=batch_size, timesteps=timesteps, transformer_seq_len=seq_len) pipe(prompt, num_images_per_prompt=batch_size, timesteps=2, transformer_seq_len=seq_len) 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}, resolution: {resolution}," f" use_xformers: {use_xformers}, use_fused_residual_norm: {use_fused_residual_norm}" ), description=model, ).blocked_autorange(min_run_time=1) return measure_max_memory_allocated(fn) def wurst_benchmark(batch_size, use_xformers): model = "warp-ai/wuerstchen" device = "cuda" dtype = torch.float16 pipe = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=dtype).to(device) if use_xformers: pipe.enable_xformers_memory_efficient_attention() def benchmark_fn(): pipe( prompt, height=1024, width=1024, prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS, prior_guidance_scale=4.0, num_images_per_prompt=batch_size, ) # warmup benchmark_fn() def fn(): return Timer( stmt="benchmark_fn()", globals={"benchmark_fn": benchmark_fn}, num_threads=num_threads, label=f"batch_size: {batch_size}, dtype: {dtype}, use_xformers: {use_xformers}", description=model, ).blocked_autorange(min_run_time=1) return measure_max_memory_allocated(fn) 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) def lcm_benchmark(batch_size, timesteps, use_xformers): model = "SimianLuo/LCM_Dreamshaper_v7" device = "cuda" dtype = torch.float16 scheduler = LCMScheduler.from_pretrained(model, subfolder="scheduler") pipe = LatentConsistencyModelPipeline.from_pretrained(model, torch_dtype=dtype, scheduler=scheduler) pipe.to(device) if use_xformers: pipe.enable_xformers_memory_efficient_attention() def benchmark_fn(): pipe(prompt, num_inference_steps=timesteps, num_images_per_prompt=batch_size) pipe(prompt, num_inference_steps=2, num_images_per_prompt=batch_size) 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) def ssd_1b_benchmark(batch_size, timesteps, use_xformers, gpu_type): model = "segmind/SSD-1B" device = "cuda" dtype = torch.float16 pipe = StableDiffusionXLPipeline.from_pretrained(model, torch_dtype=dtype, use_safetensors=True, variant="fp16") 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) def sd_benchmark(batch_size, timesteps, use_xformers): model = "runwayml/stable-diffusion-v1-5" device = "cuda" dtype = torch.float16 tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(model, subfolder="text_encoder") text_encoder.to(device=device, dtype=dtype) vae = AutoencoderKL.from_pretrained(model, subfolder="vae") vae = vae.to(device=device, dtype=dtype) unet = UNet2DConditionModel.from_pretrained(model, subfolder="unet") unet = unet.to(device=device, dtype=dtype) pipe = StableDiffusionPipeline.from_pretrained( model, vae=vae, unet=unet, text_encoder=text_encoder, tokenizer=tokenizer, safety_checker=None, ) if use_xformers: pipe.enable_xformers_memory_efficient_attention() def benchmark_fn(): pipe( prompt, num_images_per_prompt=batch_size, num_inference_steps=timesteps, ) pipe(prompt, num_images_per_prompt=batch_size, num_inference_steps=2) 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) def sd_turbo_benchmark(batch_size, timesteps, use_xformers): model = "stabilityai/sd-turbo" dtype = torch.float16 pipe = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", safety_checker=None) pipe.to("cuda") if use_xformers: pipe.enable_xformers_memory_efficient_attention() def benchmark_fn(): pipe( prompt, num_images_per_prompt=batch_size, num_inference_steps=timesteps, ) pipe(prompt, num_images_per_prompt=batch_size, num_inference_steps=2) 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) def sdxl_turbo_benchmark(batch_size, timesteps, use_xformers): model = "stabilityai/sdxl-turbo" dtype = torch.float16 pipe = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", safety_checker=None) pipe.to("cuda") if use_xformers: pipe.enable_xformers_memory_efficient_attention() def benchmark_fn(): pipe( prompt, num_images_per_prompt=batch_size, num_inference_steps=timesteps, ) pipe(prompt, num_images_per_prompt=batch_size, num_inference_steps=2) 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) def measure_max_memory_allocated(fn): torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() rv = fn() mem_bytes = torch.cuda.max_memory_allocated() return rv, mem_bytes if __name__ == "__main__": main()