Ancestral sampling with DPM-Solver++(2S) second-order steps.
(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None)
| 507 | |
| 508 | @torch.no_grad() |
| 509 | def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): |
| 510 | """Ancestral sampling with DPM-Solver++(2S) second-order steps.""" |
| 511 | extra_args = {} if extra_args is None else extra_args |
| 512 | noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler |
| 513 | s_in = x.new_ones([x.shape[0]]) |
| 514 | sigma_fn = lambda t: t.neg().exp() |
| 515 | t_fn = lambda sigma: sigma.log().neg() |
| 516 | |
| 517 | for i in trange(len(sigmas) - 1, disable=disable): |
| 518 | denoised = model(x, sigmas[i] * s_in, **extra_args) |
| 519 | sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) |
| 520 | if callback is not None: |
| 521 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) |
| 522 | if sigma_down == 0: |
| 523 | # Euler method |
| 524 | d = to_d(x, sigmas[i], denoised) |
| 525 | dt = sigma_down - sigmas[i] |
| 526 | x = x + d * dt |
| 527 | else: |
| 528 | # DPM-Solver++(2S) |
| 529 | t, t_next = t_fn(sigmas[i]), t_fn(sigma_down) |
| 530 | r = 1 / 2 |
| 531 | h = t_next - t |
| 532 | s = t + r * h |
| 533 | x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised |
| 534 | denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args) |
| 535 | x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2 |
| 536 | # Noise addition |
| 537 | if sigmas[i + 1] > 0: |
| 538 | x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up |
| 539 | return x |
| 540 | |
| 541 | |
| 542 | @torch.no_grad() |
nothing calls this directly
no test coverage detected