memef4rmer/edit_anything
0
1import torch2import torch.nn.functional as F3import math4from tqdm import tqdm5 6 7class NoiseScheduleVP:8 def __init__(9 self,10 schedule='discrete',11 betas=None,12 alphas_cumprod=None,13 continuous_beta_0=0.1,14 continuous_beta_1=20.,15 ):16 """Create a wrapper class for the forward SDE (VP type).17 ***18 Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.19 We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.20 ***21 The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).22 We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).23 Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:24 log_alpha_t = self.marginal_log_mean_coeff(t)25 sigma_t = self.marginal_std(t)26 lambda_t = self.marginal_lambda(t)27 Moreover, as lambda(t) is an invertible function, we also support its inverse function:28 t = self.inverse_lambda(lambda_t)29 ===============================================================30 We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).31 1. For discrete-time DPMs:32 For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:33 t_i = (i + 1) / N34 e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.35 We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.36 Args:37 betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)38 alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)39 Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.40 **Important**: Please pay special attention for the args for `alphas_cumprod`:41 The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that42 q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).43 Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have44 alpha_{t_n} = \sqrt{\hat{alpha_n}},45 and46 log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).47 2. For continuous-time DPMs:48 We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise49 schedule are the default settings in DDPM and improved-DDPM:50 Args:51 beta_min: A `float` number. The smallest beta for the linear schedule.52 beta_max: A `float` number. The largest beta for the linear schedule.53 cosine_s: A `float` number. The hyperparameter in the cosine schedule.54 cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.55 T: A `float` number. The ending time of the forward process.56 ===============================================================57 Args:58 schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,59 'linear' or 'cosine' for continuous-time DPMs.60 Returns:61 A wrapper object of the forward SDE (VP type).62 63 ===============================================================64 Example:65 # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):66 >>> ns = NoiseScheduleVP('discrete', betas=betas)67 # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):68 >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)69 # For continuous-time DPMs (VPSDE), linear schedule:70 >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)71 """72 73 if schedule not in ['discrete', 'linear', 'cosine']:74 raise ValueError(75 "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(76 schedule))77 78 self.schedule = schedule79 if schedule == 'discrete':80 if betas is not None:81 log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)82 else:83 assert alphas_cumprod is not None84 log_alphas = 0.5 * torch.log(alphas_cumprod)85 self.total_N = len(log_alphas)86 self.T = 1.87 self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))88 self.log_alpha_array = log_alphas.reshape((1, -1,))89 else:90 self.total_N = 100091 self.beta_0 = continuous_beta_092 self.beta_1 = continuous_beta_193 self.cosine_s = 0.00894 self.cosine_beta_max = 999.95 self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (96 1. + self.cosine_s) / math.pi - self.cosine_s97 self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))98 self.schedule = schedule99 if schedule == 'cosine':100 # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.101 # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.102 self.T = 0.9946103 else:104 self.T = 1.105 106 def marginal_log_mean_coeff(self, t):107 """108 Compute log(alpha_t) of a given continuous-time label t in [0, T].109 """110 if self.schedule == 'discrete':111 return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),112 self.log_alpha_array.to(t.device)).reshape((-1))113 elif self.schedule == 'linear':114 return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0115 elif self.schedule == 'cosine':116 log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))117 log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0118 return log_alpha_t119 120 def marginal_alpha(self, t):121 """122 Compute alpha_t of a given continuous-time label t in [0, T].123 """124 return torch.exp(self.marginal_log_mean_coeff(t))125 126 def marginal_std(self, t):127 """128 Compute sigma_t of a given continuous-time label t in [0, T].129 """130 return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))131 132 def marginal_lambda(self, t):133 """134 Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].135 """136 log_mean_coeff = self.marginal_log_mean_coeff(t)137 log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))138 return log_mean_coeff - log_std139 140 def inverse_lambda(self, lamb):141 """142 Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.143 """144 if self.schedule == 'linear':145 tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))146 Delta = self.beta_0 ** 2 + tmp147 return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)148 elif self.schedule == 'discrete':149 log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)150 t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),151 torch.flip(self.t_array.to(lamb.device), [1]))152 return t.reshape((-1,))153 else:154 log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))155 t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (156 1. + self.cosine_s) / math.pi - self.cosine_s157 t = t_fn(log_alpha)158 return t159 160 161def model_wrapper(162 model,163 noise_schedule,164 model_type="noise",165 model_kwargs={},166 guidance_type="uncond",167 condition=None,168 unconditional_condition=None,169 guidance_scale=1.,170 classifier_fn=None,171 classifier_kwargs={},172):173 """Create a wrapper function for the noise prediction model.174 DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to175 firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.176 We support four types of the diffusion model by setting `model_type`:177 1. "noise": noise prediction model. (Trained by predicting noise).178 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).179 3. "v": velocity prediction model. (Trained by predicting the velocity).180 The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].181 [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."182 arXiv preprint arXiv:2202.00512 (2022).183 [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."184 arXiv preprint arXiv:2210.02303 (2022).185 186 4. "score": marginal score function. (Trained by denoising score matching).187 Note that the score function and the noise prediction model follows a simple relationship:188 ```189 noise(x_t, t) = -sigma_t * score(x_t, t)190 ```191 We support three types of guided sampling by DPMs by setting `guidance_type`:192 1. "uncond": unconditional sampling by DPMs.193 The input `model` has the following format:194 ``195 model(x, t_input, **model_kwargs) -> noise | x_start | v | score196 ``197 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.198 The input `model` has the following format:199 ``200 model(x, t_input, **model_kwargs) -> noise | x_start | v | score201 ``202 The input `classifier_fn` has the following format:203 ``204 classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)205 ``206 [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"207 in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.208 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.209 The input `model` has the following format:210 ``211 model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score212 ``213 And if cond == `unconditional_condition`, the model output is the unconditional DPM output.214 [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."215 arXiv preprint arXiv:2207.12598 (2022).216 217 The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)218 or continuous-time labels (i.e. epsilon to T).219 We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:220 ``221 def model_fn(x, t_continuous) -> noise:222 t_input = get_model_input_time(t_continuous)223 return noise_pred(model, x, t_input, **model_kwargs)224 ``225 where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.226 ===============================================================227 Args:228 model: A diffusion model with the corresponding format described above.229 noise_schedule: A noise schedule object, such as NoiseScheduleVP.230 model_type: A `str`. The parameterization type of the diffusion model.231 "noise" or "x_start" or "v" or "score".232 model_kwargs: A `dict`. A dict for the other inputs of the model function.233 guidance_type: A `str`. The type of the guidance for sampling.234 "uncond" or "classifier" or "classifier-free".235 condition: A pytorch tensor. The condition for the guided sampling.236 Only used for "classifier" or "classifier-free" guidance type.237 unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.238 Only used for "classifier-free" guidance type.239 guidance_scale: A `float`. The scale for the guided sampling.240 classifier_fn: A classifier function. Only used for the classifier guidance.241 classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.242 Returns:243 A noise prediction model that accepts the noised data and the continuous time as the inputs.244 """245 246 def get_model_input_time(t_continuous):247 """248 Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.249 For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].250 For continuous-time DPMs, we just use `t_continuous`.251 """252 if noise_schedule.schedule == 'discrete':253 return (t_continuous - 1. / noise_schedule.total_N) * 1000.254 else:255 return t_continuous256 257 def noise_pred_fn(x, t_continuous, cond=None):258 if t_continuous.reshape((-1,)).shape[0] == 1:259 t_continuous = t_continuous.expand((x.shape[0]))260 t_input = get_model_input_time(t_continuous)261 if cond is None:262 output = model(x, t_input, **model_kwargs)263 else:264 output = model(x, t_input, cond, **model_kwargs)265 if model_type == "noise":266 return output267 elif model_type == "x_start":268 alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)269 dims = x.dim()270 return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)271 elif model_type == "v":272 alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)273 dims = x.dim()274 return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x275 elif model_type == "score":276 sigma_t = noise_schedule.marginal_std(t_continuous)277 dims = x.dim()278 return -expand_dims(sigma_t, dims) * output279 280 def cond_grad_fn(x, t_input):281 """282 Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).283 """284 with torch.enable_grad():285 x_in = x.detach().requires_grad_(True)286 log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)287 return torch.autograd.grad(log_prob.sum(), x_in)[0]288 289 def model_fn(x, t_continuous):290 """291 The noise predicition model function that is used for DPM-Solver.292 """293 if t_continuous.reshape((-1,)).shape[0] == 1:294 t_continuous = t_continuous.expand((x.shape[0]))295 if guidance_type == "uncond":296 return noise_pred_fn(x, t_continuous)297 elif guidance_type == "classifier":298 assert classifier_fn is not None299 t_input = get_model_input_time(t_continuous)300 cond_grad = cond_grad_fn(x, t_input)301 sigma_t = noise_schedule.marginal_std(t_continuous)302 noise = noise_pred_fn(x, t_continuous)303 return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad304 elif guidance_type == "classifier-free":305 if guidance_scale == 1. or unconditional_condition is None:306 return noise_pred_fn(x, t_continuous, cond=condition)307 else:308 x_in = torch.cat([x] * 2)309 t_in = torch.cat([t_continuous] * 2)310 c_in = torch.cat([unconditional_condition, condition])311 noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)312 return noise_uncond + guidance_scale * (noise - noise_uncond)313 314 assert model_type in ["noise", "x_start", "v"]315 assert guidance_type in ["uncond", "classifier", "classifier-free"]316 return model_fn317 318 319class DPM_Solver:320 def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):321 """Construct a DPM-Solver.322 We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").323 If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).324 If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).325 In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.326 The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.327 Args:328 model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):329 ``330 def model_fn(x, t_continuous):331 return noise332 ``333 noise_schedule: A noise schedule object, such as NoiseScheduleVP.334 predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.335 thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].336 max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.337 338 [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.339 """340 self.model = model_fn341 self.noise_schedule = noise_schedule342 self.predict_x0 = predict_x0343 self.thresholding = thresholding344 self.max_val = max_val345 346 def noise_prediction_fn(self, x, t):347 """348 Return the noise prediction model.349 """350 return self.model(x, t)351 352 def data_prediction_fn(self, x, t):353 """354 Return the data prediction model (with thresholding).355 """356 noise = self.noise_prediction_fn(x, t)357 dims = x.dim()358 alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)359 x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)360 if self.thresholding:361 p = 0.995 # A hyperparameter in the paper of "Imagen" [1].362 s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)363 s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)364 x0 = torch.clamp(x0, -s, s) / s365 return x0366 367 def model_fn(self, x, t):368 """369 Convert the model to the noise prediction model or the data prediction model.370 """371 if self.predict_x0:372 return self.data_prediction_fn(x, t)373 else:374 return self.noise_prediction_fn(x, t)375 376 def get_time_steps(self, skip_type, t_T, t_0, N, device):377 """Compute the intermediate time steps for sampling.378 Args:379 skip_type: A `str`. The type for the spacing of the time steps. We support three types:380 - 'logSNR': uniform logSNR for the time steps.381 - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)382 - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)383 t_T: A `float`. The starting time of the sampling (default is T).384 t_0: A `float`. The ending time of the sampling (default is epsilon).385 N: A `int`. The total number of the spacing of the time steps.386 device: A torch device.387 Returns:388 A pytorch tensor of the time steps, with the shape (N + 1,).389 """390 if skip_type == 'logSNR':391 lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))392 lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))393 logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)394 return self.noise_schedule.inverse_lambda(logSNR_steps)395 elif skip_type == 'time_uniform':396 return torch.linspace(t_T, t_0, N + 1).to(device)397 elif skip_type == 'time_quadratic':398 t_order = 2399 t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)400 return t401 else:402 raise ValueError(403 "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))404 405 def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):406 """407 Get the order of each step for sampling by the singlestep DPM-Solver.408 We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".409 Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:410 - If order == 1:411 We take `steps` of DPM-Solver-1 (i.e. DDIM).412 - If order == 2:413 - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.414 - If steps % 2 == 0, we use K steps of DPM-Solver-2.415 - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.416 - If order == 3:417 - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.418 - 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.419 - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.420 - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.421 ============================================422 Args:423 order: A `int`. The max order for the solver (2 or 3).424 steps: A `int`. The total number of function evaluations (NFE).425 skip_type: A `str`. The type for the spacing of the time steps. We support three types:426 - 'logSNR': uniform logSNR for the time steps.427 - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)428 - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)429 t_T: A `float`. The starting time of the sampling (default is T).430 t_0: A `float`. The ending time of the sampling (default is epsilon).431 device: A torch device.432 Returns:433 orders: A list of the solver order of each step.434 """435 if order == 3:436 K = steps // 3 + 1437 if steps % 3 == 0:438 orders = [3, ] * (K - 2) + [2, 1]439 elif steps % 3 == 1:440 orders = [3, ] * (K - 1) + [1]441 else:442 orders = [3, ] * (K - 1) + [2]443 elif order == 2:444 if steps % 2 == 0:445 K = steps // 2446 orders = [2, ] * K447 else:448 K = steps // 2 + 1449 orders = [2, ] * (K - 1) + [1]450 elif order == 1:451 K = 1452 orders = [1, ] * steps453 else:454 raise ValueError("'order' must be '1' or '2' or '3'.")455 if skip_type == 'logSNR':456 # To reproduce the results in DPM-Solver paper457 timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)458 else:459 timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[460 torch.cumsum(torch.tensor([0, ] + orders)).to(device)]461 return timesteps_outer, orders462 463 def denoise_to_zero_fn(self, x, s):464 """465 Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.466 """467 return self.data_prediction_fn(x, s)468 469 def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):470 """471 DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.472 Args:473 x: A pytorch tensor. The initial value at time `s`.474 s: A pytorch tensor. The starting time, with the shape (x.shape[0],).475 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).476 model_s: A pytorch tensor. The model function evaluated at time `s`.477 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.478 return_intermediate: A `bool`. If true, also return the model value at time `s`.479 Returns:480 x_t: A pytorch tensor. The approximated solution at time `t`.481 """482 ns = self.noise_schedule483 dims = x.dim()484 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)485 h = lambda_t - lambda_s486 log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)487 sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)488 alpha_t = torch.exp(log_alpha_t)489 490 if self.predict_x0:491 phi_1 = torch.expm1(-h)492 if model_s is None:493 model_s = self.model_fn(x, s)494 x_t = (495 expand_dims(sigma_t / sigma_s, dims) * x496 - expand_dims(alpha_t * phi_1, dims) * model_s497 )498 if return_intermediate:499 return x_t, {'model_s': model_s}500 else:501 return x_t502 else:503 phi_1 = torch.expm1(h)504 if model_s is None:505 model_s = self.model_fn(x, s)506 x_t = (507 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x508 - expand_dims(sigma_t * phi_1, dims) * model_s509 )510 if return_intermediate:511 return x_t, {'model_s': model_s}512 else:513 return x_t514 515 def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,516 solver_type='dpm_solver'):517 """518 Singlestep solver DPM-Solver-2 from time `s` to time `t`.519 Args:520 x: A pytorch tensor. The initial value at time `s`.521 s: A pytorch tensor. The starting time, with the shape (x.shape[0],).522 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).523 r1: A `float`. The hyperparameter of the second-order solver.524 model_s: A pytorch tensor. The model function evaluated at time `s`.525 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.526 return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).527 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.528 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.529 Returns:530 x_t: A pytorch tensor. The approximated solution at time `t`.531 """532 if solver_type not in ['dpm_solver', 'taylor']:533 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))534 if r1 is None:535 r1 = 0.5536 ns = self.noise_schedule537 dims = x.dim()538 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)539 h = lambda_t - lambda_s540 lambda_s1 = lambda_s + r1 * h541 s1 = ns.inverse_lambda(lambda_s1)542 log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(543 s1), ns.marginal_log_mean_coeff(t)544 sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)545 alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)546 547 if self.predict_x0:548 phi_11 = torch.expm1(-r1 * h)549 phi_1 = torch.expm1(-h)550 551 if model_s is None:552 model_s = self.model_fn(x, s)553 x_s1 = (554 expand_dims(sigma_s1 / sigma_s, dims) * x555 - expand_dims(alpha_s1 * phi_11, dims) * model_s556 )557 model_s1 = self.model_fn(x_s1, s1)558 if solver_type == 'dpm_solver':559 x_t = (560 expand_dims(sigma_t / sigma_s, dims) * x561 - expand_dims(alpha_t * phi_1, dims) * model_s562 - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)563 )564 elif solver_type == 'taylor':565 x_t = (566 expand_dims(sigma_t / sigma_s, dims) * x567 - expand_dims(alpha_t * phi_1, dims) * model_s568 + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (569 model_s1 - model_s)570 )571 else:572 phi_11 = torch.expm1(r1 * h)573 phi_1 = torch.expm1(h)574 575 if model_s is None:576 model_s = self.model_fn(x, s)577 x_s1 = (578 expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x579 - expand_dims(sigma_s1 * phi_11, dims) * model_s580 )581 model_s1 = self.model_fn(x_s1, s1)582 if solver_type == 'dpm_solver':583 x_t = (584 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x585 - expand_dims(sigma_t * phi_1, dims) * model_s586 - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)587 )588 elif solver_type == 'taylor':589 x_t = (590 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x591 - expand_dims(sigma_t * phi_1, dims) * model_s592 - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)593 )594 if return_intermediate:595 return x_t, {'model_s': model_s, 'model_s1': model_s1}596 else:597 return x_t598 599 def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,600 return_intermediate=False, solver_type='dpm_solver'):601 """602 Singlestep solver DPM-Solver-3 from time `s` to time `t`.603 Args:604 x: A pytorch tensor. The initial value at time `s`.605 s: A pytorch tensor. The starting time, with the shape (x.shape[0],).606 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).607 r1: A `float`. The hyperparameter of the third-order solver.608 r2: A `float`. The hyperparameter of the third-order solver.609 model_s: A pytorch tensor. The model function evaluated at time `s`.610 If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.611 model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).612 If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.613 return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).614 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.615 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.616 Returns:617 x_t: A pytorch tensor. The approximated solution at time `t`.618 """619 if solver_type not in ['dpm_solver', 'taylor']:620 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))621 if r1 is None:622 r1 = 1. / 3.623 if r2 is None:624 r2 = 2. / 3.625 ns = self.noise_schedule626 dims = x.dim()627 lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)628 h = lambda_t - lambda_s629 lambda_s1 = lambda_s + r1 * h630 lambda_s2 = lambda_s + r2 * h631 s1 = ns.inverse_lambda(lambda_s1)632 s2 = ns.inverse_lambda(lambda_s2)633 log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(634 s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)635 sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(636 s2), ns.marginal_std(t)637 alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)638 639 if self.predict_x0:640 phi_11 = torch.expm1(-r1 * h)641 phi_12 = torch.expm1(-r2 * h)642 phi_1 = torch.expm1(-h)643 phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.644 phi_2 = phi_1 / h + 1.645 phi_3 = phi_2 / h - 0.5646 647 if model_s is None:648 model_s = self.model_fn(x, s)649 if model_s1 is None:650 x_s1 = (651 expand_dims(sigma_s1 / sigma_s, dims) * x652 - expand_dims(alpha_s1 * phi_11, dims) * model_s653 )654 model_s1 = self.model_fn(x_s1, s1)655 x_s2 = (656 expand_dims(sigma_s2 / sigma_s, dims) * x657 - expand_dims(alpha_s2 * phi_12, dims) * model_s658 + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)659 )660 model_s2 = self.model_fn(x_s2, s2)661 if solver_type == 'dpm_solver':662 x_t = (663 expand_dims(sigma_t / sigma_s, dims) * x664 - expand_dims(alpha_t * phi_1, dims) * model_s665 + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)666 )667 elif solver_type == 'taylor':668 D1_0 = (1. / r1) * (model_s1 - model_s)669 D1_1 = (1. / r2) * (model_s2 - model_s)670 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)671 D2 = 2. * (D1_1 - D1_0) / (r2 - r1)672 x_t = (673 expand_dims(sigma_t / sigma_s, dims) * x674 - expand_dims(alpha_t * phi_1, dims) * model_s675 + expand_dims(alpha_t * phi_2, dims) * D1676 - expand_dims(alpha_t * phi_3, dims) * D2677 )678 else:679 phi_11 = torch.expm1(r1 * h)680 phi_12 = torch.expm1(r2 * h)681 phi_1 = torch.expm1(h)682 phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.683 phi_2 = phi_1 / h - 1.684 phi_3 = phi_2 / h - 0.5685 686 if model_s is None:687 model_s = self.model_fn(x, s)688 if model_s1 is None:689 x_s1 = (690 expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x691 - expand_dims(sigma_s1 * phi_11, dims) * model_s692 )693 model_s1 = self.model_fn(x_s1, s1)694 x_s2 = (695 expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x696 - expand_dims(sigma_s2 * phi_12, dims) * model_s697 - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)698 )699 model_s2 = self.model_fn(x_s2, s2)700 if solver_type == 'dpm_solver':701 x_t = (702 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x703 - expand_dims(sigma_t * phi_1, dims) * model_s704 - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)705 )706 elif solver_type == 'taylor':707 D1_0 = (1. / r1) * (model_s1 - model_s)708 D1_1 = (1. / r2) * (model_s2 - model_s)709 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)710 D2 = 2. * (D1_1 - D1_0) / (r2 - r1)711 x_t = (712 expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x713 - expand_dims(sigma_t * phi_1, dims) * model_s714 - expand_dims(sigma_t * phi_2, dims) * D1715 - expand_dims(sigma_t * phi_3, dims) * D2716 )717 718 if return_intermediate:719 return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}720 else:721 return x_t722 723 def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):724 """725 Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.726 Args:727 x: A pytorch tensor. The initial value at time `s`.728 model_prev_list: A list of pytorch tensor. The previous computed model values.729 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)730 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).731 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.732 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.733 Returns:734 x_t: A pytorch tensor. The approximated solution at time `t`.735 """736 if solver_type not in ['dpm_solver', 'taylor']:737 raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))738 ns = self.noise_schedule739 dims = x.dim()740 model_prev_1, model_prev_0 = model_prev_list741 t_prev_1, t_prev_0 = t_prev_list742 lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(743 t_prev_0), ns.marginal_lambda(t)744 log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)745 sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)746 alpha_t = torch.exp(log_alpha_t)747 748 h_0 = lambda_prev_0 - lambda_prev_1749 h = lambda_t - lambda_prev_0750 r0 = h_0 / h751 D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)752 if self.predict_x0:753 if solver_type == 'dpm_solver':754 x_t = (755 expand_dims(sigma_t / sigma_prev_0, dims) * x756 - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0757 - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0758 )759 elif solver_type == 'taylor':760 x_t = (761 expand_dims(sigma_t / sigma_prev_0, dims) * x762 - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0763 + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0764 )765 else:766 if solver_type == 'dpm_solver':767 x_t = (768 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x769 - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0770 - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0771 )772 elif solver_type == 'taylor':773 x_t = (774 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x775 - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0776 - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0777 )778 return x_t779 780 def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):781 """782 Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.783 Args:784 x: A pytorch tensor. The initial value at time `s`.785 model_prev_list: A list of pytorch tensor. The previous computed model values.786 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)787 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).788 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.789 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.790 Returns:791 x_t: A pytorch tensor. The approximated solution at time `t`.792 """793 ns = self.noise_schedule794 dims = x.dim()795 model_prev_2, model_prev_1, model_prev_0 = model_prev_list796 t_prev_2, t_prev_1, t_prev_0 = t_prev_list797 lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(798 t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)799 log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)800 sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)801 alpha_t = torch.exp(log_alpha_t)802 803 h_1 = lambda_prev_1 - lambda_prev_2804 h_0 = lambda_prev_0 - lambda_prev_1805 h = lambda_t - lambda_prev_0806 r0, r1 = h_0 / h, h_1 / h807 D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)808 D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)809 D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)810 D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)811 if self.predict_x0:812 x_t = (813 expand_dims(sigma_t / sigma_prev_0, dims) * x814 - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0815 + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1816 - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2817 )818 else:819 x_t = (820 expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x821 - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0822 - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1823 - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2824 )825 return x_t826 827 def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,828 r2=None):829 """830 Singlestep DPM-Solver with the order `order` from time `s` to time `t`.831 Args:832 x: A pytorch tensor. The initial value at time `s`.833 s: A pytorch tensor. The starting time, with the shape (x.shape[0],).834 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).835 order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.836 return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).837 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.838 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.839 r1: A `float`. The hyperparameter of the second-order or third-order solver.840 r2: A `float`. The hyperparameter of the third-order solver.841 Returns:842 x_t: A pytorch tensor. The approximated solution at time `t`.843 """844 if order == 1:845 return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)846 elif order == 2:847 return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,848 solver_type=solver_type, r1=r1)849 elif order == 3:850 return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,851 solver_type=solver_type, r1=r1, r2=r2)852 else:853 raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))854 855 def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):856 """857 Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.858 Args:859 x: A pytorch tensor. The initial value at time `s`.860 model_prev_list: A list of pytorch tensor. The previous computed model values.861 t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)862 t: A pytorch tensor. The ending time, with the shape (x.shape[0],).863 order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.864 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.865 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.866 Returns:867 x_t: A pytorch tensor. The approximated solution at time `t`.868 """869 if order == 1:870 return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])871 elif order == 2:872 return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)873 elif order == 3:874 return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)875 else:876 raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))877 878 def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5,879 solver_type='dpm_solver'):880 """881 The adaptive step size solver based on singlestep DPM-Solver.882 Args:883 x: A pytorch tensor. The initial value at time `t_T`.884 order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.885 t_T: A `float`. The starting time of the sampling (default is T).886 t_0: A `float`. The ending time of the sampling (default is epsilon).887 h_init: A `float`. The initial step size (for logSNR).888 atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].889 rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.890 theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].891 t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the892 current time and `t_0` is less than `t_err`. The default setting is 1e-5.893 solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.894 The type slightly impacts the performance. We recommend to use 'dpm_solver' type.895 Returns:896 x_0: A pytorch tensor. The approximated solution at time `t_0`.897 [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.898 """899 ns = self.noise_schedule900 s = t_T * torch.ones((x.shape[0],)).to(x)901 lambda_s = ns.marginal_lambda(s)902 lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))903 h = h_init * torch.ones_like(s).to(x)904 x_prev = x905 nfe = 0906 if order == 2:907 r1 = 0.5908 lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)909 higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,910 solver_type=solver_type,911 **kwargs)912 elif order == 3:913 r1, r2 = 1. / 3., 2. / 3.914 lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,915 return_intermediate=True,916 solver_type=solver_type)917 higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,918 solver_type=solver_type,919 **kwargs)920 else:921 raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))922 while torch.abs((s - t_0)).mean() > t_err:923 t = ns.inverse_lambda(lambda_s + h)924 x_lower, lower_noise_kwargs = lower_update(x, s, t)925 x_higher = higher_update(x, s, t, **lower_noise_kwargs)926 delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))927 norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))928 E = norm_fn((x_higher - x_lower) / delta).max()929 if torch.all(E <= 1.):930 x = x_higher931 s = t932 x_prev = x_lower933 lambda_s = ns.marginal_lambda(s)934 h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)935 nfe += order936 print('adaptive solver nfe', nfe)937 return x938 939 def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',940 method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',941 atol=0.0078, rtol=0.05,942 ):943 """944 Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.945 =====================================================946 We support the following algorithms for both noise prediction model and data prediction model:947 - 'singlestep':948 Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.949 We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).950 The total number of function evaluations (NFE) == `steps`.951 Given a fixed NFE == `steps`, the sampling procedure is:952 - If `order` == 1:953 - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).954 - If `order` == 2:955 - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.956 - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.957 - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.958 - If `order` == 3:959 - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.960 - 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.961 - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.962 - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.963 - 'multistep':964 Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.965 We initialize the first `order` values by lower order multistep solvers.966 Given a fixed NFE == `steps`, the sampling procedure is:967 Denote K = steps.968 - If `order` == 1:969 - We use K steps of DPM-Solver-1 (i.e. DDIM).970 - If `order` == 2:971 - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.972 - If `order` == 3:973 - 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.974 - 'singlestep_fixed':975 Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).976 We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.977 - 'adaptive':978 Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).979 We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.980 You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs981 (NFE) and the sample quality.982 - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.983 - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.984 =====================================================985 Some advices for choosing the algorithm:986 - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:987 Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.988 e.g.989 >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)990 >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,991 skip_type='time_uniform', method='singlestep')992 - For **guided sampling with large guidance scale** by DPMs:993 Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.994 e.g.995 >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)996 >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,997 skip_type='time_uniform', method='multistep')998 We support three types of `skip_type`:999 - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**1000 - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.1001 - 'time_quadratic': quadratic time for the time steps.1002 =====================================================1003 Args:1004 x: A pytorch tensor. The initial value at time `t_start`1005 e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.1006 steps: A `int`. The total number of function evaluations (NFE).1007 t_start: A `float`. The starting time of the sampling.1008 If `T` is None, we use self.noise_schedule.T (default is 1.0).1009 t_end: A `float`. The ending time of the sampling.1010 If `t_end` is None, we use 1. / self.noise_schedule.total_N.1011 e.g. if total_N == 1000, we have `t_end` == 1e-3.1012 For discrete-time DPMs:1013 - We recommend `t_end` == 1. / self.noise_schedule.total_N.1014 For continuous-time DPMs:1015 - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.1016 order: A `int`. The order of DPM-Solver.1017 skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.1018 method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.1019 denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.1020 Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).1021 This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and1022 score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID1023 for diffusion models sampling by diffusion SDEs for low-resolutional images1024 (such as CIFAR-10). However, we observed that such trick does not matter for1025 high-resolutional images. As it needs an additional NFE, we do not recommend1026 it for high-resolutional images.1027 lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.1028 Only valid for `method=multistep` and `steps < 15`. We empirically find that1029 this trick is a key to stabilizing the sampling by DPM-Solver with very few steps1030 (especially for steps <= 10). So we recommend to set it to be `True`.1031 solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.1032 atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.1033 rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.1034 Returns:1035 x_end: A pytorch tensor. The approximated solution at time `t_end`.1036 """1037 t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end1038 t_T = self.noise_schedule.T if t_start is None else t_start1039 device = x.device1040 if method == 'adaptive':1041 with torch.no_grad():1042 x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,1043 solver_type=solver_type)1044 elif method == 'multistep':1045 assert steps >= order1046 timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)1047 assert timesteps.shape[0] - 1 == steps1048 with torch.no_grad():1049 vec_t = timesteps[0].expand((x.shape[0]))1050 model_prev_list = [self.model_fn(x, vec_t)]1051 t_prev_list = [vec_t]1052 # Init the first `order` values by lower order multistep DPM-Solver.1053 for init_order in tqdm(range(1, order), desc="DPM init order"):1054 vec_t = timesteps[init_order].expand(x.shape[0])1055 x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,1056 solver_type=solver_type)1057 model_prev_list.append(self.model_fn(x, vec_t))1058 t_prev_list.append(vec_t)1059 # Compute the remaining values by `order`-th order multistep DPM-Solver.1060 for step in tqdm(range(order, steps + 1), desc="DPM multistep"):1061 vec_t = timesteps[step].expand(x.shape[0])1062 if lower_order_final and steps < 15:1063 step_order = min(order, steps + 1 - step)1064 else:1065 step_order = order1066 x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order,1067 solver_type=solver_type)1068 for i in range(order - 1):1069 t_prev_list[i] = t_prev_list[i + 1]1070 model_prev_list[i] = model_prev_list[i + 1]1071 t_prev_list[-1] = vec_t1072 # We do not need to evaluate the final model value.1073 if step < steps:1074 model_prev_list[-1] = self.model_fn(x, vec_t)1075 elif method in ['singlestep', 'singlestep_fixed']:1076 if method == 'singlestep':1077 timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,1078 skip_type=skip_type,1079 t_T=t_T, t_0=t_0,1080 device=device)1081 elif method == 'singlestep_fixed':1082 K = steps // order1083 orders = [order, ] * K1084 timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)1085 for i, order in enumerate(orders):1086 t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]1087 timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),1088 N=order, device=device)1089 lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)1090 vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])1091 h = lambda_inner[-1] - lambda_inner[0]1092 r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h1093 r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h1094 x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)1095 if denoise_to_zero:1096 x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)1097 return x1098 1099 1100#############################################################1101# other utility functions1102#############################################################1103 1104def interpolate_fn(x, xp, yp):1105 """1106 A piecewise linear function y = f(x), using xp and yp as keypoints.1107 We implement f(x) in a differentiable way (i.e. applicable for autograd).1108 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.)1109 Args:1110 x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).1111 xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.1112 yp: PyTorch tensor with shape [C, K].1113 Returns:1114 The function values f(x), with shape [N, C].1115 """1116 N, K = x.shape[0], xp.shape[1]1117 all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)1118 sorted_all_x, x_indices = torch.sort(all_x, dim=2)1119 x_idx = torch.argmin(x_indices, dim=2)1120 cand_start_idx = x_idx - 11121 start_idx = torch.where(1122 torch.eq(x_idx, 0),1123 torch.tensor(1, device=x.device),1124 torch.where(1125 torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,1126 ),1127 )1128 end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)1129 start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)1130 end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)1131 start_idx2 = torch.where(1132 torch.eq(x_idx, 0),1133 torch.tensor(0, device=x.device),1134 torch.where(1135 torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,1136 ),1137 )1138 y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)1139 start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)1140 end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)1141 cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)1142 return cand1143 1144 1145def expand_dims(v, dims):1146 """1147 Expand the tensor `v` to the dim `dims`.1148 Args:1149 `v`: a PyTorch tensor with shape [N].1150 `dim`: a `int`.1151 Returns:1152 a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.1153 """1154 return v[(...,) + (None,) * (dims - 1)]