构建自定义块

ModularPipelineBlocks 是 ModularPipeline 的基本构建模块。您可以通过定义其输入、输出和计算逻辑来创建自定义块。本指南将演示如何创建和使用自定义块。

探索 Modular Diffusers Custom Blocks 集合，以获取官方自定义块。

项目结构

您的自定义 block 项目应采用以下结构：

.
├── block.py
└── modular_config.json

block.py 包含自定义的 block 实现
modular_config.json 包含加载该块所需的元数据

使用模板快速入门

创建自定义 block 最快的方式是使用我们的模板。该模板提供了一个预配置的项目结构，包含 block.py 和 modular_config.json 文件，并附有带注释的示例，演示如何定义组件、输入、输出以及 __call__ 方法——这样你就可以专注于自定义逻辑，而无需处理样板式的设置工作。

下载模板

from diffusers import ModularPipelineBlocks

model_id = "diffusers/custom-block-template"
local_dir = model_id.split("/")[-1]

blocks = ModularPipelineBlocks.from_pretrained(
    model_id, 
    trust_remote_code=True, 
    local_dir=local_dir
)

这会将模板文件保存到本地的 custom-block-template/ 目录，或者您可以使用 local_dir 将其保存到指定位置。

在本地编辑

打开 block.py 并实现你的自定义 block。模板中包含注释示例，展示了如何定义每个属性。请参阅下方的 Florence-2 示例以查看完整实现。

测试您的 block

from diffusers import ModularPipelineBlocks

blocks = ModularPipelineBlocks.from_pretrained(local_dir, trust_remote_code=True)
pipeline = blocks.init_pipeline()
output = pipeline(...)  # your inputs here

上传到 Hub

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

示例：Florence-2 图像标注器

此示例使用 Florence-2 创建一个自定义块，用于处理输入图像并生成用于图像修复的掩码。

定义组件

定义该块所需的组件，即 Florence2ForConditionalGeneration 及其处理器。在定义组件时，需指定 name（在代码中如何访问该组件）、type_hint（模型类）以及 pretrained_model_name_or_path（加载权重的来源）。

# Inside block.py
from diffusers.modular_pipelines import ModularPipelineBlocks, ComponentSpec
from transformers import AutoProcessor, Florence2ForConditionalGeneration


class Florence2ImageAnnotatorBlock(ModularPipelineBlocks):

    @property
    def expected_components(self):
        return [
            ComponentSpec(
                name="image_annotator",
                type_hint=Florence2ForConditionalGeneration,
                pretrained_model_name_or_path="florence-community/Florence-2-base-ft",
            ),
            ComponentSpec(
                name="image_annotator_processor",
                type_hint=AutoProcessor,
                pretrained_model_name_or_path="florence-community/Florence-2-base-ft",
            ),
        ]

定义输入和输出

输入包括图像、标注任务和提示词。输出包括生成的掩码和标注。

from typing import List, Union
from PIL import Image
from diffusers.modular_pipelines import InputParam, OutputParam


class Florence2ImageAnnotatorBlock(ModularPipelineBlocks):

    # ... expected_components from above ...

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "image",
                type_hint=Union[Image.Image, List[Image.Image]],
                required=True,
                description="Image(s) to annotate",
            ),
            InputParam(
                "annotation_task",
                type_hint=str,
                default="<REFERRING_EXPRESSION_SEGMENTATION>",
                description="Annotation task to perform (e.g., <OD>, <CAPTION>, <REFERRING_EXPRESSION_SEGMENTATION>)",
            ),
            InputParam(
                "annotation_prompt",
                type_hint=str,
                required=True,
                description="Prompt to provide context for the annotation task",
            ),
            InputParam(
                "annotation_output_type",
                type_hint=str,
                default="mask_image",
                description="Output type: 'mask_image', 'mask_overlay', or 'bounding_box'",
            ),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "mask_image",
                type_hint=Image.Image,
                description="Inpainting mask for the input image",
            ),
            OutputParam(
                "annotations",
                type_hint=dict,
                description="Raw annotation predictions",
            ),
            OutputParam(
                "image",
                type_hint=Image.Image,
                description="Annotated image",
            ),
        ]

实现 call 方法

__call__ 方法包含该 block 的逻辑。通过 block_state 访问输入，执行计算，并将输出设置回 block_state。

import torch
from diffusers.modular_pipelines import PipelineState


class Florence2ImageAnnotatorBlock(ModularPipelineBlocks):

    # ... expected_components, inputs, intermediate_outputs from above ...

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        
        images, annotation_task_prompt = self.prepare_inputs(
            block_state.image, block_state.annotation_prompt
        )
        task = block_state.annotation_task
        fill = block_state.fill
        
        annotations = self.get_annotations(
            components, images, annotation_task_prompt, task
        )
        block_state.annotations = annotations
        if block_state.annotation_output_type == "mask_image":
            block_state.mask_image = self.prepare_mask(images, annotations)
        else:
            block_state.mask_image = None

        if block_state.annotation_output_type == "mask_overlay":
            block_state.image = self.prepare_mask(images, annotations, overlay=True, fill=fill)

        elif block_state.annotation_output_type == "bounding_box":
            block_state.image = self.prepare_bounding_boxes(images, annotations)

        self.set_block_state(state, block_state)

        return components, state
    
    # Helper methods for mask/bounding box generation...

完整实现请参见 diffusers/Florence2-image-Annotator。

使用自定义块

使用 from_pretrained() 加载自定义块，并设置 trust_remote_code=True。

import torch
from diffusers import ModularPipeline
from diffusers.utils import load_image

# Load the Florence-2 annotator pipeline
image_annotator = ModularPipeline.from_pretrained(
    "diffusers/Florence2-image-Annotator",
    trust_remote_code=True
)

# Check the docstring to see inputs/outputs
print(image_annotator.blocks.doc)

使用块生成掩码：

image_annotator.load_components(dtype=torch.bfloat16)
image_annotator.to("cuda")  # or "mps", "xpu", "cpu"

image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg")
image = image.resize((1024, 1024))
prompt = ["A red car"]
annotation_task = "<REFERRING_EXPRESSION_SEGMENTATION>"
annotation_prompt = ["the car"]

mask_image = image_annotator_node(
    prompt=prompt,
    image=image,
    annotation_task=annotation_task,
    annotation_prompt=annotation_prompt,
    annotation_output_type="mask_image",
).images
mask_image[0].save("car-mask.png")

将其与其他块组合，以创建新的流水线：

# Get the annotator block
annotator_block = image_annotator.blocks

# Get an inpainting workflow and insert the annotator at the beginning
inpaint_blocks = ModularPipeline.from_pretrained("Qwen/Qwen-Image").blocks.get_workflow("inpainting")
inpaint_blocks.sub_blocks.insert("image_annotator", annotator_block, 0)

# Initialize the combined pipeline
pipe = inpaint_blocks.init_pipeline()
pipe.load_components(dtype=torch.float16, device="cuda")  # or "mps", "xpu", "cpu"

# Now the pipeline automatically generates masks from prompts
output = pipe(
    prompt=prompt,
    image=image,
    annotation_task=annotation_task,
    annotation_prompt=annotation_prompt,
    annotation_output_type="mask_image",
    num_inference_steps=35,
    guidance_scale=7.5,
    strength=0.95,
    output="images"
)
output[0].save("florence-inpainting.png")

编辑自定义块

通过将其下载到本地来编辑自定义块。这与使用模板快速入门的工作流程相同，只是从现有块开始，而不是从模板开始。

使用 local_dir 参数将自定义块下载到指定文件夹：

from diffusers import ModularPipelineBlocks

# Download to a local folder for editing
annotator_block = ModularPipelineBlocks.from_pretrained(
    "diffusers/Florence2-image-Annotator",
    trust_remote_code=True,
    local_dir="./my-florence-block"
)

在此文件夹中对块文件所做的任何更改，都会在您再次加载该块时生效。当您准备好分享更改时，请上传到一个新的仓库：

pipeline = annotator_block.init_pipeline()
pipeline.save_pretrained("./my-florence-block", repo_id="your-username/my-custom-florence", push_to_hub=True)

后续步骤

了解块类型

在 Mellon 中使用

探索现有块

依赖

在自定义块中声明包依赖项可防止后续出现运行时导入错误。Diffusers 会验证这些依赖项，如果某个包缺失或不兼容，则会返回警告。

在您的 block 类中设置 _requirements 属性，将包名映射到版本说明符。

from diffusers.modular_pipelines import PipelineBlock

class MyCustomBlock(PipelineBlock):
    _requirements = {
        "transformers": ">=4.44.0",
        "sentencepiece": ">=0.2.0"
    }

当存在具有不同要求的块时，Diffusers 会合并这些要求。

from diffusers.modular_pipelines import SequentialPipelineBlocks

class BlockA(PipelineBlock):
    _requirements = {"transformers": ">=4.44.0"}
    # ...

class BlockB(PipelineBlock):
    _requirements = {"sentencepiece": ">=0.2.0"}
    # ...

pipe = SequentialPipelineBlocks.from_blocks_dict({
    "block_a": BlockA,
    "block_b": BlockB,
})

当此代码块通过 save_pretrained() 保存时，相关依赖要求会被写入 modular_config.json 文件。加载此代码块时，Diffusers 会将每项依赖要求与当前环境进行比对。如果存在不匹配或未找到相应包，Diffusers 将返回以下警告。

# missing package
xyz-package was specified in the requirements but wasn't found in the current environment.

# version mismatch
xyz requirement 'specific-version' is not satisfied by the installed version 'actual-version'. Things might work unexpected.

Update on GitHub

Diffusers

构建自定义块

项目结构

使用模板快速入门

下载模板

在本地编辑

测试您的 block

上传到 Hub

示例：Florence-2 图像标注器

定义组件

定义输入和输出

实现 __call__ 方法

使用自定义块

编辑自定义块

后续步骤

依赖

实现 call 方法