Modular Diffusers 简介 - 可组合的扩散管道构建块

来源：https://huggingface.co/blog/modular-diffusers 返回文章列表 (https://huggingface.co/blog)

• 快速开始 (https://huggingface.co/blog/modular-diffusers#quickstart)
• 自定义块 (https://huggingface.co/blog/modular-diffusers#custom-blocks)
  • 编写自定义块 (https://huggingface.co/blog/modular-diffusers#writing-a-custom-block)
  • 将块组合成工作流 (https://huggingface.co/blog/modular-diffusers#composing-blocks-into-workflows)
  • 在 Hub 上共享自定义块 (https://huggingface.co/blog/modular-diffusers#sharing-custom-blocks-on-the-hub)
• 模块化仓库 (https://huggingface.co/blog/modular-diffusers#modular-repositories)
• 社区管道 (https://huggingface.co/blog/modular-diffusers#community-pipelines)
• 与 Mellon 的集成 (https://huggingface.co/blog/modular-diffusers#integration-with-mellon)
• 总结 (https://huggingface.co/blog/modular-diffusers#conclusion)
• 资源 (https://huggingface.co/blog/modular-diffusers#resources)

Modular Diffusers 引入了一种新的扩散管道构建方式，通过组合可重用的块来实现。与其说是从头开始编写整个管道，不如说你可以混合搭配各种块来创建适合你需求的工作流！这是对现有 DiffusionPipeline 类的补充，提供了一个更灵活、可组合的替代方案。

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

目录

• 快速开始 (https://huggingface.co/blog/modular-diffusers#quickstart)
• 自定义块 (https://huggingface.co/blog/modular-diffusers#custom-blocks)
• 模块化仓库 (https://huggingface.co/blog/modular-diffusers#modular-repositories)
• 社区管道 (https://huggingface.co/blog/modular-diffusers#community-pipelines)
• 与 Mellon 的集成 (https://huggingface.co/blog/modular-diffusers#integration-with-mellon)

快速开始

以下是如何使用预构建块通过 FLUX.2 Klein 4B 运行推理的简单示例：

import torch
from diffusers import ModularPipeline

# 创建模块化管道 - 这只定义工作流，模型权重还没有被加载
pipe = ModularPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B"
)
# 现在加载模型权重 — 在这一步配置数据类型、量化等参数
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")

# 生成图像 - API 与 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() 加载模型权重。

# 获取块的副本
blocks = pipe.blocks

# 弹出文本编码器块
text_blocks = blocks.sub_blocks.pop("text_encoder")

# 将其作为独立管道运行
text_pipe = text_blocks.init_pipeline("black-forest-labs/FLUX.2-klein-4B")

# 加载 text_encoder，或重用已加载的组件：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

# 从剩余的块创建新管道
# 它现在直接接受 prompt_embeds 而不是 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 文档 (https://huggingface.co/docs/diffusers/en/modular_diffusers/overview)。

自定义块

当创建你自己的块时，Modular Diffusers 真正发光。自定义块是一个 Python 类，定义其组件、输入、输出和计算逻辑——一旦定义，你就可以将其插入任何工作流。

编写自定义块

这是一个示例块，使用 Depth Anything V2 (https://huggingface.co/depth-anything/Depth-Anything-V2-Large) 从图像中提取深度图。

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 工作流并在开始处插入深度块：

# 创建 Qwen Image 管道
pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image")

print(pipe.blocks.available_workflows)
#       支持的工作流：
#        - `text2image`: 需要 `prompt`
#        - `image2image`: 需要 `prompt`, `image`
#        - `inpainting`: 需要 `prompt`, `mask_image`, `image`
#        - `controlnet_text2image`: 需要 `prompt`, `control_image`
#        - `controlnet_image2image`: 需要 `prompt`, `image`, `control_image`

# 提取 ControlNet 工作流 — 它期望一个 control_image 输入
blocks = pipe.blocks.get_workflow("controlnet_text2image")
# 显示此工作流使用的块
print(blocks)

# 在开始处插入深度块 — 其输出 (control_image)
# 自动流向需要它的 ControlNet 块
blocks.sub_blocks.insert("depth", DepthProcessorBlock(), 0)

# 你可以用 print(blocks.doc) 检查任何块的输入和输出
blocks.sub_blocks['depth'].doc

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

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

# ComponentsManager 在多个管道之间处理内存 —
# 当不使用时，它会自动将模型卸载到 CPU
manager = ComponentsManager()

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

# 深度模型从我们在 expected_components 中设置的默认路径自动加载 —
# 即使它不是 Qwen 仓库的一部分，也不需要手动加载。
# 但 controlnet 默认不包含，所以我们需要从不同的仓库加载它
controlnet = AutoModel.from_pretrained("InstantX/Qwen-Image-ControlNet-Union", torch_dtype=torch.bfloat16)
pipeline.update_components(controlnet=controlnet)

# 管道现在接受图像作为输入
image = load_image("https://huggingface.co/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 加载它。我们创建了一个模板 (https://huggingface.co/diffusers/custom-block-template) 来帮助你入门——查看构建自定义块指南 (https://huggingface.co/docs/diffusers/main/en/modular_diffusers/custom_blocks#quick-start-with-template) 了解完整的演练。

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

本文中的 DepthProcessorBlock 发布在 diffusers/depth-processor-custom-block (https://huggingface.co/diffusers/depth-processor-custom-block)——你可以直接加载和使用它：

from diffusers import ModularPipelineBlocks

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

我们发布了一个随时可用的自定义块集合 (https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks)。

模块化仓库

ModularPipeline.from_pretrained 可以直接用于任何现有的 Diffusers 仓库，但 Modular Diffusers 还引入了一种新型仓库：模块化仓库。

模块化仓库能够从其原始模型仓库引用组件。例如，diffusers/flux2-bnb-4bit-modular (https://huggingface.co/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 代码和可视化 UI 配置托管，用于 Mellon (https://huggingface.co/docs/diffusers/main/en/modular_diffusers/mellon) 等工具——所有内容都在一个地方。

社区管道

社区已经开始使用 Modular Diffusers 构建完整的管道，并在 Hub 上发布它们，包括模型权重和随时可运行的代码。

• Krea 实时视频 (https://huggingface.co/krea/krea-realtime-video)——一个 14B 参数的实时视频生成模型，从 Wan 2.1 蒸馏而来，在单个 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 (https://huggingface.co/Overworld/Waypoint-1-Small)——来自 Overworld (https://over.world/) 的 2.3B 参数实时扩散世界模型。它从控制输入和文本提示自回归生成交互式世界——你可以在消费级硬件上实时探索和与生成的环境交互。

团队可以构建新颖的架构，将其打包为块，并在 Hub 上发布整个管道，供任何人使用 ModularPipeline.from_pretrained。

查看完整的社区管道集合 (https://huggingface.co/collections/diffusers/modular-pipelines) 了解更多。

与 Mellon 的集成

💡 Mellon 仍在早期开发阶段，还未准备好用于生产。把这看作是集成工作方式的一个抢先预览！

Mellon (https://github.com/cubiq/Mellon) 是与 Modular Diffusers 集成的节点图界面。如果你对 ComfyUI 等基于节点的工具很熟悉，你会感到如鱼得水——但有一些关键区别：

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

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

例如，diffusers/FLUX.2-klein-4B-modular (https://huggingface.co/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 (https://huggingface.co/diffusers/gemini-prompt-expander-mellon)：

1. 拖入一个动态块节点并输入 repo_id（例如 diffusers/gemini-prompt-expander-mellon）
2. 点击加载自定义块——节点会自动添加一个文本框用于输入提示和一个名为“prompt“的输出套接字，所有配置都来自仓库
3. 输入短提示，将输出连接到编码提示节点，然后运行

Gemini 在生成图像前将你的短提示扩展为详细描述。无需代码，无需