Optimizing Stable Diffusion XL (SDXL) for Inference Speed and Memory
Hugging Face has detailed a series of simple optimizations to make Stable Diffusion XL (SDXL) more practical for inference, particularly on consumer GPUs. By implementing lower precision, memory-efficient attention, and various offloading techniques, the memory footprint can be reduced from 28GB to 11.47GB and inference latency can be dropped from 72.2 seconds to roughly 10.3 seconds.
Inference Speed Optimizations
Optimizing for speed is critical in diffusion models because the iterative nature of image generation often requires multiple runs to achieve a desired result. The following techniques focus on reducing latency:
Lower Precision (fp16)
Using float16 (fp16) instead of the standard float32 (fp32) reduces memory usage by half and increases calculation speed because fp16 captures a narrower range of floating numbers and is better supported by modern GPU hardware. In the tested SDXL pipeline, switching to fp16 reduced memory usage to 21.7GB and dropped inference time from 72.2 seconds to 14.8 seconds.
Memory-Efficient Attention (SDPA)
Attention blocks in transformers can create memory bottlenecks because memory requirements increase quadratically with input sequence length. PyTorch 2.0 introduced Scaled Dot Product Attention (SDPA), which provides fused implementations of Flash Attention and memory-efficient attention (xFormers). In 🤗Diffusers, SDPA is enabled by default for PyTorch ≥ 2.0, further reducing inference time to 11.4 seconds while maintaining the 21.7GB memory footprint.
JIT Compilation with torch.compile
PyTorch 2.0's torch.compile API allows for just-in-time (JIT) compilation of PyTorch code into optimized kernels. Wrapping the SDXL UNet with torch.compile (using mode="reduce-overhead") further improves inference time to 10.2 seconds. Note that the first compilation run is slower, but subsequent calls are significantly faster.
Reducing the Model Memory Footprint
Because SDXL is approximately 3x larger than previous Stable Diffusion models (with a 3.5B parameter UNet), fitting it into VRAM is a primary challenge. Several techniques can reduce the memory footprint:
CPU Offloading
Model offloading moves components of the pipeline to the CPU when they are not actively needed on the GPU:
- Model CPU Offloading: Loads the UNet into GPU memory while keeping text encoders and the VAE on the CPU. This reduced memory usage to 20.2GB.
- Sequential CPU Offloading: Offloads weights of individual UNet submodules to the CPU, loading them onto the GPU only immediately before the forward pass. This reduces memory to 19.9GB but significantly increases latency to 67 seconds.
VAE Slicing
The Variational Autoencoder (VAE) decodes latents into images, a process that scales in memory usage with the batch size. VAE slicing splits the input tensor into smaller slices and decodes them over several steps. This optimization reduces memory usage to 15.4GB without significantly impacting latency.
Caching Computations
SDXL utilizes two text encoders to compute embeddings from prompts. Since these embeddings remain constant throughout the reverse diffusion process, they can be precomputed and cached. Once the embeddings are generated, the text encoders and tokenizers can be removed from GPU memory, resulting in a memory footprint of 21.9GB when combined with SDPA and fp16.
Tiny Autoencoder (TAESD)
Replacing the standard VAE with a distilled version, such as the Tiny Autoencoder by madebyollin (approximately 10MB), reduces memory usage to 15.6GB and decreases inference latency. However, the Tiny Autoencoder may omit fine-grained details and is primarily recommended for image previews.
Performance Summary
Tests conducted on an A100 GPU (40 GB) generating 4 images per prompt show the following trade-offs between memory and latency:
| Technique | Memory (GB) | Inference Latency (ms) |
|---|---|---|
| Unoptimized pipeline | 28.09 | 72200.5 |
| fp16 | 21.72 | 14800.9 |
| fp16 + SDPA (default) | 21.72 | 11413.0 |
default + torch.compile |
21.73 | 10296.7 |
| default + model CPU offload | 20.21 | 16082.2 |
| default + sequential CPU offload | 19.91 | 67034.0 |
| default + VAE slicing | 15.40 | 11232.2 |
| default + VAE slicing + sequential CPU offload | 11.47 | 66869.2 |
| default + precomputed text embeddings | 21.85 | 11909.0 |
| default + Tiny Autoencoder | 15.48 | 10449.7 |