Callback của Pipeline

Một callback được thực thi vào cuối mỗi bước khử nhiễu và có thể thay đổi trạng thái của pipeline hoặc các tensor cho các bước tiếp theo. Hãy sử dụng nó để điều chỉnh các thuộc tính hoặc biến tensor nhằm tạo ra hành vi mới mà không cần viết lại pipeline.

Các callback này hoạt động ngay với các vòng lặp DiffusionPipeline cổ điển. Chúng cũng hoạt động với các pipeline Modular Diffusers khi một khối thực thi chúng được chèn vào vòng lặp khử nhiễu; xem Modular pipelines.

Diffusers cung cấp một số callback trong tổng quan của pipeline.

Để kích hoạt một callback, hãy cấu hình thời điểm callback được thực thi sau một số bước khử nhiễu nhất định bằng một trong các đối số sau.

• cutoff_step_ratio chỉ định thời điểm kích hoạt callback dưới dạng phần trăm của tổng số bước khử nhiễu. Sử dụng khi ngưỡng cắt cần tỷ lệ với num_inference_steps (ví dụ: bỏ CFG sau 40% quá trình chạy).
• cutoff_step_index chỉ định số bước chính xác mà một callback được kích hoạt. Sử dụng khi bạn quan tâm đến một bước tuyệt đối (ví dụ, bước 10 trên lịch trình cố định 25 bước).

Ví dụ bên dưới sử dụng cutoff_step_ratio=0.4, có nghĩa là callback được kích hoạt khi quá trình khử nhiễu đạt 40% tổng số bước suy luận. SDXLCFGCutoffCallback vô hiệu hóa classifier-free guidance (CFG) sau một số bước nhất định, điều này có thể giúp tiết kiệm tài nguyên tính toán mà không ảnh hưởng đáng kể đến hiệu suất.

Định nghĩa một callback với một trong các đối số cutoff và truyền nó vào tham số callback_on_step_end trong pipeline.

import torch
from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
from diffusers.callbacks import SDXLCFGCutoffCallback

callback = SDXLCFGCutoffCallback(cutoff_step_ratio=0.4)
# if using cutoff_step_index
# callback = SDXLCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=10)

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_karras_sigmas=True)

prompt = "a sports car on the road, best quality, high quality, high detail, 8k resolution"
output = pipeline(
    prompt=prompt,
    negative_prompt="",
    guidance_scale=6.5,
    num_inference_steps=25,
    callback_on_step_end=callback,
)

Các callback chính thức tự thiết lập các đầu vào tensor của chúng. Đối với một hàm tùy chỉnh, hãy truyền callback_on_step_end_tensor_inputs như trong Hiển thị hình ảnh trung gian.

Nếu bạn muốn thêm một callback chính thức mới, hãy mở một yêu cầu tính năng hoặc gửi một PR. Ngoài ra, bạn cũng có thể tự tạo callback của riêng mình như minh họa bên dưới.

Dừng sớm

Early stopping hữu ích nếu bạn không hài lòng với các kết quả trung gian trong quá trình sinh. Callback này đặt một điểm dừng cố định bằng cách thiết lập thuộc tính _interrupt thành True, khiến vòng lặp khử nhiễu bỏ qua các bước còn lại.

import torch
from diffusers import DiffusionPipeline

def interrupt_callback(pipeline, i, t, callback_kwargs):
    stop_idx = 10
    if i == stop_idx:
        pipeline._interrupt = True
    return callback_kwargs

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image",
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
)
pipeline(
    prompt="A photo of a cat",
    num_inference_steps=50,
    callback_on_step_end=interrupt_callback,
)

Hiển thị hình ảnh trung gian

Trực quan hóa các hình ảnh trung gian rất hữu ích để theo dõi tiến độ. Bản xem trước bên dưới chỉ dành cho SDXL. Nó ánh xạ các latent của SDXL sang RGB bằng một phép biến đổi tuyến tính để xem nhanh trong quá trình khử nhiễu. Các trọng số đó không áp dụng được cho các mô hình khác. Đối với Qwen-Image và các checkpoint tương tự, hãy giải mã bằng VAE của mô hình thay vì dùng tiện ích này.

Chuyển đổi latents của Stable Diffusion XL (4 kênh) sang tensor RGB (3 kênh).

import torch
from PIL import Image
from diffusers import AutoPipelineForText2Image

def latents_to_rgb(latents):
    weights = (
        (60, -60, 25, -70),
        (60,  -5, 15, -50),
        (60,  10, -5, -35),
    )

    weights_tensor = torch.t(torch.tensor(weights, dtype=latents.dtype).to(latents.device))
    biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(latents.device)
    rgb_tensor = torch.einsum("...lxy,lr -> ...rxy", latents, weights_tensor) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
    image_array = rgb_tensor.clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)

    return Image.fromarray(image_array)

Trích xuất các latent và chuyển đổi hình ảnh đầu tiên trong batch sang RGB. Lưu hình ảnh dưới dạng tệp PNG với số bước.

def decode_tensors(pipe, step, timestep, callback_kwargs):
    latents = callback_kwargs["latents"]

    image = latents_to_rgb(latents[0])
    image.save(f"{step}.png")

    return callback_kwargs

Sử dụng callback_on_step_end_tensor_inputs để chọn các tensor mà callback nhận được, trong trường hợp này là các latents.

pipeline = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

image = pipeline(
    prompt="A croissant shaped like a cute bear.",
    negative_prompt="Deformed, ugly, bad anatomy",
    callback_on_step_end=decode_tensors,
    callback_on_step_end_tensor_inputs=["latents"],
).images[0]

Đường ống xử lý dạng module

Các pipeline dạng module không có cơ chế callback_on_step_end và interrupt. Thay vào đó, chúng sẽ có hỗ trợ streaming cấp một: huggingface/diffusers#14159, chưa được hợp nhất, thêm pipe.stream(), trả về trạng thái pipeline ở mỗi bước khử nhiễu và kết thúc quá trình chạy khi bạn dừng lặp. Hiện tại, để tái sử dụng các hàm callback hiện có, hãy chèn một khối vào vòng lặp khử nhiễu: khối này nhận BlockState của vòng lặp với chỉ số bước và timestep, và có thể đọc hoặc chỉnh sửa latents.

Sửa các khối trước khi gọi init_pipeline(): pipe.blocks trả về một bản sao, do đó việc thay đổi nó sau khi pipeline đã tồn tại sẽ không có hiệu quả.

Khối bên dưới nhận callback dưới dạng đầu vào của pipeline và gọi nó với cùng chữ ký như ở trên: pipeline, chỉ số bước, timestep và một dict chứa các tensor, trong đó các giá trị trả về sẽ được ghi ngược lại. walk_denoise_loops tìm mọi vòng lặp khử nhiễu trong các khối, bất kể chúng thuộc workflow nào, nên cùng một đoạn mã này hoạt động cho bất kỳ pipeline modular nào.

from typing import Callable

import torch
from diffusers import ModularPipeline
from diffusers.modular_pipelines import InputParam, LoopSequentialPipelineBlocks, ModularPipelineBlocks


class CallbackOnStepEnd(ModularPipelineBlocks):
    @property
    def inputs(self):
        return [
            InputParam(
                "callback_on_step_end",
                type_hint=Callable,
                description="Called after each scheduler step as `callback(pipe, i, t, {'latents': ...})`; the dict it returns is written back.",
            ),
            InputParam("latents", required=True),
        ]

    def __call__(self, components, block_state, i, t):
        if block_state.callback_on_step_end is not None:
            outputs = block_state.callback_on_step_end(components, i, t, {"latents": block_state.latents})
            block_state.latents = outputs.get("latents", block_state.latents)
        return components, block_state


def walk_denoise_loops(block):
    """Every denoising loop in the blocks, whichever workflow it belongs to."""
    if isinstance(block, LoopSequentialPipelineBlocks) and "denoiser" in block.sub_blocks:
        yield block
    for sub in getattr(block, "sub_blocks", {}).values():
        yield from walk_denoise_loops(sub)


repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
blocks = ModularPipeline.from_pretrained(repo_id).blocks
for loop in walk_denoise_loops(blocks):
    loop.sub_blocks["callback"] = CallbackOnStepEnd()  # appended after the last step of the loop
pipe = blocks.init_pipeline(repo_id)
pipe.load_components(dtype=torch.float16)
pipe.to("cuda")

callback_on_step_end hiện là một đầu vào của pipeline, và một callback được viết cho DiffusionPipeline, chẳng hạn như decode_tensors ở trên, hoạt động như bình thường:

image = pipe(prompt="A croissant shaped like a cute bear.", callback_on_step_end=decode_tensors, output="images")[0]

Dừng sớm hoạt động khác: một khối không thể bỏ qua vòng lặp mà nó đang chạy, nên việc đặt _interrupt sẽ không có tác dụng. Thay vào đó, hãy ném ngoại lệ từ callback: ngoại lệ sẽ lan ra khỏi pipe(...) (không có hình ảnh nào được giải mã), và pipeline vẫn có thể sử dụng được sau đó.

Cách này hoạt động với mọi pipeline modular ngoại trừ vòng lặp bên trong của các vòng lặp lồng nhau (các mô hình video tự hồi quy như Helios và Wan-Animate-2 khử nhiễu từng chunk bên trong một khối của vòng lặp bên ngoài). Tại đó, hãy chèn khối vào vòng lặp bên ngoài thay thế: nó chạy một lần cho mỗi chunk và nhận chỉ số chunk k thay cho i và t.