Creación de bloques personalizados

Los ModularPipelineBlocks son los bloques de construcción fundamentales de un ModularPipeline. Puede crear bloques personalizados definiendo sus entradas, salidas y lógica de cálculo. Esta guía demuestra cómo crear y usar un bloque personalizado.

Explora la colección Modular Diffusers Custom Blocks para ver los bloques personalizados oficiales.

Estructura del proyecto

Su proyecto de bloque personalizado debe usar la siguiente estructura:

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

block.py contiene la implementación personalizada del bloque
modular_config.json contiene los metadatos necesarios para cargar el bloque

Inicio rápido con Plantilla

La forma más rápida de crear un bloque personalizado es partir de nuestra plantilla. La plantilla proporciona una estructura de proyecto preconfigurada con los archivos block.py y modular_config.json, además de ejemplos comentados que muestran cómo definir componentes, entradas, salidas y el método __call__, de modo que puedas centrarte en tu lógica personalizada en lugar de en la configuración inicial.

Descargar la plantilla

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
)

Esto guarda los archivos de plantilla en custom-block-template/ localmente, o bien puedes usar local_dir para guardarlos en una ubicación específica.

Editar localmente

Abre block.py e implementa tu bloque personalizado. La plantilla incluye ejemplos comentados que muestran cómo definir cada propiedad. Consulta el ejemplo de Florence-2 a continuación para una implementación completa.

Prueba tu bloque

from diffusers import ModularPipelineBlocks

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

Subir al Hub

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

Ejemplo: Anotador de imágenes Florence-2

Este ejemplo crea un bloque personalizado con Florence-2 para procesar una imagen de entrada y generar una máscara para inpainting.

Definir componentes

Define los componentes que necesita el bloque, Florence2ForConditionalGeneration y su procesador. Al definir los componentes, especifica name (cómo lo accederás en el código), type_hint (la clase del modelo) y pretrained_model_name_or_path (de dónde cargar los pesos).

# 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",
            ),
        ]

Definir entradas y salidas

Las entradas incluyen la imagen, la tarea de anotación y el prompt. Las salidas incluyen la máscara generada y las anotaciones.

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",
            ),
        ]

Implementar el método call

El método __call__ contiene la lógica del bloque. Acceda a las entradas a través de block_state, ejecute su cálculo y establezca las salidas de nuevo en 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...

Vea la implementación completa en diffusers/Florence2-image-Annotator.

Uso de bloques personalizados

Cargue un bloque personalizado con from_pretrained() y establezca 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)

Use el bloque para generar una máscara:

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")

Componga con otros bloques para crear una nueva pipeline:

# 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")

Edición de bloques personalizados

Edite los bloques personalizados descargándolos localmente. Este es el mismo flujo de trabajo que el Quick Start with Template, pero comenzando desde un bloque existente en lugar de la plantilla.

Use el argumento local_dir para descargar un bloque personalizado a una carpeta específica:

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"
)

Cualquier cambio realizado en los archivos de bloques de esta carpeta se reflejará al volver a cargar el bloque. Cuando estés listo para compartir tus cambios, súbelos a un nuevo repositorio:

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

Próximos pasos

Conocer los tipos de bloques

Uso en Mellon

Explorar bloques existentes

Dependencias

Declarar las dependencias de paquetes en bloques personalizados evita errores de importación en tiempo de ejecución. Diffusers valida las dependencias y devuelve una advertencia si falta un paquete o es incompatible.

Establezca un atributo _requirements en su clase de bloque, que mapee los nombres de paquetes a especificadores de versión.

from diffusers.modular_pipelines import PipelineBlock

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

Cuando hay bloques con requisitos diferentes, Diffusers fusiona sus requisitos.

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,
})

Cuando este bloque se guarda con save_pretrained(), los requisitos se guardan en el archivo modular_config.json. Cuando este bloque se carga, Diffusers comprueba cada requisito contra el entorno actual. Si hay una discrepancia o no se encuentra un paquete, Diffusers devuelve la siguiente advertencia.

# 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

Creación de bloques personalizados

Estructura del proyecto

Inicio rápido con Plantilla

Descargar la plantilla

Editar localmente

Prueba tu bloque

Subir al Hub

Ejemplo: Anotador de imágenes Florence-2

Definir componentes

Definir entradas y salidas

Implementar el método __call__

Uso de bloques personalizados

Edición de bloques personalizados

Próximos pasos

Dependencias

Implementar el método call