| 51 | |
| 52 | |
| 53 | class GaussianDiffusion: |
| 54 | def __init__(self,betas, loss_type, model_mean_type, model_var_type): |
| 55 | self.loss_type = loss_type |
| 56 | self.model_mean_type = model_mean_type |
| 57 | self.model_var_type = model_var_type |
| 58 | assert isinstance(betas, np.ndarray) |
| 59 | self.np_betas = betas = betas.astype(np.float64) # computations here in float64 for accuracy |
| 60 | assert (betas > 0).all() and (betas <= 1).all() |
| 61 | timesteps, = betas.shape |
| 62 | self.num_timesteps = int(timesteps) |
| 63 | |
| 64 | # initialize twice the actual length so we can keep running for eval |
| 65 | # betas = np.concatenate([betas, np.full_like(betas[:int(0.2*len(betas))], betas[-1])]) |
| 66 | |
| 67 | alphas = 1. - betas |
| 68 | alphas_cumprod = torch.from_numpy(np.cumprod(alphas, axis=0)).float() |
| 69 | alphas_cumprod_prev = torch.from_numpy(np.append(1., alphas_cumprod[:-1])).float() |
| 70 | |
| 71 | self.betas = torch.from_numpy(betas).float() |
| 72 | self.alphas_cumprod = alphas_cumprod.float() |
| 73 | self.alphas_cumprod_prev = alphas_cumprod_prev.float() |
| 74 | |
| 75 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 76 | self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod).float() |
| 77 | self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod).float() |
| 78 | self.log_one_minus_alphas_cumprod = torch.log(1. - alphas_cumprod).float() |
| 79 | self.sqrt_recip_alphas_cumprod = torch.sqrt(1. / alphas_cumprod).float() |
| 80 | self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1. / alphas_cumprod - 1).float() |
| 81 | |
| 82 | betas = torch.from_numpy(betas).float() |
| 83 | alphas = torch.from_numpy(alphas).float() |
| 84 | # calculations for posterior q(x_{t-1} | x_t, x_0) |
| 85 | posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) |
| 86 | # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) |
| 87 | self.posterior_variance = posterior_variance |
| 88 | # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain |
| 89 | self.posterior_log_variance_clipped = torch.log(torch.max(posterior_variance, 1e-20 * torch.ones_like(posterior_variance))) |
| 90 | self.posterior_mean_coef1 = betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod) |
| 91 | self.posterior_mean_coef2 = (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod) |
| 92 | |
| 93 | @staticmethod |
| 94 | def _extract(a, t, x_shape): |
| 95 | """ |
| 96 | Extract some coefficients at specified timesteps, |
| 97 | then reshape to [batch_size, 1, 1, 1, 1, ...] for broadcasting purposes. |
| 98 | """ |
| 99 | bs, = t.shape |
| 100 | assert x_shape[0] == bs |
| 101 | out = torch.gather(a, 0, t) |
| 102 | assert out.shape == torch.Size([bs]) |
| 103 | return torch.reshape(out, [bs] + ((len(x_shape) - 1) * [1])) |
| 104 | |
| 105 | |
| 106 | |
| 107 | def q_mean_variance(self, x_start, t): |
| 108 | mean = self._extract(self.sqrt_alphas_cumprod.to(x_start.device), t, x_start.shape) * x_start |
| 109 | variance = self._extract(1. - self.alphas_cumprod.to(x_start.device), t, x_start.shape) |
| 110 | log_variance = self._extract(self.log_one_minus_alphas_cumprod.to(x_start.device), t, x_start.shape) |