Diffusers documentation

Membangun Blok Kustom

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Membangun Blok Kustom

ModularPipelineBlocks adalah blok pembangun dasar dari sebuah ModularPipeline. Anda dapat membuat blok kustom dengan mendefinisikan input, output, dan logika komputasinya. Panduan ini menunjukkan cara membuat dan menggunakan blok kustom.

Jelajahi koleksi Modular Diffusers Custom Blocks untuk blok kustom resmi.

Struktur Proyek

Proyek blok kustom Anda harus menggunakan struktur berikut:

.
├── block.py
└── modular_config.json
  • block.py berisi implementasi blok kustom
  • modular_config.json berisi metadata yang diperlukan untuk memuat blok

Mulai Cepat dengan Template

Cara tercepat untuk membuat blok kustom adalah memulai dari template kami. Template menyediakan struktur proyek yang sudah dikonfigurasi sebelumnya dengan file block.py dan modular_config.json, serta contoh berkode yang menunjukkan cara mendefinisikan komponen, input, output, dan metode __call__—sehingga Anda dapat fokus pada logika kustom Anda alih-alih setup boilerplate.

Unduh templat

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
)

Ini menyimpan file template ke custom-block-template/ secara lokal, atau Anda bisa menggunakan local_dir untuk menyimpannya ke lokasi tertentu.

Edit secara lokal

Buka block.py dan implementasikan blok kustom Anda. Template mencakup contoh yang dikomentari yang menunjukkan cara mendefinisikan setiap properti. Lihat contoh Florence-2 di bawah ini untuk implementasi lengkap.

Uji blok Anda

from diffusers import ModularPipelineBlocks

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

Unggah ke Hub

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

Contoh: Florence-2 Image Annotator

Contoh ini membuat blok kustom dengan Florence-2 untuk memproses gambar input dan menghasilkan mask untuk inpainting.

Definisikan komponen

Definisikan komponen yang dibutuhkan blok, yaitu Florence2ForConditionalGeneration dan prosesor-nya. Saat mendefinisikan komponen, tentukan name (cara Anda mengaksesnya dalam kode), type_hint (kelas model), dan pretrained_model_name_or_path (dari mana bobot dimuat).

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

Definisikan input dan output

Input mencakup gambar, tugas anotasi, dan prompt. Output mencakup mask yang dihasilkan dan anotasi.

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

Implementasikan metode __call__

Metode __call__ berisi logika blok. Akses input melalui block_state, jalankan komputasi Anda, dan tetapkan output kembali ke 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...

Lihat implementasi lengkap di diffusers/Florence2-image-Annotator.

Menggunakan Blok Kustom

Muat blok kustom dengan from_pretrained() dan atur 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)

Gunakan blok untuk menghasilkan mask:

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

Gabungkan dengan blok lain untuk membuat pipeline baru:

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

Mengedit blok kustom

Edit blok kustom dengan mengunduhnya secara lokal. Ini adalah alur kerja yang sama dengan Quick Start with Template, tetapi dimulai dari blok yang sudah ada, bukan dari template.

Gunakan argumen local_dir untuk mengunduh blok kustom ke folder tertentu:

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

Semua perubahan yang dilakukan pada file blok di folder ini akan tercermin saat Anda memuat ulang blok tersebut. Saat Anda siap untuk membagikan perubahan Anda, unggah ke repositori baru:

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

Langkah Selanjutnya

Pelajari jenis blok
Gunakan di Mellon
Jelajahi blok yang ada

Panduan ini membahas pembuatan satu blok kustom. Pelajari cara menggabungkan beberapa blok:

Dependensi

Mendeklarasikan dependensi paket di blok kustom mencegah kesalahan impor saat runtime di kemudian hari. Diffusers memvalidasi dependensi dan mengembalikan peringatan jika sebuah paket hilang atau tidak kompatibel.

Setel atribut _requirements di kelas blok Anda, yang memetakan nama paket ke spesifikasi versi.

from diffusers.modular_pipelines import PipelineBlock

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

Ketika ada blok dengan persyaratan yang berbeda, Diffusers menggabungkan persyaratan mereka.

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

Ketika blok ini disimpan dengan save_pretrained(), persyaratan disimpan ke file modular_config.json. Ketika blok ini dimuat, Diffusers memeriksa setiap persyaratan terhadap lingkungan saat ini. Jika ada ketidakcocokan atau paket tidak ditemukan, Diffusers mengembalikan peringatan berikut.

# 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