Implements Algorithm 2 (Euler 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.)
| 116 | |
| 117 | @torch.no_grad() |
| 118 | def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): |
| 119 | """Implements Algorithm 2 (Euler steps) from Karras et al. (2022).""" |
| 120 | extra_args = {} if extra_args is None else extra_args |
| 121 | s_in = x.new_ones([x.shape[0]]) |
| 122 | for i in trange(len(sigmas) - 1, disable=disable): |
| 123 | gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0. |
| 124 | eps = torch.randn_like(x) * s_noise |
| 125 | sigma_hat = sigmas[i] * (gamma + 1) |
| 126 | if gamma > 0: |
| 127 | x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5 |
| 128 | denoised = model(x, sigma_hat * s_in, **extra_args) |
| 129 | d = to_d(x, sigma_hat, denoised) |
| 130 | if callback is not None: |
| 131 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) |
| 132 | dt = sigmas[i + 1] - sigma_hat |
| 133 | # Euler method |
| 134 | x = x + d * dt |
| 135 | return x |
| 136 | |
| 137 | |
| 138 | @torch.no_grad() |