| 317 | |
| 318 | |
| 319 | class 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 noise |
| 332 | `` |
| 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_fn |
| 341 | self.noise_schedule = noise_schedule |
| 342 | self.predict_x0 = predict_x0 |
| 343 | self.thresholding = thresholding |
| 344 | self.max_val = max_val |
| 345 | |
| 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) / s |
| 365 | return x0 |
| 366 | |
| 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): |