Callbacks de Pipeline

Un callback se ejecuta al final de un paso de desruido y puede cambiar el estado del pipeline o los tensores para los pasos posteriores. Úselo para ajustar atributos o variables de tensor para un nuevo comportamiento sin reescribir el pipeline.

Estos callbacks funcionan con los bucles clásicos de DiffusionPipeline sin configuración adicional. También funcionan con las pipelines de Modular Diffusers una vez que se inserta un bloque que los ejecuta en el bucle de desruido; consulte Modular pipelines.

Diffusers ofrece varios callbacks en el resumen del pipeline.

Para habilitar un callback, configure cuándo se ejecuta el callback después de un número determinado de pasos de desruido con uno de los siguientes argumentos.

cutoff_step_ratio especifica cuándo se activa un callback como porcentaje de los pasos totales de desruido. Úsalo cuando el punto de corte debe escalar con num_inference_steps (por ejemplo, eliminar CFG después del 40% de la ejecución).
cutoff_step_index especifica el número exacto de paso en el que se activa un callback. Úsalo cuando te importe un paso absoluto (por ejemplo, el paso 10 en un plan fijo de 25 pasos).

El ejemplo siguiente utiliza cutoff_step_ratio=0.4, lo que significa que el callback se activa una vez que el desruido alcanza el 40 % de los pasos totales de inferencia. SDXLCFGCutoffCallback desactiva la guía libre de clasificador (CFG) después de un cierto número de pasos, lo que puede ayudar a ahorrar cómputo sin afectar significativamente el rendimiento.

Defina un callback con uno de los argumentos cutoff y páselo al parámetro callback_on_step_end en el 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,
)

Los callbacks oficiales definen sus propias entradas de tensor. Para una función personalizada, pase callback_on_step_end_tensor_inputs como se muestra en Mostrar imágenes intermedias.

Si deseas agregar un nuevo callback oficial, no dudes en abrir una solicitud de funcionalidad o enviar un PR. De lo contrario, también puedes crear tu propio callback, como se muestra a continuación.

Detención temprana

El early stopping es útil si no estás satisfecho con los resultados intermedios durante la generación. Este callback establece un punto de parada fijo al configurar el atributo _interrupt en True, lo que hace que el bucle de denoising omita los pasos 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,
)

Mostrar imágenes intermedias

Visualizar las imágenes intermedias es útil para monitorear el progreso. La vista previa de abajo es exclusiva de SDXL. Mapea los latentes de SDXL a RGB mediante una transformación lineal para obtener una vista rápida durante el desruido. Esos pesos no se transfieren a otros modelos. Para Qwen-Image y puntos de control similares, decodifica con el VAE del modelo en lugar de este helper.

Convertir los latentes de Stable Diffusion XL (4 canales) a tensores RGB (3 canales).

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)

Extrae los latentes y convierte la primera imagen del lote a RGB. Guarda la imagen como un archivo PNG con el número de paso.

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 elegir qué tensores recibe el callback, que en este caso son los 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

Los pipelines modulares no incluyen el mecanismo callback_on_step_end y interrupt. En su lugar, se incorporará soporte de streaming de primera clase: huggingface/diffusers#14159, que aún no se ha fusionado, añade pipe.stream(), que devuelve el estado del pipeline en cada paso de desruido y finaliza la ejecución cuando dejas de iterar. Por ahora, para reutilizar tus funciones de callback existentes, inserta un bloque en el bucle de desruido: recibe el BlockState del bucle con el índice del paso y el timestep, y puede leer o editar latents.

Edita los bloques antes de init_pipeline(): pipe.blocks devuelve una copia, por lo que modificarlo después de que exista el pipeline no tiene efecto.

El bloque siguiente toma el callback como entrada de la pipeline y lo invoca con la misma firma que arriba: la pipeline, el índice del paso, el timestep y un diccionario de tensores, cuyos valores devueltos se escriben de vuelta. walk_denoise_loops encuentra todos los bucles de desruido en los bloques, sin importar a qué flujo de trabajo pertenezcan, por lo que el mismo fragmento de código funciona para cualquier 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 ahora es una entrada del pipeline, y un callback escrito para un DiffusionPipeline, como decode_tensors arriba, funciona tal cual:

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

Early stopping funciona de manera diferente: un bloque no puede omitir el bucle en el que se ejecuta, por lo que establecer _interrupt no tiene efecto. En su lugar, lanza una excepción desde el callback: la excepción se propaga fuera de pipe(...) (no se decodifica ninguna imagen) y el pipeline sigue siendo utilizable después.

Esto funciona para cada pipeline modular, excepto en el bucle interno de bucles anidados (los modelos de video autoregresivos como Helios y Wan-Animate-2 desruidifican cada fragmento dentro de un bloque del bucle externo). Allí, inserta el bloque en el bucle externo en su lugar: se ejecuta una vez por fragmento y recibe el índice del fragmento k en lugar de i y t.

Update on GitHub