加速推理

有几种方法可以优化 Diffusers 的推理速度，例如通过降低数据精度或使用轻量级蒸馏模型来减少计算负担。PyTorch 2.0 中还有一些内存高效的注意力实现，例如 xFormers 和 缩放点积注意力，它们减少了内存使用，也间接加快了推理速度。不同的速度优化可以叠加在一起，以获得最快的推理时间。

TIP

优化推理速度或减少内存使用可能会导致另一类性能的提升，因此你应该尽可能地尝试同时优化这两者。本指南侧重于推理速度，但你可以在 减少内存使用 指南中了解更多关于降低内存使用的信息。

以下推理时间是从使用 50 个 DDIM 步在 NVIDIA A100 上从提示“宇航员骑着马在火星上的照片”生成单个 512x512 图像获得的。

setup	latency	speed-up
baseline	5.27s	x1
tf32	4.14s	x1.27
fp16	3.51s	x1.50
combined	3.41s	x1.54

TensorFloat-32

在 Ampere 及更高版本的 CUDA 设备上，矩阵乘法和卷积可以使用 TensorFloat-32 (tf32) 模式进行更快的计算，但精度略低。默认情况下，PyTorch 为卷积启用了 tf32 模式，但未为矩阵乘法启用。除非你的网络需要完整的 float32 精度，否则我们建议为矩阵乘法启用 tf32。它可以显著加快计算速度，通常不会造成明显的数值精度损失。

import torch

torch.backends.cuda.matmul.allow_tf32 = True

在 混合精度训练 指南中了解更多关于 tf32 的信息。

半精度权重

为了节省 GPU 内存并获得更高的速度，请设置 torch_dtype=torch.float16 以直接使用半精度权重加载和运行模型权重。

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    use_safetensors=True,
)
pipe = pipe.to("cuda")

蒸馏模型

你也可以使用蒸馏的 Stable Diffusion 模型和自动编码器来加速推理。在蒸馏过程中，UNet 的许多残差和注意力块被移除，从而将模型大小减少了 51%，并将 CPU/GPU 上的延迟提高了 43%。蒸馏模型更快，使用更少的内存，同时生成与完整 Stable Diffusion 模型质量相当的图像。

TIP

阅读 开源 SD-Small 和 SD-Tiny 的知识蒸馏代码和权重 博客文章，了解知识蒸馏训练如何生成更快、更小、更便宜的生成模型。

以下推理时间是从使用 NVIDIA A100 在 25 个 PNDM 步长下从提示“宇航员骑着马在火星上的照片”生成 4 张图像获得的。每个生成使用 Nota AI 的蒸馏 Stable Diffusion v1.4 模型重复 3 次。

setup	latency	speed-up
baseline	6.37s	x1
distilled	4.18s	x1.52
distilled + tiny autoencoder	3.83s	x1.66

让我们加载蒸馏后的 Stable Diffusion 模型，并将其与原始 Stable Diffusion 模型进行比较。

from diffusers import StableDiffusionPipeline
import torch

distilled = StableDiffusionPipeline.from_pretrained(
    "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")
prompt = "a golden vase with different flowers"
generator = torch.manual_seed(2023)
image = distilled("a golden vase with different flowers", num_inference_steps=25, generator=generator).images[0]
image

original Stable Diffusion

distilled Stable Diffusion

小型自动编码器

为了进一步加快推理速度，请用其蒸馏版本替换自动编码器。

import torch
from diffusers import AutoencoderTiny, StableDiffusionPipeline

distilled = StableDiffusionPipeline.from_pretrained(
    "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")
distilled.vae = AutoencoderTiny.from_pretrained(
    "sayakpaul/taesd-diffusers", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")

prompt = "a golden vase with different flowers"
generator = torch.manual_seed(2023)
image = distilled("a golden vase with different flowers", num_inference_steps=25, generator=generator).images[0]
image

distilled Stable Diffusion + Tiny AutoEncoder

其他 Stable Diffusion 模型（如 Stable Diffusion 3）的更多小型自动编码器模型可从madebyollin获得。