(self, beta_schedule="linear", timesteps=1000,
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3,
zero_terminal_snr=False)
| 127 | self.loss_type = loss_type |
| 128 | |
| 129 | def register_schedule(self, beta_schedule="linear", timesteps=1000, |
| 130 | linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3, |
| 131 | zero_terminal_snr=False): |
| 132 | betas = make_beta_schedule( |
| 133 | beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s |
| 134 | ) |
| 135 | if zero_terminal_snr: |
| 136 | betas = enforce_zero_terminal_snr(betas) |
| 137 | alphas = 1. - betas |
| 138 | alphas_cumprod = np.cumprod(alphas, axis=0) |
| 139 | alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) |
| 140 | alphas_cumprod_next = np.append(alphas_cumprod[1:], 0.0) |
| 141 | |
| 142 | timesteps, = betas.shape |
| 143 | self.num_timesteps = int(timesteps) |
| 144 | self.linear_start = linear_start |
| 145 | self.linear_end = linear_end |
| 146 | assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' |
| 147 | |
| 148 | to_torch = partial(torch.tensor, dtype=torch.float32) |
| 149 | |
| 150 | self.register_buffer('betas', to_torch(betas)) |
| 151 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 152 | self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) |
| 153 | self.register_buffer('alphas_cumprod_next', to_torch(alphas_cumprod_next)) |
| 154 | |
| 155 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 156 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) |
| 157 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) |
| 158 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) |
| 159 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) |
| 160 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) |
| 161 | |
| 162 | # calculations for posterior q(x_{t-1} | x_t, x_0) |
| 163 | posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( |
| 164 | 1. - alphas_cumprod) + self.v_posterior * betas |
| 165 | # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) |
| 166 | self.register_buffer('posterior_variance', to_torch(posterior_variance)) |
| 167 | # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain |
| 168 | self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) |
| 169 | self.register_buffer('posterior_mean_coef1', to_torch( |
| 170 | betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) |
| 171 | self.register_buffer('posterior_mean_coef2', to_torch( |
| 172 | (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) |
| 173 | |
| 174 | if self.parameterization == "eps": |
| 175 | lvlb_weights = self.betas ** 2 / ( |
| 176 | 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) |
| 177 | elif self.parameterization == "x0": |
| 178 | lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) |
| 179 | elif self.parameterization == "v": |
| 180 | lvlb_weights = torch.ones_like(self.betas ** 2 / ( |
| 181 | 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))) |
| 182 | else: |
| 183 | raise NotImplementedError(f"Parameterization {self.parameterization} not supported yet") |
| 184 | lvlb_weights[0] = lvlb_weights[1] |
| 185 | self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) |
| 186 | assert not torch.isnan(self.lvlb_weights).all() |
no test coverage detected