(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True)
| 21 | setattr(self, name, attr) |
| 22 | |
| 23 | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): |
| 24 | self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, |
| 25 | num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose) |
| 26 | alphas_cumprod = self.model.alphas_cumprod |
| 27 | assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep' |
| 28 | to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device) |
| 29 | |
| 30 | self.register_buffer('betas', to_torch(self.model.betas)) |
| 31 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 32 | self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev)) |
| 33 | |
| 34 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 35 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu()))) |
| 36 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu()))) |
| 37 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu()))) |
| 38 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu()))) |
| 39 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1))) |
| 40 | |
| 41 | # ddim sampling parameters |
| 42 | ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(), |
| 43 | ddim_timesteps=self.ddim_timesteps, |
| 44 | eta=ddim_eta,verbose=verbose) |
| 45 | self.register_buffer('ddim_sigmas', ddim_sigmas) |
| 46 | self.register_buffer('ddim_alphas', ddim_alphas) |
| 47 | self.register_buffer('ddim_alphas_prev', ddim_alphas_prev) |
| 48 | self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas)) |
| 49 | sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt( |
| 50 | (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * ( |
| 51 | 1 - self.alphas_cumprod / self.alphas_cumprod_prev)) |
| 52 | self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps) |
| 53 | |
| 54 | @torch.no_grad() |
| 55 | def sample(self, |
no test coverage detected