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