Обратные вызовы Pipeline

Обратный вызов выполняется в конце шага денойзинга и может изменять состояние пайплайна или тензоры для последующих шагов. Используйте его для настройки атрибутов или тензорных переменных с целью изменения поведения без переписывания пайплайна.

Эти колбэки работают с классическими циклами DiffusionPipeline без дополнительной настройки. Они также работают с конвейерами Modular Diffusers, если в цикл денойзинга вставлен блок, который их выполняет; см. Modular pipelines.

Diffusers предоставляет несколько колбэков в обзоре пайплайна.

Чтобы включить обратный вызов, укажите, после какого количества шагов шумоподавления он должен выполняться, с помощью одного из следующих аргументов.

• cutoff_step_ratio определяет момент активации обратного вызова в процентах от общего числа шагов денойзинга. Используйте, когда точка отсечения должна масштабироваться вместе с num_inference_steps (например, отключить CFG после 40% выполнения).
• cutoff_step_index задаёт точный номер шага, на котором активируется callback. Используйте, когда важна абсолютная позиция шага (например, шаг 10 в фиксированном расписании из 25 шагов).

Приведённый ниже пример использует cutoff_step_ratio=0.4, что означает, что колбэк активируется, когда процесс денойзинга достигает 40% от общего числа шагов инференса. SDXLCFGCutoffCallback отключает классификаторно-свободное наведение (CFG) после определённого числа шагов, что позволяет сэкономить вычислительные ресурсы, не оказывая значительного влияния на качество результата.

Определите колбэк с одним из аргументов cutoff и передайте его в параметр callback_on_step_end в конвейере.

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

Официальные колбэки задают собственные тензорные входные данные. Для пользовательской функции передайте callback_on_step_end_tensor_inputs, как показано в разделе Отображение промежуточных изображений.

Если вы хотите добавить новый официальный callback, не стесняйтесь открыть feature request или submit a PR. В противном случае вы также можете создать собственный callback, как показано ниже.

Ранняя остановка

Ранняя остановка полезна, если промежуточные результаты генерации вас не устраивают. Этот callback устанавливает фиксированную точку остановки, присваивая атрибуту _interrupt значение True, что заставляет цикл денойзинга пропустить оставшиеся шаги.

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

Отображение промежуточных изображений

Визуализация промежуточных изображений полезна для отслеживания прогресса. Приведённый ниже предпросмотр работает только с SDXL. Он отображает латенты SDXL в RGB с помощью линейного преобразования, чтобы дать быстрый взгляд на процесс во время денойзинга. Эти веса не переносятся на другие модели. Для Qwen-Image и аналогичных чекпоинтов декодируйте изображения с помощью VAE модели, а не с помощью этого вспомогательного средства.

Конвертация латентов Stable Diffusion XL (4 канала) в RGB-тензоры (3 канала).

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)

Извлеките латенты и преобразуйте первое изображение в батче в RGB. Сохраните изображение в файл PNG с номером шага.

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

Используйте callback_on_step_end_tensor_inputs, чтобы выбрать, какие тензоры получает колбэк, в данном случае это латенты.

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]

Модульные конвейеры

Модульные конвейеры не поддерживают механизмы callback_on_step_end и interrupt. Вместо этого будет реализована поддержка потоковой передачи первого класса: huggingface/diffusers#14159 (ещё не объединён) добавляет pipe.stream(), который возвращает состояние конвейера на каждом шаге денойзинга и завершает выполнение, когда вы прекращаете итерацию. Пока что, чтобы переиспользовать ваши существующие функции обратного вызова, вставьте блок в цикл денойзинга: он получает BlockState цикла с индексом шага и временным шагом и может читать или изменять latents.

Измените блоки до вызова init_pipeline(): pipe.blocks возвращает копию, поэтому изменения после создания конвейера не имеют эффекта.

Блок ниже принимает колбэк в качестве входа в пайплайн и вызывает его с той же сигнатурой, что и выше: пайплайн, индекс шага, временной шаг и словарь тензоров, возвращаемые значения которых записываются обратно. walk_denoise_loops находит все циклы денойзинга в блоках, независимо от того, к какому рабочему процессу они относятся, поэтому один и тот же фрагмент кода работает для любого модульного пайплайна.

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 теперь является входным параметром пайплайна, и колбэк, написанный для DiffusionPipeline, например decode_tensors выше, работает без изменений:

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

Ранняя остановка работает иначе: блок не может пропустить цикл, в котором он выполняется, поэтому установка _interrupt не имеет эффекта. Вместо этого вызывайте исключение из колбэка: оно распространится наружу из pipe(...) (изображение не будет декодировано), и пайплайн останется пригодным для дальнейшего использования.

Это работает для каждого модульного конвейера, кроме внутреннего цикла вложенных циклов (авто регрессионные видеомодели, такие как Helios и Wan-Animate-2, выполняют денойзинг каждого чанка внутри блока внешнего цикла). Там блок нужно вставить во внешний цикл: он выполняется один раз на чанк и получает индекс чанка k вместо i и t.