A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).
(
denoiser,
x,
sigmas,
progress=False,
s_churn=0.0,
s_tmin=0.0,
s_tmax=float("inf"),
s_noise=1.0,
)
| 281 | |
| 282 | @th.no_grad() |
| 283 | def sample_dpm( |
| 284 | denoiser, |
| 285 | x, |
| 286 | sigmas, |
| 287 | progress=False, |
| 288 | s_churn=0.0, |
| 289 | s_tmin=0.0, |
| 290 | s_tmax=float("inf"), |
| 291 | s_noise=1.0, |
| 292 | ): |
| 293 | """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).""" |
| 294 | s_in = x.new_ones([x.shape[0]]) |
| 295 | indices = range(len(sigmas) - 1) |
| 296 | if progress: |
| 297 | from tqdm.auto import tqdm |
| 298 | |
| 299 | indices = tqdm(indices) |
| 300 | |
| 301 | for i in indices: |
| 302 | gamma = ( |
| 303 | min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0 |
| 304 | ) |
| 305 | eps = th.randn_like(x) * s_noise |
| 306 | sigma_hat = sigmas[i] * (gamma + 1) |
| 307 | if gamma > 0: |
| 308 | x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5 |
| 309 | denoised = denoiser(x, sigma_hat * s_in) |
| 310 | d = to_d(x, sigma_hat, denoised) |
| 311 | yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "denoised": denoised} |
| 312 | # Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule |
| 313 | sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2) ** 3 |
| 314 | dt_1 = sigma_mid - sigma_hat |
| 315 | dt_2 = sigmas[i + 1] - sigma_hat |
| 316 | x_2 = x + d * dt_1 |
| 317 | denoised_2 = denoiser(x_2, sigma_mid * s_in) |
| 318 | d_2 = to_d(x_2, sigma_mid, denoised_2) |
| 319 | x = x + d_2 * dt_2 |
| 320 | yield {"x": x, "pred_xstart": denoised} |
| 321 | |
| 322 | |
| 323 | def append_dims(x, target_dims): |