Diffusers documentation
Pipeline 回调
Pipeline 回调
回调在去噪步骤结束时运行，可以修改管道状态或张量，以影响后续步骤。使用它来调整属性或张量变量，从而获得新的行为，而无需重写管道。
这些回调函数开箱即可与经典的 DiffusionPipeline 循环配合使用。在去噪循环中插入一个执行它们的块之后，它们也能与 Modular Diffusers 流水线配合使用；请参阅 Modular pipelines。
Diffusers 在流水线概览中提供了多种回调。
要启用回调，请使用以下参数之一，配置在特定数量的去噪步骤后执行该回调的时机。
cutoff_step_ratio以总去噪步数的百分比形式指定回调的激活时机。当截断点需要随num_inference_steps缩放时使用（例如，在运行 40% 后丢弃 CFG）。cutoff_step_index指定回调被激活的确切步数。当您关注绝对步数时（例如，在固定的 25 步计划中的第10步），请使用此参数。
以下示例使用了 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。
如果您想添加一个新的官方回调，请随时提交一个功能请求或提交 PR。否则，您也可以按照下文所示创建自己的回调。
早停
如果您对生成过程中的中间结果不满意，早停（Early stopping）会很有用。此回调通过将 _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() 之前编辑 blocks：
pipe.blocks返回的是副本，因此在 pipeline 创建之后对其进行修改不会生效。
下方的代码块将回调函数作为 pipeline 的输入，并以与上文相同的签名调用它：即 pipeline、步骤索引、时间步以及一个张量字典，其返回值会被写回。walk_denoise_loops 会找到所有块中的去噪循环，无论其属于哪个工作流，因此同一段代码适用于任何模块化 pipeline。
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 现在是 pipeline 的输入参数，为 DiffusionPipeline 编写的回调函数（例如上面的 decode_tensors）可以直接使用：
image = pipe(prompt="A croissant shaped like a cute bear.", callback_on_step_end=decode_tensors, output="images")[0]Early stopping 的工作方式有所不同：一个 block 无法跳过其所在的循环，因此设置 _interrupt 不会产生任何效果。请改为在回调中抛出异常：该异常会从 pipe(...) 中传播出来（不会解码图像），并且管道在之后仍可继续使用。
这适用于所有模块化流水线，但嵌套循环的内层循环除外（例如 Helios 和 Wan-Animate-2 等自回归视频模型，它们在外层循环的一个块内对每个块进行去噪）。在这种情况下，请将块插入到外层循环中：它每个块运行一次，并接收块索引 k 来代替 i 和 t。