Hugging Face 发布模块化扩散模型库

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

Integration with Mellon

下面是一个使用预构建块运行 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 类，它定义了其组件、输入、输出和计算逻辑 —— 一旦定义，你就可以将其插入任何工作流中。

Writing a Custom Block

这是一个示例块，它使用 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 是计算逻辑所在的位置。

Composing Blocks into Workflows

让我们将这个块与 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://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]

Sharing Custom Blocks on the Hub

你可以将自定义块发布到 Hub，这样任何人都可以使用 trust_remote_code=True 加载它。我们创建了一个模板来帮助你开始 —— 查看 Building Custom Blocks 指南了解完整的演练。

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 还引入了一种新的仓库类型：Modular Repository。

一个模块化仓库能够引用来自其原始模型仓库的组件。例如，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 等工具使用 —— 所有内容都在一个地方。

社区已经开始使用 Modular Diffusers 构建完整的流水线并在 Hub 上发布，包括模型权重和现成可运行的代码。

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 —— 来自 Overworld 的 2.3B 参数实时扩散世界模型。它从控制输入和文本提示自回归生成交互式世界 —— 你可以在消费级硬件上实时探索和交互生成的环境。

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

查看完整的社区流水线集合了解更多。

Integration with Mellon

💡 Mellon 仍处于早期开发阶段，尚未准备好投入生产使用。将其视为集成工作方式的预览！

Mellon 是与 Modular Diffusers 集成的节点图界面。如果你熟悉 ComfyUI 等基于节点的工具，你会感到如鱼得水 —— 但有一些关键差异：

Dynamic nodes —— 不是数十个特定模型的节点，我们有一小套节点，它们根据你选择的模型自动调整其界面。学习一次，使用任何模型。

Single-node workflows —— 感谢 Modular Diffusers 的可组合块系统，你可以将整个流水线折叠成单个节点。在同一画布上运行多个工作流，不会杂乱。

Hub integration out of the box —— 发布到 Hugging Face 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 x Modular Diffusers

自定义块集合

社区流水线集合与 Modular Diffusers

感谢 Chun Te Lee 提供的缩略图，感谢 Poli、Pedro、Lysandre、Linoy、Aritra 和 Steven 的深思熟虑的审查。

本文提到的模型 8 个

本文提到的集合 2 个

我们博客的更多文章

深入了解文本到视频模型

使用 Gradio 的可见水印

这是写得非常好的博文！感谢 diffusers 团队给我们带来了这个。