Accelerated PyTorch 2.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 accelerated transformers implementation with memory-efficient attention – no extra dependencies required.
2. torch.compile support for extra performance boost when individual models are compiled.

Installation

pip install --upgrade torch torchvision diffusers

Using accelerated transformers and torch.compile.

1. Accelerated Transformers implementation

   PyTorch 2.0 includes an optimized and memory-efficient attention implementation through the torch.nn.functional.scaled_dot_product_attention function, which automatically enables several optimizations depending on the inputs and the GPU type. This is similar to the memory_efficient_attention from xFormers, but built natively into PyTorch.

   These optimizations will be enabled by default in Diffusers if PyTorch 2.0 is installed and if torch.nn.functional.scaled_dot_product_attention is available. To use it, just install torch 2.0 as suggested above and simply use the pipeline. For example:

   Copied

   import torch
   from diffusers import DiffusionPipeline
   
   pipe = DiffusionPipeline.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 DiffusionPipeline
   from diffusers.models.attention_processor import AttnProcessor2_0
   
   pipe = DiffusionPipeline.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. More details in our benchmark.

2. torch.compile

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

   Copied

   import torch
   from diffusers import DiffusionPipeline
   
   pipe = DiffusionPipeline.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, compile() can yield between 2-9% of additional speed-up over the accelerated transformer optimizations. Note, however, that compilation is able to squeeze more performance improvements in more recent GPU architectures such as Ampere (A100, 3090), Ada (4090) and Hopper (H100).

   Compilation takes 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 stable-diffusion-v1-4 model with 50 steps. The xFormers benchmark is done using the torch==1.13.1 version, while the accelerated transformers optimizations are tested using nightly versions of PyTorch 2.0. The tables below summarize the results we got.

Please refer to our featured blog post in the PyTorch site for more details.

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 of up of 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. In this case, torch.nn.functional.scaled_dot_product_attention is faster than xFormers on all the GPUs we tested.

Using torch.compile in addition to the accelerated transformers implementation can yield up to 19% performance improvement over xFormers in Ampere and Ada cards, and up to 20% (Ampere) or 28% (Ada) 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 large batch sizes.

For more details about how this benchmark was run, please refer to this PR and to the blog post.