跳到正文
原文
Hugging Face Blog·· 2026-03-05精选AI 评分62

Hugging Face 推出 Modular Diffusers,可组合的扩散流水线构建模块

Introducing Modular Diffusers - Composable Building Blocks for Diffusion Pipelines

AI 导读

Hugging Face 发布 Modular Diffusers,用可复用 block 组合扩散流水线,作为 DiffusionPipeline 的灵活替代方案,支持添加、删除和替换 block,并可用 .init_pipeline() 把 block 转为可运行流水线。

推荐理由

原文给出可组合模块的 API 写法与自定义 block 示例,读者可据此判断现有扩散流水线能否改造成模块化流程。

正文 · AI 翻译

Modular Diffusers 引入了一种通过组合可复用模块来构建扩散管道的新方式。无需从头编写整个管道,你可以混合搭配各个模块,创建适合自己需求的工作流!它与现有的 DiffusionPipeline 类相辅相成,提供了一种更灵活、可组合的替代方案。

在本文中,我们将带你了解 Modular Diffusers 的工作原理——从使用熟悉的 API 运行模块化管道,到构建完全自定义的模块并将它们组合成你自己的工作流。我们还将展示它如何与 Mellon 集成,Mellon 是一个基于节点的可视化工作流界面,你可以用它把 Modular Diffusers 的模块连接起来。

目录

快速开始

下面是一个使用预构建模块通过 FLUX.2 Klein 4B 运行推理的简单示例:

import torch
from diffusers import ModularPipeline

# Create a modular pipeline - this only defines the workflow, model weights have not been loaded yet
pipe = ModularPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B"
)
#  Now load the model weights — configure dtype, quantization, etc in this step
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")

# Generate an image - API remains the same as DiffusionPipeline
image = pipe(
    prompt="a serene landscape at sunset",
    num_inference_steps=4,
).images[0]

image.save("output.png")

你会得到与标准 DiffusionPipeline 相同的结果,但管道在底层却截然不同:它由灵活的模块组成——文本编码、图像编码、去噪和解码——你可以直接查看它们:

print(pipe.blocks)
Flux2KleinAutoBlocks(
  ...
  Sub-Blocks:
    [0] text_encoder (Flux2KleinTextEncoderStep)
    [1] vae_encoder (Flux2KleinAutoVaeEncoderStep)
    [2] denoise (Flux2KleinCoreDenoiseStep)
    [3] decode (Flux2DecodeStep)
)

每个模块都是自包含的,拥有自己的输入和输出。你可以将任意模块作为独立管道运行,也可以自由地添加、移除和替换模块——它们会动态重组以适配保留下来的模块。使用 .init_pipeline() 将模块转换为可运行的管道,使用 .load_components() 加载模型权重。

# get a copy of the blocks
blocks = pipe.blocks

# pop out the text_encoder block
text_blocks = blocks.sub_blocks.pop("text_encoder")

# run it as its own pipeline
text_pipe = text_blocks.init_pipeline("black-forest-labs/FLUX.2-klein-4B")

# load the text_encoder, or reuse already loaded components: text_pipe.update_components(text_encoder=pipe.text_encoder)
text_pipe.load_components(torch_dtype=torch.bfloat16)
text_pipe.to("cuda")
prompt_embeds = text_pipe(prompt="a serene landscape at sunset").prompt_embeds

# create a new pipeline from the remaining blocks
# it now accepts prompt_embeds directly instead of prompt
remaining_pipe = blocks.init_pipeline("black-forest-labs/FLUX.2-klein-4B")
remaining_pipe.load_components(torch_dtype=torch.bfloat16)
remaining_pipe.to("cuda")
image = remaining_pipe(prompt_embeds=prompt_embeds, num_inference_steps=4).images[0]

关于模块类型、组合模式、懒加载以及使用 ComponentsManager 进行内存管理的更多内容,请查看 Modular Diffusers 文档。

自定义模块

Modular Diffusers 在创建你自己的模块时真正大放异彩。自定义模块是一个 Python 类,定义了它的组件、输入、输出和计算逻辑——一旦定义好,你就可以把它接入任何工作流。

编写自定义模块

下面是一个使用 Depth Anything V2 从图像中提取深度图的示例模块。

class DepthProcessorBlock(ModularPipelineBlocks):
    @property
    def expected_components(self):
        return [
            ComponentSpec("depth_processor", DepthPreprocessor,
                          pretrained_model_name_or_path="depth-anything/Depth-Anything-V2-Large-hf")
        ]

    @property
    def inputs(self):
        return [
            InputParam("image", required=True,
                       description="Image(s) to extract depth maps from"),
        ]

    @property
    def intermediate_outputs(self):
        return [
            OutputParam("control_image", type_hint=torch.Tensor,
                        description="Depth map(s) of input image(s)"),
        ]

    @torch.no_grad()
    def __call__(self, components, state):
        block_state = self.get_block_state(state)
        depth_map = components.depth_processor(block_state.image)
        block_state.control_image = depth_map.to(block_state.device)
        self.set_block_state(state, block_state)
        return components, state
  • expected_components 定义了该模块需要哪些模型——在本例中是一个深度估计模型。pretrained_model_name_or_path 参数设置了一个默认的 Hub 仓库用于加载,因此 load_components 会自动获取深度模型,除非你在 modular_model_index.json 中覆盖它。
  • inputs 和 intermediate_outputs 定义了输入和输出。
  • __call__ 是计算逻辑所在之处。

将模块组合成工作流

让我们把这个模块用在 Qwen 的 ControlNet 工作流中。提取 ControlNet 工作流,并在开头插入深度模块:

# Create Qwen Image pipeline
pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image")

print(pipe.blocks.available_workflows)
#       Supported workflows:
#        - `text2image`: requires `prompt`
#        - `image2image`: requires `prompt`, `image`
#        - `inpainting`: requires `prompt`, `mask_image`, `image`
#        - `controlnet_text2image`: requires `prompt`, `control_image`
#        - `controlnet_image2image`: requires `prompt`, `image`, `control_image`

# Extract the ControlNet workflow — it expects a control_image input
blocks = pipe.blocks.get_workflow("controlnet_text2image")
# Show the blocks this workflow uses
print(blocks)

# Insert depth block at the beginning — its output (control_image)
# automatically flows to the ControlNet block that needs it
blocks.sub_blocks.insert("depth", DepthProcessorBlock(), 0)

# You can inspect any block's inputs and outputs with print(blocks.doc)
blocks.sub_blocks['depth'].doc

序列中的模块会自动共享数据:深度模块的 control_image 输出会流向需要它的下游模块,而它的 image 输入会成为管道输入,因为没有更早的模块提供它。

blocks_composed

from diffusers import ComponentsManager, AutoModel
from diffusers.utils import load_image

# ComponentsManager handles memory across multiple pipelines —
# it automatically offloads models to CPU when not in use
manager = ComponentsManager()

pipeline = blocks.init_pipeline("Qwen/Qwen-Image", components_manager=manager)
pipeline.load_components(torch_dtype=torch.bfloat16)

# The depth model loads automatically from the default path we set in expected_components —
# no need to load it manually even though it's not part of the Qwen repo.
# But controlnet is not included by default, so we do need to load it from a different repo
controlnet = AutoModel.from_pretrained("InstantX/Qwen-Image-ControlNet-Union", torch_dtype=torch.bfloat16)
pipeline.update_components(controlnet=controlnet)

# pipeline now takes image as input
image = load_image("https://hf-mirror.com/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg")
output = pipeline(
    prompt="an astronaut hatching from an egg, detailed, fantasy, Pixar, Disney",
    image=image,
).images[0]

在 Hub 上分享自定义模块

你可以将自定义模块发布到 Hub,这样任何人都可以用 trust_remote_code=True 加载它。我们创建了一个模板帮助你入门——查看构建自定义模块指南获取完整教程。

pipeline.save_pretrained(local_dir, repo_id="your-username/your-block-name", push_to_hub=True)

本文中的 DepthProcessorBlock 发布在 diffusers/depth-processor-custom-block——你可以直接加载并使用它:

from diffusers import ModularPipelineBlocks

depth_block = ModularPipelineBlocks.from_pretrained(
    "diffusers/depth-processor-custom-block", trust_remote_code=True
)

我们已在此处发布了一系列即用型自定义模块。

模块化仓库

ModularPipeline.from_pretrained 开箱即用,可与任何现有的 Diffusers 仓库配合,但 Modular Diffusers 还引入了一种新型仓库:模块化仓库。

模块化仓库可以从其原始模型仓库引用组件。例如,diffusers/flux2-bnb-4bit-modular 包含一个量化 transformer,并从原始仓库加载其余组件。

// diffusers/flux2-bnb-4bit-modular/modular_model_index.json
{
    "transformer": [
        "diffusers", 
        "Flux2Transformer2DModel", 
        {
            "pretrained_model_name_or_path": "diffusers/flux2-bnb-4bit-modular",
            "subfolder": "transformer",
            "type_hint": ["diffusers", "Flux2Transformer2DModel"]
        }
    ],
    "vae": [
        "diffusers", 
        "AutoencoderKLFlux2", 
        {
            "pretrained_model_name_or_path": "black-forest-labs/FLUX.2-dev",
            "subfolder": "vae",
            "type_hint": ["diffusers", "AutoencoderKLFlux2"]
        }
    ],
    ...
}

模块化仓库还可以托管自定义管道块,以 Python 代码形式提供,并为类似 Mellon 的工具提供可视化 UI 配置——全部集中在一处。

社区管道

社区已经开始使用 Modular Diffusers 构建完整的管道,并将其发布到 Hub 上,附带模型权重和可直接运行的代码。

  • Krea Realtime Video —— 一个从 Wan 2.1 蒸馏而来的 14B 参数实时视频生成模型,在单张 B200 GPU 上达到 11fps。它支持文本生成视频、视频生成视频和流式视频生成视频——全部构建为模块化块。用户可以在生成过程中修改提示词、即时重新设计视频风格,并在 1 秒内看到首帧。
import torch
from diffusers import ModularPipeline

pipe = ModularPipeline.from_pretrained("krea/krea-realtime-video", trust_remote_code=True)
pipe.load_components(
    trust_remote_code=True, 
    device_map="cuda",
    torch_dtype={"default": torch.bfloat16, "vae": torch.float16}
)
  • Waypoint-1 —— 来自 Overworld 的 2.3B 参数实时扩散世界模型。它根据控制输入和文本提示自回归地生成交互式世界——你可以在消费级硬件上实时探索并与生成的环境互动。

团队可以构建新颖的架构,将其打包为块,并将整个管道发布到 Hub 上,让任何人都能通过 ModularPipeline.from_pretrained 使用。

查看完整的社区管道集合以了解更多。

与 Mellon 集成

💡 Mellon 仍处于早期开发阶段,尚未准备好投入生产使用。请将此视为集成工作原理的预览!

Mellon 是一个与 Modular Diffusers 集成的节点图界面。如果你熟悉 ComfyUI 等基于节点的工具,你会感到非常亲切——但也有一些关键区别:

  • 动态节点 —— 我们不再使用数十个特定于模型的节点,而是使用一小组节点,它们会根据你选择的模型自动调整其界面。学一次,即可用于任何模型。
  • 单节点工作流 —— 得益于 Modular Diffusers 的可组合块系统,你可以将整个管道折叠为单个节点。在同一画布上运行多个工作流,而不会显得杂乱。
  • 开箱即用的 Hub 集成 —— 发布到 HF Mirror Hub 的自定义块可立即在 Mellon 中工作。我们提供了一个实用函数,可根据你的块定义自动生成节点界面——无需 UI 代码。

这种集成之所以可行,是因为每个块都暴露相同的属性(inputs、intermediate_outputs、expected_components)。这种一致的 API 意味着 Mellon 可以从任何块定义自动生成节点的 UI,并将块组合为更高级别的节点。

例如,diffusers/FLUX.2-klein-4B-modular 包含管道定义、组件引用和 mellon_pipeline_config.json——全部在一个仓库中。在 Python 中使用 ModularPipeline.from_pretrained("diffusers/FLUX.2-klein-4B-modular") 加载它,或在 Mellon 中加载它以创建单节点或多节点工作流。

这里有一个快速示例。我们向现有的文本生成图像工作流添加一个 Gemini 提示词扩展节点——托管为 diffusers/gemini-prompt-expander-mellon 处的模块化仓库:

  1. 拖入一个 Dynamic Block 节点并输入 repo_id(即 diffusers/gemini-prompt-expander-mellon)
  2. 点击 LOAD CUSTOM BLOCK——该节点会自动为你的提示词输入生成一个文本框,以及一个名为 "prompt" 的输出套接字,全部根据仓库配置
  3. 输入一段简短提示词,将输出连接到 Encode Prompt 节点,然后运行

Gemini 会在生成图像前,将你的简短提示词扩展为详细描述。无需代码,无需配置——只需一个 Hub 仓库 id。

这只是一个示例。如需详细的分步指南,请查看 Mellon x Modular Diffusers 指南。你也可以查看这个视频!

结论

Modular Diffusers 带来了社区一直期待的可组合性与灵活性,同时不牺牲让 Diffusers 强大的各项功能。目前仍处于早期阶段——我们希望听到你的反馈,来塑造接下来的方向。试一试,告诉我们哪些做得好、哪些不好、还缺什么。

资源

感谢 Chun Te Lee 制作缩略图,感谢 Poli、Pedro、Lysandre、Linoy、Aritra 和 Steven 的悉心审阅。

来源:Hugging Face Blog · hf-mirror.com