Add Consistency Models Pipeline

[dg845]

This PR implements a pipeline for Consistency Models as discussed in #3073.

Model/Pipeline Description

Consistency Models (paper, code) are a new family of generative models similar to continuous-time diffusion models which support fast one-step generation. Diffusion models can be distilled into a consistency model for faster sampling, and consistency models can also be trained from scratch. From the paper abstract:

Diffusion models have made significant breakthroughs in image, audio, and video generation, but they depend on an iterative generation process that causes slow sampling speed and caps their potential for real-time applications. To overcome this limitation, we propose consistency models, a new family of generative models that achieve high sample quality without adversarial training. They support fast one-step generation by design, while still allowing for few-step sampling to trade compute for sample quality. They also support zero-shot data editing, like image inpainting, colorization, and super-resolution, without requiring explicit training on these tasks. Consistency models can be trained either as a way to distill pre-trained diffusion models, or as standalone generative models. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one- and few-step generation. For example, we achieve the new state-of-the-art FID of 3.55 on CIFAR-10 and 6.20 on ImageNet 64x64 for one-step generation. When trained as standalone generative models, consistency models also outperform single-step, non-adversarial generative models on standard benchmarks like CIFAR-10, ImageNet 64x64 and LSUN 256x256.

In this PR, we implement a consistency model sampling pipeline as described in the paper.

Usage Examples

import torch

from diffusers import ConsistencyModelPipeline

device = "cuda"
# Load the cd_imagenet64_l2 checkpoint.
model_id_or_path = "dg845/consistency-model-pipelines"
pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

# Onestep Sampling
image = pipe(num_inference_steps=1).images[0]
image.save("consistency_model_onestep_sample.png")

# Onestep sampling, class-conditional image generation
# ImageNet-64 class label 145 corresponds to king penguins
image = pipe(num_inference_steps=1, class_labels=145).images[0]
image.save("consistency_model_onestep_sample_penguin.png")

# Multistep sampling, class-conditional image generation
# Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo.
image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]
image.save("consistency_model_multistep_sample_penguin.png")

TODO

• Implement Consistency Models pipeline (ConsistencyModelPipeline)
• Script to convert Consistency Model checkpoints to diffusers checkpoints
• Upload pre-trained Consistency Model checkpoints [original checkpoints]
• Add model cards for checkpoints and move checkpoints to appropriate repositories
• Create tests for Consistency Models
• Add docstrings for new classes
• Create documentation for Consistency Models
• Add usage example(s)

Discussion

• (TBD)

CC

@patrickvonplaten
@yang-song (author of original paper and code)
@tyshiwo1
@ayushtues
@xenova

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[dg845]

The sampling code is still very much a work in progress right now, I think I need to understand Karras schedulers better 😅 .

@ayushtues, not sure what the best way to collaborate is. One suggestion is that I can continue to work on the sampling code and you can start work on scripts/examples to distill diffusion models to consistency models, which I think people will probably want alongside a consistency models pipeline. Also open to other suggestions :).

[ayushtues]

Hey @dg845 if you're working on sampling, I can pick up on replicating the model architecture and porting the weights from the original checkpoints to diffusers for now, does that sound okay?

[ayushtues]

As a side note, have you been able to sample any good images from the original repo? I was trying to sample some Imagenet images, but could only get random garbage.

I was not able to install Flash Attention, but even using normal attention should not give such different results

I created a branch with some changes to the original code, removing distributed training/argparse and stuff to make it easily usable here: https://github.com/ayushtues/consistency_models/tree/local, feel free to use it.

Also set use_fp16 to false, it was leading to numerical unstability and giving NaNs

[ayushtues]

Also for the sampler, can we not use HeunDiscreteScheduler or EulerDiscreteScheduler already present in Diffusers?

[dg845]

Hey @dg845 if you're working on sampling, I can pick up on replicating the model architecture and porting the weights from the original checkpoints to diffusers for now, does that sound okay?

Sounds good :).

As a side note, have you been able to sample any good images from the original repo? I was trying to sample some Imagenet images, but could only get random garbage.

I ran into a problem installing the dependencies (I think it was related to distributed training/sampling) and haven't fixed it yet. I might try out your branch.

Also for the sampler, can we not use HeunDiscreteScheduler or EulerDiscreteScheduler already present in Diffusers?

In my understanding (admittedly a little shaky), the original implementation supports two types of schedulers:

1. Karras-style schedulers like Euler or Heun.
2. The multi-step scheduler (e.g. stochastic_iterative_sampler in the original code), of which the one-step scheduler could be considered a special case. This implements Algorithm 1 in the paper and doesn't correspond to any currently implemented scheduler, although KarrasVeScheduler is similar (to the best of my knowledge).

I think what I have currently supports Karras-style schedulers (and leverages already implemented Karras schedulers like HeunDiscreteScheduler), but I'm not sure about the best way to implement the multi-step scheduler.

[dg845]

Hi, would someone be able to do a preliminary design review of ConsistencyModelPipeline and CMStochasticIterativeScheduler? In particular, I have a few questions:

1. Should CMStochasticIterativeScheduler.__init__ configure a beta schedule like most of the samplers supported in the original implementation (e.g. Heun, DPM)? Right now it takes arguments in "original" sigma space, like KarrasDenoiser in the original implementation or KarrasVeScheduler.
2. Does it make sense to support schedulers like KarrasVeScheduler or ScoreSdeVeScheduler (since their implementation is different from the KarrasDiffusionSchedulers)? These schedulers aren't supported in the original code but I'm not sure if they're theoretically compatible with consistency models. Right now I'm trying to support them.

[ayushtues]

Hey @dg845 I'll also help in reviewing the schedulers, wanted to understand how they work.

Wasn't able to get much work done over the weekend, will update soon in the upcoming week!

[ayushtues]

Started a PR for adding the Unet and conversion scripts dg845#1.

[patrickvonplaten]

Also cc @williamberman here

[williamberman]

Awesome, thanks for getting this started! My high level understanding of the paper is that the consistency model directly outputs the unnoised result. If my understanding is correct, do we need a separate consistency model scheduler? If we want to do an iterative denoising process, we can re-use an existing scheduler then?

[dg845]

I'm not sure about what the best design should be. Currently, the new scheduler CMStochasticIterativeScheduler's step() method is trivial because (unless I'm misunderstanding) after we add_noise_to_input, all we need to do is to return the consistency model evaluated at the noised sample sample_hat:

diffusers/src/diffusers/schedulers/scheduling_consistency_models.py

Lines 197 to 201 in f2e53da

	# Assume model output is the consistency model evaluated at sample_hat.
	sample_prev = model_output
	if not return_dict:
	return (sample_prev,)

An alternative design would be to put the get_sigmas_karras/add_noise_to_input methods into the ConsistencyModelPipeline pipeline and then use those to run multi-step sampling. Under this design the pipeline kind of takes on some of the functions that a scheduler normally might, which is why I'm not sure about it.

To the best of my knowledge, I don't think there is a currently-implemented scheduler which is equivalent to the multi-step sampler in the original code; KarrasVeScheduler is probably the most similar, but the add_noise_to_input and step methods are not the same. For other samplers used by the original implementation, such as Euler and Heun, the existing diffusers implementations (e.g. EulerDiscreteScheduler, etc.) should be reused.

[ayushtues]

Was able to convert the Unet into diffusers, load pretrained checkpoints and get same outputs from forward passes:

ckpt : https://huggingface.co/ayushtues/consistency_models/tree/main

unet = UNet2DModel.from_pretrained("ayushtues/consistency_models", subfolder="diffusers_cd_imagenet64_l2")

[patrickvonplaten]

Was able to convert the Unet into diffusers, load pretrained checkpoints and get same outputs from forward passes:

ckpt : https://huggingface.co/ayushtues/consistency_models/tree/main

unet = UNet2DModel.from_pretrained("ayushtues/consistency_models", subfolder="diffusers_cd_imagenet64_l2")

Wow great work @ayushtues ! Exciting!

[williamberman]

I'm not sure about what the best design should be. Currently, the new scheduler CMStochasticIterativeScheduler's step() method is trivial because (unless I'm misunderstanding) after we add_noise_to_input, all we need to do is to return the consistency model evaluated at the noised sample sample_hat:

diffusers/src/diffusers/schedulers/scheduling_consistency_models.py

Lines 197 to 201 in f2e53da

	# Assume model output is the consistency model evaluated at sample_hat.
	sample_prev = model_output
	if not return_dict:
	return (sample_prev,)

An alternative design would be to put the get_sigmas_karras/add_noise_to_input methods into the ConsistencyModelPipeline pipeline and then use those to run multi-step sampling. Under this design the pipeline kind of takes on some of the functions that a scheduler normally might, which is why I'm not sure about it.

To the best of my knowledge, I don't think there is a currently-implemented scheduler which is equivalent to the multi-step sampler in the original code; KarrasVeScheduler is probably the most similar, but the add_noise_to_input and step methods are not the same. For other samplers used by the original implementation, such as Euler and Heun, the existing diffusers implementations (e.g. EulerDiscreteScheduler, etc.) should be reused.

I'd avoid putting them on the pipeline, can you state how the two methods are different? I'm not familiar off the top of my head. If the only way the methods are different is in the configured noise schedule but the formulas are the same, we can re-use the existing scheduler

[dg845]

For the add_noise_to_input method, KarrasVeScheduler and the consistency models implementation both make a noise update of the form

$$\hat{x}_n \leftarrow x + \sqrt{\tau_n^2 - \epsilon^2}\textbf{z}$$

where $\textbf{z} \sim \mathcal{N}(0, s_{noise}^2\textbf{I})$, but in KarrasVeScheduler $\epsilon = \sigma_{n}$ while in the consistency models implementation $\epsilon = \sigma_{min}$ (not in Karras space). Although $\tau_{n} = \hat{\sigma}_{n}$ in both implementations, $\hat{\sigma}_{n}$ is calculated differently, with KarrasVeScheduler using a gamma factor: $\hat{\sigma}_{n} = \sigma_n + \gamma \cdot \sigma_n$ while the consistency models implementation takes the next Karras sigma $\sigma_{n + 1}$ in the schedule as $\hat{\sigma}_n$.

For the step method, KarrasVeScheduler takes a Euler step from $\sigma_n$ to $\sigma_{n + 1}$, while the consistency model implementation simply evaluates the consistency model at $\hat{x}_n$ and $\hat{\sigma}_n$.

[williamberman]

Ok great super helpful! I think I have my head wrapped around it.

1. Just to note, the karras_ve and sde_ve are deprecated interfaces for interacting with diffeq based schedulers. We don't anymore explicitly call the higher order solving steps or noising in the pipelines. This has been encapsulated through multiple calls to call which tracks the current correction step and applies noise if needed. This means that the don't want to use the add_noise_to_input method.

2. This does need a new scheduler as existing schedulers which noise a prediction_type of sample do not do so according to a sigma based noise schedule (ddpm is the only noising a sample prediction type I know off the top of my head and it uses the particular alpha beta parameterization). However, we can re-use components (sigma schedule, adding noise) from existing schedulers pretty easily as what the code you linked is doing is straightforward.

Sigma schedule: We can re-use all existing components for constructing the sigma schedule, see scheduling_dpmsolver_sde.py for an example of how it copy's the methods from the euler scheduler

Adding noise: I believe the noise adding step to update the prediction of x_0 to x_{t-1} is just our standard add_noise method for a sigma based scheduler but updated for a starting noise amount which in the case of the code you linked is always t_min. So for this, I'd just adapt the add_noise method from the euler scheduler with an additional parameter for the starting noise amount. (I'm guessing a bit here that it is a starting noise amount, if you could confirm that would be helpful).

[patrickvonplaten]

We don't have to tackle it now but I think it should be a simple fix. Please correct me if I am wrong.

I was trying to get the UNet to compile with torch.compile() and when running inference it errors out:

Unsupported: call_function BuiltinVariable(iadd) [ConstantVariable(tuple), TupleVariable()] {}

from user code:
   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unet_2d.py", line 283, in forward
    sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unet_2d_blocks.py", line 935, in forward
    output_states += (hidden_states,)

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

Full code:

import torch
from diffusers import ConsistencyModelPipeline

device = "cuda"
model_id_or_path = "dg845/diffusers-cd_bedroom256_l2"
pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

_ = pipe(num_inference_steps=5)

Happens when num_inference_steps is set to 1 too.

This is now fixed - torch compile works!

[patrickvonplaten]

@patrickvonplaten for now I have modified the tests in SchedulerCommonTest so that the tests for the new scheduler pass. Not super confident in the design, would appreciate if you could review :).

That works!

[patrickvonplaten]

Make sure contributors are mentioned

[patrickvonplaten]

Good to be merged from my side. I pushed all checkpoints here:
https://huggingface.co/models?other=consistency-model

@ayushtues and @dg845 - really amazing work here 🚀

I've made sure to mention you on every model card here: https://huggingface.co/openai/diffusers-ct_cat256

as well as on the official docs: https://github.com/huggingface/diffusers/pull/3492/files#r1253304997

Hope that is ok?

[dg845]

I've made sure to mention you on every model card here: https://huggingface.co/openai/diffusers-ct_cat256

as well as on the official docs: https://github.com/huggingface/diffusers/pull/3492/files#r1253304997

Hope that is ok?

Sounds good! Thanks for helping out with the PR :).

I think one additional thing that might need to be updated is scripts/convert_consistency_to_diffusers.py, which still uses old UNet2DModel configs with AttnDownsampleBlock2D/AttnUpsampleBlock2D instead of AttnDownBlock2D/AttnUpBlock2D. The conversion logic might also need to be changed slightly to match the new block keys.

[williamberman]

Do we need a specific test for this? It seems outside of the scope of testing if this specific model works

[patrickvonplaten]

Think it won't really hurt tbh no?

[williamberman]

Nice! one quick question about if we need the flash attention tests but looks great!

[svjack]

The output seems bad, Can you release a stable config and image output for generation