降低内存使用量
使用扩散模型的一个障碍是需要大量的内存。为了克服这一挑战，您可以使用多种减少内存的技术来在免费层或消费级 GPU 上运行一些最大的模型。其中一些技术甚至可以结合使用以进一步减少内存使用量。
在许多情况下，针对内存或速度进行优化会导致另一方面的性能提升，因此您应该尽可能地同时针对两者进行优化。本指南侧重于最大程度地减少内存使用量，但您还可以了解有关如何加快推理速度的更多信息。
以下结果是在 Nvidia Titan RTX 上使用 50 个 DDIM 步长从提示“火星上骑马的宇航员的照片”生成单个 512x512 图像获得的，展示了由于减少内存消耗而可以预期的加速效果。
| 延迟 | 加速 | |
|---|---|---|
| 原始 | 9.50 秒 | x1 |
| fp16 | 3.61 秒 | x2.63 |
| 通道最后 | 3.30 秒 | x2.88 |
| 追踪 UNet | 3.21 秒 | x2.96 |
| 内存高效注意力 | 2.63 秒 | x3.61 |
分片 VAE
分片 VAE 通过一次解码一批潜变量中的一个图像，从而能够在 VRAM 有限的情况下解码大量图像批次或包含 32 个或更多图像的批次。如果您安装了 xFormers，则可能需要将其与enable_xformers_memory_efficient_attention()结合使用以进一步减少内存使用量。
要使用分片 VAE，请在推理之前在您的管道上调用enable_vae_slicing()
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
pipe = pipe.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_vae_slicing()
#pipe.enable_xformers_memory_efficient_attention()
images = pipe([prompt] * 32).images您可能会在多图像批次的 VAE 解码中看到少量性能提升，并且在单图像批次中不应产生性能影响。
平铺 VAE
平铺 VAE 处理还能够在 VRAM 有限的情况下处理大型图像（例如，在 8GB VRAM 上生成 4k 图像），方法是将图像分割成重叠的图块，解码这些图块，然后将输出混合在一起以构成最终图像。如果您安装了 xFormers，则还应将平铺 VAE 与enable_xformers_memory_efficient_attention()结合使用以进一步减少内存使用量。
要使用平铺 VAE 处理，请在推理之前在您的管道上调用enable_vae_tiling()
import torch
from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")
prompt = "a beautiful landscape photograph"
pipe.enable_vae_tiling()
#pipe.enable_xformers_memory_efficient_attention()
image = pipe([prompt], width=3840, height=2224, num_inference_steps=20).images[0]输出图像存在一些图块到图块之间的色调变化，因为图块是分别解码的，但您不应在图块之间看到任何尖锐且明显的接缝。对于大小为 512x512 或更小的图像，将关闭平铺功能。
CPU 卸载
将权重卸载到 CPU 并仅在执行前向传递时将它们加载到 GPU 上也可以节省内存。通常，此技术可以将内存消耗减少到不到 3GB。
要执行 CPU 卸载，请调用enable_sequential_cpu_offload()
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_sequential_cpu_offload()
image = pipe(prompt).images[0]CPU 卸载适用于子模块而不是整个模型。这是最大程度地减少内存消耗的最佳方法，但由于扩散过程的迭代性质，推理速度会慢得多。管道的 UNet 组件运行多次（多达 num_inference_steps）；每次，不同的 UNet 子模块都会根据需要依次加载和卸载，从而导致大量的内存传输。
如果您想优化速度，请考虑使用模型卸载，因为它速度快得多。权衡是您的内存节省不会那么大。
使用enable_sequential_cpu_offload()时，不要事先将管道移动到 CUDA，否则内存消耗的增益将很小（请参阅此问题以获取更多信息）。
enable_sequential_cpu_offload() 是一个有状态操作，它在模型上安装挂钩。
模型卸载
模型卸载需要 🤗 Accelerate 版本 0.17.0 或更高版本。
顺序 CPU 卸载 保留了大量内存，但它使推理速度变慢，因为子模块会根据需要移动到 GPU，并在新模块运行时立即返回到 CPU。
全模型卸载是一种替代方法，它将整个模型移动到 GPU，而不是处理每个模型的组成子模块。对推理时间的影响可以忽略不计（与将管道移动到 cuda 相比），并且它仍然可以节省一些内存。
在模型卸载期间，管道的主要组件之一（通常是文本编码器、UNet 和 VAE）位于 GPU 上，而其他组件则等待在 CPU 上。像 UNet 这样的运行多次迭代的组件会保留在 GPU 上，直到不再需要它们为止。
通过在管道上调用enable_model_cpu_offload()来启用模型卸载
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
)
prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_model_cpu_offload()
image = pipe(prompt).images[0]为了在调用模型后正确卸载模型，需要运行整个管道，并且模型按管道预期顺序调用。如果在安装挂钩后在管道上下文之外重用模型，请谨慎操作。有关更多信息，请参阅删除挂钩。
enable_model_cpu_offload() 是一个有状态操作，它在模型上安装挂钩并在管道上安装状态。
通道后置内存格式
通道后置内存格式是一种用于重新排序 NCHW 张量在内存中的存储方式，以保持维度顺序。通道后置张量按像素逐像素存储通道的方式进行排序，使通道成为最密集的维度。由于目前并非所有算子都支持通道后置格式，因此可能会导致性能下降，但您仍应尝试并查看它是否适用于您的模型。
例如，要将管道的 UNet 设置为使用通道后置格式
print(pipe.unet.conv_out.state_dict()["weight"].stride()) # (2880, 9, 3, 1)
pipe.unet.to(memory_format=torch.channels_last) # in-place operation
print(
pipe.unet.conv_out.state_dict()["weight"].stride()
) # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works追踪
追踪会将示例输入张量通过模型，并捕获在该输入通过模型的各层时执行的操作。返回的可执行文件或ScriptFunction会通过即时编译进行优化。
要追踪 UNet
import time
import torch
from diffusers import StableDiffusionPipeline
import functools
# torch disable grad
torch.set_grad_enabled(False)
# set variables
n_experiments = 2
unet_runs_per_experiment = 50
# load inputs
def generate_inputs():
sample = torch.randn((2, 4, 64, 64), device="cuda", dtype=torch.float16)
timestep = torch.rand(1, device="cuda", dtype=torch.float16) * 999
encoder_hidden_states = torch.randn((2, 77, 768), device="cuda", dtype=torch.float16)
return sample, timestep, encoder_hidden_states
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
unet = pipe.unet
unet.eval()
unet.to(memory_format=torch.channels_last) # use channels_last memory format
unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default
# warmup
for _ in range(3):
with torch.inference_mode():
inputs = generate_inputs()
orig_output = unet(*inputs)
# trace
print("tracing..")
unet_traced = torch.jit.trace(unet, inputs)
unet_traced.eval()
print("done tracing")
# warmup and optimize graph
for _ in range(5):
with torch.inference_mode():
inputs = generate_inputs()
orig_output = unet_traced(*inputs)
# benchmarking
with torch.inference_mode():
for _ in range(n_experiments):
torch.cuda.synchronize()
start_time = time.time()
for _ in range(unet_runs_per_experiment):
orig_output = unet_traced(*inputs)
torch.cuda.synchronize()
print(f"unet traced inference took {time.time() - start_time:.2f} seconds")
for _ in range(n_experiments):
torch.cuda.synchronize()
start_time = time.time()
for _ in range(unet_runs_per_experiment):
orig_output = unet(*inputs)
torch.cuda.synchronize()
print(f"unet inference took {time.time() - start_time:.2f} seconds")
# save the model
unet_traced.save("unet_traced.pt")将管道中的unet属性替换为已追踪的模型
from diffusers import StableDiffusionPipeline
import torch
from dataclasses import dataclass
@dataclass
class UNet2DConditionOutput:
sample: torch.Tensor
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
# use jitted unet
unet_traced = torch.jit.load("unet_traced.pt")
# del pipe.unet
class TracedUNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.in_channels = pipe.unet.config.in_channels
self.device = pipe.unet.device
def forward(self, latent_model_input, t, encoder_hidden_states):
sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
return UNet2DConditionOutput(sample=sample)
pipe.unet = TracedUNet()
with torch.inference_mode():
image = pipe([prompt] * 1, num_inference_steps=50).images[0]内存高效注意力机制
最近在优化注意力块带宽方面的工作取得了巨大的速度提升和 GPU 内存使用量的降低。最新的内存高效注意力机制类型是Flash Attention（您可以在HazyResearch/flash-attention查看原始代码）。
如果您安装了 PyTorch >= 2.0，则启用xformers时不应期望推理速度加快。
要使用 Flash Attention，请安装以下内容：
- PyTorch > 1.12
- CUDA 可用
- xFormers
然后在管道上调用enable_xformers_memory_efficient_attention()
from diffusers import DiffusionPipeline
import torch
pipe = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
pipe.enable_xformers_memory_efficient_attention()
with torch.inference_mode():
sample = pipe("a small cat")
# optional: You can disable it via
# pipe.disable_xformers_memory_efficient_attention()使用xformers时的迭代速度应与此处描述的 PyTorch 2.0 的迭代速度匹配此处。