Diffusers 文档
加速推理
并获得增强的文档体验
开始使用
加速推理
有几种方法可以优化 Diffusers 以提高推理速度，例如通过降低数据精度或使用轻量级蒸馏模型来减少计算负担。还有内存高效的注意力机制实现，xFormers 和 PyTorch 2.0 中的 scaled dot product attention，它们减少内存使用，这也间接加快了推理速度。不同的速度优化可以叠加在一起以获得最快的推理时间。
优化推理速度或减少内存使用可以提高另一类别的性能，因此您应该尽可能尝试同时优化两者。本指南侧重于推理速度，但您可以在减少内存使用指南中了解更多关于降低内存使用的信息。
下面的推理时间是通过在 NVIDIA A100 上使用 50 个 DDIM 步骤，从提示“火星上宇航员骑马的照片”生成单张 512x512 图像获得的。
| 设置 | 延迟 | 加速 |
|---|---|---|
| 基线 | 5.27秒 | x1 |
| tf32 | 4.14秒 | x1.27 |
| fp16 | 3.51秒 | x1.50 |
| 组合 | 3.41秒 | 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")不要在任何 pipelines 中使用 torch.autocast，因为它可能导致黑图像，并且总是比纯 float16 精度慢。
蒸馏模型
您还可以使用蒸馏的 Stable Diffusion 模型和自动编码器来加速推理。在蒸馏过程中，许多 UNet 的残差和注意力模块被去除，以将模型尺寸减少 51%，并将 CPU/GPU 上的延迟提高 43%。蒸馏模型速度更快，使用的内存更少，同时生成的图像质量与完整的 Stable Diffusion 模型相当。
阅读开源 SD-Small 和 SD-Tiny 的知识蒸馏代码和权重博客文章，以了解更多关于知识蒸馏训练如何工作以产生更快、更小、更便宜的生成模型的信息。
下面的推理时间是通过在 NVIDIA A100 上使用 25 个 PNDM 步骤，从提示“火星上宇航员骑马的照片”生成 4 张图像获得的。每次生成重复 3 次，使用 Nota AI 的蒸馏 Stable Diffusion v1.4 模型。
| 设置 | 延迟 | 加速 |
|---|---|---|
| 基线 | 6.37秒 | x1 |
| 蒸馏模型 | 4.18秒 | x1.52 |
| 蒸馏模型 + 微型自动编码器 | 3.83秒 | 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
微型自动编码器
为了进一步加速推理，请将自动编码器替换为它的蒸馏版本。
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
来自 madebyollin 提供了适用于其他 Stable Diffusion 模型（如 Stable Diffusion 3）的更多微型自动编码器模型。
< > 在 GitHub 上更新