DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. s: A pytorch tensor. The starting time, with the shape (1,). t: A pytorch tensor. The ending time, with the shape (1,).
(self, x, s, t, model_s=None, return_intermediate=False)
| 553 | return self.data_prediction_fn(x, s) |
| 554 | |
| 555 | def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False): |
| 556 | """ |
| 557 | DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. |
| 558 | |
| 559 | Args: |
| 560 | x: A pytorch tensor. The initial value at time `s`. |
| 561 | s: A pytorch tensor. The starting time, with the shape (1,). |
| 562 | t: A pytorch tensor. The ending time, with the shape (1,). |
| 563 | model_s: A pytorch tensor. The model function evaluated at time `s`. |
| 564 | If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. |
| 565 | return_intermediate: A `bool`. If true, also return the model value at time `s`. |
| 566 | Returns: |
| 567 | x_t: A pytorch tensor. The approximated solution at time `t`. |
| 568 | """ |
| 569 | ns = self.noise_schedule |
| 570 | dims = x.dim() |
| 571 | lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) |
| 572 | h = lambda_t - lambda_s |
| 573 | log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t) |
| 574 | sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t) |
| 575 | alpha_t = torch.exp(log_alpha_t) |
| 576 | |
| 577 | if self.algorithm_type == "dpmsolver++": |
| 578 | phi_1 = torch.expm1(-h) |
| 579 | if model_s is None: |
| 580 | model_s = self.model_fn(x, s) |
| 581 | x_t = ( |
| 582 | sigma_t / sigma_s * x |
| 583 | - alpha_t * phi_1 * model_s |
| 584 | ) |
| 585 | if return_intermediate: |
| 586 | return x_t, {'model_s': model_s} |
| 587 | else: |
| 588 | return x_t |
| 589 | else: |
| 590 | phi_1 = torch.expm1(h) |
| 591 | if model_s is None: |
| 592 | model_s = self.model_fn(x, s) |
| 593 | x_t = ( |
| 594 | torch.exp(log_alpha_t - log_alpha_s) * x |
| 595 | - (sigma_t * phi_1) * model_s |
| 596 | ) |
| 597 | if return_intermediate: |
| 598 | return x_t, {'model_s': model_s} |
| 599 | else: |
| 600 | return x_t |
| 601 | |
| 602 | def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpmsolver'): |
| 603 | """ |
no test coverage detected