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