커스텀 블록 만들기

ModularPipelineBlocks는 ModularPipeline의 기본 구성 요소입니다. 입력, 출력 및 계산 로직을 정의하여 사용자 정의 블록을 만들 수 있습니다. 이 가이드에서는 사용자 정의 블록을 만들고 사용하는 방법을 보여줍니다.

공식 커스텀 블록을 확인하려면 Modular Diffusers Custom Blocks 컬렉션을 탐색해 보세요.

프로젝트 구조

사용자 정의 블록 프로젝트는 다음 구조를 사용해야 합니다:

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

block.py에는 커스텀 블록 구현이 포함되어 있습니다
modular_config.json에는 블록을 로드하는 데 필요한 메타데이터가 포함되어 있습니다

Template로 빠른 시작

커스텀 블록을 만드는 가장 빠른 방법은 템플릿을 기반으로 시작하는 것입니다. 템플릿은 block.py와 modular_config.json 파일이 포함된 사전 구성된 프로젝트 구조를 제공하며, 컴포넌트, 입력, 출력 및 __call__ 메서드를 정의하는 방법에 대한 주석 달린 예제를 보여줍니다. 이를 통해 보일러플레이트 설정 대신 커스텀 로직에 집중할 수 있습니다.

Download the template

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 이미지 어노테이터

이 예제는 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")

커스텀 블록 편집

커스텀 블록을 로컬로 다운로드하여 편집합니다. 이 워크플로우는 Quick Start with Template와 동일하지만, 템플릿 대신 기존 블록에서 시작합니다.

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에서 사용

기존 블록 탐색

의존성

커스텀 블록에서 패키지 의존성을 선언하면 이후 런타임 import 오류를 방지할 수 있습니다. 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

Diffusers

커스텀 블록 만들기

프로젝트 구조

Template로 빠른 시작

Download the template

로컬에서 편집

블록 테스트

Hub에 업로드

예: Florence-2 이미지 어노테이터

컴포넌트 정의

입력과 출력을 정의합니다

__call__ 메서드를 구현합니다

사용자 정의 블록 사용

커스텀 블록 편집

다음 단계

의존성

call 메서드를 구현합니다