Diffusers enable_vae_slicing 怎样降低多图解码的峰值显存

在 Diffusers 中，enable_vae_slicing() 适合解决“同一次任务要解码多张图片，VAE 阶段出现显存峰值”的问题。它把 VAE 解码的批次拆成较小的切片，重点降低解码瞬间同时驻留在显存中的数据量；代价是解码可能需要更多轮次，整体速度不一定更快。

官方地址：https://huggingface.co/docs/diffusers/en/optimization/memory

本文只讨论多图解码的峰值显存。示例是教学代码，不代表已经在某个本机模型上执行；模型标识、显存数字和最终速度应以自己的设备测量结果为准。

先区分生成批次与 VAE 解码批次

多图生成通常会把多张潜变量一起送入管线。前面的去噪阶段、潜变量张量和最后的 VAE 解码都可能产生显存压力，但 enable_vae_slicing() 针对的是 VAE 解码这一段。它不是把生成任务改成单张生成，也不会自动改变 num_images_per_prompt 的语义。

因此，排查时先问两个问题：一次调用到底要输出多少张图，以及峰值是否出现在潜变量转换为像素图的阶段。如果只是单张输出，切片通常没有值得期待的批量收益；如果一次要输出多张，VAE 阶段更容易成为需要关注的峰值点。

图中的“切片”是解码输入的组织方式，不应理解为把最终图片裁成若干块。输出仍然对应原来的多张图片，变化主要发生在 VAE 同时处理的输入规模。

在管线级启用 enable_vae_slicing

如果使用的是 Diffusers 管线，可以在开始生成前调用 pipe.enable_vae_slicing()。管线会把这个设置传给内部 VAE；生成参数仍由调用方控制。之后若要恢复普通解码，可调用对应的 pipe.disable_vae_slicing()。

import torch
from diffusers import DiffusionPipeline

# 把这里替换成自己有权限使用的模型标识；半精度只是示例设备配置
model_id = "your-model-id"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# 只改变 VAE 的解码组织方式，不改变 prompt 或输出张数的含义
pipe.enable_vae_slicing()

prompt = "一间有柔和天光的阅读室，插画风格"
with torch.inference_mode():
    # 多图输出更容易暴露 VAE 阶段的峰值显存问题
    result = pipe(prompt, num_images_per_prompt=4)
    images = result.images

# 如果后续场景更看重解码速度，可以显式恢复普通模式
# pipe.disable_vae_slicing()

这里的开关放在管线准备阶段，便于把策略与一次完整的生成任务绑定。它不替代模型加载、半精度、CPU offload 或其它内存策略，也不会替调用方处理显存不足后的错误恢复。

用受控的多图任务测量峰值

不要用“能生成成功”作为唯一结论。更有用的对比是：固定模型、提示词、推理步数、输出分辨率和输出张数，只切换 VAE slicing，然后读取 CUDA 的峰值分配值。下面的函数只展示测量方法，不写入任何虚构结果。

def run_once(pipe, prompt, image_count, sliced):
    # 每轮测试前清除上一次的峰值统计，避免两个模式互相污染
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    if sliced:
        # 切片模式降低 VAE 解码时的并行输入规模
        pipe.enable_vae_slicing()
    else:
        # 对照组明确关闭切片，确保比较条件可复现
        pipe.disable_vae_slicing()

    with torch.inference_mode():
        pipe(prompt, num_images_per_prompt=image_count)

    # 返回峰值分配值；不要把它写成所有显卡都相同的固定数字
    peak_mib = torch.cuda.max_memory_allocated() / (1024 ** 2)
    return peak_mib


prompt = "一座雾中的木桥，清晨自然光"
normal_peak = run_once(pipe, prompt, image_count=4, sliced=False)
sliced_peak = run_once(pipe, prompt, image_count=4, sliced=True)
print({"normal_mib": normal_peak, "sliced_mib": sliced_peak})

测量时要保持同一个进程和同一组输入，至少分别比较单图与多图。若每次测试都换了分辨率、精度或步数，最后看到的变化就不能归因于 enable_vae_slicing()。峰值统计也只描述本轮 CUDA 分配情况，不等于整台机器的总内存占用。

按场景决定是否长期打开

我的取舍是把 slicing 当成“多图解码的内存旋钮”，而不是默认的性能加速开关。线上服务如果请求张数变化很大，可以把它作为显存预算不足时的固定策略；如果单图延迟是首要指标，则应先用同样的输入测量解码耗时，再决定是否在所有请求上开启。

常见问题

enable_vae_slicing 会把四张图变成四次独立生成吗？不会。它主要改变 VAE 解码阶段对批次输入的处理方式，生成调用仍然可以请求多张图片。

打开 slicing 后显存一定减半吗？不能这样推断。峰值取决于分辨率、批量、精度、模型结构和其它驻留张量，官方开关没有给出适用于所有设备的固定比例。

VAE slicing 和 VAE tiling 是一回事吗？不是。前者面向批次切片，后者面向更大空间尺寸下的分块解码；两者解决的压力来源不同，不能只凭名称互换。

如果你的问题是多图任务在 VAE 解码瞬间冲高显存，先固定输入并比较 enable_vae_slicing() 前后的峰值。确认收益后，再结合延迟要求决定是否长期启用；如果 OOM 出现在其它阶段，则继续针对对应阶段排查。