Diffusers documentation

Memory and speed

You are viewing v0.9.0 version. A newer version v0.14.0 is available.
Join the Hugging Face community

to get started

# Memory and speed

We present some techniques and ideas to optimize 🤗 Diffusers inference for memory or speed.

Latency Speedup
original 9.50s x1
cuDNN auto-tuner 9.37s x1.01
autocast (fp16) 5.47s x1.74
fp16 3.61s x2.63
channels last 3.30s x2.88
traced UNet 3.21s x2.96
memory efficient attention 2.63s x3.61
obtained on NVIDIA TITAN RTX by generating a single image of size 512x512 from the prompt "a photo of an astronaut riding a horse on mars" with 50 DDIM steps.

## Enable cuDNN auto-tuner

NVIDIA cuDNN supports many algorithms to compute a convolution. Autotuner runs a short benchmark and selects the kernel with the best performance on a given hardware for a given input size.

Since we’re using convolutional networks (other types currently not supported), we can enable cuDNN autotuner before launching the inference by setting:

import torch

torch.backends.cudnn.benchmark = True

### Use tf32 instead of fp32 (on Ampere and later CUDA devices)

On Ampere and later CUDA devices matrix multiplications and convolutions can use the TensorFloat32 (TF32) mode for faster but slightly less accurate computations. By default PyTorch enables TF32 mode for convolutions but not matrix multiplications, and unless a network requires full float32 precision we recommend enabling this setting for matrix multiplications, too. It can significantly speed up computations with typically negligible loss of numerical accuracy. You can read more about it here. All you need to do is to add this before your inference:

import torch

torch.backends.cuda.matmul.allow_tf32 = True

## Automatic mixed precision (AMP)

If you use a CUDA GPU, you can take advantage of torch.autocast to perform inference roughly twice as fast at the cost of slightly lower precision. All you need to do is put your inference call inside an autocast context manager. The following example shows how to do it using Stable Diffusion text-to-image generation as an example:

from torch import autocast
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
image = pipe(prompt).images[0]

Despite the precision loss, in our experience the final image results look the same as the float32 versions. Feel free to experiment and report back!

## Half precision weights

To save more GPU memory and get even more speed, you can load and run the model weights directly in half precision. This involves loading the float16 version of the weights, which was saved to a branch named fp16, and telling PyTorch to use the float16 type when loading them:

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

## Sliced attention for additional memory savings

For even additional memory savings, you can use a sliced version of attention that performs the computation in steps instead of all at once.

Attention slicing is useful even if a batch size of just 1 is used - as long as the model uses more than one attention head. If there is more than one attention head the *QK^T* attention matrix can be computed sequentially for each head which can save a significant amount of memory.

To perform the attention computation sequentially over each head, you only need to invoke enable_attention_slicing() in your pipeline before inference, like here:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_attention_slicing()
image = pipe(prompt).images[0]

There’s a small performance penalty of about 10% slower inference times, but this method allows you to use Stable Diffusion in as little as 3.2 GB of VRAM!

## Sliced VAE decode for larger batches

To decode large batches of images with limited VRAM, or to enable batches with 32 images or more, you can use sliced VAE decode that decodes the batch latents one image at a time.

You likely want to couple this with enable_attention_slicing() or enable_xformers_memory_efficient_attention() to further minimize memory use.

To perform the VAE decode one image at a time, invoke enable_vae_slicing() in your pipeline before inference. For example:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_vae_slicing()
images = pipe([prompt] * 32).images

You may see a small performance boost in VAE decode on multi-image batches. There should be no performance impact on single-image batches.

For additional memory savings, you can offload the weights to CPU and load them to GPU when performing the forward pass.

To perform CPU offloading, all you have to do is invoke enable_sequential_cpu_offload():

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

And you can get the memory consumption to < 2GB.

If is also possible to chain it with attention slicing for minimal memory consumption, running it in as little as < 800mb of GPU vRAM:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_attention_slicing(1)

image = pipe(prompt).images[0]

## Using Channels Last memory format

Channels last memory format is an alternative way of ordering NCHW tensors in memory preserving dimensions ordering. Channels last tensors ordered in such a way that channels become the densest dimension (aka storing images pixel-per-pixel). Since not all operators currently support channels last format it may result in a worst performance, so it’s better to try it and see if it works for your model.

For example, in order to set the UNet model in our pipeline to use channels last format, we can use the following:

print(pipe.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipe.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
pipe.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

## Tracing

Tracing runs an example input tensor through your model, and captures the operations that are invoked as that input makes its way through the model’s layers so that an executable or ScriptFunction is returned that will be optimized using just-in-time compilation.

To trace our UNet model, we can use the following:

import time
import torch
from diffusers import StableDiffusionPipeline
import functools

# set variables
n_experiments = 2
unet_runs_per_experiment = 50

def generate_inputs():
sample = torch.randn(2, 4, 64, 64).half().cuda()
timestep = torch.rand(1).half().cuda() * 999
encoder_hidden_states = torch.randn(2, 77, 768).half().cuda()
return sample, timestep, encoder_hidden_states

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
).to("cuda")
unet = pipe.unet
unet.eval()
unet.to(memory_format=torch.channels_last)  # use channels_last memory format
unet.forward = functools.partial(unet.forward, return_dict=False)  # set return_dict=False as default

# warmup
for _ in range(3):
with torch.inference_mode():
inputs = generate_inputs()
orig_output = unet(*inputs)

# trace
print("tracing..")
unet_traced = torch.jit.trace(unet, inputs)
unet_traced.eval()
print("done tracing")

# warmup and optimize graph
for _ in range(5):
with torch.inference_mode():
inputs = generate_inputs()
orig_output = unet_traced(*inputs)

# benchmarking
with torch.inference_mode():
for _ in range(n_experiments):
torch.cuda.synchronize()
start_time = time.time()
for _ in range(unet_runs_per_experiment):
orig_output = unet_traced(*inputs)
torch.cuda.synchronize()
print(f"unet traced inference took {time.time() - start_time:.2f} seconds")
for _ in range(n_experiments):
torch.cuda.synchronize()
start_time = time.time()
for _ in range(unet_runs_per_experiment):
orig_output = unet(*inputs)
torch.cuda.synchronize()
print(f"unet inference took {time.time() - start_time:.2f} seconds")

# save the model
unet_traced.save("unet_traced.pt")

Then we can replace the unet attribute of the pipeline with the traced model like the following

from diffusers import StableDiffusionPipeline
import torch
from dataclasses import dataclass

@dataclass
class UNet2DConditionOutput:
sample: torch.FloatTensor

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
).to("cuda")

# use jitted unet
# del pipe.unet
class TracedUNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.in_channels = pipe.unet.in_channels
self.device = pipe.unet.device

def forward(self, latent_model_input, t, encoder_hidden_states):
sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
return UNet2DConditionOutput(sample=sample)

pipe.unet = TracedUNet()

with torch.inference_mode():
image = pipe([prompt] * 1, num_inference_steps=50).images[0]

## Memory Efficient Attention

Recent work on optimizing the bandwitdh in the attention block have generated huge speed ups and gains in GPU memory usage. The most recent being Flash Attention (from @tridao, [code](https://github.com/HazyResearch/flash-attention), [paper](https://arxiv.org/pdf/2205.14135.pdf)) . Here are the speedups we obtain on a few Nvidia GPUs when running the inference at 512x512 with a batch size of 1 (one prompt):
GPU Base Attention FP16 Memory Efficient Attention FP16
NVIDIA Tesla T4 3.5it/s 5.5it/s
NVIDIA 3060 RTX 4.6it/s 7.8it/s
NVIDIA A10G 8.88it/s 15.6it/s
NVIDIA RTX A6000 11.7it/s 21.09it/s
NVIDIA TITAN RTX 12.51it/s 18.22it/s
A100-SXM4-40GB 18.6it/s 29.it/s
A100-SXM-80GB 18.7it/s 29.5it/s

To leverage it just make sure you have:

• PyTorch > 1.12
• Cuda available
• Installed the xformers library
from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16,
).to("cuda")

pipe.enable_xformers_memory_efficient_attention()

with torch.inference_mode():
sample = pipe("a small cat")

# optional: You can disable it via
# pipe.disable_xformers_memory_efficient_attention()