| 10 | |
| 11 | |
| 12 | class PLMSSampler(object): |
| 13 | def __init__(self, model, schedule="linear", **kwargs): |
| 14 | super().__init__() |
| 15 | self.model = model |
| 16 | self.ddpm_num_timesteps = model.num_timesteps |
| 17 | self.schedule = schedule |
| 18 | |
| 19 | def register_buffer(self, name, attr): |
| 20 | if type(attr) == torch.Tensor: |
| 21 | if attr.device != torch.device("cuda"): |
| 22 | attr = attr.to(torch.device("cuda")) |
| 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, |
| 60 | S, |
| 61 | batch_size, |
| 62 | shape, |
| 63 | conditioning=None, |
| 64 | callback=None, |
| 65 | normals_sequence=None, |
| 66 | img_callback=None, |
| 67 | quantize_x0=False, |
| 68 | eta=0., |
| 69 | mask=None, |
nothing calls this directly
no outgoing calls
no test coverage detected