Mac SDXL计算时间及内存占用优化方法

在制作AI写真整合包时，发现使用SDXL内存占用非常大，因此在网上找到一些解决方案，分享给大家！

Stable Diffusion XL (SDXL)是由Stability AI推出的最新潜在扩散模型，用于生成高质量超现实图像。它克服了以前Stable Diffusion模型的挑战，如正确处理手部和文本以及空间上正确的构图。此外，SDXL还更具上下文感知能力，需要较少的词语即可生成更好看的图像。

然而，所有这些改进都是以显著增加模型大小为代价的。有多大呢？基础的SDXL模型有35亿个参数（特别是UNet），大约是之前Stable Diffusion模型的3倍。

为了探索如何优化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这样的内存密集型模型，我们将向你展示如何做到！我们将关注的两个方面是推理速度和内存。

推理速度

扩散是一个随机过程，所以你没有保证会得到一个你喜欢的图像。通常情况下，你需要多次运行推理并迭代，这就是为什么优化速度至关重要。本节重点介绍使用较低精度权重、结合内存高效的注意力机制和torch.compile（来自PyTorch 2.0）来提高速度并减少推理时间。

较低精度

模型权重以特定精度存储，表示为浮点数据类型。标准浮点数据类型是float32（fp32），可以准确表示广泛的浮点数。对于推理来说，你通常不需要那么精确，所以你应该使用float16（fp16），它捕获的浮点数范围更窄。这意味着与fp32相比，fp16只需要一半的内存来存储，并且因为计算更容易，所以速度是两倍。此外，现代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秒。你几乎将推理速度提高了整整一分钟！

内存高效注意力

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

内存高效的注意力算法旨在减少计算注意力所需的内存负担，无论是通过利用稀疏性还是平铺。这些优化算法过去主要作为需要单独安装的第三方库提供。但从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,
).to("cuda")
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,
).to("cuda")
pipe.enable_sequential_cpu_offload()

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

切片

在SDXL中，变分编码器（VAE）将UNet预测的精炼潜在变量（latents）解码为逼真的图像。这一步骤的内存需求随着被预测图像的数量（批量大小）而增加。根据图像分辨率和可用的GPU VRAM，它可能相当占用内存。

这就是“切片”派上用场的地方。要解码的输入张量被分成切片，计算在几个步骤中完成。这样可以节省内存并允许更大的批量大小。

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张（1024×1024）图像。然而，使用顺序卸载，推理延迟也会增加。

缓存计算

任何基于文本条件的图像生成模型通常使用文本编码器从输入提示中计算嵌入。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的Tiny Autoencoder，可在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，同时也减少了推理延迟。

结论

总结我们的优化节省：

我们希望这些优化能让你轻松运行你最喜爱的管道。尝试这些技巧并与我们分享你的图像！🤗