Callbacks de pipeline

Un rappel s’exécute à la fin d’une étape de débruitage et peut modifier l’état du pipeline ou les tenseurs pour les étapes ultérieures. Utilisez-le pour ajuster des attributs ou des variables de tenseurs afin d’obtenir un nouveau comportement sans réécrire le pipeline.

Ces callbacks fonctionnent avec les boucles classiques de DiffusionPipeline sans configuration supplémentaire. Ils fonctionnent également avec les pipelines Modular Diffusers dès qu’un bloc qui les exécute est inséré dans la boucle de débruitage ; voir Modular pipelines.

Diffusers fournit plusieurs callbacks dans l'aperçu du pipeline.

Pour activer un rappel, configurez le moment où il est exécuté après un certain nombre d'étapes de débruitage à l'aide de l'un des arguments suivants.

• cutoff_step_ratio spécifie le moment d’activation d’un rappel, exprimé en pourcentage du nombre total d’étapes de débruitage. À utiliser lorsque le seuil doit s’adapter à num_inference_steps (par exemple, désactiver le CFG après 40 % de l’exécution).
• cutoff_step_index spécifie le numéro exact de l'étape à laquelle un rappel est activé. À utiliser lorsque vous vous souciez d'une étape absolue (par exemple, l'étape 10 sur un planning fixe de 25 étapes).

L’exemple ci-dessous utilise cutoff_step_ratio=0.4, ce qui signifie que le rappel est activé une fois que le débruitage atteint 40 % des étapes d’inférence totales. SDXLCFGCutoffCallback désactive le guidage sans classifieur (CFG) après un certain nombre d’étapes, ce qui peut aider à économiser des ressources de calcul sans affecter significativement les performances.

Définissez un rappel avec l'un des arguments cutoff et passez-le au paramètre callback_on_step_end du 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,
)

Les callbacks officiels définissent leurs propres entrées tensorielles. Pour une fonction personnalisée, passez callback_on_step_end_tensor_inputs comme indiqué dans Affichage d’images intermédiaires.

Si vous souhaitez ajouter un nouveau callback officiel, n'hésitez pas à ouvrir une demande de fonctionnalité ou à soumettre une PR. Sinon, vous pouvez également créer votre propre callback, comme indiqué ci-dessous.

Arrêt anticipé

L’arrêt anticipé est utile si vous n’êtes pas satisfait des résultats intermédiaires pendant la génération. Ce rappel définit un point d’arrêt en dur en définissant l’attribut _interrupt sur True, ce qui permet à la boucle de débruitage d’ignorer les étapes 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,
)

Afficher les images intermédiaires

La visualisation des images intermédiaires est utile pour le suivi de la progression. L’aperçu ci-dessous est réservé à SDXL. Il mappe les latents de SDXL vers le RGB à l’aide d’une transformation linéaire pour un coup d’œil rapide pendant le débruitage. Ces poids ne sont pas transférables à d’autres modèles. Pour Qwen-Image et des points de contrôle similaires, décodez avec le VAE du modèle plutôt qu’avec cet utilitaire.

Convertir les latents Stable Diffusion XL (4 canaux) en tenseurs RGB (3 canaux).

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)

Extrait les latents et convertit la première image du lot en RGB. Enregistre l'image sous forme de fichier PNG avec le numéro d'étape.

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

Utilisez callback_on_step_end_tensor_inputs pour choisir les tenseurs que le rappel reçoit, qui dans ce cas sont les 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]

Pipelines modulaires

Les pipelines modulaires ne prennent pas en charge les mécanismes callback_on_step_end et interrupt. Ils seront remplacés par un support du streaming de premier ordre : huggingface/diffusers#14159, pas encore fusionné, ajoute pipe.stream(), qui renvoie l’état du pipeline à chaque étape de débruitage et met fin à l’exécution lorsque vous cessez d’itérer. En attendant, pour réutiliser vos fonctions de rappel existantes, insérez un bloc dans la boucle de débruitage : il reçoit l’BlockState de la boucle avec l’indice d’étape et le timestep, et peut lire ou modifier les latents.

Modifiez les blocs avant init_pipeline() : pipe.blocks renvoie une copie, de sorte que toute modification effectuée après l’existence du pipeline n’a aucun effet.

Le bloc ci-dessous prend le rappel en tant qu'entrée du pipeline et l'appelle avec la même signature que ci-dessus : le pipeline, l'index de l'étape, le timestep et un dictionnaire de tenseurs, dont les valeurs retournées sont réécrites. walk_denoise_loops trouve toutes les boucles de débruitage dans les blocs, quel que soit le workflow auquel ils appartiennent, de sorte que le même extrait fonctionne pour n'importe quel pipeline modulaire.

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 est désormais une entrée du pipeline, et un callback écrit pour un DiffusionPipeline, comme decode_tensors ci-dessus, fonctionne tel quel :

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

L’arrêt anticipé fonctionne différemment : un bloc ne peut pas ignorer la boucle dans laquelle il s’exécute, donc définir _interrupt n’a aucun effet. Il faut plutôt lever une exception depuis le rappel : l’exception se propage hors de pipe(...) (aucune image n’est décodée) et le pipeline reste utilisable par la suite.

Cela fonctionne pour chaque pipeline modulaire, à l’exception de la boucle interne des boucles imbriquées (les modèles vidéo autorégressifs tels que Helios et Wan-Animate-2 débruitent chaque bloc à l’intérieur d’un bloc de la boucle externe). Dans ce cas, insérez le bloc dans la boucle externe à la place : il s’exécute une fois par bloc et reçoit l’indice du bloc k à la place de i et t.