Ancestral sampling with DPM-Solver second-order steps.
(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None)
| 217 | |
| 218 | @torch.no_grad() |
| 219 | def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): |
| 220 | """Ancestral sampling with DPM-Solver second-order steps.""" |
| 221 | extra_args = {} if extra_args is None else extra_args |
| 222 | noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler |
| 223 | s_in = x.new_ones([x.shape[0]]) |
| 224 | for i in trange(len(sigmas) - 1, disable=disable): |
| 225 | denoised = model(x, sigmas[i] * s_in, **extra_args) |
| 226 | sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) |
| 227 | if callback is not None: |
| 228 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) |
| 229 | d = to_d(x, sigmas[i], denoised) |
| 230 | if sigma_down == 0: |
| 231 | # Euler method |
| 232 | dt = sigma_down - sigmas[i] |
| 233 | x = x + d * dt |
| 234 | else: |
| 235 | # DPM-Solver-2 |
| 236 | sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp() |
| 237 | dt_1 = sigma_mid - sigmas[i] |
| 238 | dt_2 = sigma_down - sigmas[i] |
| 239 | x_2 = x + d * dt_1 |
| 240 | denoised_2 = model(x_2, sigma_mid * s_in, **extra_args) |
| 241 | d_2 = to_d(x_2, sigma_mid, denoised_2) |
| 242 | x = x + d_2 * dt_2 |
| 243 | x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up |
| 244 | return x |
| 245 | |
| 246 | |
| 247 | def linear_multistep_coeff(order, t, i, j): |
nothing calls this directly
no test coverage detected