Création de blocs personnalisés

Les ModularPipelineBlocks sont les briques de base d'un ModularPipeline. Vous pouvez créer des blocs personnalisés en définissant leurs entrées, leurs sorties et leur logique de calcul. Ce guide montre comment créer et utiliser un bloc personnalisé.

Explorez la collection Modular Diffusers Custom Blocks pour les blocs personnalisés officiels.

Structure du projet

Votre projet de bloc personnalisé doit utiliser la structure suivante :

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

block.py contient l'implémentation du bloc personnalisé
modular_config.json contient les métadonnées nécessaires pour charger le bloc

Démarrage rapide avec un modèle

La méthode la plus rapide pour créer un bloc personnalisé consiste à partir de notre modèle. Le modèle fournit une structure de projet préconfigurée avec les fichiers block.py et modular_config.json, ainsi que des exemples commentés montrant comment définir les composants, les entrées, les sorties et la méthode __call__—vous permettant ainsi de vous concentrer sur votre logique personnalisée plutôt que sur la configuration de base.

Télécharger le modèle

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
)

Cela enregistre les fichiers de modèle dans custom-block-template/ localement, ou vous pouvez utiliser local_dir pour les enregistrer à un emplacement spécifique.

Modifier localement

Ouvrez block.py et implémentez votre bloc personnalisé. Le modèle contient des exemples commentés montrant comment définir chaque propriété. Consultez l'exemple Florence-2 ci-dessous pour une implémentation complète.

Testez votre bloc

from diffusers import ModularPipelineBlocks

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

Téléverser sur le Hub

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

Exemple : Annotateur d'images Florence-2

Cet exemple crée un bloc personnalisé avec Florence-2 pour traiter une image d'entrée et générer un masque pour l'inpainting.

Définir des composants

Définissez les composants dont le bloc a besoin, Florence2ForConditionalGeneration et son processeur. Lors de la définition des composants, spécifiez le name (la manière dont vous y accéderez dans le code), le type_hint (la classe du modèle) et le pretrained_model_name_or_path (l’emplacement de chargement des poids).

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

Définir les entrées et les sorties

Les entrées comprennent l'image, la tâche d'annotation et le prompt. Les sorties comprennent le masque généré et les annotations.

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

Implémenter la méthode call

La méthode __call__ contient la logique du bloc. Accédez aux entrées via block_state, exécutez votre calcul, puis réaffectez les sorties à 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...

Consultez l'implémentation complète sur diffusers/Florence2-image-Annotator.

Utilisation de blocs personnalisés

Chargez un bloc personnalisé avec from_pretrained() et définissez 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)

Utilisez le bloc pour générer un masque :

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

Composez-le avec d'autres blocs pour créer un nouveau 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")

Édition de blocs personnalisés

Modifier des blocs personnalisés en les téléchargeant localement. Il s’agit du même flux de travail que le Démarrage rapide avec un modèle, mais en partant d’un bloc existant plutôt que du modèle.

Utilisez l'argument local_dir pour télécharger un bloc personnalisé dans un dossier spécifique :

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

Toute modification apportée aux fichiers de blocs de ce dossier sera prise en compte lors du rechargement du bloc. Lorsque vous êtes prêt à partager vos modifications, téléversez-les dans un nouveau dépôt :

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

Étapes suivantes

Découvrir les types de blocs

Utiliser dans Mellon

Explorer les blocs existants

Dépendances

Déclarer les dépendances de paquets dans les blocs personnalisés permet d'éviter les erreurs d'importation au moment de l'exécution. Diffusers valide les dépendances et renvoie un avertissement si un paquet est manquant ou incompatible.

Définissez un attribut _requirements dans votre classe de bloc, associant les noms de packages aux spécificateurs de version.

from diffusers.modular_pipelines import PipelineBlock

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

Lorsqu'il existe des blocs aux exigences différentes, Diffusers fusionne leurs exigences.

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

Lorsque ce bloc est enregistré avec save_pretrained(), les dépendances sont enregistrées dans le fichier modular_config.json. Lors du chargement de ce bloc, Diffusers vérifie chaque dépendance par rapport à l’environnement actuel. En cas d’incompatibilité ou si un package n’est pas trouvé, Diffusers renvoie l’avertissement suivant.

# 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

Création de blocs personnalisés

Structure du projet

Démarrage rapide avec un modèle

Télécharger le modèle

Modifier localement

Testez votre bloc

Téléverser sur le Hub

Exemple : Annotateur d'images Florence-2

Définir des composants

Définir les entrées et les sorties

Implémenter la méthode __call__

Utilisation de blocs personnalisés

Édition de blocs personnalisés

Étapes suivantes

Dépendances

Implémenter la méthode call