稳定扩散 XL（SDXL） 是 Stability AI 最新的潜在扩散模型，用于生成高质量、逼真的图像。它解决了以前稳定扩散模型的一些挑战，例如处理手部和文本的正确性以及空间上正确的构图。此外，SDXL 还更具上下文意识，并且在生成更好的图像时需要较少的提示词。
然而，所有这些改进都以模型更大的代价为代价。有多大呢？基本的 SDXL 模型有 35 亿个参数（特别是 UNet），比以前的稳定扩散模型大约大了 3 倍。
为了探索如何优化 SDXL 的推理速度和内存使用，我们在 A100 GPU（40 GB）上进行了一些测试。对于每次推理运行，我们会生成 4 张图像，并重复 3 次。在计算推理延迟时，我们只考虑 3 次迭代中的最后一次。
因此，如果您直接使用默认精度和默认的注意力机制运行 SDXL，它将消耗 28GB 的内存并花费 72.2 秒！
from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")pipeline.unet.set_default_attn_processor()这并不太实际，而且可能会减慢您的速度，因为您通常会生成多于 4 张的图像。如果您没有更强大的 GPU，那么可能会遇到令人沮丧的内存不足错误消息。那么，我们如何优化 SDXL 以提高推理速度并减少其内存使用呢？
在🤗 Diffusers，我们有一些优化技巧和技术，可以帮助您运行内存密集型模型，如SDXL，并向您展示如何做到！我们将重点关注推理速度和内存这两个方面。
推理速度
扩散是一个随机过程，所以不能保证您会得到您喜欢的图像。通常情况下，您需要多次运行推理并进行迭代，这就是为什么优化速度至关重要的原因。本节重点介绍使用较低精度权重和结合 PyTorch 2.0 中的内存高效注意力和torch.compile来提升速度和降低推理时间。
降低精度
模型权重以某种精度存储，该精度用浮点数据类型表示。标准的浮点数据类型是 float32 (fp32)，它可以准确表示各种浮点数。在推理中，您通常不需要那么精确，所以应该使用 float16 (fp16)，它代表了较窄范围的浮点数。这意味着与 fp32 相比，fp16 只需一半的内存来存储，而且计算更容易，速度提升了一倍。此外，现代的 GPU 卡具有优化的硬件来运行 fp16 计算，使其速度更快。
使用🤗 Diffusers，您可以通过将 torch.dtype 参数指定为加载模型时要转换权重的数据类型来在推理中使用 fp16：
from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16,).to("cuda")pipeline.unet.set_default_attn_processor()与完全未经过优化的 SDXL 管道相比，使用 fp16 需要花费 21.7GB 的内存，仅需 14.8 秒。您几乎节省了一整分钟的推理时间！
内存高效注意力
在 transformer 模块中使用的注意力块可能成为巨大的瓶颈，因为随着输入序列变得越来越长，内存增加呈二次方增长。这可能很快占用大量内存，并导致内存不足的错误消息。
内存高效的注意力算法旨在减少计算注意力的内存负担，无论是通过利用稀疏性还是 tile 的方式。这些优化算法曾经主要作为第三方库来使用，需要单独安装。但是从 PyTorch 2.0 开始，情况已经改变了。PyTorch 2 引入了缩放点积注意力（SDPA），它提供了Flash Attention、内存高效的注意力（xFormers）的融合实现以及 C++ 中的 PyTorch 实现。SDPA 可能是加速推理的最简单方法：如果您使用的是 PyTorch ≥ 2.0，并且使用了🤗 Diffusers，则它会自动启用！
from diffusers import StableDiffusionXLPipelinepipeline = 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代码即时编译为更优化的推理内核。与其他编译器解决方案不同，torch.compile对现有代码的更改要求很少，只需将模型包装在该函数中即可。
通过mode参数，您可以在编译过程中优化内存开销或推理速度，这给您提供了更大的灵活性。
from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16).to("cuda")pipeline.unet = torch.compile(pipeline.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 StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipeline.enable_model_cpu_offload()与基准相比，现在它需要20.2GB的内存，节省了1.5GB的内存。
顺序CPU卸载
另一种可以节省更多内存但推理速度较慢的卸载方式是顺序CPU卸载。与将整个模型（如UNet）卸载不同，存储在不同UNet子模块中的模型权重被卸载到CPU上，并且仅在正向传递之前加载到GPU上。实质上，您每次只加载模型的部分，这样可以节省更多内存。唯一的缺点是速度明显变慢，因为您需要多次加载和卸载子模块。
from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipeline.enable_sequential_cpu_offload()与基准相比，这需要19.9GB的内存，但推理时间增加到67秒。
切片
在SDXL中，变分编码器（VAE）将由UNet预测的精炼潜变量解码为逼真的图像。此步骤的内存需求与要预测的图像数量（批大小）成比例。根据图像分辨率和可用的GPU VRAM，它可能需要相当多的内存。
这就是“切片”有用的地方。要解码的输入张量被分割成多个片段，并在几个步骤中完成解码计算。这样可以节省内存并允许更大的批大小。
pipe = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipe = pipe.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_2flush()现在嵌入向量准备好进入SDXL流水线：
from diffusers import StableDiffusionXLPipelinepipe = 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,)pipe = pipe.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大小的限制。所以，让我们使用一个较小的自动编码器！由可在the Hub上获取的madebyollin推出的小型自动编码器，只有10MB，并且它是由SDXL使用的原始VAE蒸馏而来。
from diffusers import AutoencoderTinypipe = 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 = pipe.to("cuda")通过这个设置，我们将内存需求减少到15.6GB，同时减少推理延迟。
结论
总结并归纳我们优化所带来的节省：
我们希望这些优化技术能帮助您轻松运行您喜爱的流水线。尝试这些技术，并与我们分享您的图片！🤗