Spaces:
Runtime error
Runtime error
| # -------------------------------------------------------- | |
| # ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621) | |
| # Github source: https://github.com/mbzuai-nlp/ArTST | |
| # Based on speecht5, fairseq and espnet code bases | |
| # https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet | |
| # -------------------------------------------------------- | |
| import contextlib | |
| import torch | |
| import torch.nn as nn | |
| from espnet.nets.pytorch_backend.tacotron2.decoder import Postnet | |
| class SpeechDecoderPostnet(nn.Module): | |
| """ | |
| Args: | |
| in_channels (int): the number of input channels | |
| mid_channels (int): the number of intermediate channels | |
| out_channels (int): the number of output channels | |
| kernel_sizes (List[int]): the kernel size for each convolutional layer | |
| """ | |
| def __init__( | |
| self, | |
| odim, | |
| args, | |
| ): | |
| super(SpeechDecoderPostnet, self).__init__() | |
| # define decoder postnet | |
| # define final projection | |
| self.feat_out = torch.nn.Linear(args.decoder_embed_dim, odim * args.reduction_factor) | |
| self.prob_out = torch.nn.Linear(args.decoder_embed_dim, args.reduction_factor) | |
| # define postnet | |
| self.postnet = ( | |
| None | |
| if args.postnet_layers == 0 | |
| else Postnet( | |
| idim=0, | |
| odim=odim, | |
| n_layers=args.postnet_layers, | |
| n_chans=args.postnet_chans, | |
| n_filts=args.postnet_filts, | |
| use_batch_norm=args.use_batch_norm, | |
| dropout_rate=args.postnet_dropout_rate, | |
| ) | |
| ) | |
| self.odim = odim | |
| self.num_updates = 0 | |
| self.freeze_decoder_updates = args.freeze_decoder_updates | |
| def forward(self, zs): | |
| ft = self.freeze_decoder_updates <= self.num_updates | |
| with torch.no_grad() if not ft else contextlib.ExitStack(): | |
| # (B, Lmax//r, odim * r) -> (B, Lmax//r * r, odim) | |
| before_outs = self.feat_out(zs).view(zs.size(0), -1, self.odim) | |
| # (B, Lmax//r, r) -> (B, Lmax//r * r) | |
| logits = self.prob_out(zs).view(zs.size(0), -1) | |
| # postnet -> (B, Lmax//r * r, odim) | |
| if self.postnet is None: | |
| after_outs = before_outs | |
| else: | |
| after_outs = before_outs + self.postnet( | |
| before_outs.transpose(1, 2) | |
| ).transpose(1, 2) | |
| return before_outs, after_outs, logits | |
| def set_num_updates(self, num_updates): | |
| """Set the number of parameters updates.""" | |
| self.num_updates = num_updates | |