| 354 | |
| 355 | |
| 356 | class OfflineGaussianDiffusion(GaussianDiffusion): |
| 357 | def forward(self, txt_tokens, mel2ph=None, spk_embed=None, |
| 358 | ref_mels=None, f0=None, uv=None, energy=None, infer=False): |
| 359 | b, *_, device = *txt_tokens.shape, txt_tokens.device |
| 360 | |
| 361 | ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy, |
| 362 | skip_decoder=True, infer=True) |
| 363 | cond = ret['decoder_inp'].transpose(1, 2) |
| 364 | fs2_mels = ref_mels[1] |
| 365 | ref_mels = ref_mels[0] |
| 366 | |
| 367 | if not infer: |
| 368 | t = torch.randint(0, self.K_step, (b,), device=device).long() |
| 369 | x = ref_mels |
| 370 | x = self.norm_spec(x) |
| 371 | x = x.transpose(1, 2)[:, None, :, :] # [B, 1, M, T] |
| 372 | ret['diff_loss'] = self.p_losses(x, t, cond) |
| 373 | else: |
| 374 | t = self.K_step |
| 375 | fs2_mels = self.norm_spec(fs2_mels) |
| 376 | fs2_mels = fs2_mels.transpose(1, 2)[:, None, :, :] |
| 377 | |
| 378 | x = self.q_sample(x_start=fs2_mels, t=torch.tensor([t - 1], device=device).long()) |
| 379 | |
| 380 | if hparams.get('gaussion_start') is not None and hparams['gaussion_start']: |
| 381 | print('===> gaussion start.') |
| 382 | shape = (cond.shape[0], 1, self.mel_bins, cond.shape[2]) |
| 383 | x = torch.randn(shape, device=device) |
| 384 | for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t): |
| 385 | x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond) |
| 386 | x = x[:, 0].transpose(1, 2) |
| 387 | ret['mel_out'] = self.denorm_spec(x) |
| 388 | return ret |