Implements Algorithm 2 (Heun steps) 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.)
| 157 | |
| 158 | @torch.no_grad() |
| 159 | def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): |
| 160 | """Implements Algorithm 2 (Heun steps) from Karras et al. (2022).""" |
| 161 | extra_args = {} if extra_args is None else extra_args |
| 162 | s_in = x.new_ones([x.shape[0]]) |
| 163 | for i in trange(len(sigmas) - 1, disable=disable): |
| 164 | gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0. |
| 165 | eps = torch.randn_like(x) * s_noise |
| 166 | sigma_hat = sigmas[i] * (gamma + 1) |
| 167 | if gamma > 0: |
| 168 | x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5 |
| 169 | denoised = model(x, sigma_hat * s_in, **extra_args) |
| 170 | d = to_d(x, sigma_hat, denoised) |
| 171 | if callback is not None: |
| 172 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) |
| 173 | dt = sigmas[i + 1] - sigma_hat |
| 174 | if sigmas[i + 1] == 0: |
| 175 | # Euler method |
| 176 | x = x + d * dt |
| 177 | else: |
| 178 | # Heun's method |
| 179 | x_2 = x + d * dt |
| 180 | denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args) |
| 181 | d_2 = to_d(x_2, sigmas[i + 1], denoised_2) |
| 182 | d_prime = (d + d_2) / 2 |
| 183 | x = x + d_prime * dt |
| 184 | return x |
| 185 | |
| 186 | |
| 187 | @torch.no_grad() |