ding.torch_utils.diffusion_SDE.dpm_solver_pytorch

ding.torch_utils.diffusion_SDE.dpm_solver_pytorch

NoiseScheduleVP

__init__(schedule='discrete', betas=None, alphas_cumprod=None, continuous_beta_0=0.1, continuous_beta_1=20.0)

Create a wrapper class for the forward SDE (VP type).

---

Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t. We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.

---

The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper). Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have: log_alpha_t = self.marginal_log_mean_coeff(t) sigma_t = self.marginal_std(t) lambda_t = self.marginal_lambda(t)

Moreover, as lambda(t) is an invertible function, we also support its inverse function:

t = self.inverse_lambda(lambda_t)

===============================================================

We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).

1. For discrete-time DPMs:

   For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by: t_i = (i + 1) / N e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.

   Args: betas: A torch.Tensor. The beta array for the discrete-time DPM. (See the original DDPM paper for details) alphas_cumprod: A torch.Tensor. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)

   Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of betas and alphas_cumprod.

   Important: Please pay special attention for the args for alphas_cumprod: The alphas_cumprod is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have alpha_{t_n} = \sqrt{\hat{alpha_n}}, and log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).

2. For continuous-time DPMs:

   We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise schedule are the default settings in DDPM and improved-DDPM:

   Args: beta_min: A float number. The smallest beta for the linear schedule. beta_max: A float number. The largest beta for the linear schedule. cosine_s: A float number. The hyperparameter in the cosine schedule. cosine_beta_max: A float number. The hyperparameter in the cosine schedule. T: A float number. The ending time of the forward process.

===============================================================

Parameters:

Name	Type	Description	Default
schedule		A str. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs, 'linear' or 'cosine' for continuous-time DPMs.	'discrete'

Returns: A wrapper object of the forward SDE (VP type).

===============================================================

Example:

For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):

ns = NoiseScheduleVP('discrete', betas=betas)

For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):

ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)

For continuous-time DPMs (VPSDE), linear schedule:

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

marginal_log_mean_coeff(t)

Compute log(alpha_t) of a given continuous-time label t in [0, T].

marginal_alpha(t)

Compute alpha_t of a given continuous-time label t in [0, T].

marginal_std(t)

Compute sigma_t of a given continuous-time label t in [0, T].

marginal_lambda(t)

Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].

inverse_lambda(lamb)

Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.

DPM_Solver

__init__(model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.0)

Construct a DPM-Solver.

We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0"). If predict_x0 is False, we use the solver for the noise prediction model (DPM-Solver). If predict_x0 is True, we use the solver for the data prediction model (DPM-Solver++). In such case, we further support the "dynamic thresholding" in [1] when thresholding is True. The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.

Parameters:

Name	Type	Description	Default
model_fn		A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]): def model_fn(x, t_continuous): return noise	required
noise_schedule		A noise schedule object, such as NoiseScheduleVP.	required
predict_x0		A bool. If true, use the data prediction model; else, use the noise prediction model.	False
thresholding		A bool. Valid when predict_x0 is True. Whether to use the "dynamic thresholding" in [1].	False
max_val		A float. Valid when both predict_x0 and thresholding are True. The max value for thresholding.	1.0

[1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.

noise_prediction_fn(x, t)

Return the noise prediction model.

data_prediction_fn(x, t)

Return the data prediction model (with thresholding).

model_fn(x, t)

Convert the model to the noise prediction model or the data prediction model.

get_time_steps(skip_type, t_T, t_0, N, device)

Compute the intermediate time steps for sampling.

Parameters:

Name	Type	Description	Default
skip_type		A str. The type for the spacing of the time steps. We support three types: - 'logSNR': uniform logSNR for the time steps. - 'time_uniform': uniform time for the time steps. (Recommended for high-resolutional data.) - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)	required
t_T		A float. The starting time of the sampling (default is T).	required
t_0		A float. The ending time of the sampling (default is epsilon).	required
N		A int. The total number of the spacing of the time steps.	required
device		A torch device.	required

Returns: A pytorch tensor of the time steps, with the shape (N + 1,).

get_orders_for_singlestep_solver(steps, order)

Get the order of each step for sampling by the singlestep DPM-Solver.

We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". Given a fixed number of function evaluations by steps, the sampling procedure by DPM-Solver-fast is: - If order == 1: We take steps of DPM-Solver-1 (i.e. DDIM). - If order == 2: - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. - If steps % 2 == 0, we use K steps of DPM-Solver-2. - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. - If order == 3: - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.

============================================ Args: order: A int. The max order for the solver (2 or 3). steps: A int. The total number of function evaluations (NFE). Returns: orders: A list of the solver order of each step.

denoise_fn(x, s)

Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.

dpm_solver_first_update(x, s, t, model_s=None, return_intermediate=False)

DPM-Solver-1 (equivalent to DDIM) from time s to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
s		A pytorch tensor. The starting time, with the shape (x.shape[0],).	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
model_s		A pytorch tensor. The model function evaluated at time s. If model_s is None, we evaluate the model by x and s; otherwise we directly use it.	None
return_intermediate		A bool. If true, also return the model value at time s.	False

Returns: x_t: A pytorch tensor. The approximated solution at time t.

singlestep_dpm_solver_second_update(x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpm_solver')

Singlestep solver DPM-Solver-2 from time s to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
s		A pytorch tensor. The starting time, with the shape (x.shape[0],).	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
r1		A float. The hyperparameter of the second-order solver.	0.5
model_s		A pytorch tensor. The model function evaluated at time s. If model_s is None, we evaluate the model by x and s; otherwise we directly use it.	None
return_intermediate		A bool. If true, also return the model value at time s and s1 (the intermediate time).	False
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_t: A pytorch tensor. The approximated solution at time t.

singlestep_dpm_solver_third_update(x, s, t, r1=1.0 / 3.0, r2=2.0 / 3.0, model_s=None, model_s1=None, return_intermediate=False, solver_type='dpm_solver')

Singlestep solver DPM-Solver-3 from time s to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
s		A pytorch tensor. The starting time, with the shape (x.shape[0],).	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
r1		A float. The hyperparameter of the third-order solver.	1.0 / 3.0
r2		A float. The hyperparameter of the third-order solver.	2.0 / 3.0
model_s		A pytorch tensor. The model function evaluated at time s. If model_s is None, we evaluate the model by x and s; otherwise we directly use it.	None
model_s1		A pytorch tensor. The model function evaluated at time s1 (the intermediate time given by r1). If model_s1 is None, we evaluate the model at s1; otherwise we directly use it.	None
return_intermediate		A bool. If true, also return the model value at time s, s1 and s2 (the intermediate times).	False
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_t: A pytorch tensor. The approximated solution at time t.

multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type='dpm_solver')

Multistep solver DPM-Solver-2 from time t_prev_list[-1] to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
model_prev_list		A list of pytorch tensor. The previous computed model values.	required
t_prev_list		A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_t: A pytorch tensor. The approximated solution at time t.

multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type='dpm_solver')

Multistep solver DPM-Solver-3 from time t_prev_list[-1] to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
model_prev_list		A list of pytorch tensor. The previous computed model values.	required
t_prev_list		A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_t: A pytorch tensor. The approximated solution at time t.

singlestep_dpm_solver_update(x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None, r2=None)

Singlestep DPM-Solver with the order order from time s to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
s		A pytorch tensor. The starting time, with the shape (x.shape[0],).	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
order		A int. The order of DPM-Solver. We only support order == 1 or 2 or 3.	required
return_intermediate		A bool. If true, also return the model value at time s, s1 and s2 (the intermediate times).	False
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'
r1		A float. The hyperparameter of the second-order or third-order solver.	None
r2		A float. The hyperparameter of the third-order solver.	None

Returns: x_t: A pytorch tensor. The approximated solution at time t.

multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver')

Multistep DPM-Solver with the order order from time t_prev_list[-1] to time t.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time s.	required
model_prev_list		A list of pytorch tensor. The previous computed model values.	required
t_prev_list		A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)	required
t		A pytorch tensor. The ending time, with the shape (x.shape[0],).	required
order		A int. The order of DPM-Solver. We only support order == 1 or 2 or 3.	required
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_t: A pytorch tensor. The approximated solution at time t.

dpm_solver_adaptive(x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-05, solver_type='dpm_solver')

The adaptive step size solver based on singlestep DPM-Solver.

Parameters:

Name	Type	Description	Default
x		A pytorch tensor. The initial value at time t_T.	required
order		A int. The (higher) order of the solver. We only support order == 2 or 3.	required
t_T		A float. The starting time of the sampling (default is T).	required
t_0		A float. The ending time of the sampling (default is epsilon).	required
h_init		A float. The initial step size (for logSNR).	0.05
atol		A float. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].	0.0078
rtol		A float. The relative tolerance of the solver. The default setting is 0.05.	0.05
theta		A float. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].	0.9
t_err		A float. The tolerance for the time. We solve the diffusion ODE until the absolute error between the current time and t_0 is less than t_err. The default setting is 1e-5.	1e-05
solver_type		either 'dpm_solver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpm_solver' type.	'dpm_solver'

Returns: x_0: A pytorch tensor. The approximated solution at time t_0.

[1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.

sample(x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform', method='singlestep', denoise=False, solver_type='dpm_solver', atol=0.0078, rtol=0.05)

Compute the sample at time t_end by DPM-Solver, given the initial x at time t_start.

=====================================================

We support the following algorithms for both noise prediction model and data prediction model

• 'singlestep': Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of 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. Given a fixed NFE == steps, the sampling procedure is: - If order == 1: - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM). - If order == 2: - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling. - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2. - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - If order == 3: - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1. - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
• 'multistep': Multistep DPM-Solver with the order of order. The total number of function evaluations (NFE) == steps. We initialize the first order values by lower order multistep solvers. Given a fixed NFE == steps, the sampling procedure is: Denote K = steps. - If order == 1: - We use K steps of DPM-Solver-1 (i.e. DDIM). - If order == 2: - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2. - If order == 3: - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
• 'singlestep_fixed': Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3). We use singlestep DPM-Solver-order for order=1 or 2 or 3, with total [steps // order] * order NFE.
• 'adaptive': Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper). 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. - If order == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2. - If order == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.

=====================================================

Some advices for choosing the algorithm

• For unconditional sampling or guided sampling with small guidance scale by DPMs: Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with order = 3. e.g. >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False) >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, skip_type='time_uniform', method='singlestep')
• For guided sampling with large guidance scale by DPMs: Use multistep DPM-Solver with predict_x0 = True and order = 2. e.g. >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True) >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2, skip_type='time_uniform', method='multistep')

We support three types of skip_type: - 'logSNR': uniform logSNR for the time steps. Recommended for low-resolutional images - 'time_uniform': uniform time for the time steps. Recommended for high-resolutional images. - 'time_quadratic': quadratic time for the time steps.

===================================================== Args: x: A pytorch tensor. The initial value at time t_start e.g. if t_start == T, then x is a sample from the standard normal distribution. steps: A int. The total number of function evaluations (NFE). t_start: A float. The starting time of the sampling. If T is None, we use self.noise_schedule.T (default is 1.0). t_end: A float. The ending time of the sampling. If t_end is None, we use 1. / self.noise_schedule.total_N. e.g. if total_N == 1000, we have t_end == 1e-3. For discrete-time DPMs: - We recommend t_end == 1. / self.noise_schedule.total_N. For continuous-time DPMs: - We recommend t_end == 1e-3 when steps <= 15; and t_end == 1e-4 when steps > 15. order: A int. The order of DPM-Solver. skip_type: A str. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'. method: A str. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'. denoise: A bool. Whether to denoise at the final step. Default is False. If denoise is True, the total NFE is (steps + 1). solver_type: A str. The taylor expansion type for the solver. dpm_solver or taylor. We recommend dpm_solver. atol: A float. The absolute tolerance of the adaptive step size solver. Valid when method == 'adaptive'. rtol: A float. The relative tolerance of the adaptive step size solver. Valid when method == 'adaptive'. Returns: x_end: A pytorch tensor. The approximated solution at time t_end.

model_wrapper(model, noise_schedule, model_type='noise', model_kwargs={}, guidance_type='uncond', condition=None, unconditional_condition=None, guidance_scale=1.0, classifier_fn=None, classifier_kwargs={})

Create a wrapper function for the noise prediction model.

DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.

We support four types of the diffusion model by setting model_type:

1. "noise": noise prediction model. (Trained by predicting noise).

2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).

3. "v": velocity prediction model. (Trained by predicting the velocity).
    The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].

    [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
        arXiv preprint arXiv:2202.00512 (2022).
    [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
        arXiv preprint arXiv:2210.02303 (2022).

4. "score": marginal score function. (Trained by denoising score matching).
    Note that the score function and the noise prediction model follows a simple relationship:
    ```
        noise(x_t, t) = -sigma_t * score(x_t, t)
    ```

We support three types of guided sampling by DPMs by setting guidance_type: 1. "uncond": unconditional sampling by DPMs. The input model has the following format: model(x, t_input, **model_kwargs) -> noise | x_start | v | score

2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
    The input `model` has the following format:
    ``
        model(x, t_input, **model_kwargs) -> noise | x_start | v | score
    ``

    The input `classifier_fn` has the following format:
    ``
        classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
    ``

    [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
        in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.

3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
    The input `model` has the following format:
    ``
        model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
    ``
    And if cond == `unconditional_condition`, the model output is the unconditional DPM output.

    [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
        arXiv preprint arXiv:2207.12598 (2022).

The t_input is the time label of the model, which may be discrete-time labels (i.e. 0 to 999) or continuous-time labels (i.e. epsilon to T).

We wrap the model function to accept only x and t_continuous as inputs, and outputs the predicted noise: def model_fn(x, t_continuous) -> noise: t_input = get_model_input_time(t_continuous) return noise_pred(model, x, t_input, **model_kwargs) where t_continuous is the continuous time labels (i.e. epsilon to T). And we use model_fn for DPM-Solver.

===============================================================

Parameters:

Name	Type	Description	Default
model		A diffusion model with the corresponding format described above.	required
noise_schedule		A noise schedule object, such as NoiseScheduleVP.	required
model_type		A str. The parameterization type of the diffusion model. "noise" or "x_start" or "v" or "score".	'noise'
model_kwargs		A dict. A dict for the other inputs of the model function.	{}
guidance_type		A str. The type of the guidance for sampling. "uncond" or "classifier" or "classifier-free".	'uncond'
condition		A pytorch tensor. The condition for the guided sampling. Only used for "classifier" or "classifier-free" guidance type.	None
unconditional_condition		A pytorch tensor. The condition for the unconditional sampling. Only used for "classifier-free" guidance type.	None
guidance_scale		A float. The scale for the guided sampling.	1.0
classifier_fn		A classifier function. Only used for the classifier guidance.	None
classifier_kwargs		A dict. A dict for the other inputs of the classifier function.	{}

Returns: A noise prediction model that accepts the noised data and the continuous time as the inputs.

interpolate_fn(x, xp, yp)

A piecewise linear function y = f(x), using xp and yp as keypoints. We implement f(x) in a differentiable way (i.e. applicable for autograd). The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)

Parameters:

Name	Type	Description	Default
x		PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels	required
xp		PyTorch tensor with shape [C, K], where K is the number of keypoints.	required
yp		PyTorch tensor with shape [C, K].	required

Returns: The function values f(x), with shape [N, C].

expand_dims(v, dims)

Expand the tensor v to the dim dims.

Parameters:

Name	Type	Description	Default
`v`		a PyTorch tensor with shape [N].	required
`dim`		a int.	required

Returns: a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is dims.

Full Source Code

../ding/torch_utils/diffusion_SDE/dpm_solver_pytorch.py

1############################################################# 2# This DPM-Solver snippet is from https://github.com/ChenDRAG/CEP-energy-guided-diffusion 3# wich is based on https://github.com/LuChengTHU/dpm-solver 4############################################################# 5 6import torch 7import math 8 9 10class NoiseScheduleVP: 11 12 def __init__( 13 self, 14 schedule='discrete', 15 betas=None, 16 alphas_cumprod=None, 17 continuous_beta_0=0.1, 18 continuous_beta_1=20., 19 ): 20 """Create a wrapper class for the forward SDE (VP type). 21 22 *** 23 Update: We support discrete-time diffusion models by implementing a picewise linear interpolation \ 24 for log_alpha_t. We recommend to use schedule='discrete' for the discrete-time diffusion models, \ 25 especially for high-resolution images. 26 *** 27 28 The forward SDE ensures that the condition distribution \ 29 q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). 30 We further define lambda_t = log(alpha_t) - log(sigma_t), \ 31 which is the half-logSNR (described in the DPM-Solver paper). 32 Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. 33 For t in [0, T], we have: 34 log_alpha_t = self.marginal_log_mean_coeff(t) 35 sigma_t = self.marginal_std(t) 36 lambda_t = self.marginal_lambda(t) 37 38 Moreover, as lambda(t) is an invertible function, we also support its inverse function: 39 40 t = self.inverse_lambda(lambda_t) 41 42 =============================================================== 43 44 We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and \ 45 continuous-time DPMs (trained on t in [t_0, T]). 46 47 1. For discrete-time DPMs: 48 49 For discrete-time DPMs trained on n = 0, 1, ..., N-1, \ 50 we convert the discrete steps to continuous time steps by: 51 t_i = (i + 1) / N 52 e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. 53 We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3. 54 55 Args: 56 betas: A `torch.Tensor`. The beta array for the discrete-time DPM. \ 57 (See the original DDPM paper for details) 58 alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. \ 59 (See the original DDPM paper for details) 60 61 Note that we always have alphas_cumprod = cumprod(betas). \ 62 Therefore, we only need to set one of `betas` and `alphas_cumprod`. 63 64 **Important**: Please pay special attention for the args for `alphas_cumprod`: 65 The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. \ 66 Specifically, DDPMs assume that 67 q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). 68 Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. 69 In fact, we have 70 alpha_{t_n} = \sqrt{\hat{alpha_n}}, 71 and 72 log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}). 73 74 75 2. For continuous-time DPMs: 76 77 We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). 78 The hyperparameters for the noise schedule are the default settings in DDPM and improved-DDPM: 79 80 Args: 81 beta_min: A `float` number. The smallest beta for the linear schedule. 82 beta_max: A `float` number. The largest beta for the linear schedule. 83 cosine_s: A `float` number. The hyperparameter in the cosine schedule. 84 cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule. 85 T: A `float` number. The ending time of the forward process. 86 87 =============================================================== 88 89 Args: 90 schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs, 91 'linear' or 'cosine' for continuous-time DPMs. 92 Returns: 93 A wrapper object of the forward SDE (VP type). 94 95 =============================================================== 96 97 Example: 98 99 # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1): 100 >>> ns = NoiseScheduleVP('discrete', betas=betas) 101 102 # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array \ 103 for n = 0, 1, ..., N - 1): 104 >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod) 105 106 # For continuous-time DPMs (VPSDE), linear schedule: 107 >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.) 108 109 """ 110 111 if schedule not in ['discrete', 'linear', 'cosine']: 112 raise ValueError( 113 "Unsupported noise schedule {}. \ 114 The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule) 115 ) 116 117 self.schedule = schedule 118 if schedule == 'discrete': 119 if betas is not None: 120 log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0) 121 else: 122 assert alphas_cumprod is not None 123 log_alphas = 0.5 * torch.log(alphas_cumprod) 124 self.total_N = len(log_alphas) 125 self.T = 1. 126 self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)) 127 self.log_alpha_array = log_alphas.reshape(( 128 1, 129 -1, 130 )) 131 else: 132 self.total_N = 1000 133 self.beta_0 = continuous_beta_0 134 self.beta_1 = continuous_beta_1 135 self.cosine_s = 0.008 136 self.cosine_beta_max = 999. 137 self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi 138 ) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s 139 self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.)) 140 self.schedule = schedule 141 if schedule == 'cosine': 142 # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T. 143 # Note that T = 0.9946 may be not the optimal setting. However, we find it works well. 144 self.T = 0.9946 145 else: 146 self.T = 1. 147 148 def marginal_log_mean_coeff(self, t): 149 """ 150 Compute log(alpha_t) of a given continuous-time label t in [0, T]. 151 """ 152 if self.schedule == 'discrete': 153 return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), 154 self.log_alpha_array.to(t.device)).reshape((-1)) 155 elif self.schedule == 'linear': 156 return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0 157 elif self.schedule == 'cosine': 158 log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.)) 159 log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0 160 return log_alpha_t 161 162 def marginal_alpha(self, t): 163 """ 164 Compute alpha_t of a given continuous-time label t in [0, T]. 165 """ 166 return torch.exp(self.marginal_log_mean_coeff(t)) 167 168 def marginal_std(self, t): 169 """ 170 Compute sigma_t of a given continuous-time label t in [0, T]. 171 """ 172 return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t))) 173 174 def marginal_lambda(self, t): 175 """ 176 Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T]. 177 """ 178 log_mean_coeff = self.marginal_log_mean_coeff(t) 179 log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff)) 180 return log_mean_coeff - log_std 181 182 def inverse_lambda(self, lamb): 183 """ 184 Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t. 185 """ 186 if self.schedule == 'linear': 187 tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1, )).to(lamb)) 188 Delta = self.beta_0 ** 2 + tmp 189 return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0) 190 elif self.schedule == 'discrete': 191 log_alpha = -0.5 * torch.logaddexp(torch.zeros((1, )).to(lamb.device), -2. * lamb) 192 t = interpolate_fn( 193 log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), 194 torch.flip(self.t_array.to(lamb.device), [1]) 195 ) 196 return t.reshape((-1, )) 197 else: 198 log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1, )).to(lamb)) 199 t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0) 200 ) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s 201 t = t_fn(log_alpha) 202 return t 203 204 205def model_wrapper( 206 model, 207 noise_schedule, 208 model_type="noise", 209 model_kwargs={}, 210 guidance_type="uncond", 211 condition=None, 212 unconditional_condition=None, 213 guidance_scale=1., 214 classifier_fn=None, 215 classifier_kwargs={}, 216): 217 """Create a wrapper function for the noise prediction model. 218 219 DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to 220 firstly wrap the model function to a noise prediction model that accepts the continuous time as the input. 221 222 We support four types of the diffusion model by setting `model_type`: 223 224 1. "noise": noise prediction model. (Trained by predicting noise). 225 226 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0). 227 228 3. "v": velocity prediction model. (Trained by predicting the velocity). 229 The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2]. 230 231 [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." 232 arXiv preprint arXiv:2202.00512 (2022). 233 [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models." 234 arXiv preprint arXiv:2210.02303 (2022). 235 236 4. "score": marginal score function. (Trained by denoising score matching). 237 Note that the score function and the noise prediction model follows a simple relationship: 238 ``` 239 noise(x_t, t) = -sigma_t * score(x_t, t) 240 ``` 241 242 We support three types of guided sampling by DPMs by setting `guidance_type`: 243 1. "uncond": unconditional sampling by DPMs. 244 The input `model` has the following format: 245 `` 246 model(x, t_input, **model_kwargs) -> noise | x_start | v | score 247 `` 248 249 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier. 250 The input `model` has the following format: 251 `` 252 model(x, t_input, **model_kwargs) -> noise | x_start | v | score 253 `` 254 255 The input `classifier_fn` has the following format: 256 `` 257 classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond) 258 `` 259 260 [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis," 261 in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794. 262 263 3. "classifier-free": classifier-free guidance sampling by conditional DPMs. 264 The input `model` has the following format: 265 `` 266 model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score 267 `` 268 And if cond == `unconditional_condition`, the model output is the unconditional DPM output. 269 270 [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." 271 arXiv preprint arXiv:2207.12598 (2022). 272 273 274 The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999) 275 or continuous-time labels (i.e. epsilon to T). 276 277 We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise: 278 `` 279 def model_fn(x, t_continuous) -> noise: 280 t_input = get_model_input_time(t_continuous) 281 return noise_pred(model, x, t_input, **model_kwargs) 282 `` 283 where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver. 284 285 =============================================================== 286 287 Args: 288 model: A diffusion model with the corresponding format described above. 289 noise_schedule: A noise schedule object, such as NoiseScheduleVP. 290 model_type: A `str`. The parameterization type of the diffusion model. 291 "noise" or "x_start" or "v" or "score". 292 model_kwargs: A `dict`. A dict for the other inputs of the model function. 293 guidance_type: A `str`. The type of the guidance for sampling. 294 "uncond" or "classifier" or "classifier-free". 295 condition: A pytorch tensor. The condition for the guided sampling. 296 Only used for "classifier" or "classifier-free" guidance type. 297 unconditional_condition: A pytorch tensor. The condition for the unconditional sampling. 298 Only used for "classifier-free" guidance type. 299 guidance_scale: A `float`. The scale for the guided sampling. 300 classifier_fn: A classifier function. Only used for the classifier guidance. 301 classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function. 302 Returns: 303 A noise prediction model that accepts the noised data and the continuous time as the inputs. 304 """ 305 306 def get_model_input_time(t_continuous): 307 """ 308 Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time. 309 For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N]. 310 For continuous-time DPMs, we just use `t_continuous`. 311 """ 312 if noise_schedule.schedule == 'discrete': 313 return (t_continuous - 1. / noise_schedule.total_N) * 1000. 314 else: 315 return t_continuous 316 317 def noise_pred_fn(x, t_continuous, cond=None): 318 if t_continuous.reshape((-1, )).shape[0] == 1: 319 t_continuous = t_continuous.expand((x.shape[0])) 320 t_input = get_model_input_time(t_continuous) 321 if cond is None: 322 output = model(x, t_input, **model_kwargs) 323 else: 324 output = model(x, t_input, cond, **model_kwargs) 325 if model_type == "noise": 326 return output 327 elif model_type == "x_start": 328 alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) 329 dims = x.dim() 330 return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims) 331 elif model_type == "v": 332 alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) 333 dims = x.dim() 334 return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x 335 elif model_type == "score": 336 sigma_t = noise_schedule.marginal_std(t_continuous) 337 dims = x.dim() 338 return -expand_dims(sigma_t, dims) * output 339 340 def cond_grad_fn(x, t_input): 341 """ 342 Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t). 343 """ 344 with torch.enable_grad(): 345 x_in = x.detach().requires_grad_(True) 346 log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs) 347 return torch.autograd.grad(log_prob.sum(), x_in)[0] 348 349 def model_fn(x, t_continuous): 350 """ 351 The noise predicition model function that is used for DPM-Solver. 352 """ 353 if t_continuous.reshape((-1, )).shape[0] == 1: 354 t_continuous = t_continuous.expand((x.shape[0])) 355 if guidance_type == "uncond": 356 return noise_pred_fn(x, t_continuous) 357 elif guidance_type == "classifier": 358 assert classifier_fn is not None 359 t_input = get_model_input_time(t_continuous) 360 cond_grad = cond_grad_fn(x, t_input) 361 sigma_t = noise_schedule.marginal_std(t_continuous) 362 noise = noise_pred_fn(x, t_continuous) 363 return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad 364 elif guidance_type == "classifier-free": 365 if guidance_scale == 1. or unconditional_condition is None: 366 return noise_pred_fn(x, t_continuous, cond=condition) 367 else: 368 x_in = torch.cat([x] * 2) 369 t_in = torch.cat([t_continuous] * 2) 370 c_in = torch.cat([unconditional_condition, condition]) 371 noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2) 372 return noise_uncond + guidance_scale * (noise - noise_uncond) 373 374 assert model_type in ["noise", "x_start", "v"] 375 assert guidance_type in ["uncond", "classifier", "classifier-free"] 376 return model_fn 377 378 379class DPM_Solver: 380 381 def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.): 382 """Construct a DPM-Solver. 383 384 We support both the noise prediction model ("predicting epsilon") and \ 385 the data prediction model ("predicting x0"). 386 If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver). 387 If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++). 388 In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True. 389 The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs \ 390 with large guidance scales. 391 392 Args: 393 model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]): 394 `` 395 def model_fn(x, t_continuous): 396 return noise 397 `` 398 noise_schedule: A noise schedule object, such as NoiseScheduleVP. 399 predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model. 400 thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1]. 401 max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. \ 402 The max value for thresholding. 403 404 [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, \ 405 Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. \ 406 Photorealistic text-to-image diffusion models with deep language understanding. \ 407 arXiv preprint arXiv:2205.11487, 2022b. 408 """ 409 self.model = model_fn 410 self.noise_schedule = noise_schedule 411 self.predict_x0 = predict_x0 412 self.thresholding = thresholding 413 self.max_val = max_val 414 415 def noise_prediction_fn(self, x, t): 416 """ 417 Return the noise prediction model. 418 """ 419 return self.model(x, t) 420 421 def data_prediction_fn(self, x, t): 422 """ 423 Return the data prediction model (with thresholding). 424 """ 425 noise = self.noise_prediction_fn(x, t) 426 dims = x.dim() 427 alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t) 428 x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims) 429 if self.thresholding: 430 p = 0.995 # A hyperparameter in the paper of "Imagen" [1]. 431 s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1) 432 s = expand_dims(torch.maximum(s, torch.ones_like(s).to(s.device)), dims) 433 x0 = torch.clamp(x0, -s, s) / (s / self.max_val) 434 return x0 435 436 def model_fn(self, x, t): 437 """ 438 Convert the model to the noise prediction model or the data prediction model. 439 """ 440 if self.predict_x0: 441 return self.data_prediction_fn(x, t) 442 else: 443 return self.noise_prediction_fn(x, t) 444 445 def get_time_steps(self, skip_type, t_T, t_0, N, device): 446 """Compute the intermediate time steps for sampling. 447 448 Args: 449 skip_type: A `str`. The type for the spacing of the time steps. We support three types: 450 - 'logSNR': uniform logSNR for the time steps. 451 - 'time_uniform': uniform time for the time steps. \ 452 (**Recommended for high-resolutional data**.) 453 - 'time_quadratic': quadratic time for the time steps. \ 454 (Used in DDIM for low-resolutional data.) 455 t_T: A `float`. The starting time of the sampling (default is T). 456 t_0: A `float`. The ending time of the sampling (default is epsilon). 457 N: A `int`. The total number of the spacing of the time steps. 458 device: A torch device. 459 Returns: 460 A pytorch tensor of the time steps, with the shape (N + 1,). 461 """ 462 if skip_type == 'logSNR': 463 lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device)) 464 lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device)) 465 logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device) 466 return self.noise_schedule.inverse_lambda(logSNR_steps) 467 elif skip_type == 'time_uniform': 468 return torch.linspace(t_T, t_0, N + 1).to(device) 469 elif skip_type == 'time_quadratic': 470 t_order = 2 471 t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device) 472 return t 473 else: 474 raise ValueError( 475 "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type) 476 ) 477 478 def get_orders_for_singlestep_solver(self, steps, order): 479 """ 480 Get the order of each step for sampling by the singlestep DPM-Solver. 481 482 We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". 483 Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is: 484 - If order == 1: 485 We take `steps` of DPM-Solver-1 (i.e. DDIM). 486 - If order == 2: 487 - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. 488 - If steps % 2 == 0, we use K steps of DPM-Solver-2. 489 - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. 490 - If order == 3: 491 - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. 492 - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, \ 493 and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. 494 - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. 495 - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2. 496 497 ============================================ 498 Args: 499 order: A `int`. The max order for the solver (2 or 3). 500 steps: A `int`. The total number of function evaluations (NFE). 501 Returns: 502 orders: A list of the solver order of each step. 503 """ 504 if order == 3: 505 K = steps // 3 + 1 506 if steps % 3 == 0: 507 orders = [ 508 3, 509 ] * (K - 2) + [2, 1] 510 elif steps % 3 == 1: 511 orders = [ 512 3, 513 ] * (K - 1) + [1] 514 else: 515 orders = [ 516 3, 517 ] * (K - 1) + [2] 518 return orders 519 elif order == 2: 520 K = steps // 2 521 if steps % 2 == 0: 522 # orders = [2,] * K 523 K = steps // 2 + 1 524 orders = [ 525 2, 526 ] * (K - 2) + [ 527 1, 528 ] * 2 529 else: 530 orders = [ 531 2, 532 ] * K + [1] 533 return orders 534 elif order == 1: 535 return [ 536 1, 537 ] * steps 538 else: 539 raise ValueError("'order' must be '1' or '2' or '3'.") 540 541 def denoise_fn(self, x, s): 542 """ 543 Denoise at the final step, which is equivalent to solve the ODE \ 544 from lambda_s to infty by first-order discretization. 545 """ 546 return self.data_prediction_fn(x, s) 547 548 def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False): 549 """ 550 DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. 551 552 Args: 553 x: A pytorch tensor. The initial value at time `s`. 554 s: A pytorch tensor. The starting time, with the shape (x.shape[0],). 555 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 556 model_s: A pytorch tensor. The model function evaluated at time `s`. 557 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. 558 return_intermediate: A `bool`. If true, also return the model value at time `s`. 559 Returns: 560 x_t: A pytorch tensor. The approximated solution at time `t`. 561 """ 562 ns = self.noise_schedule 563 dims = x.dim() 564 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) 565 h = lambda_t - lambda_s 566 log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t) 567 sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t) 568 alpha_t = torch.exp(log_alpha_t) 569 570 if self.predict_x0: 571 phi_1 = torch.expm1(-h) 572 if model_s is None: 573 model_s = self.model_fn(x, s) 574 x_t = (expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s) 575 if return_intermediate: 576 return x_t, {'model_s': model_s} 577 else: 578 return x_t 579 else: 580 phi_1 = torch.expm1(h) 581 if model_s is None: 582 model_s = self.model_fn(x, s) 583 x_t = ( 584 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - 585 expand_dims(sigma_t * phi_1, dims) * model_s 586 ) 587 if return_intermediate: 588 return x_t, {'model_s': model_s} 589 else: 590 return x_t 591 592 def singlestep_dpm_solver_second_update( 593 self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpm_solver' 594 ): 595 """ 596 Singlestep solver DPM-Solver-2 from time `s` to time `t`. 597 598 Args: 599 x: A pytorch tensor. The initial value at time `s`. 600 s: A pytorch tensor. The starting time, with the shape (x.shape[0],). 601 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 602 r1: A `float`. The hyperparameter of the second-order solver. 603 model_s: A pytorch tensor. The model function evaluated at time `s`. 604 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. 605 return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` \ 606 (the intermediate time). 607 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 608 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 609 Returns: 610 x_t: A pytorch tensor. The approximated solution at time `t`. 611 """ 612 if solver_type not in ['dpm_solver', 'taylor']: 613 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) 614 if r1 is None: 615 r1 = 0.5 616 ns = self.noise_schedule 617 dims = x.dim() 618 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) 619 h = lambda_t - lambda_s 620 lambda_s1 = lambda_s + r1 * h 621 s1 = ns.inverse_lambda(lambda_s1) 622 log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff( 623 s1 624 ), ns.marginal_log_mean_coeff(t) 625 sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t) 626 alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t) 627 628 if self.predict_x0: 629 phi_11 = torch.expm1(-r1 * h) 630 phi_1 = torch.expm1(-h) 631 632 if model_s is None: 633 model_s = self.model_fn(x, s) 634 x_s1 = (expand_dims(sigma_s1 / sigma_s, dims) * x - expand_dims(alpha_s1 * phi_11, dims) * model_s) 635 model_s1 = self.model_fn(x_s1, s1) 636 if solver_type == 'dpm_solver': 637 x_t = ( 638 expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s - 639 (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s) 640 ) 641 elif solver_type == 'taylor': 642 x_t = ( 643 expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + 644 (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (model_s1 - model_s) 645 ) 646 else: 647 phi_11 = torch.expm1(r1 * h) 648 phi_1 = torch.expm1(h) 649 650 if model_s is None: 651 model_s = self.model_fn(x, s) 652 x_s1 = ( 653 expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x - 654 expand_dims(sigma_s1 * phi_11, dims) * model_s 655 ) 656 model_s1 = self.model_fn(x_s1, s1) 657 if solver_type == 'dpm_solver': 658 x_t = ( 659 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - 660 expand_dims(sigma_t * phi_1, dims) * model_s - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * 661 (model_s1 - model_s) 662 ) 663 elif solver_type == 'taylor': 664 x_t = ( 665 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - 666 expand_dims(sigma_t * phi_1, dims) * model_s - 667 (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s) 668 ) 669 if return_intermediate: 670 return x_t, {'model_s': model_s, 'model_s1': model_s1} 671 else: 672 return x_t 673 674 def singlestep_dpm_solver_third_update( 675 self, 676 x, 677 s, 678 t, 679 r1=1. / 3., 680 r2=2. / 3., 681 model_s=None, 682 model_s1=None, 683 return_intermediate=False, 684 solver_type='dpm_solver' 685 ): 686 """ 687 Singlestep solver DPM-Solver-3 from time `s` to time `t`. 688 689 Args: 690 x: A pytorch tensor. The initial value at time `s`. 691 s: A pytorch tensor. The starting time, with the shape (x.shape[0],). 692 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 693 r1: A `float`. The hyperparameter of the third-order solver. 694 r2: A `float`. The hyperparameter of the third-order solver. 695 model_s: A pytorch tensor. The model function evaluated at time `s`. 696 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. 697 model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`). 698 If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it. 699 return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` \ 700 (the intermediate times). 701 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 702 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 703 Returns: 704 x_t: A pytorch tensor. The approximated solution at time `t`. 705 """ 706 if solver_type not in ['dpm_solver', 'taylor']: 707 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) 708 if r1 is None: 709 r1 = 1. / 3. 710 if r2 is None: 711 r2 = 2. / 3. 712 ns = self.noise_schedule 713 dims = x.dim() 714 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) 715 h = lambda_t - lambda_s 716 lambda_s1 = lambda_s + r1 * h 717 lambda_s2 = lambda_s + r2 * h 718 s1 = ns.inverse_lambda(lambda_s1) 719 s2 = ns.inverse_lambda(lambda_s2) 720 log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff( 721 s 722 ), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t) 723 sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std( 724 s2 725 ), ns.marginal_std(t) 726 alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t) 727 728 if self.predict_x0: 729 phi_11 = torch.expm1(-r1 * h) 730 phi_12 = torch.expm1(-r2 * h) 731 phi_1 = torch.expm1(-h) 732 phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1. 733 phi_2 = phi_1 / h + 1. 734 phi_3 = phi_2 / h - 0.5 735 736 if model_s is None: 737 model_s = self.model_fn(x, s) 738 if model_s1 is None: 739 x_s1 = (expand_dims(sigma_s1 / sigma_s, dims) * x - expand_dims(alpha_s1 * phi_11, dims) * model_s) 740 model_s1 = self.model_fn(x_s1, s1) 741 x_s2 = ( 742 expand_dims(sigma_s2 / sigma_s, dims) * x - expand_dims(alpha_s2 * phi_12, dims) * model_s + 743 r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s) 744 ) 745 model_s2 = self.model_fn(x_s2, s2) 746 if solver_type == 'dpm_solver': 747 x_t = ( 748 expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + 749 (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s) 750 ) 751 elif solver_type == 'taylor': 752 D1_0 = (1. / r1) * (model_s1 - model_s) 753 D1_1 = (1. / r2) * (model_s2 - model_s) 754 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) 755 D2 = 2. * (D1_1 - D1_0) / (r2 - r1) 756 x_t = ( 757 expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + 758 expand_dims(alpha_t * phi_2, dims) * D1 - expand_dims(alpha_t * phi_3, dims) * D2 759 ) 760 else: 761 phi_11 = torch.expm1(r1 * h) 762 phi_12 = torch.expm1(r2 * h) 763 phi_1 = torch.expm1(h) 764 phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1. 765 phi_2 = phi_1 / h - 1. 766 phi_3 = phi_2 / h - 0.5 767 768 if model_s is None: 769 model_s = self.model_fn(x, s) 770 if model_s1 is None: 771 x_s1 = ( 772 expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x - 773 expand_dims(sigma_s1 * phi_11, dims) * model_s 774 ) 775 model_s1 = self.model_fn(x_s1, s1) 776 x_s2 = ( 777 expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x - 778 expand_dims(sigma_s2 * phi_12, dims) * model_s - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * 779 (model_s1 - model_s) 780 ) 781 model_s2 = self.model_fn(x_s2, s2) 782 if solver_type == 'dpm_solver': 783 x_t = ( 784 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - 785 expand_dims(sigma_t * phi_1, dims) * model_s - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * 786 (model_s2 - model_s) 787 ) 788 elif solver_type == 'taylor': 789 D1_0 = (1. / r1) * (model_s1 - model_s) 790 D1_1 = (1. / r2) * (model_s2 - model_s) 791 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) 792 D2 = 2. * (D1_1 - D1_0) / (r2 - r1) 793 x_t = ( 794 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - 795 expand_dims(sigma_t * phi_1, dims) * model_s - expand_dims(sigma_t * phi_2, dims) * D1 - 796 expand_dims(sigma_t * phi_3, dims) * D2 797 ) 798 799 if return_intermediate: 800 return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2} 801 else: 802 return x_t 803 804 def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"): 805 """ 806 Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`. 807 808 Args: 809 x: A pytorch tensor. The initial value at time `s`. 810 model_prev_list: A list of pytorch tensor. The previous computed model values. 811 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) 812 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 813 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 814 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 815 Returns: 816 x_t: A pytorch tensor. The approximated solution at time `t`. 817 """ 818 if solver_type not in ['dpm_solver', 'taylor']: 819 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) 820 ns = self.noise_schedule 821 dims = x.dim() 822 model_prev_1, model_prev_0 = model_prev_list 823 t_prev_1, t_prev_0 = t_prev_list 824 lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda( 825 t_prev_0 826 ), ns.marginal_lambda(t) 827 log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) 828 sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) 829 alpha_t = torch.exp(log_alpha_t) 830 831 h_0 = lambda_prev_0 - lambda_prev_1 832 h = lambda_t - lambda_prev_0 833 r0 = h_0 / h 834 D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1) 835 if self.predict_x0: 836 if solver_type == 'dpm_solver': 837 x_t = ( 838 expand_dims(sigma_t / sigma_prev_0, dims) * x - 839 expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0 - 840 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0 841 ) 842 elif solver_type == 'taylor': 843 x_t = ( 844 expand_dims(sigma_t / sigma_prev_0, dims) * x - 845 expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0 + 846 expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0 847 ) 848 else: 849 if solver_type == 'dpm_solver': 850 x_t = ( 851 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - 852 expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - 853 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0 854 ) 855 elif solver_type == 'taylor': 856 x_t = ( 857 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - 858 expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - 859 expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0 860 ) 861 return x_t 862 863 def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'): 864 """ 865 Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`. 866 867 Args: 868 x: A pytorch tensor. The initial value at time `s`. 869 model_prev_list: A list of pytorch tensor. The previous computed model values. 870 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) 871 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 872 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 873 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 874 Returns: 875 x_t: A pytorch tensor. The approximated solution at time `t`. 876 """ 877 ns = self.noise_schedule 878 dims = x.dim() 879 model_prev_2, model_prev_1, model_prev_0 = model_prev_list 880 t_prev_2, t_prev_1, t_prev_0 = t_prev_list 881 lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda( 882 t_prev_1 883 ), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t) 884 log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) 885 sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) 886 alpha_t = torch.exp(log_alpha_t) 887 888 h_1 = lambda_prev_1 - lambda_prev_2 889 h_0 = lambda_prev_0 - lambda_prev_1 890 h = lambda_t - lambda_prev_0 891 r0, r1 = h_0 / h, h_1 / h 892 D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1) 893 D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2) 894 D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1) 895 D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1) 896 if self.predict_x0: 897 x_t = ( 898 expand_dims(sigma_t / sigma_prev_0, dims) * x - expand_dims(alpha_t * 899 (torch.exp(-h) - 1.), dims) * model_prev_0 + 900 expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1 - 901 expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2 902 ) 903 else: 904 x_t = ( 905 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - 906 expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - 907 expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1 - 908 expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2 909 ) 910 return x_t 911 912 def singlestep_dpm_solver_update( 913 self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None, r2=None 914 ): 915 """ 916 Singlestep DPM-Solver with the order `order` from time `s` to time `t`. 917 918 Args: 919 x: A pytorch tensor. The initial value at time `s`. 920 s: A pytorch tensor. The starting time, with the shape (x.shape[0],). 921 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 922 order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. 923 return_intermediate: A `bool`. If true, also return the model value at time `s`, \ 924 `s1` and `s2` (the intermediate times). 925 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 926 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 927 r1: A `float`. The hyperparameter of the second-order or third-order solver. 928 r2: A `float`. The hyperparameter of the third-order solver. 929 Returns: 930 x_t: A pytorch tensor. The approximated solution at time `t`. 931 """ 932 if order == 1: 933 return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate) 934 elif order == 2: 935 return self.singlestep_dpm_solver_second_update( 936 x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1 937 ) 938 elif order == 3: 939 return self.singlestep_dpm_solver_third_update( 940 x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1, r2=r2 941 ) 942 else: 943 raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) 944 945 def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'): 946 """ 947 Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`. 948 949 Args: 950 x: A pytorch tensor. The initial value at time `s`. 951 model_prev_list: A list of pytorch tensor. The previous computed model values. 952 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) 953 t: A pytorch tensor. The ending time, with the shape (x.shape[0],). 954 order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. 955 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 956 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 957 Returns: 958 x_t: A pytorch tensor. The approximated solution at time `t`. 959 """ 960 if order == 1: 961 return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1]) 962 elif order == 2: 963 return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) 964 elif order == 3: 965 return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) 966 else: 967 raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) 968 969 def dpm_solver_adaptive( 970 self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5, solver_type='dpm_solver' 971 ): 972 """ 973 The adaptive step size solver based on singlestep DPM-Solver. 974 975 Args: 976 x: A pytorch tensor. The initial value at time `t_T`. 977 order: A `int`. The (higher) order of the solver. We only support order == 2 or 3. 978 t_T: A `float`. The starting time of the sampling (default is T). 979 t_0: A `float`. The ending time of the sampling (default is epsilon). 980 h_init: A `float`. The initial step size (for logSNR). 981 atol: A `float`. The absolute tolerance of the solver. \ 982 For image data, the default setting is 0.0078, followed [1]. 983 rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05. 984 theta: A `float`. The safety hyperparameter for adapting the step size. \ 985 The default setting is 0.9, followed [1]. 986 t_err: A `float`. The tolerance for the time. \ 987 We solve the diffusion ODE until the absolute error between the \ 988 current time and `t_0` is less than `t_err`. The default setting is 1e-5. 989 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. 990 The type slightly impacts the performance. We recommend to use 'dpm_solver' type. 991 Returns: 992 x_0: A pytorch tensor. The approximated solution at time `t_0`. 993 994 [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, \ 995 "Gotta go fast when generating data with score-based models," \ 996 arXiv preprint arXiv:2105.14080, 2021. 997 """ 998 ns = self.noise_schedule 999 s = t_T * torch.ones((x.shape[0], )).to(x)1000 lambda_s = ns.marginal_lambda(s)1001 lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))1002 h = h_init * torch.ones_like(s).to(x)1003 x_prev = x1004 nfe = 01005 if order == 2:1006 r1 = 0.51007 lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)1008 higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(1009 x, s, t, r1=r1, solver_type=solver_type, **kwargs1010 )1011 elif order == 3:1012 r1, r2 = 1. / 3., 2. / 3.1013 lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(1014 x, s, t, r1=r1, return_intermediate=True, solver_type=solver_type1015 )1016 higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(1017 x, s, t, r1=r1, r2=r2, solver_type=solver_type, **kwargs1018 )1019 else:1020 raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))1021 while torch.abs((s - t_0)).mean() > t_err:1022 t = ns.inverse_lambda(lambda_s + h)1023 x_lower, lower_noise_kwargs = lower_update(x, s, t)1024 x_higher = higher_update(x, s, t, **lower_noise_kwargs)1025 delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))1026 norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))1027 E = norm_fn((x_higher - x_lower) / delta).max()1028 if torch.all(E <= 1.):1029 x = x_higher1030 s = t1031 x_prev = x_lower1032 lambda_s = ns.marginal_lambda(s)1033 h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)1034 nfe += order1035 print('adaptive solver nfe', nfe)1036 return x10371038 def sample(1039 self,1040 x,1041 steps=20,1042 t_start=None,1043 t_end=None,1044 order=3,1045 skip_type='time_uniform',1046 method='singlestep',1047 denoise=False,1048 solver_type='dpm_solver',1049 atol=0.0078,1050 rtol=0.05,1051 ):1052 """1053 Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.10541055 =====================================================10561057 We support the following algorithms for both noise prediction model and data prediction model:1058 - 'singlestep':1059 Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), \1060 which combines different orders of singlestep DPM-Solver.1061 We combine all the singlestep solvers with order <= `order` \1062 to use up all the function evaluations (steps).1063 The total number of function evaluations (NFE) == `steps`.1064 Given a fixed NFE == `steps`, the sampling procedure is:1065 - If `order` == 1:1066 - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).1067 - If `order` == 2:1068 - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.1069 - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.1070 - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 \1071 and 1 step of DPM-Solver-1.1072 - If `order` == 3:1073 - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.1074 - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, \1075 and 1 step of singlestep DPM-Solver-2 \1076 and 1 step of DPM-Solver-1.1077 - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 \1078 and 1 step of DPM-Solver-1.1079 - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 \1080 and 1 step of singlestep DPM-Solver-2.1081 - 'multistep':1082 Multistep DPM-Solver with the order of `order`.1083 The total number of function evaluations (NFE) == `steps`.1084 We initialize the first `order` values by lower order multistep solvers.1085 Given a fixed NFE == `steps`, the sampling procedure is:1086 Denote K = steps.1087 - If `order` == 1:1088 - We use K steps of DPM-Solver-1 (i.e. DDIM).1089 - If `order` == 2:1090 - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.1091 - If `order` == 3:1092 - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, \1093 then (K - 2) step of multistep DPM-Solver-3.1094 - 'singlestep_fixed':1095 Fixed order singlestep DPM-Solver \1096 (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).1097 We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, \1098 with total [`steps` // `order`] * `order` NFE.1099 - 'adaptive':1100 Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).1101 We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.1102 You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` \1103 to balance the computatation costs (NFE) and the sample quality.1104 - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.1105 - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 \1106 and singlestep DPM-Solver-3.11071108 =====================================================11091110 Some advices for choosing the algorithm:1111 - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:1112 Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.1113 e.g.1114 >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)1115 >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,1116 skip_type='time_uniform', method='singlestep')1117 - For **guided sampling with large guidance scale** by DPMs:1118 Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.1119 e.g.1120 >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)1121 >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,1122 skip_type='time_uniform', method='multistep')11231124 We support three types of `skip_type`:1125 - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**1126 - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.1127 - 'time_quadratic': quadratic time for the time steps.11281129 =====================================================1130 Args:1131 x: A pytorch tensor. The initial value at time `t_start`1132 e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.1133 steps: A `int`. The total number of function evaluations (NFE).1134 t_start: A `float`. The starting time of the sampling.1135 If `T` is None, we use self.noise_schedule.T (default is 1.0).1136 t_end: A `float`. The ending time of the sampling.1137 If `t_end` is None, we use 1. / self.noise_schedule.total_N.1138 e.g. if total_N == 1000, we have `t_end` == 1e-3.1139 For discrete-time DPMs:1140 - We recommend `t_end` == 1. / self.noise_schedule.total_N.1141 For continuous-time DPMs:1142 - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.1143 order: A `int`. The order of DPM-Solver.1144 skip_type: A `str`. The type for the spacing of the time steps. \1145 'time_uniform' or 'logSNR' or 'time_quadratic'.1146 method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.1147 denoise: A `bool`. Whether to denoise at the final step. Default is False.1148 If `denoise` is True, the total NFE is (`steps` + 1).1149 solver_type: A `str`. The taylor expansion type for the solver. \1150 `dpm_solver` or `taylor`. We recommend `dpm_solver`.1151 atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.1152 rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.1153 Returns:1154 x_end: A pytorch tensor. The approximated solution at time `t_end`.11551156 """1157 t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end1158 t_T = self.noise_schedule.T if t_start is None else t_start1159 device = x.device1160 if method == 'adaptive':1161 with torch.no_grad():1162 x = self.dpm_solver_adaptive(1163 x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol, solver_type=solver_type1164 )1165 elif method == 'multistep':1166 assert steps >= order1167 timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)1168 assert timesteps.shape[0] - 1 == steps1169 with torch.no_grad():1170 vec_t = timesteps[0].expand((x.shape[0]))1171 model_prev_list = [self.model_fn(x, vec_t)]1172 t_prev_list = [vec_t]1173 # Init the first `order` values by lower order multistep DPM-Solver.1174 for init_order in range(1, order):1175 vec_t = timesteps[init_order].expand(x.shape[0])1176 x = self.multistep_dpm_solver_update(1177 x, model_prev_list, t_prev_list, vec_t, init_order, solver_type=solver_type1178 )1179 model_prev_list.append(self.model_fn(x, vec_t))1180 t_prev_list.append(vec_t)1181 # Compute the remaining values by `order`-th order multistep DPM-Solver.1182 for step in range(order, steps + 1):1183 vec_t = timesteps[step].expand(x.shape[0])1184 x = self.multistep_dpm_solver_update(1185 x, model_prev_list, t_prev_list, vec_t, order, solver_type=solver_type1186 )1187 for i in range(order - 1):1188 t_prev_list[i] = t_prev_list[i + 1]1189 model_prev_list[i] = model_prev_list[i + 1]1190 t_prev_list[-1] = vec_t1191 # We do not need to evaluate the final model value.1192 if step < steps:1193 model_prev_list[-1] = self.model_fn(x, vec_t)1194 elif method in ['singlestep', 'singlestep_fixed']:1195 if method == 'singlestep':1196 orders = self.get_orders_for_singlestep_solver(steps=steps, order=order)1197 timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)1198 elif method == 'singlestep_fixed':1199 K = steps // order1200 orders = [1201 order,1202 ] * K1203 timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=(K * order), device=device)1204 with torch.no_grad():1205 i = 01206 for order in orders:1207 vec_s, vec_t = timesteps[i].expand(x.shape[0]), timesteps[i + order].expand(x.shape[0])1208 h = self.noise_schedule.marginal_lambda(timesteps[i + order]1209 ) - self.noise_schedule.marginal_lambda(timesteps[i])1210 r1 = None if order <= 1 else (1211 self.noise_schedule.marginal_lambda(timesteps[i + 1]) -1212 self.noise_schedule.marginal_lambda(timesteps[i])1213 ) / h1214 r2 = None if order <= 2 else (1215 self.noise_schedule.marginal_lambda(timesteps[i + 2]) -1216 self.noise_schedule.marginal_lambda(timesteps[i])1217 ) / h1218 x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)1219 i += order1220 if denoise:1221 x = self.denoise_fn(x, torch.ones((x.shape[0], )).to(device) * t_0)1222 return x122312241225#############################################################1226# other utility functions1227#############################################################122812291230def interpolate_fn(x, xp, yp):1231 """1232 A piecewise linear function y = f(x), using xp and yp as keypoints.1233 We implement f(x) in a differentiable way (i.e. applicable for autograd).1234 The function f(x) is well-defined for all x-axis. \1235 (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)12361237 Args:1238 x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels1239 (we use C = 1 for DPM-Solver).1240 xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.1241 yp: PyTorch tensor with shape [C, K].1242 Returns:1243 The function values f(x), with shape [N, C].1244 """1245 N, K = x.shape[0], xp.shape[1]1246 all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)1247 sorted_all_x, x_indices = torch.sort(all_x, dim=2)1248 x_idx = torch.argmin(x_indices, dim=2)1249 cand_start_idx = x_idx - 11250 start_idx = torch.where(1251 torch.eq(x_idx, 0),1252 torch.tensor(1, device=x.device),1253 torch.where(1254 torch.eq(x_idx, K),1255 torch.tensor(K - 2, device=x.device),1256 cand_start_idx,1257 ),1258 )1259 end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)1260 start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)1261 end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)1262 start_idx2 = torch.where(1263 torch.eq(x_idx, 0),1264 torch.tensor(0, device=x.device),1265 torch.where(1266 torch.eq(x_idx, K),1267 torch.tensor(K - 2, device=x.device),1268 cand_start_idx,1269 ),1270 )1271 y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)1272 start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)1273 end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)1274 cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)1275 return cand127612771278def expand_dims(v, dims):1279 """1280 Expand the tensor `v` to the dim `dims`.12811282 Args:1283 `v`: a PyTorch tensor with shape [N].1284 `dim`: a `int`.1285 Returns:1286 a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.1287 """1288 return v[(..., ) + (None, ) * (dims - 1)]