Callbacks de Pipeline

Um callback é executado no final de uma etapa de desruidamento e pode alterar o estado do pipeline ou os tensores para as etapas subsequentes. Use-o para ajustar atributos ou variáveis de tensor para um novo comportamento sem reescrever o pipeline.

Esses callbacks funcionam com os loops clássicos de DiffusionPipeline sem necessidade de configuração adicional. Eles também funcionam com os pipelines de Modular Diffusers assim que um bloco que os executa é inserido no loop de remoção de ruído; consulte Modular pipelines.

O Diffusers fornece vários callbacks no visão geral do pipeline.

Para habilitar um callback, configure quando o callback é executado após um determinado número de etapas de remoção de ruído com um dos seguintes argumentos.

• cutoff_step_ratio especifica quando um callback é ativado como uma porcentagem do total de etapas de desruidamento. Use quando o corte deve escalar com num_inference_steps (por exemplo, descartar o CFG após 40% da execução).
• cutoff_step_index especifica o número exato de etapa em que um callback é ativado. Use quando você se importar com uma etapa absoluta (por exemplo, a etapa 10 em um cronograma fixo de 25 etapas).

O exemplo abaixo usa cutoff_step_ratio=0.4, o que significa que o callback é ativado quando o desruidamento atinge 40% das etapas totais de inferência. SDXLCFGCutoffCallback desativa o classifier-free guidance (CFG) após um certo número de etapas, o que pode ajudar a economizar computação sem afetar significativamente o desempenho.

Defina um callback com um dos argumentos cutoff e passe-o para o parâmetro callback_on_step_end no 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,
)

Os callbacks oficiais definem seus próprios inputs de tensor. Para uma função personalizada, passe callback_on_step_end_tensor_inputs conforme mostrado em Exibir imagens intermediárias.

Se você deseja adicionar um novo callback oficial, sinta-se à vontade para abrir uma solicitação de recurso ou enviar um PR. Caso contrário, você também pode criar seu próprio callback, conforme mostrado abaixo.

Parada antecipada

O early stopping é útil se você não estiver satisfeito com os resultados intermediários durante a geração. Este callback define um ponto de parada fixo ao definir o atributo _interrupt como True, o que faz com que o loop de desruido pule as etapas restantes.

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

Exibir imagens intermediárias

Visualizar imagens intermediárias é útil para monitorar o progresso. A pré-visualização abaixo é exclusiva do SDXL. Ela mapeia os latentes do SDXL para RGB com uma transformação linear para uma visão rápida durante o desruidamento. Esses pesos não se transferem para outros modelos. Para Qwen-Image e checkpoints semelhantes, decodifique com o VAE do modelo em vez deste helper.

Converte latentes do Stable Diffusion XL (4 canais) em tensores RGB (3 canais).

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)

Extraia os latentes e converta a primeira imagem do lote para RGB. Salve a imagem como um arquivo PNG com o número da etapa.

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

Use callback_on_step_end_tensor_inputs para escolher quais tensores o callback recebe, que, neste caso, são os latentes.

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]

Pipelines modulares

Os pipelines modulares não possuem o mecanismo callback_on_step_end e interrupt. Em vez disso, virão com suporte de streaming de primeira classe: huggingface/diffusers#14159, ainda não mesclado, adiciona pipe.stream(), que fornece o estado do pipeline em cada etapa de desruidamento e encerra a execução quando você para de iterar. Por enquanto, para reutilizar suas funções de callback existentes, insira um bloco no loop de desruidamento: ele recebe o BlockState do loop com o índice da etapa e o timestep, e pode ler ou editar latents.

Edite os blocos antes de init_pipeline(): pipe.blocks retorna uma cópia, portanto, alterá-lo após a existência do pipeline não tem efeito.

O bloco abaixo recebe o callback como entrada do pipeline e o chama com a mesma assinatura acima: o pipeline, o índice da etapa, o timestep e um dict de tensores, cujos valores retornados são gravados de volta. walk_denoise_loops encontra todos os loops de desruidamento nos blocos, independentemente do workflow ao qual pertencem, de modo que o mesmo trecho de código funciona para qualquer pipeline modular.

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 agora é uma entrada do pipeline, e um callback escrito para um DiffusionPipeline, como decode_tensors acima, funciona sem alterações:

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

Early stopping funciona de forma diferente: um bloco não pode pular o loop em que está sendo executado, portanto, definir _interrupt não tem efeito. Em vez disso, lance uma exceção a partir do callback: a exceção se propaga para fora de pipe(...) (nenhuma imagem é decodificada) e o pipeline continua utilizável em seguida.

Isso funciona para todos os pipelines modulares, exceto para o laço interno de laços aninhados (modelos de vídeo autoregressivos, como Helios e Wan-Animate-2, desruidam cada bloco dentro de um bloco do laço externo). Lá, insira o bloco no laço externo em vez disso: ele é executado uma vez por bloco e recebe o índice do bloco k no lugar de i e t.