Hugging Face's logo

Diffusers documentation

Создание пользовательских блоков

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Создание пользовательских блоков

ModularPipelineBlocks — это базовые строительные блоки ModularPipeline. Вы можете создавать собственные блоки, определяя их входные данные, выходные данные и логику вычислений. В этом руководстве показано, как создать и использовать собственный блок.

Ознакомьтесь с коллекцией Modular Diffusers Custom Blocks, чтобы найти официальные пользовательские блоки.

Структура проекта

Проект вашего пользовательского блока должен использовать следующую структуру:

.
├── block.py
└── modular_config.json
  • block.py содержит реализацию пользовательского блока
  • modular_config.json содержит метаданные, необходимые для загрузки блока

Быстрый старт с шаблоном

Самый быстрый способ создать собственный блок — начать с нашего шаблона. Шаблон предоставляет заранее настроенную структуру проекта с файлами 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 и реализуйте собственный блок. В шаблоне приведены закомментированные примеры, показывающие, как определить каждое свойство. Полную реализацию можно посмотреть в примере Florence-2 ниже.

Проверьте ваш блок

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 Image Annotator

В этом примере создаётся пользовательский блок с 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_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
Изучение существующих блоков

В этом руководстве рассматривалось создание одного пользовательского блока. Узнайте, как объединять несколько блоков:

  • SequentialPipelineBlocks: цепочка блоков для последовательного выполнения
  • ConditionalPipelineBlocks: создание условных блоков, выбирающих различные пути выполнения
  • LoopSequentialPipelineBlocks: определение итеративных рабочих процессов, таких как цикл денойзинга

Зависимости

Объявление зависимостей пакетов в пользовательских блоках предотвращает ошибки импорта во время выполнения. Diffusers проверяет зависимости и возвращает предупреждение, если пакет отсутствует или несовместим.

Установите атрибут _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

←Автоматические шаблоны docstring и параметров Использование пользовательских блоков с Mellon→