使用大型模型

像 Stable Diffusion XL (SDXL) 这样的现代扩散模型，不仅仅是一个单一模型，而是一组多个模型的集合。SDXL 有四个不同的模型级组件：

变分自动编码器 (VAE)
两个文本编码器
用于去噪的 UNet

通常，文本编码器和去噪器比 VAE 大得多。

随着模型越来越大，性能也越来越好，你的模型可能大到连单个副本都无法放入内存。但这并不意味着它无法加载。如果你有多个 GPU，则有更多内存可用于存储你的模型。在这种情况下，最好将你的模型检查点拆分成几个较小的 检查点分片。

当文本编码器检查点有多个分片时，例如用于 SD3 的 T5-xxl，它会由 Transformers 库自动处理，因为它是使用 [StableDiffusion3Pipeline] 时 Diffusers 的必要依赖项。更具体地说，Transformers 会自动处理在请求的模型类中加载多个分片，并将其准备好以便可以执行推理。

去噪器检查点也可以有多个分片，并支持推理，这得益于 Accelerate 库。

TIP

请参考处理大型模型以进行推理指南，了解处理难以放入内存的大型模型的一般指南。

例如，让我们为 SDXL UNet 保存一个分片检查点：

python

from diffusers import UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

SDXL UNet 检查点的 fp32 变体的尺寸约为 10.4GB。将 max_shard_size 参数设置为 5GB 以创建 3 个分片。保存后，你可以在 [StableDiffusionXLPipeline] 中加载它们：

python

from diffusers import UNet2DConditionModel, StableDiffusionXLPipeline
import torch

unet = UNet2DConditionModel.from_pretrained(
    "sayakpaul/sdxl-unet-sharded", torch_dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16
).to("cuda")

image = pipeline("a cute dog running on the grass", num_inference_steps=30).images[0]
image.save("dog.png")

如果无法将所有模型级组件一次性放置在 GPU 上，请使用 [~DiffusionPipeline.enable_model_cpu_offload] 来帮助你：

diff

- pipeline.to("cuda")
+ pipeline.enable_model_cpu_offload()

一般来说，我们建议在检查点超过 5GB（fp32）时进行分片。

设备放置

在分布式设置中，你可以使用 Accelerate 在多个 GPU 上运行推理。

WARNING

此功能处于实验阶段，其 API 可能会在将来发生变化。

使用 Accelerate，你可以使用 device_map 来确定如何在多个设备上分布管道模型。这在你有不止一个 GPU 的情况下很有用。

例如，如果你有两个 8GB 的 GPU，那么使用 [~DiffusionPipeline.enable_model_cpu_offload] 可能效果不佳，因为：

它只在一个 GPU 上工作
单个模型可能无法容纳在一个 GPU 上（[~DiffusionPipeline.enable_sequential_cpu_offload] 可能有效，但速度会非常慢，而且它也仅限于单个 GPU）

为了利用两个 GPU，你可以使用“平衡”设备放置策略，该策略将模型分布到所有可用的 GPU 上。

WARNING

目前只支持“平衡”策略，我们计划在将来支持其他映射策略。

diff

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
-    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
+    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True, device_map="balanced"
)
image = pipeline("a dog").images[0]
image

你也可以传递一个字典来强制执行每个设备上可使用的最大 GPU 内存：

diff

from diffusers import DiffusionPipeline
import torch

max_memory = {0:"1GB", 1:"1GB"}
pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    use_safetensors=True,
    device_map="balanced",
+   max_memory=max_memory
)
image = pipeline("a dog").images[0]
image

如果设备不在 max_memory 中，则它将被完全忽略，不会参与设备放置。

默认情况下，Diffusers 使用所有设备的最大内存。如果模型不适合 GPU，它们将被卸载到 CPU。如果 CPU 没有足够的内存，你可能会看到错误。在这种情况下，你可以选择使用 [~DiffusionPipeline.enable_sequential_cpu_offload] 和 [~DiffusionPipeline.enable_model_cpu_offload]。

调用 [~DiffusionPipeline.reset_device_map] 重置管道的 device_map。如果你想在设备映射的管道上使用 to()、[~DiffusionPipeline.enable_sequential_cpu_offload] 和 [~DiffusionPipeline.enable_model_cpu_offload] 等方法，这也是必要的。

pipeline.reset_device_map()

管道被设备映射后，你也可以通过 hf_device_map 访问它的设备映射：

print(pipeline.hf_device_map)

一个示例设备映射如下所示：

bash

{'unet': 1, 'vae': 1, 'safety_checker': 0, 'text_encoder': 0}

使用大型模型 ​

设备放置 ​

使用大型模型

设备放置