DiffusionPipeline

扩散模型由多个组件构成，例如 UNet 或扩散 Transformer（DiT）、文本编码器、变分自编码器（VAE）以及调度器。DiffusionPipeline 将所有这些组件封装为一个易于使用的统一 API，同时保留了修改其内部组件的灵活性。

本指南将向您展示如何加载 DiffusionPipeline。

加载 pipeline

DiffusionPipeline 是一个基础流水线类，它会扫描 model_index.json 文件中的类名，自动选择并返回模型流水线子类的实例，例如 QwenImagePipeline。

将模型 ID 传递给 from_pretrained() 以加载流水线。

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

每个模型都有一个特定的流水线子类，该子类继承自 DiffusionPipeline。子类通常具有较窄的适用范围，并针对特定任务进行设计。下表给出了一个示例。

pipeline 子类	任务
QwenImagePipeline	文本到图像
QwenImageImg2ImgPipeline	图像到图像
QwenImageInpaintPipeline	inpaint

您可以直接通过向 from_pretrained() 传递模型 ID 来使用该子类。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

请参阅单文件格式文档，了解如何加载单文件模型。

本地流水线

Pipelines 也可以在本地运行。使用 snapshot_download 下载模型仓库。

from huggingface_hub import snapshot_download

local_dir = snapshot_download(repo_id="Qwen/Qwen-Image")

将该路径传递给 from_pretrained() 以加载它。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  local_dir, dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

当 from_pretrained() 方法检测到本地路径时，不会从 Hub 下载文件。但这也意味着，它同样不会下载并缓存对该模型所做的任何更新。

Pipeline 数据类型

使用 from_pretrained() 中的 dtype 参数，可以以特定的数据类型加载模型。这允许您以不同的精度加载不同的模型。例如，以半精度加载大型 transformer 模型可以减少所需的内存。

将每个模型的数据类型以字典形式传递给 dtype。使用 default 键来设置默认数据类型。如果某个模型不在字典中且未提供 default，则该模型将以全精度（torch.float32）加载。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image",
  dtype={"transformer": torch.bfloat16, "default": torch.float16},
)
print(pipeline.transformer.dtype, pipeline.vae.dtype)

如果您以相同的数据类型加载所有模型，则无需使用字典。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16
)
print(pipeline.transformer.dtype, pipeline.vae.dtype)

设备放置

device_map 参数用于将流水线放置到设备上。在运行时，Diffusers 接受 "balanced"、"cpu" 以及当前机器对应的加速器字符串（通常为 "cuda"、"mps" 或 "xpu"）。

参数	描述
accelerator 字符串	将 pipeline 放置到该设备上。请传入你的机器支持的加速器，例如 `"cuda"`、`"mps"` 或 `"xpu"`。
`"balanced"`	将流水线组件分散到可见的 GPU 上。当您需要设置每个设备的内存上限时，请配合使用 `max_memory`。
`"cpu"`	将流水线放置在 CPU 上。

使用 from_pretrained() 中的 max_memory 参数，为每个设备分配可使用的最大内存量。默认情况下，Diffusers 会使用可用的最大内存。

import torch
from diffusers import DiffusionPipeline

max_memory = {0: "16GB", 1: "16GB"}
pipeline = DiffusionPipeline.from_pretrained(
  "Qwen/Qwen-Image",
  dtype=torch.bfloat16,
  device_map="balanced",
  max_memory=max_memory,
)

hf_device_map 属性允许您访问并查看 device_map。

print(pipeline.hf_device_map)
# {'transformer': 1, 'vae': 1, 'text_encoder': 0}

使用 reset_device_map() 方法重置 pipeline 的 device_map。如果您希望使用 .to()、enable_sequential_cpu_offload() 和 enable_model_cpu_offload() 等方法，则必须执行此操作。

pipeline.reset_device_map()

并行加载

大型模型通常会被分片为较小的文件，以便于加载。Diffusers 支持并行加载分片，以加快加载过程。

将 HF_ENABLE_PARALLEL_LOADING 设置为 "YES" 以启用分片的并行加载。

device_map 参数应设置为 "cuda"，以便根据模型大小预先分配一大块内存。这能显著缩短模型加载时间，因为内存分配器的预热过程避免了后续大量较小的分配器调用。

import os
import torch
from diffusers import DiffusionPipeline

os.environ["HF_ENABLE_PARALLEL_LOADING"] = "YES"

pipeline = DiffusionPipeline.from_pretrained(
  "Wan-AI/Wan2.2-I2V-A14B-Diffusers", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

在流水线中替换模型

DiffusionPipeline 具有灵活性，支持加载其他模型或调度器。您可以尝试不同的调度器，以优化生成速度或质量，也可以将模型替换为性能更优的模型。

以下示例使用了更稳定的 VAE 版本。

import torch
from diffusers import DiffusionPipeline, AutoModel

vae = AutoModel.from_pretrained(
  "madebyollin/sdxl-vae-fp16-fix", dtype=torch.float16
)

pipeline = DiffusionPipeline.from_pretrained(
  "stabilityai/stable-diffusion-xl-base-1.0",
  vae=vae,
  dtype=torch.float16,
  device_map="cuda"  # or "mps", "xpu", "cpu"
)

在多个流水线中复用模型

在处理使用同一模型的多个流水线时，from_pipe() 方法允许复用模型，而无需每次重新加载。这样，您可以在不增加内存占用的情况下使用多个流水线。

内存使用量由内存需求最高的流水线决定，与流水线的数量无关。

以下示例加载一个 pipeline，然后使用 from_pipe() 加载第二个 pipeline，以利用 perturbed-attention guidance (PAG) 来提升生成质量。

请使用 AutoPipelineForText2Image，因为 DiffusionPipeline 不支持 PAG。有关更多信息，请参阅 AutoPipeline 文档。

import torch
from diffusers import AutoPipelineForText2Image

pipeline_sdxl = AutoPipelineForText2Image.from_pretrained(
  "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
image = pipeline_sdxl(prompt).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Max memory reserved: 10.47 GB

在第二个 pipeline 中设置 enable_pag=True 以启用 PAG。第二个 pipeline 使用的内存量与第一个相同，因为它与第一个共享模型权重。

pipeline = AutoPipelineForText2Image.from_pipe(
  pipeline_sdxl, enable_pag=True
)
prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
image = pipeline(prompt).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Max memory reserved: 10.47 GB

由 from_pipe() 创建的管道共享相同的模型和 state。在一个管道中修改某个模型的状态，会影响所有共享该模型的其他管道。

Some methods may not work correctly on pipelines created with from_pipe(). For example, enable_model_cpu_offload() relies on a unique model execution order, which may differ in the new pipeline. To ensure proper functionality, reapply these methods on the new pipeline.

后续步骤

加载流水线后，您通常会调整去噪调度、更换权重格式或附加适配器。

Schedulers 涵盖去噪算法的切换与配置。
模型格式涵盖了 GGUF、单文件检查点以及其他权重布局。
减少内存使用涵盖了卸载及其他内存工具。
LoRA 涵盖了在管道上加载适配器的内容。
旧版检查点介绍了用于旧版 Stable Diffusion 检查点的安全检查器。

Update on GitHub