A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).
(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.)
| 186 | |
| 187 | @torch.no_grad() |
| 188 | def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): |
| 189 | """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).""" |
| 190 | extra_args = {} if extra_args is None else extra_args |
| 191 | s_in = x.new_ones([x.shape[0]]) |
| 192 | for i in trange(len(sigmas) - 1, disable=disable): |
| 193 | gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0. |
| 194 | eps = torch.randn_like(x) * s_noise |
| 195 | sigma_hat = sigmas[i] * (gamma + 1) |
| 196 | if gamma > 0: |
| 197 | x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5 |
| 198 | denoised = model(x, sigma_hat * s_in, **extra_args) |
| 199 | d = to_d(x, sigma_hat, denoised) |
| 200 | if callback is not None: |
| 201 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) |
| 202 | if sigmas[i + 1] == 0: |
| 203 | # Euler method |
| 204 | dt = sigmas[i + 1] - sigma_hat |
| 205 | x = x + d * dt |
| 206 | else: |
| 207 | # DPM-Solver-2 |
| 208 | sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp() |
| 209 | dt_1 = sigma_mid - sigma_hat |
| 210 | dt_2 = sigmas[i + 1] - sigma_hat |
| 211 | x_2 = x + d * dt_1 |
| 212 | denoised_2 = model(x_2, sigma_mid * s_in, **extra_args) |
| 213 | d_2 = to_d(x_2, sigma_mid, denoised_2) |
| 214 | x = x + d_2 * dt_2 |
| 215 | return x |
| 216 | |
| 217 | |
| 218 | @torch.no_grad() |