Singlestep solver DPM-Solver-2 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, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpmsolver')
| 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 | """ |
| 604 | Singlestep solver DPM-Solver-2 from time `s` to time `t`. |
| 605 | |
| 606 | Args: |
| 607 | x: A pytorch tensor. The initial value at time `s`. |
| 608 | s: A pytorch tensor. The starting time, with the shape (1,). |
| 609 | t: A pytorch tensor. The ending time, with the shape (1,). |
| 610 | r1: A `float`. The hyperparameter of the second-order solver. |
| 611 | model_s: A pytorch tensor. The model function evaluated at time `s`. |
| 612 | If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. |
| 613 | return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time). |
| 614 | solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. |
| 615 | The type slightly impacts the performance. We recommend to use 'dpmsolver' type. |
| 616 | Returns: |
| 617 | x_t: A pytorch tensor. The approximated solution at time `t`. |
| 618 | """ |
| 619 | if solver_type not in ['dpmsolver', 'taylor']: |
| 620 | raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type)) |
| 621 | if r1 is None: |
| 622 | r1 = 0.5 |
| 623 | ns = self.noise_schedule |
| 624 | lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) |
| 625 | h = lambda_t - lambda_s |
| 626 | lambda_s1 = lambda_s + r1 * h |
| 627 | s1 = ns.inverse_lambda(lambda_s1) |
| 628 | log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(t) |
| 629 | sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t) |
| 630 | alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t) |
| 631 | |
| 632 | if self.algorithm_type == "dpmsolver++": |
| 633 | phi_11 = torch.expm1(-r1 * h) |
| 634 | phi_1 = torch.expm1(-h) |
| 635 | |
| 636 | if model_s is None: |
| 637 | model_s = self.model_fn(x, s) |
| 638 | x_s1 = ( |
| 639 | (sigma_s1 / sigma_s) * x |
| 640 | - (alpha_s1 * phi_11) * model_s |
| 641 | ) |
| 642 | model_s1 = self.model_fn(x_s1, s1) |
| 643 | if solver_type == 'dpmsolver': |
| 644 | x_t = ( |
| 645 | (sigma_t / sigma_s) * x |
| 646 | - (alpha_t * phi_1) * model_s |
| 647 | - (0.5 / r1) * (alpha_t * phi_1) * (model_s1 - model_s) |
| 648 | ) |
| 649 | elif solver_type == 'taylor': |
| 650 | x_t = ( |
| 651 | (sigma_t / sigma_s) * x |
| 652 | - (alpha_t * phi_1) * model_s |
| 653 | + (1. / r1) * (alpha_t * (phi_1 / h + 1.)) * (model_s1 - model_s) |
| 654 | ) |
| 655 | else: |
| 656 | phi_11 = torch.expm1(r1 * h) |
| 657 | phi_1 = torch.expm1(h) |
| 658 | |
| 659 | if model_s is None: |
no test coverage detected