Diffusers 文档
Pipeline 回调
加入 Hugging Face 社区
并获得增强的文档体验
开始使用
Pipeline 回调
pipeline 的去噪循环可以使用自定义函数通过 callback_on_step_end 参数进行修改。回调函数在每个步骤结束时执行，并修改 pipeline 属性和变量以供下一步使用。这对于动态调整某些 pipeline 属性或修改张量变量非常有用。这种多功能性允许有趣的用例，例如在每个时间步更改 prompt embeddings，为 prompt embeddings 分配不同的权重，以及编辑 guidance scale。通过回调，您无需修改底层代码即可实现新功能！
🤗 Diffusers 当前仅支持 callback_on_step_end，但如果您有很酷的用例并需要具有不同执行点的回调函数，请随时打开功能请求！
本指南将通过您可以使用回调实现的一些功能来演示回调的工作原理。
官方回调
我们提供了一个回调列表，您可以将其插入到现有的 pipeline 中并修改去噪循环。以下是当前官方回调列表
SDCFGCutoffCallback：在一定步数后禁用所有 SD 1.5 pipeline 的 CFG，包括文本到图像、图像到图像、图像修复和 controlnet。SDXLCFGCutoffCallback：在一定步数后禁用所有 SDXL pipeline 的 CFG，包括文本到图像、图像到图像、图像修复和 controlnet。IPAdapterScaleCutoffCallback：在一定步数后禁用所有支持 IP-Adapter 的 pipeline 的 IP Adapter。
要设置回调，您需要指定回调生效的去噪步数。您可以使用以下两个参数之一来执行此操作
cutoff_step_ratio：浮点数，表示步数的比率。cutoff_step_index：整数，表示步数的精确数量。
import torch
from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
from diffusers.callbacks import SDXLCFGCutoffCallback
callback = SDXLCFGCutoffCallback(cutoff_step_ratio=0.4)
# can also be used with cutoff_step_index
# callback = SDXLCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=10)
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_karras_sigmas=True)
prompt = "a sports car at the road, best quality, high quality, high detail, 8k resolution"
generator = torch.Generator(device="cpu").manual_seed(2628670641)
out = pipeline(
prompt=prompt,
negative_prompt="",
guidance_scale=6.5,
num_inference_steps=25,
generator=generator,
callback_on_step_end=callback,
)
out.images[0].save("official_callback.png")
动态无分类器指导
动态无分类器指导 (CFG) 是一项功能，允许您在一定数量的推理步骤后禁用 CFG，这可以帮助您节省计算资源，而性能损失极小。此功能的回调函数应具有以下参数
pipeline（或 pipeline 实例）提供对重要属性（例如num_timesteps和guidance_scale）的访问。您可以通过更新底层属性来修改这些属性。对于此示例，您将通过设置pipeline._guidance_scale=0.0来禁用 CFG。step_index和timestep告诉您在去噪循环中的位置。使用step_index在达到num_timesteps的 40% 后关闭 CFG。callback_kwargs是一个字典，其中包含您可以在去噪循环期间修改的张量变量。它仅包含在callback_on_step_end_tensor_inputs参数中指定的变量，该参数传递给 pipeline 的__call__方法。不同的 pipeline 可能会使用不同的变量集，因此请检查 pipeline 的_callback_tensor_inputs属性以获取您可以修改的变量列表。一些常见变量包括latents和prompt_embeds。对于此函数，在设置guidance_scale=0.0后更改prompt_embeds的批量大小，以便使其正常工作。
您的回调函数应如下所示
def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
# adjust the batch_size of prompt_embeds according to guidance_scale
if step_index == int(pipeline.num_timesteps * 0.4):
prompt_embeds = callback_kwargs["prompt_embeds"]
prompt_embeds = prompt_embeds.chunk(2)[-1]
# update guidance_scale and prompt_embeds
pipeline._guidance_scale = 0.0
callback_kwargs["prompt_embeds"] = prompt_embeds
return callback_kwargs现在，您可以将回调函数传递给 callback_on_step_end 参数，并将 prompt_embeds 传递给 callback_on_step_end_tensor_inputs。
import torch
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
generator = torch.Generator(device="cuda").manual_seed(1)
out = pipeline(
prompt,
generator=generator,
callback_on_step_end=callback_dynamic_cfg,
callback_on_step_end_tensor_inputs=['prompt_embeds']
)
out.images[0].save("out_custom_cfg.png")中断扩散过程
对于 StableDiffusionPipeline 和 StableDiffusionXLPipeline，文本到图像、图像到图像和图像修复都支持中断回调。
在构建使用 Diffusers 的 UI 时，提前停止扩散过程非常有用，因为它允许用户在对中间结果不满意时停止生成过程。您可以使用回调将此功能集成到您的 pipeline 中。
此回调函数应接受以下参数：pipeline、i、t 和 callback_kwargs（必须返回）。将 pipeline 的 _interrupt 属性设置为 True 以在一定步数后停止扩散过程。您也可以自由地在回调中实现您自己的自定义停止逻辑。
在此示例中，即使 num_inference_steps 设置为 50，扩散过程也会在 10 步后停止。
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
pipeline.enable_model_cpu_offload()
num_inference_steps = 50
def interrupt_callback(pipeline, i, t, callback_kwargs):
stop_idx = 10
if i == stop_idx:
pipeline._interrupt = True
return callback_kwargs
pipeline(
"A photo of a cat",
num_inference_steps=num_inference_steps,
callback_on_step_end=interrupt_callback,
)在每个生成步骤后显示图像
此技巧由 asomoza 贡献。
通过在每个步骤后访问和转换 latents 为图像，在每个生成步骤后显示图像。潜在空间被压缩为 128x128，因此图像也为 128x128，这对于快速预览非常有用。
- 使用以下函数将 SDXL latents（4 个通道）转换为 RGB 张量（3 个通道），如 Explaining the SDXL latent space 博客文章中所述。
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)- 创建一个函数来解码 latents 并将其保存为图像。
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- 将
decode_tensors函数传递给callback_on_step_end参数，以便在每个步骤后解码张量。您还需要在callback_on_step_end_tensor_inputs参数中指定要修改的内容，在本例中为 latents。
from diffusers import AutoPipelineForText2Image
import torch
from PIL import Image
pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
).to("cuda")
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]
