search

LuChengTHU / dpm-solver

Official code for "DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps" (Neurips 2022 Oral)
MIT License
1.56k stars 121 forks source link
diffusion-models machine-learning score-based-generative-models stable-diffusion

readme

DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps

Hugging Face Spaces

The official code for the paper DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps (Neurips 2022 Oral) and DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models by Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li and Jun Zhu.

DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with the convergence order guarantee. DPM-Solver is suitable for both discrete-time and continuous-time diffusion models without any further training. Experimental results show that DPM-Solver can generate high-quality samples in only 10 to 20 function evaluations on various datasets.

Guided-Diffusion with DPM-Solver:

DPM-Solver

Stable-Diffusion with DPM-Solver++:

sdm

DiffEdit with DPM-Solver++:

inpainting

Use the SOTA Multistep DPM-Solver++ with Diffusers

🤗 Diffusers is a fantastic library for diffusion models. It supports both DPM-Solver and DPM-Solver++. The multistep DPM-Solver++ is the fastest solver currently.

Stable-Diffusion

The second-order multistep DPM-Solver++ is the default solver for Stable-Diffusion online demos (e.g., see example) and can also be used in LoRA (e.g., see example). Here is an example:

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "stabilityai/stable-diffusion-2-1"

# Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

image.save("astronaut_rides_horse.png")

DeepFloyd-IF

We recommend the SDE version DPM-Solver++ for the stage-1, and the ODE version DPM-Solver++ for the upscaling stages (both stage-2 and 3).

from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import pt_to_pil
import torch

# stage 1
stage_1 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16)
stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
stage_1.enable_model_cpu_offload()

# stage 2
stage_2 = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16
)
stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
stage_2.enable_model_cpu_offload()

# stage 3
safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker, "watermarker": stage_1.watermarker}
stage_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", **safety_modules, torch_dtype=torch.float16)
stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
stage_3.enable_model_cpu_offload()

def set_scheduler(stage):
    if scheduler_name == 'dpm++':
        scheduler = DPMSolverMultistepScheduler.from_config(stage.scheduler.config)
        scheduler.config.algorithm_type = 'dpmsolver++'
    elif scheduler_name == 'sde-dpm++':
        scheduler = DPMSolverMultistepScheduler.from_config(stage.scheduler.config)
        scheduler.config.algorithm_type = 'sde-dpmsolver++'
    stage.scheduler = scheduler
    return stage

upscale_steps = 25
stage_1 = set_scheduler(stage_1, 'sde-dpm++')
stage_2 = set_scheduler(stage_2, 'dpm++')
stage_3 = set_scheduler(stage_3, 'dpm++')

prompt = "casual photo of a leaf maple syrup glass container sitting on a wooden table in a log cabin, high depth of field during golden hour as the sunlight shines through the windows, dusty air"

# text embeds
prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)

generator = torch.manual_seed(0)

# stage 1
image = stage_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt").images
pt_to_pil(image)[0].save("./if_stage_I.png")

# stage 2
image = stage_2(
    image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt", num_inference_steps=upscale_steps
).images
pt_to_pil(image)[0].save("./if_stage_II.png")

# stage 3
image = stage_3(prompt=prompt, image=image, generator=generator, noise_level=100, num_inference_steps=upscale_steps).images
image[0].save("./if_stage_III.png")


News


Supported Models and Algorithms

Models

We support the following four types of diffusion models. You can set the model type by the argument model_type in the function model_wrapper.

Model Type Training Objective Example Paper
"noise": noise prediction model $\epsilon_\theta$ $E{x{0},\epsilon,t}\left[\omega1(t)||\epsilon\theta(x_t,t)-\epsilon||_2^2\right]$ DDPM, Stable-Diffusion
"xstart": data prediction model $x\theta$ $E_{x_0,\epsilon,t}\left[\omega2(t)||x\theta(x_t,t)-x_0||_2^2\right]$ DALL·E 2
"v": velocity prediction model $v_\theta$ $E_{x_0,\epsilon,t}\left[\omega3(t)||v\theta(x_t,t)-(\alpha_t\epsilon - \sigma_t x_0)||_2^2\right]$ Imagen Video
"score": marginal score function $s_\theta$ $E_{x_0,\epsilon,t}\left[\omega_4(t)||\sigmat s\theta(x_t,t)+\epsilon||_2^2\right]$ ScoreSDE

Sampling Types

We support the following three types of sampling by diffusion models. You can set the argument guidance_type in the function model_wrapper.

Sampling Type Equation for Noise Prediction Model Example Paper
"uncond": unconditional sampling $\tilde\epsilon_\theta(xt,t)=\epsilon\theta(x_t,t)$ DDPM
"classifier": classifier guidance $\tilde\epsilon_\theta(xt,t,c)=\epsilon\theta(x_t,t)-s\cdot\sigmat\nabla{xt}\log q\phi(x_t,t,c)$ ADM, GLIDE
"classifier-free": classifier-free guidance $\tilde\epsilon_\theta(xt,t,c)=s\cdot \epsilon\theta(xt,t,c)+(1-s)\cdot\epsilon\theta(x_t,t)$ DALL·E 2, Imagen, Stable-Diffusion

Algorithms in DPM-Solver

We support the following four algorithms. The algorithms are DPM-Solver and DPM-Solver++.

We also support the dynamic thresholding introduced by Imagen for algorithms with data-prediction. The dynamic thresholding method can further improve the sample quality by pixel-space DPMs with large guidance scales.

Note that the model_fn for initializing DPM-Solver is always the noise prediction model. The setting for algorithm_type is for the algorithm (DPM-Solver or DPM-Solver++), not for the model. In other words, both DPM-Solver and DPM-Solver++ is suitable for all the four model types.

The performance of singlestep solvers (i.e. Runge-Kutta-like solvers) and the multistep solvers (i.e. Adams-Bashforth-like solvers) are different. We recommend to use different solvers for different tasks.

Method Supported Orders Supporting Thresholding Remark
DPM-Solver, singlestep 1, 2, 3 No
DPM-Solver, multistep 1, 2, 3 No
DPM-Solver++, singlestep 1, 2, 3 Yes
DPM-Solver++, multistep 1, 2, 3 Yes Recommended for guided sampling with order = 2, and for unconditional sampling with order = 3.


Code Examples

DDPM and Guided-Diffusion with DPM-Solver

We provide an example of guided-diffusion with DPM-Solver in examples/ddpm_and_guided-diffusion.

Text-to-Image by Stable-Diffusion with DPM-Solver

We provide an example of stable diffusion with DPM-Solver in examples/stable-diffusion. DPM-Solver can greatly accelerate the sampling speed of the original stable-diffusion.

Image Editing (DiffEdit) by Stable-Diffusion with DPM-Solver

We provide an example of DiffEdit with DPM-Solver, which can be used for image editing. The idea of DiffEdit can be general decribe as, using DDIM to get a invertable latent serise, then apply different prompt for inpainting (controled by auto generated mask).

We could easily accelerate such editing / inpainting by DPM-Solver in only 20 steps.

ScoreSDE with DPM-Solver

We provide a pytorch example and a JAX example in examples/ which apply DPM-Solver for Yang Song's score_sde repo on CIFAR-10.


Use DPM-Solver in your own code

It is very easy to combine DPM-Solver with your own diffusion models. We support both Pytorch and JAX code. You can just copy the file dpm_solver_pytorch.py or dpm_solver_jax.py to your own code files and import it.

In each step, DPM-Solver needs to compute the corresponding $\alpha_t$, $\sigma_t$ and $\lambda_t$ of the noise schedule. We support the commonly-used variance preserving (VP) noise schedule for both discrete-time and continuous-time DPMs:

Moreover, DPM-Solver is designed for the continuous-time diffusion ODEs. For discrete-time diffusion models, we also implement a wrapper function to convert the discrete-time diffusion models to the continuous-time diffusion models in the model_wrapper function.


Suggestions for Choosing the Hyperparameters

If you want to find the best setting for accelerating the sampling procedure by your own diffusion models, we provide a reference guide here:

  1. IMPORTANT: First run 1000-step DDIM to check the sample quality of your model. If the sample quality is poor, then DPM-Solver cannot improve it. Please further check your model defination or training process.

    Reason: DDIM is the first-order special case of DPM-Solver (proved in our paper). So given the same noise sample at time $T$, the converged samples of DDIM and DPM-Solver are the same. DPM-Solver can accelerate the convergence, but cannot improve the converged sample quality.

  2. If 1000-step DDIM can generate quite good samples, then DPM-Solver can achieve a quite good sample quality within very few steps because it can greatly accelerate the convergence. You may want to further choose the detailed hyperparameters of DPM-Solver. Here we provide a comprehensive searching routine:

    • Comparing algorithm_type="dpmsolver" and algorithm_type="dpmsolver++". Note that these settings are for the algorithm, not for the model. In other words, even for algorithm_type="dpmsolver++, you can still use the noise prediction model (such as stable-diffusion) and the algorithm can work well.

    • (Optional) Comparing with / without dynamic thresholding.

      IMPORTANT: our supported dynamic thresholding method is only valid for pixel-space diffusion models with algorithm_type="dpmsolver++. For example, Imagen uses the dynamic thresholding method and greatly improves the sample quality. The thresholding method pushes the pixel-space samples into the bounded area, so it can generate reasonable images. However, for latent-space diffusion models (such as stable-diffusion), the thresholding method is unsuitable because the $x_0$ at time $0$ of the diffusion model is in fact the "latent variable" in the latent space and it is unbounded.

    • Comparing singlestep or multistep methods.

    • Comparing order = 2, 3. Note that the all the first-order versions are equivalent to DDIM, so you do not need to try it.

    • Comparing steps = 10, 15, 20, 25, 50, 100. It depends on your computation resources and the need of sample quality.

    • (Optional) Comparing the time_uniform, logSNR and time_quadratic for the skip type.

      We empirically find that for high-resolutional images, the best setting is the time_uniform. So we recommend this setting and there is no need for extra searching. However, for low-resolutional images such as CIFAR-10, we empirically find that logSNR is the best setting.

    • (Optional) Comparing denoise_to_zero=True or denoise_to_zero=False.

      Empirically, the denoise_to_zero=True can improve the FID for low-resolutional images such as CIFAR-10. However, the influence of this method for high-resolutional images seem to be small. As the denoise_to_zero method needs one additional function evaluation (i.e. one additional step), we do not recommend to use the denoise_to_zero method for high-resolutional images.

    The detailed pseudo code is like:

    for algorithm_type in ["dpmsolver", "dpmsolver++"]:
# Optional, for correcting_x0_fn in [None, "dynamic_thresholding"]:
    dpm_solver = DPM_Solver(..., algorithm_type=algorithm_type) # ... means other arguments
    for method in ['singlestep', 'multistep']:
        for order in [2, 3]:
            for steps in [10, 15, 20, 25, 50, 100]:
                sample = dpm_solver.sample(
                    ..., # ... means other arguments
                    method=method,
                    order=order,
                    steps=steps,
                    # optional: skip_type='time_uniform' or 'logSNR' or 'time_quadratic',
                    # optional: denoise_to_zero=True or False
                )

    And then compare the samples to choose the best setting.

Moreover, for unconditional sampling and guided sampling, we have some recommendation settings and code examples, which are listed in the following section.


Suggestions for the Detailed Settings

We recommend to use the following two types of solvers for different tasks:

Specifically, we have the following suggestions:


Example: Unconditional Sampling by DPM-Solver

We recommend to use the 3rd-order (dpmsolver or dpmsolver++, multistep) DPM-Solver. Here is an example for discrete-time DPMs:

from dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver

## You need to firstly define your model and the extra inputs of your model,
## And initialize an `x_T` from the standard normal distribution.
## `model` has the format: model(x_t, t_input, **model_kwargs).
## If your model has no extra inputs, just let model_kwargs = {}.

## If you use discrete-time DPMs, you need to further define the
## beta arrays for the noise schedule.

# model = ....
# model_kwargs = {...}
# x_T = ...
# betas = ....

## 1. Define the noise schedule.
noise_schedule = NoiseScheduleVP(schedule='discrete', betas=betas)

## 2. Convert your discrete-time `model` to the continuous-time
## noise prediction model. Here is an example for a diffusion model
## `model` with the noise prediction type ("noise") .
model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type="noise",  # or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
)

## 3. Define dpm-solver and sample by singlestep DPM-Solver.
## (We recommend singlestep DPM-Solver for unconditional sampling)
## You can adjust the `steps` to balance the computation
## costs and the sample quality.
dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
## Can also try
# dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver")

## You can use steps = 10, 12, 15, 20, 25, 50, 100.
## Empirically, we find that steps in [10, 20] can generate quite good samples.
## And steps = 20 can almost converge.
x_sample = dpm_solver.sample(
    x_T,
    steps=20,
    order=3,
    skip_type="time_uniform",
    method="multistep",
)


Example: Classifier Guidance Sampling by DPM-Solver

We recommend to use the 2nd-order (dpmsolver++, multistep) DPM-Solver, especially for large guidance scales. Here is an example for discrete-time DPMs:

from dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver

## You need to firstly define your model and the extra inputs of your model,
## And initialize an `x_T` from the standard normal distribution.
## `model` has the format: model(x_t, t_input, **model_kwargs).
## If your model has no extra inputs, just let model_kwargs = {}.

## If you use discrete-time DPMs, you need to further define the
## beta arrays for the noise schedule.

## For classifier guidance, you need to further define a classifier function,
## a guidance scale and a condition variable.

# model = ....
# model_kwargs = {...}
# x_T = ...
# condition = ...
# betas = ....
# classifier = ...
# classifier_kwargs = {...}
# guidance_scale = ...

## 1. Define the noise schedule.
noise_schedule = NoiseScheduleVP(schedule='discrete', betas=betas)

## 2. Convert your discrete-time `model` to the continuous-time
## noise prediction model. Here is an example for a diffusion model
## `model` with the noise prediction type ("noise") .
model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type="noise",  # or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
    guidance_type="classifier",
    condition=condition,
    guidance_scale=guidance_scale,
    classifier_fn=classifier,
    classifier_kwargs=classifier_kwargs,
)

## 3. Define dpm-solver and sample by multistep DPM-Solver.
## (We recommend multistep DPM-Solver for conditional sampling)
## You can adjust the `steps` to balance the computation
## costs and the sample quality.

dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")

## If the DPM is defined on pixel-space images, you can further
## set `correcting_x0_fn="dynamic_thresholding"`. e.g.:

# dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++",
#   correcting_x0_fn="dynamic_thresholding")

## You can use steps = 10, 12, 15, 20, 25, 50, 100.
## Empirically, we find that steps in [10, 20] can generate quite good samples.
## And steps = 20 can almost converge.
x_sample = dpm_solver.sample(
    x_T,
    steps=20,
    order=2,
    skip_type="time_uniform",
    method="multistep",
)

Example: Classifier-Free Guidance Sampling by DPM-Solver

We recommend to use the 2nd-order (dpmsolver++, multistep) DPM-Solver, especially for large guidance scales. Here is an example for discrete-time DPMs:

from dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver

## You need to firstly define your model and the extra inputs of your model,
## And initialize an `x_T` from the standard normal distribution.
## `model` has the format: model(x_t, t_input, cond, **model_kwargs).
## If your model has no extra inputs, just let model_kwargs = {}.

## If you use discrete-time DPMs, you need to further define the
## beta arrays for the noise schedule.

## For classifier-free guidance, you need to further define a guidance scale,
## a condition variable and an unconditioanal condition variable.

# model = ....
# model_kwargs = {...}
# x_T = ...
# condition = ...
# unconditional_condition = ...
# betas = ....
# guidance_scale = ...

## 1. Define the noise schedule.
noise_schedule = NoiseScheduleVP(schedule='discrete', betas=betas)

## 2. Convert your discrete-time `model` to the continuous-time
## noise prediction model. Here is an example for a diffusion model
## `model` with the noise prediction type ("noise") .
model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type="noise",  # or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
    guidance_type="classifier-free",
    condition=condition,
    unconditional_condition=unconditional_condition,
    guidance_scale=guidance_scale,
)

## 3. Define dpm-solver and sample by multistep DPM-Solver.
## (We recommend multistep DPM-Solver for conditional sampling)
## You can adjust the `steps` to balance the computation
## costs and the sample quality.

dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")

## If the DPM is defined on pixel-space images, you can further
## set `correcting_x0_fn="dynamic_thresholding"`. e.g.:

# dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++",
#   correcting_x0_fn="dynamic_thresholding")

## You can use steps = 10, 12, 15, 20, 25, 50, 100.
## Empirically, we find that steps in [10, 20] can generate quite good samples.
## And steps = 20 can almost converge.
x_sample = dpm_solver.sample(
    x_T,
    steps=20,
    order=2,
    skip_type="time_uniform",
    method="multistep",
)


Documentation

1. Define the noise schedule.

We support the commonly-used variance preserving (VP) noise schedule for both discrete-time and continuous-time DPMs:

1.1. Discrete-time DPMs

We support a picewise linear interpolation of $\log\alpha_{t}$ in the NoiseScheduleVP class. It can support all types of VP noise schedules.

We need either the $\beta_i$ array or the $\bar{\alpha}_i$ array (see DDPM for details) to define the noise schedule. Note that the $\bar{\alpha}_i$ in DDPM is different from the $\alpha_t$ in DPM-Solver, and the detailed relationship is:

$$ \bar{\alpha}_i = \prod (1 - \beta_k) $$

$$ \alpha_{t_i} = \sqrt{\bar{\alpha}_i} $$

Define the discrete-time noise schedule by the $\beta_i$ array:

noise_schedule = NoiseScheduleVP(schedule='discrete', betas=betas)

Or define the discrete-time noise schedule by the $\bar{\alpha}_i$ array:

noise_schedule = NoiseScheduleVP(schedule='discrete', alphas_cumprod=alphas_cumprod)


1.2. Continuous-time DPMs

We support both linear schedule (as used in DDPM and ScoreSDE) and cosine schedule (as used in improved-DDPM) for the continuous-time DPMs.

Define the continuous-time linear noise schedule:

noise_schedule = NoiseScheduleVP(schedule='linear', continuous_beta_0=0.1, continuous_beta_1=20.)

Define the continuous-time cosine noise schedule:

noise_schedule = NoiseScheduleVP(schedule='cosine')


2. Wrap your model to a continuous-time noise predicition model.

For a given diffusion model with an input of the time label (may be discrete-time labels (i.e. 0 to 999) or continuous-time times (i.e. 0 to 1)), and the output type of the model may be "noise" or "x_start" or "v" or "score", we wrap the model function to the following format:

model_fn(x, t_continuous) -> noise

where t_continuous is the continuous time labels (i.e. 0 to 1), and the output type of the model is "noise", i.e. a noise prediction model. And we use the continuous-time noise prediction model model_fn for DPM-Solver.

Note that DPM-Solver only needs the noise prediction model (the $\epsilon_\theta(x_t, t)$ model, also as known as the "mean" model), so for diffusion models which predict both "mean" and "variance" (such as improved-DDPM), you need to firstly define another function by yourself to only output the "mean".


2.1. Sampling without Guidance

After defining the noise schedule, we need to further wrap the model to a continuous-time noise prediction model. The given model has the following format:

model(x_t, t_input, **model_kwargs) -> noise | x_start | v | score

And we wrap the model by:

model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type=model_type, # "noise" or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
)

where model_kwargs is the additional inputs of the model, and the model_type can be "noise" or "x_start" or "v" or "score".


2.2. Sampling with Classifier Guidance

After defining the noise schedule, we need to further wrap the model to a continuous-time noise prediction model. The given model has the following format:

model(x_t, t_input, **model_kwargs) -> noise | x_start | v | score

For DPMs with classifier guidance, we also combine the model output with the classifier gradient. We need to specify the classifier function and the guidance scale. The classifier function has the following format:

classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)

where t_input is the same time label as in the original diffusion model model, and cond is the condition variable, and classifier_kwargs is the other inputs of the classifier function.

And we wrap the model by:

model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type=model_type,  # "noise" or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
    guidance_type="classifier",
    condition=condition,
    guidance_scale=guidance_scale,
    classifier_fn=classifier,
    classifier_kwargs=classifier_kwargs,
)

where model_kwargs is the additional inputs of the model, and the model_type can be "noise" or "x_start" or "v" or "score", and guidance_scale is the classifier guidance scale, and condition is the conditional input of the classifier.


2.3. Sampling with Classifier-free Guidance

After defining the noise schedule, we need to further wrap the model to a continuous-time noise prediction model. The given model has the following format:

model(x_t, t_input, cond, **model_kwargs) -> noise | x_start | v | score

Note that for classifier-free guidance, the model needs another input cond. And if cond is a special variable unconditional_condition, the model output is the unconditional DPM output.

And we wrap the model by:

model_fn = model_wrapper(
    model,
    noise_schedule,
    model_type=model_type,  # "noise" or "x_start" or "v" or "score"
    model_kwargs=model_kwargs,
    guidance_type="classifier-free",
    condition=condition,
    unconditional_condition=unconditional_condition,
    guidance_scale=guidance_scale,
)

where model_kwargs is the additional inputs of the model, and the model_type can be "noise" or "x_start" or "v" or "score", and guidance_scale is the classifier guidance scale, and condition is the conditional input, and unconditional_condition is the special unconditional condition variable for the unconditional model.


2.4. Implementation Details for the Time Inputs

Below we introduce the detailed mapping between the discrete-time labels and the continuous-time times. However, to use DPM-Solver, it is not necessary to understand the following details.

Discrete-time DPM-Solver

For discrete-time DPMs, the noise prediction model noise-prediction is trained for the discrete-time labels from $0$ to $N-1$. Therefore, we sample from the discrete-time label $N - 1$ (e.g. 999) to the discrete-time label $0$. We convert the discrete-time labels in $[0, N-1]$ to the continuous-time times $(0,1]$ by

$$ t{\text{discrete}} = 1000 * \left(t{\text{continuous}} - \frac{1}{N}\right), $$

i.e. we map the discrete-time label $0$ to the continuous-time time $\frac{1}{N}$, and the discrete-time label $N-1$ to the continuous-time time $1$. Therefore, sampling from the discrete time from $N-1$ to $0$ is corresponding to sampling from the continuous time from $1$ to $\frac{1}{N}$.

Continuous-time DPM-Solver

For continuous-time DPMs from defined by $t \in [0,1]$, we simply wrap the model to accept only $x_t$ and $t$. Note that for continuous-time DPMs, we do not modify the time inputs.


3. Define DPM-Solver

After defining the model_fn by the function model_wrapper, we can further use model_fn to define DPM-Solver and compute samples.

We support the following four algorithms. The algorithms are DPM-Solver and DPM-Solver++.

We also support the dynamic thresholding introduced by Imagen for algorithms with data-prediction. The dynamic thresholding method can further improve the sample quality by pixel-space DPMs with large guidance scales.

Note that the model_fn for initializing DPM-Solver is always the noise prediction model. The setting for algorithm_type is for the algorithm (DPM-Solver or DPM-Solver++), not for the model. In other words, both DPM-Solver and DPM-Solver++ is suitable for all the four model types.

The performance of singlestep solvers (i.e. Runge-Kutta-like solvers) and the multistep solvers (i.e. Adams-Bashforth-like solvers) are different. We recommend to use different solvers for different tasks.

Method Supported Orders Supporting Thresholding Remark
DPM-Solver, singlestep 1, 2, 3 No
DPM-Solver, multistep 1, 2, 3 No
DPM-Solver++, singlestep 1, 2, 3 Yes
DPM-Solver++, multistep 1, 2, 3 Yes Recommended for guided sampling with order = 2, and for unconditional sampling with order = 3.

You can use dpm_solver.sample to quickly sample from DPMs. This function computes the ODE solution at time t_end by DPM-Solver, given the initial x at time t_start.

We support the following algorithms:

We support three types of skip_type for the choice of intermediate time steps:


3.1. Sampling by Singlestep DPM-Solver

We combine all the singlestep solvers with order <= order to use up all the function evaluations (steps). The total number of function evaluations (NFE) == steps.

For discrete-time DPMs, we do not need to specify the t_start and t_end. The default setting is to sample from the discrete-time label $N-1$ to the discrete-time label $0$. For example,

## discrete-time DPMs
x_sample = dpm_solver.sample(
    x_T,
    steps=20,
    order=3,
    skip_type="time_uniform",
    method="singlestep",
)


For continuous-time DPMs, we sample from t_start=1.0 (the default setting) to t_end. We recommend t_end=1e-3 for steps <= 15, and t_end=1e-4 for steps > 15. For example:

x_sample = dpm_solver.sample(
    x_T,
    t_end=1e-3,
    steps=12,
    order=3,
    skip_type="time_uniform",
    method="singlestep",
)
## continuous-time DPMs
x_sample = dpm_solver.sample(
    x_T,
    t_end=1e-4,
    steps=20,
    order=3,
    skip_type="time_uniform",
    method="singlestep",
)

Implementation Details of Singlestep DPM-Solver

Given a fixed NFE == steps, the sampling procedure is:


3.2. Sampling by multistep DPM-Solver

For discrete-time DPMs, we do not need to specify the t_start and t_end. The default setting is to sample from the discrete-time label $N-1$ to the discrete-time label $0$. For example,

## discrete-time DPMs
x_sample = dpm_solver.sample(
    x_T,
    steps=20,
    order=2,
    skip_type="time_uniform",
    method="multistep",
)


For continuous-time DPMs, we sample from t_start=1.0 (the default setting) to t_end. We recommend t_end=1e-3 for steps <= 15, and t_end=1e-4 for steps > 15. For example:

x_sample = dpm_solver.sample(
    x_T,
    t_end=1e-3,
    steps=10,
    order=2,
    skip_type="time_uniform",
    method="multistep",
)
## continuous-time DPMs
x_sample = dpm_solver.sample(
    x_T,
    t_end=1e-4,
    steps=20,
    order=3,
    skip_type="time_uniform",
    method="multistep",
)

Implementation Details of Multistep DPM-Solver

We initialize the first order values by lower order multistep solvers.

Given a fixed NFE == steps, the sampling procedure is:

3.3. Sampling by adaptive step size DPM-Solver

For continuous-time DPMs, we recommend t_end=1e-4 for better sample quality.

We ignore steps and use adaptive step size DPM-Solver with a higher order of order. You can adjust the absolute tolerance atol and the relative tolerance rtol to balance the computatation costs (NFE) and the sample quality. For image data, we recommend atol=0.0078 (the default setting).

For example, to sample by DPM-Solver-12:

x_sample = dpm_solver.sample(
    x_T,
    t_end=1e-4,
    order=2,
    method="adaptive",
    rtol=0.05,
)


3.4. Sampling by Singlestep DPM-Solver-k for k = 1, 2, 3

We use DPM-Solver-order for order = 1 or 2 or 3, with total [steps // order] * order NFE.

For example, to sample by DPM-Solver-3:

x_sample = dpm_solver.sample(
    x_T,
    steps=30,
    order=3,
    skip_type="time_uniform",
    method="singlestep_fixed",
)


TODO List


References

If you find the code useful for your research, please consider citing

@article{lu2022dpm,
  title={DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps},
  author={Lu, Cheng and Zhou, Yuhao and Bao, Fan and Chen, Jianfei and Li, Chongxuan and Zhu, Jun},
  journal={arXiv preprint arXiv:2206.00927},
  year={2022}
}