Torch2.0 support in Diffusers

Starting from version 0.13.0, Diffusers supports the latest optimization from the upcoming PyTorch 2.0 release. These include:

1. Support for native flash and memory-efficient attention without any extra dependencies.
2. torch.compile support for compiling individual models for extra performance boost.

Installation

pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu117

Using efficient attention and torch.compile.

1. Efficient Attention

   Efficient attention is implemented via the torch.nn.functional.scaled_dot_product_attention function, which automatically enables flash/memory efficient attention, depending on the input and the GPU type. This is the same as the memory_efficient_attention from xFormers but built natively into PyTorch.

   Efficient attention will be enabled by default in Diffusers if torch2.0 is installed and if torch.nn.functional.scaled_dot_product_attention is available. To use it, you can install torch2.0 as suggested above and use the pipeline. For example:

   Copied

   import torch
   from diffusers import StableDiffusionPipeline
   
   pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
   pipe = pipe.to("cuda")
   
   prompt = "a photo of an astronaut riding a horse on mars"
   image = pipe(prompt).images[0]

   If you want to enable it explicitly (which is not required), you can do so as shown below.

   Copied

   import torch
   from diffusers import StableDiffusionPipeline
   from diffusers.models.cross_attention import AttnProcessor2_0
   
   pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
   pipe.unet.set_attn_processor(AttnProcessor2_0())
   
   prompt = "a photo of an astronaut riding a horse on mars"
   image = pipe(prompt).images[0]

   This should be as fast and memory efficient as xFormers.

2. torch.compile

   To get an additional speedup, we can use the new torch.compile feature. To do so, we wrap our unet with torch.compile. For more information and different options, refer to the torch compile docs.

   Copied

   import torch
   from diffusers import StableDiffusionPipeline
   
   pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(
       "cuda"
   )
   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

   Depending on the type of GPU it can give between 2-9% speed-up over efficient attention. But note that as of now the speed-up is mostly noticeable on the more recent GPU architectures, such as in the A100.

   Note that compilation will also take some time to complete, so it is best suited for situations where you need to prepare your pipeline once and then perform the same type of inference operations multiple times.

Benchmark

We conducted a simple benchmark on different GPUs to compare vanilla attention, xFormers, torch.nn.functional.scaled_dot_product_attention and torch.compile+torch.nn.functional.scaled_dot_product_attention. For the benchmark we used the the stable-diffusion-v1-4 model with 50 steps. xFormers benchmark is done using the torch==1.13.1 version. The table below summarizes the result that we got. The Speed over xformers columns denotes the speed-up gained over xFormers using the torch.compile+torch.nn.functional.scaled_dot_product_attention.

FP16 benchmark

The table below shows the benchmark results for inference using fp16. As we can see, torch.nn.functional.scaled_dot_product_attention is as fast as xFormers (sometimes slightly faster/slower) on all the GPUs we tested. And using torch.compile gives further speed-up up to 10% over xFormers, but it’s mostly noticeable on the A100 GPU.

The time reported is in seconds.

FP32 benchmark

The table below shows the benchmark results for inference using fp32. As we can see, torch.nn.functional.scaled_dot_product_attention is as fast as xFormers (sometimes slightly faster/slower) on all the GPUs we tested. Using torch.compile with efficient attention gives up to 18% performance improvement over xFormers in Ampere cards, and up to 20% over vanilla attention.

(1) Batch Size >= 32 requires enable_vae_slicing() because of https://github.com/pytorch/pytorch/issues/81665 This is required for PyTorch 1.13.1, and also for PyTorch 2.0 and batch size of 64

For more details about how this benchmark was run, please refer to this PR.