Hugging Face 推出 Modular Diffusers,可组合的扩散流水线构建模块
Introducing Modular Diffusers - Composable Building Blocks for Diffusion Pipelines
Hugging Face 发布 Modular Diffusers,用可复用 block 组合扩散流水线,作为 DiffusionPipeline 的灵活替代方案,支持添加、删除和替换 block,并可用 .init_pipeline() 把 block 转为可运行流水线。
原文给出可组合模块的 API 写法与自定义 block 示例,读者可据此判断现有扩散流水线能否改造成模块化流程。
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 输入会成为管道输入,因为没有更早的模块提供它。
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 处的模块化仓库:
- 拖入一个 Dynamic Block 节点并输入
repo_id(即diffusers/gemini-prompt-expander-mellon) - 点击 LOAD CUSTOM BLOCK——该节点会自动为你的提示词输入生成一个文本框,以及一个名为 "prompt" 的输出套接字,全部根据仓库配置
- 输入一段简短提示词,将输出连接到 Encode Prompt 节点,然后运行
Gemini 会在生成图像前,将你的简短提示词扩展为详细描述。无需代码,无需配置——只需一个 Hub 仓库 id。
这只是一个示例。如需详细的分步指南,请查看 Mellon x Modular Diffusers 指南。你也可以查看这个视频!
结论
Modular Diffusers 带来了社区一直期待的可组合性与灵活性,同时不牺牲让 Diffusers 强大的各项功能。目前仍处于早期阶段——我们希望听到你的反馈,来塑造接下来的方向。试一试,告诉我们哪些做得好、哪些不好、还缺什么。
资源
- Modular Diffusers 的概览
- Mellon
- Mellon x Modular Diffusers
- 自定义块的合集
- 使用 Modular Diffusers 的社区流水线合集
感谢 Chun Te Lee 制作缩略图,感谢 Poli、Pedro、Lysandre、Linoy、Aritra 和 Steven 的悉心审阅。
来源:Hugging Face Blog · hf-mirror.com