| 33 | |
| 34 | |
| 35 | class FFT(FastspeechDecoder): |
| 36 | def __init__(self, hidden_size=None, num_layers=None, kernel_size=None, num_heads=None): |
| 37 | super().__init__(hidden_size, num_layers, kernel_size, num_heads=num_heads) |
| 38 | dim = hparams['residual_channels'] |
| 39 | self.input_projection = Conv1d(hparams['audio_num_mel_bins'], dim, 1) |
| 40 | self.diffusion_embedding = SinusoidalPosEmb(dim) |
| 41 | self.mlp = nn.Sequential( |
| 42 | nn.Linear(dim, dim * 4), |
| 43 | Mish(), |
| 44 | nn.Linear(dim * 4, dim) |
| 45 | ) |
| 46 | self.get_mel_out = Linear(hparams['hidden_size'], 80, bias=True) |
| 47 | self.get_decode_inp = Linear(hparams['hidden_size'] + dim + dim, |
| 48 | hparams['hidden_size']) # hs + dim + 80 -> hs |
| 49 | |
| 50 | def forward(self, spec, diffusion_step, cond, padding_mask=None, attn_mask=None, return_hiddens=False): |
| 51 | """ |
| 52 | :param spec: [B, 1, 80, T] |
| 53 | :param diffusion_step: [B, 1] |
| 54 | :param cond: [B, M, T] |
| 55 | :return: |
| 56 | """ |
| 57 | x = spec[:, 0] |
| 58 | x = self.input_projection(x).permute([0, 2, 1]) # [B, T, residual_channel] |
| 59 | diffusion_step = self.diffusion_embedding(diffusion_step) |
| 60 | diffusion_step = self.mlp(diffusion_step) # [B, dim] |
| 61 | cond = cond.permute([0, 2, 1]) # [B, T, M] |
| 62 | |
| 63 | seq_len = cond.shape[1] # [T_mel] |
| 64 | time_embed = diffusion_step[:, None, :] # [B, 1, dim] |
| 65 | time_embed = time_embed.repeat([1, seq_len, 1]) # # [B, T, dim] |
| 66 | |
| 67 | decoder_inp = torch.cat([x, cond, time_embed], dim=-1) # [B, T, dim + H + dim] |
| 68 | decoder_inp = self.get_decode_inp(decoder_inp) # [B, T, H] |
| 69 | x = decoder_inp |
| 70 | |
| 71 | ''' |
| 72 | Required x: [B, T, C] |
| 73 | :return: [B, T, C] or [L, B, T, C] |
| 74 | ''' |
| 75 | padding_mask = x.abs().sum(-1).eq(0).data if padding_mask is None else padding_mask |
| 76 | nonpadding_mask_TB = 1 - padding_mask.transpose(0, 1).float()[:, :, None] # [T, B, 1] |
| 77 | if self.use_pos_embed: |
| 78 | positions = self.pos_embed_alpha * self.embed_positions(x[..., 0]) |
| 79 | x = x + positions |
| 80 | x = F.dropout(x, p=self.dropout, training=self.training) |
| 81 | # B x T x C -> T x B x C |
| 82 | x = x.transpose(0, 1) * nonpadding_mask_TB |
| 83 | hiddens = [] |
| 84 | for layer in self.layers: |
| 85 | x = layer(x, encoder_padding_mask=padding_mask, attn_mask=attn_mask) * nonpadding_mask_TB |
| 86 | hiddens.append(x) |
| 87 | if self.use_last_norm: |
| 88 | x = self.layer_norm(x) * nonpadding_mask_TB |
| 89 | if return_hiddens: |
| 90 | x = torch.stack(hiddens, 0) # [L, T, B, C] |
| 91 | x = x.transpose(1, 2) # [L, B, T, C] |
| 92 | else: |