| 50 | |
| 51 | |
| 52 | class DDIMSampler: |
| 53 | def __init__(self, ddpm: GaussianDiffusion, schedule="linear"): |
| 54 | super().__init__() |
| 55 | self.ddpm = ddpm |
| 56 | self.ddpm_num_timesteps = ddpm.num_timesteps |
| 57 | self.schedule = schedule |
| 58 | |
| 59 | def register_buffer(self, name, attr): |
| 60 | if type(attr) == torch.Tensor: |
| 61 | if attr.device != self.device: |
| 62 | attr = attr.to(self.device) |
| 63 | setattr(self, name, attr) |
| 64 | |
| 65 | def make_schedule(self, ddim_num_steps, device, ddim_discretize="uniform", ddim_eta=0., verbose=False): |
| 66 | self.device = device |
| 67 | self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, |
| 68 | num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose) |
| 69 | alphas_cumprod = self.ddpm.alphas_cumprod |
| 70 | assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep' |
| 71 | to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.device) |
| 72 | |
| 73 | self.register_buffer('betas', to_torch(self.ddpm.betas)) |
| 74 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 75 | self.register_buffer('alphas_cumprod_prev', to_torch(self.ddpm.alphas_cumprod_prev)) |
| 76 | |
| 77 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 78 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu()))) |
| 79 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu()))) |
| 80 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu()))) |
| 81 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu()))) |
| 82 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1))) |
| 83 | |
| 84 | # ddim sampling parameters |
| 85 | ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters( |
| 86 | alphacums=alphas_cumprod.cpu(), ddim_timesteps=self.ddim_timesteps, eta=ddim_eta, verbose=verbose |
| 87 | ) |
| 88 | self.register_buffer('ddim_sigmas', ddim_sigmas) |
| 89 | self.register_buffer('ddim_alphas', ddim_alphas) |
| 90 | self.register_buffer('ddim_alphas_prev', ddim_alphas_prev) |
| 91 | self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas)) |
| 92 | sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt( |
| 93 | (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (1 - self.alphas_cumprod / self.alphas_cumprod_prev) |
| 94 | ) |
| 95 | self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps) |
| 96 | |
| 97 | @torch.no_grad() |
| 98 | def sample( |
| 99 | self, |
| 100 | model, |
| 101 | noise, |
| 102 | ddim_steps, |
| 103 | eta=0., |
| 104 | model_kwargs=None, |
| 105 | progress=True, |
| 106 | temperature=1., |
| 107 | noise_dropout=0., |
| 108 | clip_denoised=False, |
| 109 | log_every_t=100, |