探索 SDXL 的简单优化方法

Stable Diffusion XL (SDXL) 是 Stability AI 推出的最新潜在扩散模型，用于生成高质量、超逼真的图像。它克服了以往 Stable Diffusion 模型在处理手部、文本以及空间构图正确性等方面的挑战。此外，SDXL 的上下文感知能力更强，只需较少的提示词就能生成更美观的图像。

然而，所有这些改进都是以模型体积显著增大为代价的。到底大了多少？SDXL 基础模型有 35 亿个参数（尤其是 UNet 部分），大约是之前 Stable Diffusion 模型的三倍大。

为了探索如何优化 SDXL 的推理速度和内存使用，我们在 A100 GPU (40 GB) 上进行了一些测试。每次推理运行，我们生成 4 张图像并重复 3 次。在计算推理延迟时，我们只考虑 3 次迭代中的最后一次。

所以，如果你直接开箱即用地运行全精度的 SDXL，并使用默认的注意力机制，它将消耗 28GB 内存，耗时 72.2 秒！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")
pipe.unet.set_default_attn_processor()

这很不实用，并且可能会拖慢你的速度，因为你通常需要生成不止 4 张图像。而且，如果你没有更强大的 GPU，就会遇到那个令人沮丧的内存不足错误信息。那么，我们该如何优化 SDXL 以提高推理速度并减少其内存使用呢？

在 🤗 Diffusers 中，我们有许多优化技巧和技术来帮助你运行像 SDXL 这样内存密集型的模型，接下来我们将向你展示如何做！我们将重点关注两个方面：推理速度和内存。

推理速度

扩散是一个随机过程，所以不能保证你一次就能得到满意的图像。通常情况下，你需要多次运行推理并进行迭代，这就是为什么优化速度至关重要。本节将重点介绍如何使用更低精度的权重，并结合内存高效的注意力和 PyTorch 2.0 中的 torch.compile 来提升速度并减少推理时间。

更低精度

模型权重以特定的精度存储，该精度表示为一种浮点数据类型。标准的浮点数据类型是 float32 (fp32)，它可以精确表示大范围的浮点数。对于推理而言，通常不需要那么高的精度，所以你可以使用 float16 (fp16)，它能表示的浮点数范围较窄。这意味着 fp16 占用的存储空间只有 fp32 的一半，而且由于计算更简单，速度是 fp32 的两倍。此外，现代 GPU 卡有专门优化的硬件来运行 fp16 计算，使其速度更快。

在 🤗 Diffusers 中，你可以通过指定 torch.dtype 参数在加载模型时转换权重，从而在推理时使用 fp16。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet.set_default_attn_processor()

与完全未经优化的 SDXL 流水线相比，使用 fp16 仅需 21.7GB 内存，耗时 14.8 秒。你几乎将推理速度提升了整整一分钟！

内存高效注意力

Transformer 模块中使用的注意力块可能是一个巨大的瓶颈，因为随着输入序列变长，内存会呈二次方增长。这会迅速占用大量内存，并导致你收到内存不足的错误信息。😬

内存高效的注意力算法旨在减轻计算注意力的内存负担，无论是通过利用稀疏性还是分块技术。这些优化算法过去主要以需要单独安装的第三方库的形式提供。但从 PyTorch 2.0 开始，情况不再如此。PyTorch 2 引入了缩放点积注意力 (SDPA)，它提供了Flash Attention、内存高效注意力 (xFormers) 以及一个 C++ 实现的 PyTorch 融合实现。SDPA 可能是加速推理最简单的方法：如果你正在使用 PyTorch ≥ 2.0 和 🤗 Diffusers，它默认是自动启用的！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")

与完全未经优化的 SDXL 流水线相比，使用 fp16 和 SDPA 所需的内存量相同，但推理时间缩短至 11.4 秒。让我们以此为新的基准，来比较其他优化方法。

torch.compile

PyTorch 2.0 还引入了 torch.compile API，用于将你的 PyTorch 代码即时 (JIT) 编译成更优化的推理内核。与其他编译器解决方案不同，torch.compile 对现有代码的改动极小，只需用该函数包装你的模型即可。

通过 mode 参数，你可以在编译时针对内存开销或推理速度进行优化，这为你提供了更大的灵活性。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

与之前的基准 (fp16 + SDPA) 相比，用 torch.compile 包装 UNet 后，推理时间缩短至 10.2 秒。

模型内存占用

如今的模型越来越大，要将它们装入内存成为一项挑战。本节重点介绍如何减少这些庞大模型的内存占用，以便你可以在消费级 GPU 上运行它们。这些技术包括 CPU 卸载、分步解码潜在表示为图像而非一次性完成，以及使用蒸馏版的自动编码器。

模型 CPU 卸载

模型卸载通过将 UNet 加载到 GPU 内存中，而将扩散模型的其他组件（文本编码器、VAE）加载到 CPU 上来节省内存。这样，UNet 可以在 GPU 上进行多次迭代，直到不再需要它为止。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_model_cpu_offload()

与基准相比，现在需要 20.2GB 内存，为你节省了 1.5GB 内存。

顺序 CPU 卸载

另一种可以为你节省更多内存但会降低推理速度的卸载方式是顺序 CPU 卸载。它不是卸载整个模型（如 UNet），而是将存储在不同 UNet 子模块中的模型权重卸载到 CPU，并且仅在前向传播之前才加载到 GPU。本质上，你每次只加载模型的一部分，从而可以节省更多内存。唯一的缺点是它明显更慢，因为你需要多次加载和卸载子模块。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_sequential_cpu_offload()

与基准相比，这需要 19.9GB 内存，但推理时间增加到 67 秒。

切片

在 SDXL 中，变分自编码器 (VAE) 将精炼后的潜在表示（由 UNet 预测）解码为逼真的图像。这一步的内存需求与预测的图像数量（批量大小）成正比。根据图像分辨率和可用的 GPU显存，它可能会非常占用内存。

这就是“切片”技术发挥作用的地方。待解码的输入张量被分割成多个切片，解码计算分几步完成。这节省了内存并允许使用更大的批量大小。

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.enable_vae_slicing()

通过切片计算，我们将内存减少到 15.4GB。如果再加入顺序 CPU 卸载，内存可以进一步减少到 11.45GB，这让你可以在每个提示词下生成 4 张 (1024x1024) 图像。然而，使用顺序卸载时，推理延迟也会增加。

缓存计算

任何文本条件的图像生成模型通常都使用文本编码器来从输入提示中计算嵌入。SDXL 使用了两个文本编码器！这对推理延迟有相当大的影响。然而，由于这些嵌入在整个反向扩散过程中保持不变，我们可以预先计算它们并在后续过程中重复使用。这样，在计算完文本嵌入后，我们就可以将文本编码器从内存中移除。

首先，加载文本编码器及其对应的分词器，并从输入提示中计算嵌入

tokenizers = [tokenizer, tokenizer_2]
text_encoders = [text_encoder, text_encoder_2]

(
    prompt_embeds,
    negative_prompt_embeds,
    pooled_prompt_embeds,
    negative_pooled_prompt_embeds
) = encode_prompt(tokenizers, text_encoders, prompt)

接下来，清空 GPU 内存以移除文本编码器

del text_encoder, text_encoder_2, tokenizer, tokenizer_2
flush()

现在，这些嵌入可以直接用于 SDXL 流水线了

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    torch_dtype=torch.float16,
).to("cuda")

call_args = dict(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        num_images_per_prompt=num_images_per_prompt,
        num_inference_steps=num_inference_steps,
)
image = pipe(**call_args).images[0]

结合 SDPA 和 fp16，我们可以将内存减少到 21.9GB。上面讨论的其他用于优化内存的技术也可以与缓存计算一起使用。

微型自动编码器

如前所述，VAE 将潜在表示解码为图像。很自然地，这一步直接受限于 VAE 的大小。所以，让我们用一个更小的自动编码器吧！由 madebyollin 制作的微型自动编码器，在 Hub 上可用，大小仅为 10MB，它是从 SDXL 使用的原始 VAE 蒸馏而来的。

from diffusers import AutoencoderTiny

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
pipe.to("cuda")

通过这种设置，我们将内存需求减少到 15.6GB，同时还降低了推理延迟。

结论

总结一下我们优化所带来的节省：

我们希望这些优化能让你轻松运行自己喜欢的流水线。快来试试这些技术，并与我们分享你的图像吧！🤗

---

致谢：感谢 Pedro Cuenca 对草稿提出的有益审阅。