| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from transformer.Models import Encoder, Decoder |
| from transformer.Layers import PostNet |
| from modules import VarianceAdaptor |
| from utils import get_mask_from_lengths |
| import hparams as hp |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| class FastSpeech2(nn.Module): |
| """ FastSpeech2 """ |
|
|
| def __init__(self, use_postnet=True): |
| super(FastSpeech2, self).__init__() |
|
|
| self.encoder = Encoder() |
| self.variance_adaptor = VarianceAdaptor() |
|
|
| self.decoder = Decoder() |
| self.mel_linear = nn.Linear(hp.decoder_hidden, hp.n_mel_channels) |
| |
| self.use_postnet = use_postnet |
| if self.use_postnet: |
| self.postnet = PostNet() |
|
|
| def forward(self, src_seq, src_len, mel_len=None, d_target=None, p_target=None, e_target=None, max_src_len=None, max_mel_len=None): |
| src_mask = get_mask_from_lengths(src_len, max_src_len) |
| mel_mask = get_mask_from_lengths(mel_len, max_mel_len) if mel_len is not None else None |
| |
| encoder_output = self.encoder(src_seq, src_mask) |
| if d_target is not None: |
| variance_adaptor_output, d_prediction, p_prediction, e_prediction, _, _ = self.variance_adaptor( |
| encoder_output, src_mask, mel_mask, d_target, p_target, e_target, max_mel_len) |
| else: |
| variance_adaptor_output, d_prediction, p_prediction, e_prediction, mel_len, mel_mask = self.variance_adaptor( |
| encoder_output, src_mask, mel_mask, d_target, p_target, e_target, max_mel_len) |
| |
| decoder_output = self.decoder(variance_adaptor_output, mel_mask) |
| mel_output = self.mel_linear(decoder_output) |
| |
| if self.use_postnet: |
| mel_output_postnet = self.postnet(mel_output) + mel_output |
| else: |
| mel_output_postnet = mel_output |
|
|
| return mel_output, mel_output_postnet, d_prediction, p_prediction, e_prediction, src_mask, mel_mask, mel_len |
|
|
|
|
| if __name__ == "__main__": |
| |
| model = FastSpeech2(use_postnet=False) |
| print(model) |
| print(sum(param.numel() for param in model.parameters())) |
|
|
