Callback pipeline

Callback dijalankan di akhir langkah denoising dan dapat mengubah status pipeline atau tensor untuk langkah-langkah berikutnya. Gunakan callback untuk menyesuaikan atribut atau variabel tensor demi perilaku baru tanpa menulis ulang pipeline.

Callback ini berfungsi dengan loop DiffusionPipeline klasik secara bawaan. Mereka juga berfungsi dengan pipeline Modular Diffusers setelah blok yang menjalankannya disisipkan ke dalam loop denoising; lihat Modular pipelines.

Diffusers menyediakan beberapa callback dalam ringkasan pipeline.

Untuk mengaktifkan callback, konfigurasikan kapan callback dieksekusi setelah sejumlah langkah denoising tertentu dengan salah satu argumen berikut.

• cutoff_step_ratio menentukan kapan callback diaktifkan sebagai persentase dari total langkah denoising. Gunakan ketika cutoff harus berskala dengan num_inference_steps (misalnya, hapus CFG setelah 40% dari eksekusi).
• cutoff_step_index menentukan nomor langkah persis di mana callback diaktifkan. Gunakan ini ketika Anda peduli pada langkah absolut (misalnya, langkah 10 pada jadwal tetap 25 langkah).

Contoh di bawah menggunakan cutoff_step_ratio=0.4, yang berarti callback diaktifkan ketika denoising mencapai 40% dari total langkah inferensi. SDXLCFGCutoffCallback menonaktifkan classifier-free guidance (CFG) setelah sejumlah langkah tertentu, yang dapat membantu menghemat komputasi tanpa secara signifikan memengaruhi kinerja.

Definisikan callback dengan salah satu argumen cutoff dan teruskan ke parameter callback_on_step_end di 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,
)

Callback resmi mengatur input tensor mereka sendiri. Untuk fungsi kustom, berikan callback_on_step_end_tensor_inputs seperti pada Tampilkan gambar antara.

Jika Anda ingin menambahkan callback resmi baru, silakan buka permintaan fitur atau kirim PR. Selain itu, Anda juga dapat membuat callback Anda sendiri seperti yang ditunjukkan di bawah ini.

Penghentian dini

Early stopping berguna jika Anda tidak puas dengan hasil antara selama proses generasi. Callback ini menetapkan titik henti yang di-hardcode dengan mengatur atribut _interrupt menjadi True, yang membuat loop denoising melewati langkah-langkah yang tersisa.

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,
)

Tampilkan gambar perantara

Visualisasi gambar perantara berguna untuk memantau kemajuan. Pratinjau di bawah ini hanya untuk SDXL. Pratinjau ini memetakan latens SDXL ke RGB dengan transformasi linear untuk melihat cepat selama denoising. Bobot tersebut tidak dapat dipindahkan ke model lain. Untuk Qwen-Image dan checkpoint serupa, dekodekan dengan VAE model alih-alih helper ini.

Konversi latens Stable Diffusion XL (4 saluran) menjadi tensor RGB (3 saluran).

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)

Ekstrak latents dan konversi gambar pertama dalam batch ke RGB. Simpan gambar sebagai file PNG dengan nomor langkah.

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

Gunakan callback_on_step_end_tensor_inputs untuk memilih tensor mana yang diterima oleh callback, yang dalam hal ini adalah 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]

Pipeline modular

Pipeline modular tidak memiliki mekanisme callback_on_step_end dan interrupt. Sebagai gantinya, dukungan streaming kelas satu akan hadir: huggingface/diffusers#14159, yang belum digabungkan, menambahkan pipe.stream(), yang menghasilkan status pipeline pada setiap langkah denoising dan mengakhiri eksekusi ketika Anda berhenti mengiterasi. Untuk saat ini, agar dapat menggunakan kembali fungsi callback Anda yang ada, sisipkan blok ke dalam loop denoising: blok ini menerima BlockState dari loop dengan indeks langkah dan timestep, serta dapat membaca atau mengedit latents.

Edit blok sebelum init_pipeline(): pipe.blocks mengembalikan salinan, sehingga mengubahnya setelah pipeline ada tidak akan berdampak apa pun.

Blok di bawah ini mengambil callback sebagai input pipeline dan memanggilnya dengan tanda tangan yang sama seperti di atas: pipeline, indeks langkah, timestep, dan dict tensor, di mana nilai yang dikembalikan akan ditulis kembali. walk_denoise_loops menemukan setiap loop denoising di dalam blok, terlepas dari workflow apa pun yang dimilikinya, sehingga snippet yang sama dapat digunakan untuk pipeline modular apa pun.

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 kini menjadi input dari pipeline, dan callback yang ditulis untuk DiffusionPipeline, seperti decode_tensors di atas, berfungsi apa adanya:

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

Early stopping bekerja dengan cara berbeda: sebuah blok tidak dapat melewati loop tempatnya berjalan, sehingga pengaturan _interrupt tidak berpengaruh. Alih-alih, lempar pengecualian dari callback: pengecualian tersebut akan menyebar keluar dari pipe(...) (tidak ada gambar yang didekode), dan pipeline tetap dapat digunakan setelahnya.

Ini berlaku untuk setiap pipeline modular kecuali loop internal dari loop bersarang (model video autoregresif seperti Helios dan Wan-Animate-2 melakukan denoising pada setiap chunk di dalam blok loop eksternal). Di sana, sisipkan blok ke dalam loop eksternal: blok ini berjalan sekali per chunk dan menerima indeks chunk k sebagai pengganti i dan t.