Pipeline callbacks

The denoising loop of a pipeline can be modified with custom defined functions using the callback_on_step_end parameter. This can be really useful for dynamically adjusting certain pipeline attributes, or modifying tensor variables. The flexibility of callbacks opens up some interesting use-cases such as changing the prompt embeddings at each timestep, assigning different weights to the prompt embeddings, and editing the guidance scale.

This guide will show you how to use the callback_on_step_end parameter to disable classifier-free guidance (CFG) after 40% of the inference steps to save compute with minimal cost to performance.

The callback function should have the following arguments:

• pipe (or the pipeline instance) provides access to useful properties such as num_timestep and guidance_scale. You can modify these properties by updating the underlying attributes. For this example, you’ll disable CFG by setting pipe._guidance_scale=0.0.
• step_index and timestep tell you where you are in the denoising loop. Use step_index to turn off CFG after reaching 40% of num_timestep.
• callback_kwargs is a dict that contains tensor variables you can modify during the denoising loop. It only includes variables specified in the callback_on_step_end_tensor_inputs argument, which is passed to the pipeline’s __call__ method. Different pipelines may use different sets of variables, so please check a pipeline’s _callback_tensor_inputs attribute for the list of variables you can modify. Some common variables include latents and prompt_embeds. For this function, change the batch size of prompt_embeds after setting guidance_scale=0.0 in order for it to work properly.

Your callback function should look something like this:

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(pipe.num_timestep * 0.4):
                prompt_embeds = callback_kwargs["prompt_embeds"]
                prompt_embeds = prompt_embeds.chunk(2)[-1]

        # update guidance_scale and prompt_embeds
        pipe._guidance_scale = 0.0
        callback_kwargs["prompt_embeds"] = prompt_embeds
        return callback_kwargs

Now, you can pass the callback function to the callback_on_step_end parameter and the prompt_embeds to callback_on_step_end_tensor_inputs.

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"

generator = torch.Generator(device="cuda").manual_seed(1)
out = pipe(prompt, generator=generator, callback_on_step_end=callback_custom_cfg, callback_on_step_end_tensor_inputs=['prompt_embeds'])

out.images[0].save("out_custom_cfg.png")

The callback function is executed at the end of each denoising step, and modifies the pipeline attributes and tensor variables for the next denoising step.

With callbacks, you can implement features such as dynamic CFG without having to modify the underlying code at all!