Implements Algorithm 2 (Heun steps) 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,
)
| 237 | |
| 238 | @th.no_grad() |
| 239 | def sample_heun( |
| 240 | denoiser, |
| 241 | x, |
| 242 | sigmas, |
| 243 | progress=False, |
| 244 | s_churn=0.0, |
| 245 | s_tmin=0.0, |
| 246 | s_tmax=float("inf"), |
| 247 | s_noise=1.0, |
| 248 | ): |
| 249 | """Implements Algorithm 2 (Heun steps) from Karras et al. (2022).""" |
| 250 | s_in = x.new_ones([x.shape[0]]) |
| 251 | indices = range(len(sigmas) - 1) |
| 252 | if progress: |
| 253 | from tqdm.auto import tqdm |
| 254 | |
| 255 | indices = tqdm(indices) |
| 256 | |
| 257 | for i in indices: |
| 258 | gamma = ( |
| 259 | min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0 |
| 260 | ) |
| 261 | eps = th.randn_like(x) * s_noise |
| 262 | sigma_hat = sigmas[i] * (gamma + 1) |
| 263 | if gamma > 0: |
| 264 | x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5 |
| 265 | denoised = denoiser(x, sigma_hat * s_in) |
| 266 | d = to_d(x, sigma_hat, denoised) |
| 267 | yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "pred_xstart": denoised} |
| 268 | dt = sigmas[i + 1] - sigma_hat |
| 269 | if sigmas[i + 1] == 0: |
| 270 | # Euler method |
| 271 | x = x + d * dt |
| 272 | else: |
| 273 | # Heun's method |
| 274 | x_2 = x + d * dt |
| 275 | denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in) |
| 276 | d_2 = to_d(x_2, sigmas[i + 1], denoised_2) |
| 277 | d_prime = (d + d_2) / 2 |
| 278 | x = x + d_prime * dt |
| 279 | yield {"x": x, "pred_xstart": denoised} |
| 280 | |
| 281 | |
| 282 | @th.no_grad() |