Diffusers documentation

ModularPipeline

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

ModularPipeline

ModularPipeline 将 ModularPipelineBlocks 转换为可执行的流水线,负责加载模型并执行块中定义的各个计算步骤。它是运行流水线的主要接口,其 API 与 DiffusionPipeline 非常相似,但存在几处关键差异。

  • Loading is lazy. With DiffusionPipeline, from_pretrained() creates the pipeline and loads all models at the same time. With ModularPipeline, creating and loading are two separate steps: from_pretrained() reads the configuration and knows where to load each component from, but doesn’t actually load the model weights. You load the models later with load_components(), which is where you pass loading arguments like dtype and quantization_config.

  • 创建 pipeline 的两种方式。 你可以使用 from_pretrained() 加载现有的 diffusers 模型仓库——它会自动映射到默认的 pipeline 模块,然后转换为 ModularPipeline,无需额外配置。你可以查看 modular_pipelines_directory 以了解当前支持的模型。你也可以使用 ModularPipelineBlocks 自行组装 pipeline,并通过 init_pipeline() 方法进行转换(有关更多详细信息,请参阅 创建 pipeline)。

  • 运行 pipeline 的方式相同。加载完成后,你可以使用熟悉的参数调用 pipeline。当 pipeline 块使用 AutoPipelineBlocks 根据输入自动选择工作流时,单个 ModularPipeline 可以支持多种工作流(文生图、图生图、图像修复等)。

以下是使用 SDXL 进行文生图、图生图和图像修复的完整示例。

文本到图像
图像到图像
inpainting
import torch
from diffusers import ModularPipeline

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

image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k").images[0]
image.save("modular_t2i_out.png")

本指南将向您展示如何创建 ModularPipeline、管理其组件并运行该流水线。

创建流水线

There are two ways to create a ModularPipeline. Assemble and create a pipeline from ModularPipelineBlocks with init_pipeline(), or load an existing pipeline with from_pretrained().

您还可以初始化一个 ComponentsManager 来处理设备放置和内存管理。如果不需要自动卸载,您可以跳过此步骤,并使用 pipeline.to("cuda") 手动将流水线移动到目标设备。

有关 ComponentsManager 如何帮助在不同工作流中管理组件的更多详细信息,请参阅其文档。

init_pipeline

init_pipeline() 会将任意 ModularPipelineBlocks 转换为 ModularPipeline。

让我们定义一个最小化的块,来看看它是如何工作的:

from transformers import CLIPTextModel
from diffusers.modular_pipelines import (
    ComponentSpec,
    ModularPipelineBlocks,
    PipelineState,
)

class MyBlock(ModularPipelineBlocks):
    @property
    def expected_components(self):
        return [
            ComponentSpec(
                name="text_encoder",
                type_hint=CLIPTextModel,
                pretrained_model_name_or_path="openai/clip-vit-large-patch14",
            ),
        ]

    def __call__(self, components, state: PipelineState) -> PipelineState:
        return components, state

调用 init_pipeline() 将其转换为 pipeline。pipeline 上的 blocks 属性即为其创建时所用的 blocks——它决定了预期的输入、输出以及计算逻辑。

block = MyBlock()
pipe = block.init_pipeline()
pipe.blocks
MyBlock {
  "_class_name": "MyBlock",
  "_diffusers_version": "0.37.0.dev0"
}

Blocks 是可变的——在创建 pipeline 之前,您可以自由地添加、删除或替换 blocks。但是,一旦创建了 pipeline,修改 pipeline.blocks 不会影响该 pipeline,因为它返回的是一个副本。如果您希望使用不同的 block 结构,请在修改 blocks 后创建一个新的 pipeline。

调用 init_pipeline() 时如果不指定仓库,它会使用块中 ComponentSpec 定义的 pretrained_model_name_or_path 来确定从何处加载各个组件。打印 pipeline 可以查看组件的加载配置。

pipe
ModularPipeline {
  "_blocks_class_name": "MyBlock",
  "_class_name": "ModularPipeline",
  "_diffusers_version": "0.37.0.dev0",
  "text_encoder": [
    null,
    null,
    {
      "pretrained_model_name_or_path": "openai/clip-vit-large-patch14",
      "revision": null,
      "subfolder": "",
      "type_hint": [
        "transformers",
        "CLIPTextModel"
      ],
      "variant": null
    }
  ]
}

如果向 init_pipeline() 传入一个仓库,它会通过将该模块的组件与仓库中的流水线配置(model_index.json 或 modular_model_index.json)进行匹配,来覆盖加载路径。有关加载规范的记录与保存方式,请参阅 Modular repository。

在下面的示例中,pretrained_model_name_or_path 将被更新为 "stabilityai/stable-diffusion-xl-base-1.0"。

pipe = block.init_pipeline("stabilityai/stable-diffusion-xl-base-1.0")
pipe
ModularPipeline {
  "_blocks_class_name": "MyBlock",
  "_class_name": "ModularPipeline",
  "_diffusers_version": "0.37.0.dev0",
  "text_encoder": [
    null,
    null,
    {
      "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
      "revision": null,
      "subfolder": "text_encoder",
      "type_hint": [
        "transformers",
        "CLIPTextModel"
      ],
      "variant": null
    }
  ]
}

如果块中的某个组件在仓库中不存在,它将保持为 null,并在 load_components() 期间被跳过。

from_pretrained

from_pretrained() 提供了一种便捷的方式来创建 ModularPipeline,而无需自行定义各个 block。

它支持三种类型的仓库。

一个常规的 diffusers 仓库。 传入任何受支持的模型仓库,它都会自动映射到默认的流水线模块。目前支持的模型包括 SDXL、Wan、Qwen、Z-Image、Flux 和 Flux2。

from diffusers import ModularPipeline, ComponentsManager

components = ComponentsManager()
pipeline = ModularPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", components_manager=components
)

模块化仓库。这些仓库包含一个 modular_model_index.json 文件,用于指定从何处加载各个组件——这些组件可以来自不同的仓库,且模块化仓库本身可能不包含任何模型权重。例如,diffusers/flux2-bnb-4bit-modular 从一个仓库加载量化后的 transformer,从另一个仓库加载其余组件。有关该格式的更多详细信息,请参阅 模块化仓库。

from diffusers import ModularPipeline, ComponentsManager

components = ComponentsManager()
pipeline = ModularPipeline.from_pretrained(
    "diffusers/flux2-bnb-4bit-modular", components_manager=components
)

包含自定义代码的模块化仓库。 部分仓库在加载配置之外还包含自定义的 pipeline 块。请添加 trust_remote_code=True 以加载它们。有关如何创建自己的块,请参阅 自定义块。

from diffusers import ModularPipeline, ComponentsManager

components = ComponentsManager()
pipeline = ModularPipeline.from_pretrained(
    "diffusers/Florence2-image-Annotator", trust_remote_code=True, components_manager=components
)

当 pipeline 块定义了工作流(参见 pipeline.blocks.available_workflows)时,传入 workflow= 以仅保留该工作流的块——这与 get_workflow() 的剪枝方式相同。随后,pipeline 仅声明该工作流所使用的组件,其 docstring 也将准确描述该工作流的输入。

pipeline = ModularPipeline.from_pretrained("Qwen/Qwen-Image", workflow="inpainting")

加载组件

ModularPipeline 不会自动实例化组件。它仅加载配置和组件规范。您可以使用 load_components() 来加载组件。

这将加载所有具有有效加载规范的组件。

import torch

pipeline.load_components(dtype=torch.float16)

您也可以通过名称加载特定组件。以下示例仅加载 text_encoder。

pipeline.load_components(names=["text_encoder"], dtype=torch.float16)

在定义了工作流的管道中,workflow= 仅加载该工作流所使用的组件。由于管道会保留其所有块,因此这是跨工作流运行同一管道的便捷方式:每次调用只会补充新工作流仍缺失的部分。

pipeline.load_components(workflow="inpainting", dtype=torch.float16)

加载完成后,打印 pipeline 会显示已加载的组件——前两个字段会从 null 变为该组件的库和类。

pipeline
# text_encoder is loaded - shows library and class
"text_encoder": [
  "transformers",
  "CLIPTextModel",
  { ... }
]

# unet is not loaded yet - still null
"unet": [
  null,
  null,
  { ... }
]

诸如 dtype、variant、revision 和 quantization_config 等关键字参数会传递给每个组件的 from_pretrained()。您可以传入单个值以应用于所有组件,或传入字典以设置每个组件的值。

# apply bfloat16 to all components
pipeline.load_components(dtype=torch.bfloat16)

# different dtypes per component
pipeline.load_components(dtype={"transformer": torch.bfloat16, "default": torch.float32})

load_components() 仅加载尚未加载且具有有效加载规范的组件。这意味着,如果您已在管道上设置了某个组件,再次调用 load_components() 不会重新加载它。

更新组件

update_components() 用于将流水线中的某个组件替换为新组件。当组件被更新时,流水线配置中的加载规范也会随之更新,load_components() 在后续调用时会跳过该组件。

从 AutoModel

您可以传入使用 AutoModel.from_pretrained() 加载的模型对象。以这种方式加载的模型会自动标记其加载信息。

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "RunDiffusion/Juggernaut-XL-v9", subfolder="unet", variant="fp16", dtype=torch.float16
)
pipeline.update_components(unet=unet)

来自 ComponentSpec

使用 get_component_spec() 获取当前组件规范的副本,对其进行修改,然后加载新的组件。

unet_spec = pipeline.get_component_spec("unet")

# modify to load from a different repository
unet_spec.pretrained_model_name_or_path = "RunDiffusion/Juggernaut-XL-v9"

# load and update
unet = unet_spec.load(dtype=torch.float16)
pipeline.update_components(unet=unet)

您也可以从零开始创建一个 ComponentSpec。

并非所有组件都从预训练权重加载——有些组件是根据配置创建的(列在 pipeline.config_component_names 下)。对于这些组件,请使用 create() 而不是 load()。

guider_spec = pipeline.get_component_spec("guider")
guider_spec.config = {"guidance_scale": 5.0}
guider = guider_spec.create()
pipeline.update_components(guider=guider)

或者直接传递该对象。

from diffusers.guiders import ClassifierFreeGuidance

guider = ClassifierFreeGuidance(guidance_scale=5.0)
pipeline.update_components(guider=guider)

有关可用 guider 及其配置方式的详细信息,请参阅 Guiders 指南。

将流水线拆分为多个阶段

由于 block 是可组合的,您可以将 pipeline 拆解,并为每个阶段重建独立的 pipeline。下面的示例展示了如何将 text encoder block 从 pipeline 的其余部分中分离出来,以便您可以独立地对 prompt 进行编码,并将 embeddings 传递给主 pipeline。

from diffusers import ModularPipeline, ComponentsManager
import torch

device = "cuda"  # or "mps", "xpu", "cpu"
dtype = torch.bfloat16
repo_id = "black-forest-labs/FLUX.2-klein-4B"

# get the blocks and separate out the text encoder
blocks = ModularPipeline.from_pretrained(repo_id).blocks
text_block = blocks.sub_blocks.pop("text_encoder")

# use ComponentsManager to handle offloading across multiple pipelines
manager = ComponentsManager()
manager.enable_auto_cpu_offload(device=device)

# create separate pipelines for each stage
text_encoder_pipeline = text_block.init_pipeline(repo_id, components_manager=manager)
pipeline = blocks.init_pipeline(repo_id, components_manager=manager)

# encode text
text_encoder_pipeline.load_components(dtype=dtype)
text_embeddings = text_encoder_pipeline(prompt="a cat").get_by_kwargs("denoiser_input_fields")

# denoise and decode
pipeline.load_components(dtype=dtype)
output = pipeline(
    **text_embeddings,
    num_inference_steps=4,
).images[0]

ComponentsManager 负责在多个流水线之间管理内存。与 DiffusionPipeline 中遵循固定顺序的卸载策略不同,ComponentsManager 会在每次模型前向传播运行时,根据当前的内存状况动态做出卸载决策。这意味着无论您创建了多少个流水线,或以何种顺序运行它们,它都能正常工作。有关更多详细信息,请参阅 ComponentsManager 指南。

如果流水线阶段共享组件(例如,用于编码和解码的同一个 VAE),您可以使用 update_components() 将已加载的组件传递给另一个流水线,而无需再次加载。

模块化仓库

如果流水线块使用了预训练组件,则必须提供一个仓库。该仓库用于提供加载规范和元数据。

ModularPipeline 开箱即可与常规的 diffusers 仓库配合使用。不过,为了获得更高的灵活性,您也可以创建一个模块化仓库。模块化仓库包含一个modular_model_index.json文件,其中包含以下 3 个元素。

  • library 和 class 分别显示该组件加载自哪个库以及其所属的类。如果为 null,则表示该组件尚未加载。
  • loading_specs_dict 包含加载组件所需的信息,例如加载该组件的仓库和子文件夹。

模块化仓库的关键优势在于,各个组件可以从不同的仓库加载。例如,diffusers/flux2-bnb-4bit-modular 从 diffusers/FLUX.2-dev-bnb-4bit 加载量化后的 transformer,同时从 black-forest-labs/FLUX.2-dev 加载其余组件。

要将普通的 diffusers 仓库转换为模块化仓库,请使用普通仓库创建 pipeline,然后将其推送到 Hub。保存的仓库将包含一个 modular_model_index.json,其中包含所有加载规范。

from diffusers import ModularPipeline

# load from a regular repo
pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")

# push as a modular repository
pipeline.save_pretrained("local/path", repo_id="my-username/sdxl-modular", push_to_hub=True)

默认情况下,save_pretrained() 会写入 Diffusers 能够序列化的每个当前已加载的组件。未加载或无法序列化的组件不会被写入,并保留其现有的加载规范。这为您提供两种保存方式,具体取决于您的需求。

保存自包含的副本

加载所有组件,然后保存。每个 spec 都指向该结果,因此它会从单一位置完整重新加载,包括离线场景。

pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3")
pipe.load_components()
pipe.save_pretrained("path/to/local-copy")

modular_model_index.json 中记录的保存位置是 save_pretrained() 写入的位置:本地保存时为 save_directory,传入 push_to_hub=True 时为 repo_id。

在 huggingface_hub>=1.32.0 中,基于 Xet 的相同文件可以在任意 Hub 仓库之间共享同一个缓存负载。这包括分片检查点中的单个文件:即使其他分片不同,具有相同 Xet 哈希的分片也可以被复用。两次下载必须使用支持符号链接的相同 Hub 缓存目录。更改 dtype、序列化方式或分片边界可能会改变文件哈希,从而阻止复用。这可以节省下载流量和磁盘空间,但不会节省已加载组件所占用的内存。

共享缓存采用尽力而为的策略。设置 HF_HUB_DISABLE_SHARED_BLOBS=1 后仍支持加载,但会缓存独立的副本。modular_model_index.json 中现有的组件引用保持不变;对于未存储在模块化仓库中的组件,请保留其外部引用。

更新现有的模块化索引

迁移是可选的:对原始仓库的引用已经会复用其缓存文件。如果某个引用仅用于避免重复下载,且你希望该组件改为从模块化仓库加载,请更新该引用。

首先,检查模块化仓库是否包含相同的组件配置以及所有必需的权重文件,包括权重索引以及分片检查点中的每个分片。匹配的 Xet 哈希值决定了哪些权重文件可以复用共享缓存。保留对缺失的组件或有意从其他检查点加载的组件的引用。

仅编辑 modular_model_index.json 本地副本中相关的加载规范。例如,要在模块化仓库中使用存储在 vae/ 下的 VAE,且保持相同的权重变体:

import json
from pathlib import Path

index_path = Path("path/to/local-copy/modular_model_index.json")
index = json.loads(index_path.read_text())
spec = index["vae"][2]
spec["pretrained_model_name_or_path"] = "my-username/my-modular-repo"
spec.pop("repo", None)
spec["subfolder"] = "vae"
spec["revision"] = None
index_path.write_text(json.dumps(index, indent=2) + "\n")

这同时移除了旧版的 repo 字段。源仓库中的某个修订版无法标识目标仓库中的修订版:对于目标仓库的默认分支,请使用 None;或者,请固定一个包含该组件的目标修订版。如果文件名不同,请调整 subfolder 和 variant。

在发布更新后的索引之前,请使用 ModularPipeline.from_pretrained("path/to/local-copy") 加载它,调用 load_components(names="vae"),并验证该组件是否以预期的参数加载。仅上传编辑后的索引;重新保存权重可能会更改其哈希值,从而阻止缓存复用。

保留对现有组件的引用

仅加载新增内容(或完全不加载)。只有已加载的组件会被保存,其余组件仍保持为指向其原始仓库的指针。当您希望替换某个组件,同时继续从原始仓库加载其他组件时,请使用此模式。例如,保存一个自定义的 transformer,而其余组件继续从基础仓库加载。

pipe = ModularPipeline.from_pretrained("black-forest-labs/FLUX.2-dev")
pipe.update_components(transformer=my_custom_transformer)  # the only component in memory
pipe.save_pretrained("local/path", repo_id="my-username/flux2-custom-transformer", push_to_hub=True)

传入 overwrite_modular_index=False 可保持 modular_model_index.json 中的加载规范不变。对于已保存的组件,如果其加载规范为空,仍会填入目标位置,因为此时没有需要保留的内容。

以其他方式移动文件不会重写索引。使用 hf download --local-dir 下载的副本仍然可用:当从本地目录加载 pipeline 时,该目录中存在的每个组件都会从该目录加载,而不是从记录的仓库加载。使用 hf upload 上传的副本仍指向旧位置,因此在这种情况下需要手动更新索引。

模块化仓库还可以包含以 Python 代码编写的自定义流水线块。这允许您共享那些并非 Diffusers 原生的专用块。例如,diffusers/Florence2-image-Annotator 在加载配置之外还包含自定义块:

Florence2-image-Annotator/
├── block.py                    # Custom pipeline blocks implementation
├── config.json                 # Pipeline configuration and auto_map
├── mellon_config.json          # UI configuration for Mellon
└── modular_model_index.json    # Component loading specifications

config.json 文件包含一个 auto_map 键,用于告诉 ModularPipeline 在哪里找到自定义块:

{
  "_class_name": "Florence2AnnotatorBlocks",
  "auto_map": {
    "ModularPipelineBlocks": "block.Florence2AnnotatorBlocks"
  }
}

如 from_pretrained 中所示,使用 trust_remote_code=True 加载自定义代码仓库。有关如何创建和共享自己的代码,请参阅 Custom blocks。

Update on GitHub