| 178 | |
| 179 | |
| 180 | class GaussianDiffusion(nn.Module): |
| 181 | def __init__(self, phone_encoder, out_dims, denoise_fn, |
| 182 | timesteps=1000, loss_type='l1', betas=None, spec_min=None, spec_max=None): |
| 183 | super().__init__() |
| 184 | self.denoise_fn = denoise_fn |
| 185 | self.fs2 = FastSpeech2(phone_encoder, out_dims) |
| 186 | self.fs2.decoder = None |
| 187 | self.mel_bins = out_dims |
| 188 | |
| 189 | if exists(betas): |
| 190 | betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas |
| 191 | else: |
| 192 | betas = cosine_beta_schedule(timesteps) |
| 193 | |
| 194 | alphas = 1. - betas |
| 195 | alphas_cumprod = np.cumprod(alphas, axis=0) |
| 196 | alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) |
| 197 | |
| 198 | timesteps, = betas.shape |
| 199 | self.num_timesteps = int(timesteps) |
| 200 | self.loss_type = loss_type |
| 201 | |
| 202 | to_torch = partial(torch.tensor, dtype=torch.float32) |
| 203 | |
| 204 | self.register_buffer('betas', to_torch(betas)) |
| 205 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 206 | self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) |
| 207 | |
| 208 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 209 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) |
| 210 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) |
| 211 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) |
| 212 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) |
| 213 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) |
| 214 | |
| 215 | # calculations for posterior q(x_{t-1} | x_t, x_0) |
| 216 | posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) |
| 217 | # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) |
| 218 | self.register_buffer('posterior_variance', to_torch(posterior_variance)) |
| 219 | # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain |
| 220 | self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) |
| 221 | self.register_buffer('posterior_mean_coef1', to_torch( |
| 222 | betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) |
| 223 | self.register_buffer('posterior_mean_coef2', to_torch( |
| 224 | (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) |
| 225 | |
| 226 | self.register_buffer('spec_min', torch.FloatTensor(spec_min)[None, None, :hparams['keep_bins']]) |
| 227 | self.register_buffer('spec_max', torch.FloatTensor(spec_max)[None, None, :hparams['keep_bins']]) |
| 228 | |
| 229 | def q_mean_variance(self, x_start, t): |
| 230 | mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start |
| 231 | variance = extract(1. - self.alphas_cumprod, t, x_start.shape) |
| 232 | log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape) |
| 233 | return mean, variance, log_variance |
| 234 | |
| 235 | def predict_start_from_noise(self, x_t, t, noise): |
| 236 | return ( |
| 237 | extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - |
no outgoing calls
no test coverage detected