Implements Algorithm 2 (Euler steps) from Karras et al. (2022).
(model, x, sigmas, extra_args=None, callback=None, disable=None)
| 226 | def perp_step_wrap(s=0.5): |
| 227 | @torch.no_grad() |
| 228 | def perp_step(model, x, sigmas, extra_args=None, callback=None, disable=None): |
| 229 | """Implements Algorithm 2 (Euler steps) from Karras et al. (2022).""" |
| 230 | extra_args = {} if extra_args is None else extra_args |
| 231 | s_in = x.new_ones([x.shape[0]]) |
| 232 | previous_step = None |
| 233 | for i in trange(len(sigmas) - 1, disable=disable): |
| 234 | sigma_hat = sigmas[i] |
| 235 | denoised = model(x, sigma_hat * s_in, **extra_args) |
| 236 | d = to_d(x, sigma_hat, denoised) |
| 237 | dt = sigmas[i + 1] - sigma_hat |
| 238 | if previous_step is not None and sigmas[i + 1] != 0: |
| 239 | d = diff_step(d, previous_step, s) |
| 240 | previous_step = d.clone() |
| 241 | if callback is not None: |
| 242 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) |
| 243 | x = x + d * dt |
| 244 | return x |
| 245 | return perp_step |
| 246 | |
| 247 | # as a reference |
nothing calls this directly
no test coverage detected