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は、一定のステップ数後にclassifier-free guidance (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 の潜在変数（latents）を線形変換で 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)latentsを抽出し、バッチ内の最初の画像を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_kwargscallback_on_step_end_tensor_inputs を使用して、コールバックが受け取るテンソルを選択します。この場合、対象となるのは潜変数（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]モジュール化されたパイプライン
Modular pipelinesには、callback_on_step_endおよびinterruptの機構は含まれていません。代わりに、第一級サポートとしてのストリーミング機能が提供されます。まだマージされていないhuggingface/diffusers#14159では、pipe.stream()が追加されており、デノイジングの各ステップでパイプラインの状態を生成し、イテレーションを停止すると実行を終了します。当面は、既存のコールバック関数を再利用するために、デノイジングループ内にブロックを挿入します。このブロックは、ステップインデックスとタイムステップを含むループのBlockStateを受け取り、latentsを読み取ったり編集したりできます。
pipe.blocksはコピーを返すため、パイプラインが存在した後にこれを修正しても効果はありません。したがって、init_pipeline() の前にブロックを編集してください。
以下のブロックは、コールバックをパイプラインの入力として受け取り、上記と同様のシグネチャ（パイプライン、ステップインデックス、タイムステップ、およびテンソルの辞書）で呼び出します。返された値は書き戻されます。walk_denoise_loops は、どのワークフローに属しているかに関係なく、ブロック内のすべてのデノイジングループを見つけます。そのため、同じスニペットが任意のモジュラーパイプラインで動作します。
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 はパイプラインの入力となり、DiffusionPipeline 用に記述されたコールバック（上記の decode_tensors など）はそのまま動作します。
image = pipe(prompt="A croissant shaped like a cute bear.", callback_on_step_end=decode_tensors, output="images")[0]Early stopping は動作が異なります。ブロックは実行中のループをスキップできないため、_interrupt を設定しても効果はありません。代わりに、コールバックから例外を送出してください。この例外は pipe(...) の外に伝播します（画像はデコードされません）が、パイプラインはその後でも引き続き使用できます。
これは、ネストされたループのインナーループ（Helios や Wan-Animate-2 などの自己回帰型動画モデルは、外側ループのブロック内で各チャンクをデノイズします）を除くすべてのモジュラーパイプラインで機能します。その場合は、ブロックを外側ループに挿入してください。ブロックはチャンクごとに1回実行され、i と t の代わりにチャンクインデックス k を受け取ります。