Modular Diffusers: Cuộc cách mạng 'lắp ghép' trong việc xây dựng pipeline AI tạo sinh

Giới thiệu Modular Diffusers

Modular Diffusers giới thiệu một phương pháp hoàn toàn mới để xây dựng các pipeline khuếch tán (diffusion pipelines) bằng cách kết hợp các khối (block) có thể tái sử dụng. Thay vì phải viết toàn bộ pipeline từ đầu, giờ đây bạn có thể trộn và ghép các khối để tạo ra những luồng công việc (workflow) phù hợp với nhu cầu cụ thể của mình. Phương pháp này bổ sung cho lớp DiffusionPipeline hiện có bằng một giải pháp thay thế linh hoạt và có khả năng kết hợp cao hơn.

Trong bài viết này, chúng ta sẽ cùng tìm hiểu cách hoạt động của Modular Diffusers — từ API quen thuộc để chạy một pipeline module hóa, cho đến việc xây dựng các khối hoàn toàn tùy chỉnh và kết hợp chúng vào luồng công việc của riêng bạn. Chúng tôi cũng sẽ trình bày cách nó tích hợp với Mellon, một giao diện luồng công việc trực quan dựa trên node mà bạn có thể sử dụng để kết nối các khối Modular Diffusers với nhau.

Bắt đầu với Modular Diffusers

Dưới đây là một ví dụ đơn giản về cách chạy suy luận (inference) với mô hình FLUX.2 Klein 4B bằng cách sử dụng các khối được xây dựng sẵn:

import torch
from diffusers import ModularPipeline

# Tạo một pipeline module hóa - bước này chỉ định nghĩa luồng công việc, trọng số mô hình chưa được tải
pipe = ModularPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B"
)

# Tải trọng số mô hình — cấu hình dtype, lượng tử hóa (quantization), v.v. trong bước này
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")

# Tạo ảnh - API vẫn giống như DiffusionPipeline
image = pipe(
    prompt="a serene landscape at sunset",
    num_inference_steps=4,
).images[0]

image.save("output.png")

Bạn sẽ nhận được kết quả tương tự như với một DiffusionPipeline tiêu chuẩn, nhưng cấu trúc bên trong của pipeline lại rất khác biệt: nó được cấu thành từ các khối linh hoạt — mã hóa văn bản (text encoding), mã hóa hình ảnh (image encoding), khử nhiễu (denoising), và giải mã (decoding) — mà bạn có thể kiểm tra trực tiếp:

print(pipe.blocks)

Flux2KleinAutoBlocks(
...
Sub-Blocks:
[0] text_encoder (Flux2KleinTextEncoderStep)
[1] vae_encoder (Flux2KleinAutoVaeEncoderStep)
[2] denoise (Flux2KleinCoreDenoiseStep)
[3] decode (Flux2DecodeStep)
)

Sức mạnh của tính linh hoạt

Mỗi khối đều độc lập với các đầu vào và đầu ra của riêng nó. Bạn có thể chạy bất kỳ khối nào một cách riêng lẻ như một pipeline độc lập, hoặc thêm, bớt, và hoán đổi các khối một cách tự do — chúng sẽ tự động tái cấu trúc để hoạt động với các khối còn lại. Sử dụng .init_pipeline() để chuyển đổi các khối thành một pipeline có thể chạy được, và .load_components() để tải trọng số mô hình.

# Lấy một bản sao của các khối
blocks = pipe.blocks

# Tách khối text_encoder ra
text_blocks = blocks.sub_blocks.pop("text_encoder")

# Chạy nó như một pipeline riêng
text_pipe = text_blocks.init_pipeline("black-forest-labs/FLUX.2-klein-4B")
text_pipe.load_components(torch_dtype=torch.bfloat16)
text_pipe.to("cuda")

prompt_embeds = text_pipe(prompt="a serene landscape at sunset").prompt_embeds

# Tạo một pipeline mới từ các khối còn lại
# pipeline này giờ đây chấp nhận trực tiếp prompt_embeds thay vì prompt
remaining_pipe = blocks.init_pipeline("black-forest-labs/FLUX.2-klein-4B")
remaining_pipe.load_components(torch_dtype=torch.bfloat16)
remaining_pipe.to("cuda")

image = remaining_pipe(prompt_embeds=prompt_embeds, num_inference_steps=4).images[0]

Để biết thêm về các loại khối, các mẫu kết hợp, tải lười (lazy loading), và quản lý bộ nhớ với ComponentsManager, hãy tham khảo tài liệu của Modular Diffusers.

Tạo Block tùy chỉnh của riêng bạn

Modular Diffusers thực sự tỏa sáng khi bạn tạo ra các khối của riêng mình. Một khối tùy chỉnh là một lớp Python định nghĩa các thành phần (components), đầu vào, đầu ra, và logic tính toán của nó — và một khi đã được định nghĩa, bạn có thể cắm nó vào bất kỳ luồng công việc nào.

Dưới đây là một ví dụ về khối trích xuất bản đồ chiều sâu (depth map) từ hình ảnh bằng Depth Anything V2.

class DepthProcessorBlock(ModularPipelineBlocks):
    @property
    def expected_components(self):
        return [
            ComponentSpec("depth_processor", DepthPreprocessor,
                        pretrained_model_name_or_path="depth-anything/Depth-Anything-V2-Large-hf")
        ]

    @property
    def inputs(self):
        return [
            InputParam("image", required=True,
                       description="Image(s) to extract depth maps from"),
        ]

    @property
    def intermediate_outputs(self):
        return [
            OutputParam("control_image", type_hint=torch.Tensor,
                        description="Depth map(s) of input image(s)"),
        ]

    @torch.no_grad()
    def __call__(self, components, state):
        block_state = self.get_block_state(state)
        depth_map = components.depth_processor(block_state.image)
        block_state.control_image = depth_map.to(block_state.device)
        self.set_block_state(state, block_state)
        return components, state

expected_components định nghĩa các mô hình mà khối cần — trong trường hợp này là một mô hình ước tính chiều sâu. inputs và intermediate_outputs định nghĩa dữ liệu đầu vào và đầu ra. __call__ là nơi chứa logic tính toán.

Tích hợp vào luồng công việc ControlNet

Hãy sử dụng khối này với luồng công việc ControlNet của Qwen. Chúng ta sẽ trích xuất luồng công việc ControlNet và chèn khối xử lý chiều sâu vào đầu:

# Trích xuất luồng công việc ControlNet — nó yêu cầu đầu vào là control_image
pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image")
blocks = pipe.blocks.get_workflow("controlnet_text2image")

# Chèn khối xử lý chiều sâu vào đầu — đầu ra của nó (control_image)
# sẽ tự động được chuyển đến khối ControlNet cần nó
blocks.sub_blocks.insert("depth", DepthProcessorBlock(), 0)

Các khối trong một chuỗi sẽ tự động chia sẻ dữ liệu: đầu ra control_image của khối xử lý chiều sâu sẽ được chuyển đến các khối phía sau cần nó, và đầu vào image của nó trở thành một đầu vào của pipeline vì không có khối nào trước đó cung cấp nó.

from diffusers import ComponentsManager, AutoModel
from diffusers.utils import load_image

# ComponentsManager xử lý bộ nhớ trên nhiều pipeline
# nó tự động chuyển các mô hình sang CPU khi không sử dụng
manager = ComponentsManager()
pipeline = blocks.init_pipeline("Qwen/Qwen-Image", components_manager=manager)
pipeline.load_components(torch_dtype=torch.bfloat16)

# Mô hình chiều sâu được tải tự động từ đường dẫn mặc định
# Nhưng controlnet không được bao gồm mặc định, nên chúng ta cần tải nó từ một repo khác
controlnet = AutoModel.from_pretrained("InstantX/Qwen-Image-ControlNet-Union", torch_dtype=torch.bfloat16)
pipeline.update_components(controlnet=controlnet)

# pipeline giờ đây nhận 'image' làm đầu vào
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg")
output = pipeline(
    prompt="an astronaut hatching from an egg, detailed, fantasy, Pixar, Disney",
    image=image,
).images[0]

Chia sẻ và Tái sử dụng

Bạn có thể xuất bản khối tùy chỉnh của mình lên Hub để bất kỳ ai cũng có thể tải nó với trust_remote_code=True. Chúng tôi đã tạo một mẫu để bạn bắt đầu — hãy xem hướng dẫn Xây dựng Khối Tùy chỉnh để biết chi tiết.

Khối DepthProcessorBlock trong bài viết này đã được xuất bản tại diffusers/depth-processor-custom-block — bạn có thể tải và sử dụng nó trực tiếp:

from diffusers import ModularPipelineBlocks

depth_block = ModularPipelineBlocks.from_pretrained(
    "diffusers/depth-processor-custom-block", trust_remote_code=True
)

Ngoài ra, Modular Diffusers cũng giới thiệu một loại kho lưu trữ (repository) mới: Kho lưu trữ Module hóa (Modular Repository), được thiết kế để tham chiếu các thành phần từ các kho lưu trữ gốc của chúng, giúp việc quản lý và chia sẻ trở nên hiệu quả hơn.