Diffusers 文档

加速文本到图像扩散模型的推理

Diffusers

加入 Hugging Face 社区

并获得增强型文档体验

在模型、数据集和 Spaces 上协作

使用加速推断获得更快的示例

切换文档主题

以开始使用

加速文本到图像扩散模型的推断

扩散模型比其 GAN 对应物更慢，因为它们使用了迭代和顺序的反向扩散过程。有一些技术可以解决这个限制，例如渐进式时间步长蒸馏 (LCM LoRA)、模型压缩 (SSD-1B) 以及重用去噪器相邻特征 (DeepCache)。

但是，您并不一定要使用这些技术来加速推断。仅使用 PyTorch 2，您就可以将文本到图像扩散管道的推断延迟最多加速 3 倍。本教程将向您展示如何逐步应用 PyTorch 2 中的优化来减少推断延迟。在本教程中，您将使用 Stable Diffusion XL (SDXL) 管道，但这些技术也适用于其他文本到图像扩散管道。

确保您使用的是最新版本的 Diffusers

pip install -U diffusers

然后也升级其他必需的库

pip install -U transformers accelerate peft

安装 PyTorch nightly 以从最新、最快的内核中受益

pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121

以下报告的结果来自一个 80GB 400W A100，其时钟频率设置为最大值。如果您有兴趣了解完整的基准测试代码，请查看 huggingface/diffusion-fast。

基线

让我们从基线开始。禁用降低精度和 scaled_dot_product_attention (SDPA) 函数，该函数由 Diffusers 自动使用

from diffusers import StableDiffusionXLPipeline

# Load the pipeline in full-precision and place its model components on CUDA.
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0"
).to("cuda")

# Run the attention ops without SDPA.
pipe.unet.set_default_attn_processor()
pipe.vae.set_default_attn_processor()

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt, num_inference_steps=30).images[0]

此默认设置需要 7.36 秒。

bfloat16

启用第一个优化，降低精度，更具体地说是 bfloat16。使用降低精度有几个好处

对推断使用降低的数值精度（例如 float16 或 bfloat16）不会影响生成质量，但会显著提高延迟。
与 float16 相比，使用 bfloat16 的好处取决于硬件，但现代 GPU 往往更倾向于使用 bfloat16。
bfloat16 在与量化一起使用时比 float16 更具弹性，但我们使用的更近期的量化库版本 (torchao) 不会遇到 float16 的数值问题。

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

# Run the attention ops without SDPA.
pipe.unet.set_default_attn_processor()
pipe.vae.set_default_attn_processor()

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt, num_inference_steps=30).images[0]

bfloat16 将延迟从 7.36 秒降低到 4.63 秒。

在我们后来使用 float16 进行的实验中，torchao 的最新版本不会遇到 float16 的数值问题。

请查看加速推断指南，以详细了解如何使用降低精度运行推断。

SDPA

注意力块运行起来很密集。但使用 PyTorch 的 scaled_dot_product_attention 函数，效率要高得多。此函数在 Diffusers 中默认使用，因此您无需对代码进行任何更改。

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt, num_inference_steps=30).images[0]

缩放点积注意力将延迟从 4.63 秒提高到 3.31 秒。

torch.compile

PyTorch 2 包含 torch.compile，它使用快速且优化的内核。在 Diffusers 中，通常会编译 UNet 和 VAE，因为它们是最密集的计算模块。首先，配置一些编译器标志（有关更多选项，请参考完整列表）

from diffusers import StableDiffusionXLPipeline
import torch

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

在编译 UNet 和 VAE 时，将它们的内存布局更改为“channels_last”也很重要，以确保最大速度。

pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)

现在编译并执行推断

# Compile the UNet and VAE.
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

# First call to `pipe` is slow, subsequent ones are faster.
image = pipe(prompt, num_inference_steps=30).images[0]

torch.compile 提供了不同的后端和模式。为了获得最大的推断速度，对 inductor 后端使用“max-autotune”。“max-autotune”使用 CUDA 图表，并专门针对延迟优化编译图。CUDA 图表通过使用机制通过单个 CPU 操作启动多个 GPU 操作，从而大大减少了启动 GPU 操作的开销。

使用 SDPA 注意力和编译 UNet 和 VAE 将延迟从 3.31 秒降低到 2.54 秒。

从 PyTorch 2.3.1 开始，您可以控制 torch.compile() 的缓存行为。这对像 "max-autotune" 这样的编译模式特别有用，它对多个编译标志执行网格搜索以找到最佳配置。在 torch.compile 中的编译时间缓存教程中了解更多信息。

防止图表中断

指定 fullgraph=True 确保底层模型中没有图表中断，以便充分利用 torch.compile 而不降低性能。对于 UNet 和 VAE，这意味着更改访问返回值的方式。

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

编译后删除 GPU 同步

在迭代反向扩散过程中，step() 函数在降噪器预测出噪声更少的潜在嵌入后，每次都会在调度器上调用。在 step() 内部，sigmas 变量会被索引，当它被放置在 GPU 上时，会导致 CPU 和 GPU 之间的通信同步。这会引入延迟，当降噪器已经编译后，这种延迟会更加明显。

但是，如果 sigmas 数组始终保留在 CPU 上，CPU 和 GPU 同步就不会发生，你也不会遇到任何延迟。通常，任何 CPU 和 GPU 通信同步应该为零或保持最低限度，因为这会影响推理延迟。

组合注意力块的投影矩阵

SDXL 中的 UNet 和 VAE 使用类似 Transformer 的块，这些块由注意力块和前馈块组成。

在注意力块中，输入使用三个不同的投影矩阵（Q、K 和 V）投影到三个子空间。这些投影是在输入上分别执行的。但是，我们可以将投影矩阵水平组合成一个单一矩阵，并一步执行投影。这增加了输入投影矩阵乘法的尺寸，并改善了量化的影响。

你可以只用一行代码组合投影矩阵

pipe.fuse_qkv_projections()

这提供了从 2.54 秒到 2.52 秒的微小改进。

fuse_qkv_projections() 的支持有限且处于实验阶段。它不适用于许多非 Stable Diffusion 管道，例如 Kandinsky。你可以参考这个 PR 了解如何为其他管道启用它。

动态量化

你还可以使用超轻量级的 PyTorch 量化库，torchao（提交 SHA 54bcd5a10d0abbe7b0c045052029257099f83fd9），对 UNet 和 VAE 应用动态 int8 量化。量化会给模型添加额外的转换开销，但希望这可以通过更快的 matmuls（动态量化）来弥补。如果 matmuls 太小，这些技术可能会降低性能。

首先，配置所有编译器标签

from diffusers import StableDiffusionXLPipeline
import torch

# Notice the two new flags at the end.
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

UNet 和 VAE 中的某些线性层无法从动态 int8 量化中获益。你可以使用下面显示的 dynamic_quant_filter_fn 过滤掉这些层。

def dynamic_quant_filter_fn(mod, *args):
    return (
        isinstance(mod, torch.nn.Linear)
        and mod.in_features > 16
        and (mod.in_features, mod.out_features)
        not in [
            (1280, 640),
            (1920, 1280),
            (1920, 640),
            (2048, 1280),
            (2048, 2560),
            (2560, 1280),
            (256, 128),
            (2816, 1280),
            (320, 640),
            (512, 1536),
            (512, 256),
            (512, 512),
            (640, 1280),
            (640, 1920),
            (640, 320),
            (640, 5120),
            (640, 640),
            (960, 320),
            (960, 640),
        ]
    )


def conv_filter_fn(mod, *args):
    return (
        isinstance(mod, torch.nn.Conv2d) and mod.kernel_size == (1, 1) and 128 in [mod.in_channels, mod.out_channels]
    )

最后，应用到目前为止讨论的所有优化

# SDPA + bfloat16.
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

# Combine attention projection matrices.
pipe.fuse_qkv_projections()

# Change the memory layout.
pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)

由于动态量化仅限于线性层，因此将适当的逐点卷积层转换为线性层，以最大程度地发挥其优势。

from torchao import swap_conv2d_1x1_to_linear

swap_conv2d_1x1_to_linear(pipe.unet, conv_filter_fn)
swap_conv2d_1x1_to_linear(pipe.vae, conv_filter_fn)

应用动态量化

from torchao import apply_dynamic_quant

apply_dynamic_quant(pipe.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipe.vae, dynamic_quant_filter_fn)

最后，编译并执行推理

pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt, num_inference_steps=30).images[0]

应用动态量化将延迟从 2.52 秒提高到 2.43 秒。

< > 在 GitHub 上更新

←加载 LoRAs 进行推理使用大型模型→