Diffusers documentation

Xây dựng các khối tùy chỉnh

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Xây dựng các khối tùy chỉnh

ModularPipelineBlocks là các khối xây dựng cơ bản của một ModularPipeline. Bạn có thể tạo các khối tùy chỉnh bằng cách xác định đầu vào, đầu ra và logic tính toán của chúng. Hướng dẫn này minh họa cách tạo và sử dụng một khối tùy chỉnh.

Khám phá bộ sưu tập Modular Diffusers Custom Blocks để xem các khối tùy chỉnh chính thức.

Cấu trúc dự án

Dự án khối tùy chỉnh của bạn nên sử dụng cấu trúc sau:

.
├── block.py
└── modular_config.json
  • block.py chứa cài đặt khối tùy chỉnh
  • modular_config.json chứa các metadata cần thiết để nạp khối

Bắt đầu nhanh với Mẫu

Cách nhanh nhất để tạo một khối tùy chỉnh là bắt đầu từ mẫu của chúng tôi. Mẫu cung cấp cấu trúc dự án đã được cấu hình sẵn với các tệp block.py và modular_config.json, cùng các ví dụ có chú thích cho thấy cách định nghĩa các thành phần, đầu vào, đầu ra và phương thức __call__—để bạn có thể tập trung vào logic tùy chỉnh của mình thay vì các bước thiết lập lặp lại.

Tải xuống mẫu

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
)

Điều này lưu các tệp mẫu vào custom-block-template/ cục bộ hoặc bạn có thể sử dụng local_dir để lưu vào một vị trí cụ thể.

Chỉnh sửa cục bộ

Mở block.py và triển khai block tùy chỉnh của bạn. Mẫu bao gồm các ví dụ được chú thích cho thấy cách định nghĩa từng thuộc tính. Xem ví dụ Florence-2 bên dưới để có một triển khai hoàn chỉnh.

Kiểm tra khối của bạn

from diffusers import ModularPipelineBlocks

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

Tải lên Hub

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

Ví dụ: Công cụ chú thích hình ảnh Florence-2

Ví dụ này tạo một khối tùy chỉnh với Florence-2 để xử lý một hình ảnh đầu vào và tạo ra một mask cho inpainting.

Định nghĩa các thành phần

Định nghĩa các thành phần mà khối cần, Florence2ForConditionalGeneration và bộ xử lý của nó. Khi định nghĩa các thành phần, hãy chỉ định name (cách bạn sẽ truy cập nó trong mã), type_hint (lớp mô hình), và pretrained_model_name_or_path (nơi để tải trọng số từ).

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

Định nghĩa đầu vào và đầu ra

Đầu vào bao gồm hình ảnh, tác vụ chú thích và prompt. Đầu ra bao gồm mask được tạo ra và các chú thích.

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

Thực hiện phương thức __call__

Phương thức __call__ chứa logic của khối. Truy cập đầu vào thông qua block_state, thực hiện phép tính của bạn và gán đầu ra trở lại 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...

Xem cài đặt hoàn chỉnh tại diffusers/Florence2-image-Annotator.

Sử dụng Khối Tùy Chỉnh

Tải một khối tùy chỉnh bằng from_pretrained() và đặt 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)

Sử dụng khối để tạo một mặt nạ:

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

Kết hợp với các khối khác để tạo một pipeline mới:

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

Chỉnh sửa các khối tùy chỉnh

Chỉnh sửa các khối tùy chỉnh bằng cách tải xuống cục bộ. Đây là quy trình làm việc giống với Bắt đầu nhanh với Mẫu, nhưng bắt đầu từ một khối có sẵn thay vì từ mẫu.

Sử dụng đối số local_dir để tải một khối tùy chỉnh xuống một thư mục cụ thể:

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

Mọi thay đổi được thực hiện trên các tệp khối trong thư mục này sẽ được phản ánh khi bạn tải khối lại. Khi bạn sẵn sàng chia sẻ các thay đổi của mình, hãy tải lên một kho lưu trữ mới:

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

Các bước tiếp theo

Tìm hiểu các kiểu khối
Sử dụng trong Mellon
Khám phá các khối có sẵn

Hướng dẫn này đã đề cập đến việc tạo một khối tùy chỉnh đơn lẻ. Tìm hiểu cách kết hợp nhiều khối lại với nhau:

Phụ thuộc

Việc khai báo các phụ thuộc gói trong các khối tùy chỉnh giúp ngăn chặn các lỗi import thời gian chạy sau này. Diffusers sẽ xác thực các phụ thuộc và trả về một cảnh báo nếu một gói bị thiếu hoặc không tương thích.

Đặt thuộc tính _requirements trong lớp block của bạn, ánh xạ tên gói đến bộ chỉ định phiên bản.

from diffusers.modular_pipelines import PipelineBlock

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

Khi có các khối có yêu cầu khác nhau, Diffusers sẽ gộp các yêu cầu của chúng lại.

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

Khi khối này được lưu bằng save_pretrained(), các yêu cầu sẽ được lưu vào tệp modular_config.json. Khi khối này được tải, Diffusers sẽ kiểm tra từng yêu cầu so với môi trường hiện tại. Nếu có sự không khớp hoặc một gói không được tìm thấy, Diffusers sẽ trả về cảnh báo sau.

# 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