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