| 77 | |
| 78 | |
| 79 | class GaussianToKarrasDenoiser: |
| 80 | def __init__(self, model, diffusion): |
| 81 | from scipy import interpolate |
| 82 | |
| 83 | self.model = model |
| 84 | self.diffusion = diffusion |
| 85 | self.alpha_cumprod_to_t = interpolate.interp1d( |
| 86 | diffusion.alphas_cumprod, np.arange(0, diffusion.num_timesteps) |
| 87 | ) |
| 88 | |
| 89 | def sigma_to_t(self, sigma): |
| 90 | alpha_cumprod = 1.0 / (sigma**2 + 1) |
| 91 | if alpha_cumprod > self.diffusion.alphas_cumprod[0]: |
| 92 | return 0 |
| 93 | elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]: |
| 94 | return self.diffusion.num_timesteps - 1 |
| 95 | else: |
| 96 | return float(self.alpha_cumprod_to_t(alpha_cumprod)) |
| 97 | |
| 98 | def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None): |
| 99 | t = th.tensor( |
| 100 | [self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()], |
| 101 | dtype=th.long, |
| 102 | device=sigmas.device, |
| 103 | ) |
| 104 | c_in = append_dims(1.0 / (sigmas**2 + 1) ** 0.5, x_t.ndim) |
| 105 | out = self.diffusion.p_mean_variance( |
| 106 | self.model, x_t * c_in, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs |
| 107 | ) |
| 108 | return None, out["pred_xstart"] |
| 109 | |
| 110 | |
| 111 | def karras_sample(*args, **kwargs): |
no outgoing calls
no test coverage detected