加速推理
有多种方法可以优化 Diffusers 的推理速度，例如通过降低数据精度或使用轻量级蒸馏模型来降低计算负担。 还有内存高效的注意力实现，例如 PyTorch 2.0 中的 xFormers 和 缩放点积注意力，可以减少内存使用，从而间接加快推理速度。 不同的速度优化可以叠加在一起，以获得最快的推理时间。
优化推理速度或减少内存使用会导致另一类别的性能提升，因此应尽可能尝试同时优化这两者。 本指南重点介绍推理速度，但您可以在 减少内存使用 指南中了解有关降低内存使用的更多信息。
以下推理时间是从在 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(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
pipe = pipe.to("cuda")不要在任何管道中使用 torch.autocast，因为它会导致黑色图像，并且始终比纯 float16 精度慢。
蒸馏模型
您也可以使用蒸馏的 Stable Diffusion 模型和自动编码器来加快推理速度。 在蒸馏过程中，UNet 的许多残差块和注意力块会被去除，从而将模型大小减少 51%，并将 CPU/GPU 上的延迟提高 43%。 蒸馏模型更快，使用更少的内存，同时生成与完整 Stable Diffusion 模型质量相当的图像。
阅读 开源 SD-Small 和 SD-Tiny 的知识蒸馏代码和权重 博客文章，详细了解知识蒸馏训练如何生成更快、更小、更便宜的生成模型。
以下推理时间是从在 NVIDIA A100 上使用 25 个 PNDM 步从提示“宇航员骑着马在火星上”生成 4 张图像获得的。 每个生成使用 Nota AI 的蒸馏 Stable Diffusion v1.4 模型重复 3 次。
| 设置 | 延迟 | 加速 |
|---|---|---|
| 基线 | 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