DPM-Solver++(2M) SDE.
(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint')
| 609 | |
| 610 | @torch.no_grad() |
| 611 | def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'): |
| 612 | """DPM-Solver++(2M) SDE.""" |
| 613 | |
| 614 | if solver_type not in {'heun', 'midpoint'}: |
| 615 | raise ValueError('solver_type must be \'heun\' or \'midpoint\'') |
| 616 | |
| 617 | sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() |
| 618 | noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler |
| 619 | extra_args = {} if extra_args is None else extra_args |
| 620 | s_in = x.new_ones([x.shape[0]]) |
| 621 | |
| 622 | old_denoised = None |
| 623 | h_last = None |
| 624 | |
| 625 | for i in trange(len(sigmas) - 1, disable=disable): |
| 626 | denoised = model(x, sigmas[i] * s_in, **extra_args) |
| 627 | if callback is not None: |
| 628 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) |
| 629 | if sigmas[i + 1] == 0: |
| 630 | # Denoising step |
| 631 | x = denoised |
| 632 | else: |
| 633 | # DPM-Solver++(2M) SDE |
| 634 | t, s = -sigmas[i].log(), -sigmas[i + 1].log() |
| 635 | h = s - t |
| 636 | eta_h = eta * h |
| 637 | |
| 638 | x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised |
| 639 | |
| 640 | if old_denoised is not None: |
| 641 | r = h_last / h |
| 642 | if solver_type == 'heun': |
| 643 | x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised) |
| 644 | elif solver_type == 'midpoint': |
| 645 | x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised) |
| 646 | |
| 647 | if eta: |
| 648 | x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise |
| 649 | |
| 650 | old_denoised = denoised |
| 651 | h_last = h |
| 652 | return x |
| 653 | |
| 654 | @torch.no_grad() |
| 655 | def sample_dpmpp_2m_sde_cfg(model, x, sigmas, cfg_val, cfg_interval, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'): |
nothing calls this directly
no test coverage detected