import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None).to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.unet = torch.compile(pipe.unet)

batch_size = 10
prompt =  "A photo of an astronaut riding a horse on marse."
images = pipe(prompt,  num_inference_steps=steps, num_images_per_prompt=batch_size).images

import torch
import torch.utils.benchmark as benchmark
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.models.cross_attention import CrossAttnProcessor, AttnProccesor2_0

def benchmark_torch_function(f, *args, **kwargs):
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
    )
    return round(t0.blocked_autorange(min_run_time=1).mean, 2)

# benchmark code
model_id = "CompVis/stable-diffusion-v1-4"
prompt = "A photo of an astronaut riding a horse on mars."
steps = 50
batch_size = 10
dtype = torch.float16

# load model
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, safety_checker=None).to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.set_progress_bar_config(disable=True)
if batch_size >= 32:
    pipe.enable_vae_slicing()

# Vanilla Cross Attention
print("Running benchmark for vanilla cross attention...")
pipe.unet.set_attn_processor(CrossAttnProcessor())
f = lambda : pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images
time_vanilla = benchmark_torch_function(f)

# PyTorch sdpa
print("Running benchmark for PyTorch SDPA...")
pipe.unet.set_attn_processor(AttnProccesor2_0())
f = lambda : pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images
time_sdpa = benchmark_torch_function(f)

# PyTorch sdpa with torch.compile
print("Running benchmark for PyTorch SDPA with torch.compile...")
pipe.unet = torch.compile(pipe.unet)
# warmup
pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images
f = lambda : pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images
time_sdpa_torch_compile = benchmark_torch_function(f)

# print results with nice formatting
print(f"Model: {model_id}, dtype: {dtype}, steps: {steps}, batch_size: {batch_size}")
print(f"Vanilla Cross Attention:         {time_vanilla} s")
print(f"PyTorch SDPA:                    {time_sdpa} s")
print(f"PyTorch SDPA with torch.compile: {time_sdpa_torch_compile} s")

import torch
import torch.utils.benchmark as benchmark
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

def benchmark_torch_function(f, *args, **kwargs):
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
    )
    return round(t0.blocked_autorange(min_run_time=1).mean, 2)

# benchmark code
model_id = "CompVis/stable-diffusion-v1-4"
prompt = "A photo of an astronaut riding a horse on mars."
steps = 50
batch_size = 16
dtype = torch.float16

# load model
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, safety_checker=None).to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.set_progress_bar_config(disable=True)

# xFormers
print("Running benchmark for xFormers...")
pipe.enable_xformers_memory_efficient_attention()
f = lambda : pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images
time_xformers = benchmark_torch_function(f)

# print results with nice formatting
print(f"Model: {model_id}, dtype: {dtype}, steps: {steps}, batch_size: {batch_size}")
print(f"xFormers: {time_xformers} s")

[2023-02-10 13:02:03,854] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)																										
function: 'forward' (/home/pedro/code/hf/diffusers/diffusers/src/diffusers/models/cross_attention.py:175)																										
reasons:  ___guarded_code.valid																										
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.