a hZ @s.ddlZddlmZmZmZmZddlZddlmZmZddl m Z dgZ de e eeejjdd d Zd e e e e eeee ee fe eeejjd d dZeedddZGdddejZGdddejZGdddejZGdddejZGdddejZGdddejZGdddejZdS)!N)ListOptionalTupleUnion)nnTensor) functional Tacotron2Tlinear)in_dimout_dimbias w_init_gainreturncCs4tjj|||d}tjjj|jtjj|d|S)aLinear layer with xavier uniform initialization. Args: in_dim (int): Size of each input sample. out_dim (int): Size of each output sample. bias (bool, optional): If set to ``False``, the layer will not learn an additive bias. (Default: ``True``) w_init_gain (str, optional): Parameter passed to ``torch.nn.init.calculate_gain`` for setting the gain parameter of ``xavier_uniform_``. (Default: ``linear``) Returns: (torch.nn.Linear): The corresponding linear layer. r Zgain)torchrLinearinitxavier_uniform_weightcalculate_gain)r r r rr rI/var/www/auris/lib/python3.9/site-packages/torchaudio/models/tacotron2.py_get_linear_layer)s r) in_channels out_channels kernel_sizestridepaddingdilationr rrc Csl|dur0|ddkrtdt||dd}tjj|||||||d}tjjj|jtjj|d|S)al1D convolution with xavier uniform initialization. Args: in_channels (int): Number of channels in the input image. out_channels (int): Number of channels produced by the convolution. kernel_size (int, optional): Number of channels in the input image. (Default: ``1``) stride (int, optional): Number of channels in the input image. (Default: ``1``) padding (str, int or tuple, optional): Padding added to both sides of the input. (Default: dilation * (kernel_size - 1) / 2) dilation (int, optional): Number of channels in the input image. (Default: ``1``) w_init_gain (str, optional): Parameter passed to ``torch.nn.init.calculate_gain`` for setting the gain parameter of ``xavier_uniform_``. (Default: ``linear``) Returns: (torch.nn.Conv1d): The corresponding Conv1D layer. Nrzkernel_size must be odd)rrr r!r r) ValueErrorintrrConv1drrrr) rrrrr r!r rZconv1drrr_get_conv1d_layer;s  r&)lengthsrcCsFt|}tjd||j|jd}||dk}t|d}|S)alReturns a binary mask based on ``lengths``. The ``i``-th row and ``j``-th column of the mask is ``1`` if ``j`` is smaller than ``i``-th element of ``lengths. Args: lengths (Tensor): The length of each element in the batch, with shape (n_batch, ). Returns: mask (Tensor): The binary mask, with shape (n_batch, max of ``lengths``). r)devicedtyper) rmaxitemZaranger(r) unsqueezebytele)r'max_lenidsmaskrrr_get_mask_from_lengthsis  r2cs:eZdZdZeeedfdd ZeedddZZS)_LocationLayeraLocation layer used in the Attention model. Args: attention_n_filter (int): Number of filters for attention model. attention_kernel_size (int): Kernel size for attention model. attention_hidden_dim (int): Dimension of attention hidden representation. )attention_n_filterattention_kernel_sizeattention_hidden_dimc sHtt|dd}td|||dddd|_t||ddd|_dS)Nrr"F)rr r rr!tanhr r)super__init__r$r& location_convrlocation_dense)selfr4r5r6r  __class__rrr:s  z_LocationLayer.__init__)attention_weights_catrcCs$||}|dd}||}|S)aLocation layer used in the Attention model. Args: attention_weights_cat (Tensor): Cumulative and previous attention weights with shape (n_batch, 2, max of ``text_lengths``). Returns: processed_attention (Tensor): Cumulative and previous attention weights with shape (n_batch, ``attention_hidden_dim``). rr")r; transposer<)r=r@Zprocessed_attentionrrrforwards   z_LocationLayer.forward __name__ __module__ __qualname____doc__r$r:rrB __classcell__rrr>rr3zs  r3csdeZdZdZeeeeeddfdd ZeeeedddZeeeeeeeefd d d Z Z S) _AttentionaLocally sensitive attention model. Args: attention_rnn_dim (int): Number of hidden units for RNN. encoder_embedding_dim (int): Number of embedding dimensions in the Encoder. attention_hidden_dim (int): Dimension of attention hidden representation. attention_location_n_filter (int): Number of filters for Attention model. attention_location_kernel_size (int): Kernel size for Attention model. N)attention_rnn_dimencoder_embedding_dimr6attention_location_n_filterattention_location_kernel_sizercs\tt||ddd|_t||ddd|_t|ddd|_t||||_td |_ dS)NFr7r8rrinf) r9r:r query_layer memory_layervr3location_layerfloatscore_mask_value)r=rJrKr6rLrMr>rrr:s z_Attention.__init__)queryprocessed_memoryr@rcCs@||d}||}|t|||}|d}|S)a=Get the alignment vector. Args: query (Tensor): Decoder output with shape (n_batch, n_mels * n_frames_per_step). processed_memory (Tensor): Processed Encoder outputs with shape (n_batch, max of ``text_lengths``, attention_hidden_dim). attention_weights_cat (Tensor): Cumulative and previous attention weights with shape (n_batch, 2, max of ``text_lengths``). Returns: alignment (Tensor): attention weights, it is a tensor with shape (batch, max of ``text_lengths``). rr")rOr,rRrQrr7squeeze)r=rUrVr@Zprocessed_queryZprocessed_attention_weightsZenergies alignmentrrr_get_alignment_energiess   z"_Attention._get_alignment_energies)attention_hidden_statememoryrVr@r1rc CsN||||}|||j}tj|dd}t|d|}|d}||fS)aPass the input through the Attention model. Args: attention_hidden_state (Tensor): Attention rnn last output with shape (n_batch, ``attention_rnn_dim``). memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). processed_memory (Tensor): Processed Encoder outputs with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``). attention_weights_cat (Tensor): Previous and cumulative attention weights with shape (n_batch, current_num_frames * 2, max of ``text_lengths``). mask (Tensor): Binary mask for padded data with shape (n_batch, current_num_frames). Returns: attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``). attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``). rZdim) rYZ masked_fillrTFZsoftmaxrZbmmr,rW) r=rZr[rVr@r1rXattention_weightsattention_contextrrrrBs  z_Attention.forward) rDrErFrGr$r:rrYrrBrHrrr>rrIs   rIcs>eZdZdZeeeddfdd ZeedddZZ S) _PrenetzPrenet Module. It is consists of ``len(output_size)`` linear layers. Args: in_dim (int): The size of each input sample. output_sizes (list): The output dimension of each linear layers. N)r out_sizesrcs<t|g|dd}tddt||D|_dS)NcSsg|]\}}t||ddqS)Fr)r).0Zin_sizeZout_sizerrr z$_Prenet.__init__..)r9r:r ModuleListziplayers)r=r raZin_sizesr>rrr: s  z_Prenet.__init__xrcCs*|jD]}tjt||ddd}q|S)zPass the input through Prenet. Args: x (Tensor): The input sequence to Prenet with shape (n_batch, in_dim). Return: x (Tensor): Tensor with shape (n_batch, sizes[-1]) ?T)ptraining)rhr]dropoutrelu)r=rjr rrrrBs z_Prenet.forward) rDrErFrGr$rr:rrBrHrrr>rr`sr`cs<eZdZdZeeeedfdd ZeedddZZS)_PostnetaPostnet Module. Args: n_mels (int): Number of mel bins. postnet_embedding_dim (int): Postnet embedding dimension. postnet_kernel_size (int): Postnet kernel size. postnet_n_convolution (int): Number of postnet convolutions. )n_melspostnet_embedding_dimpostnet_kernel_sizepostnet_n_convolutionc stt|_t|D]}|dkr,|n|}||dkr@|n|}||dkrTdnd}||dkrh|n|} |jtt|||dt |ddd|dt | qt |j|_ dS)Nrrr r7r"rrr r!r) r9r:rrf convolutionsrangeappend Sequentialr&r$ BatchNorm1dlenn_convs) r=rqrrrsrtirrZ init_gainZ num_featuresr>rrr:*s,    z_Postnet.__init__ricCsZt|jD]J\}}||jdkr>tjt||d|jd}q tj||d|jd}q |S)aPass the input through Postnet. Args: x (Tensor): The input sequence with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``). Return: x (Tensor): Tensor with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``). rrk)rm) enumeratervr|r]rnrr7rm)r=rjr}convrrrrBJs z_Postnet.forwardrCrrr>rrp s  rpcs>eZdZdZeeeddfdd ZeeedddZZS) _EncoderaEncoder Module. Args: encoder_embedding_dim (int): Number of embedding dimensions in the encoder. encoder_n_convolution (int): Number of convolution layers in the encoder. encoder_kernel_size (int): The kernel size in the encoder. Examples >>> encoder = _Encoder(3, 512, 5) >>> input = torch.rand(10, 20, 30) >>> output = encoder(input) # shape: (10, 30, 512) N)rKencoder_n_convolutionencoder_kernel_sizerc stt|_t|D]@}tt|||dt|dddddt |}|j |qtj |t|ddddd|_ |j dS)Nrr"roruT) batch_first bidirectional)r9r:rrfrvrwryr&r$rzrxZLSTMlstmZflatten_parameters)r=rKrr_Z conv_layerr>rrr:ks0      z_Encoder.__init__)rj input_lengthsrcCsv|jD]}tt||d|j}q|dd}|}tjj j ||dd}| |\}}tjj j |dd\}}|S)a_Pass the input through the Encoder. Args: x (Tensor): The input sequences with shape (n_batch, encoder_embedding_dim, n_seq). input_lengths (Tensor): The length of each input sequence with shape (n_batch, ). Return: x (Tensor): A tensor with shape (n_batch, n_seq, encoder_embedding_dim). rkrr"T)r) rvr]rnrormrAcpurutilsZrnnZpack_padded_sequencerZpad_packed_sequence)r=rjrroutputsrrrrrBs  z_Encoder.forwardrCrrr>rr]s!rcs2eZdZdZeeeeeeeeeeeeeeddfdd ZeedddZ ee eeeeeeeefdd d Z eed d d Z eeee eeefdddZ eeeeeeeeeeee eeeeeeeeef d ddZeeee eeefdddZeedddZejjeee eeeefdddZZS)_Decodera,Decoder with Attention model. Args: n_mels (int): number of mel bins n_frames_per_step (int): number of frames processed per step, only 1 is supported encoder_embedding_dim (int): the number of embedding dimensions in the encoder. decoder_rnn_dim (int): number of units in decoder LSTM decoder_max_step (int): maximum number of output mel spectrograms decoder_dropout (float): dropout probability for decoder LSTM decoder_early_stopping (bool): stop decoding when all samples are finished attention_rnn_dim (int): number of units in attention LSTM attention_hidden_dim (int): dimension of attention hidden representation attention_location_n_filter (int): number of filters for attention model attention_location_kernel_size (int): kernel size for attention model attention_dropout (float): dropout probability for attention LSTM prenet_dim (int): number of ReLU units in prenet layers gate_threshold (float): probability threshold for stop token N)rqn_frames_per_steprKdecoder_rnn_dimdecoder_max_stepdecoder_dropoutdecoder_early_stoppingrJr6rLrMattention_dropout prenet_dimgate_thresholdrcst||_||_||_||_||_| |_||_||_ | |_ ||_ ||_ t ||| | g|_t| |||_t||| | | |_t|||d|_t|||||_t||dddd|_dS)NTrsigmoidr8)r9r:rqrrKrJrrrrrrrr`prenetrZLSTMCell attention_rnnrIattention_layer decoder_rnnrlinear_projection gate_layer)r=rqrrKrrrrrJr6rLrMrrrr>rrr:s4  z_Decoder.__init__)r[rcCs4|d}|j}|j}tj||j|j||d}|S)amGets all zeros frames to use as the first decoder input. Args: memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). Returns: decoder_input (Tensor): all zeros frames with shape (n_batch, max of ``text_lengths``, ``n_mels * n_frames_per_step``). rr)r(sizer)r(rzerosrqrr=r[n_batchr)r( decoder_inputrrr_get_initial_frames z_Decoder._get_initial_framecCs|d}|d}|j}|j}tj||j||d}tj||j||d}tj||j||d}tj||j||d} tj||||d} tj||||d} tj||j||d} |j |} |||| | | | | fS)aInitializes attention rnn states, decoder rnn states, attention weights, attention cumulative weights, attention context, stores memory and stores processed memory. Args: memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). Returns: attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``). attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``). attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``). processed_memory (Tensor): Processed encoder outputs with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``). rrr) rr)r(rrrJrrKrrP)r=r[rZmax_timer)r(attention_hiddenattention_celldecoder_hidden decoder_cellr^attention_weights_cumr_rVrrr_initialize_decoder_statess*   z#_Decoder._initialize_decoder_states)decoder_inputsrcCs@|dd}||dt|d|jd}|dd}|S)akPrepares decoder inputs. Args: decoder_inputs (Tensor): Inputs used for teacher-forced training, i.e. mel-specs, with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``) Returns: inputs (Tensor): Processed decoder inputs with shape (max of ``mel_specgram_lengths``, n_batch, ``n_mels``). rr"rrb)rAviewrr$r)r=rrrr_parse_decoder_inputs.s  z_Decoder._parse_decoder_inputs) mel_specgram gate_outputs alignmentsrcCsb|dd}|dd}|dd}|jdd|jf}|j|}|dd}|||fS)aqPrepares decoder outputs for output Args: mel_specgram (Tensor): mel spectrogram with shape (max of ``mel_specgram_lengths``, n_batch, ``n_mels``) gate_outputs (Tensor): predicted stop token with shape (max of ``mel_specgram_lengths``, n_batch) alignments (Tensor): sequence of attention weights from the decoder with shape (max of ``mel_specgram_lengths``, n_batch, max of ``text_lengths``) Returns: mel_specgram (Tensor): mel spectrogram with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``) gate_outputs (Tensor): predicted stop token with shape (n_batch, max of ``mel_specgram_lengths``) alignments (Tensor): sequence of attention weights from the decoder with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``) rrrbr")rA contiguousshaperqr)r=rrrrrrr_parse_decoder_outputsCs  z_Decoder._parse_decoder_outputs) rrrrrr^rr_r[rVr1rc  Cst||fd} || ||f\}}t||j|j}tj|d|dfdd} ||| | | | \}}||7}t||fd}| |||f\}}t||j |j}tj||fdd}| |}| |}|||||||||f S)a& Decoder step using stored states, attention and memory Args: decoder_input (Tensor): Output of the Prenet with shape (n_batch, ``prenet_dim``). attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``). attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``). attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``). memory (Tensor): Encoder output with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). processed_memory (Tensor): Processed Encoder outputs with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``). mask (Tensor): Binary mask for padded data with shape (n_batch, current_num_frames). Returns: decoder_output: Predicted mel spectrogram for the current frame with shape (n_batch, ``n_mels``). gate_prediction (Tensor): Prediction of the stop token with shape (n_batch, ``1``). attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``). decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``). attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``). attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``). attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``). rbrr\) rcatrr]rnrrmr,rrrrr)r=rrrrrr^rr_r[rVr1Z cell_inputr@Z decoder_hidden_attention_contextZdecoder_outputZgate_predictionrrrdecodecs0)   z_Decoder.decode)r[mel_specgram_truthmemory_lengthsrc Cs||d}||}tj||fdd}||}t|}||\}}} } } } } }ggg}}}t|| ddkr|t|}| |||| | | | | ||| \ }}}}} } } } } || dg7}|| dg7}|| g7}qh| t |t |t |\}}}|||fS)aDecoder forward pass for training. Args: memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). mel_specgram_truth (Tensor): Decoder ground-truth mel-specs for teacher forcing with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``). memory_lengths (Tensor): Encoder output lengths for attention masking (the same as ``text_lengths``) with shape (n_batch, ). Returns: mel_specgram (Tensor): Predicted mel spectrogram with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``). gate_outputs (Tensor): Predicted stop token for each timestep with shape (n_batch, max of ``mel_specgram_lengths``). alignments (Tensor): Sequence of attention weights from the decoder with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``). rr\r)rr,rrrrr2rr{rrrWrstack)r=r[rrrrr1rrrrr^rr_rVZ mel_outputsrrZ mel_output gate_outputrrrrrBs`      z_Decoder.forwardcCs4|d}|j}|j}tj||j|j||d}|S)aUGets all zeros frames to use as the first decoder input args: memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). returns: decoder_input (Tensor): All zeros frames with shape(n_batch, ``n_mels`` * ``n_frame_per_step``). rrrrrrr _get_go_frames z_Decoder._get_go_frame)r[rrcCs|d|j}}||}t|}||\}}} } } } } }tj|gtj|d}tj|gtj|d}g}g}g}t |j D]}| |}| |||| | | | | ||| \ }}}}} } } } } | |d| |dd| | ||d7<|t|d|jkO}|jr,t|r,q2|}q|t||j krLtdtj|dd}tj|dd}tj|dd}||||\}}}||||fS)aDecoder inference Args: memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``). memory_lengths (Tensor): Encoder output lengths for attention masking (the same as ``text_lengths``) with shape (n_batch, ). Returns: mel_specgram (Tensor): Predicted mel spectrogram with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``). mel_specgram_lengths (Tensor): the length of the predicted mel spectrogram (n_batch, )) gate_outputs (Tensor): Predicted stop token for each timestep with shape (n_batch, max of ``mel_specgram_lengths``). alignments (Tensor): Sequence of attention weights from the decoder with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``). rrrzZReached max decoder steps. The generated spectrogram might not cover the whole transcript.r\)rr(rr2rrrint32boolrwrrrrxr,rArrWrrallr{warningswarnrr)r=r[rZ batch_sizer(rr1rrrrr^rr_rVmel_specgram_lengthsfinishedZ mel_specgramsrrrrrrrrinfer sx    z_Decoder.infer)rDrErFrGr$rSrr:rrrrrrrrBrrjitexportrrHrrr>rrsX3 1 " K LrcseZdZdZdeeeeeeeeeeeeeeeeeeeeeeddfdd Zeeeee eeeefdddZ e j j deeee eeefdddZZS)r a Tacotron2 model from *Natural TTS Synthesis by Conditioning WaveNet on Mel Spectrogram Predictions* :cite:`shen2018natural` based on the implementation from `Nvidia Deep Learning Examples `_. See Also: * :class:`torchaudio.pipelines.Tacotron2TTSBundle`: TTS pipeline with pretrained model. Args: mask_padding (bool, optional): Use mask padding (Default: ``False``). n_mels (int, optional): Number of mel bins (Default: ``80``). n_symbol (int, optional): Number of symbols for the input text (Default: ``148``). n_frames_per_step (int, optional): Number of frames processed per step, only 1 is supported (Default: ``1``). symbol_embedding_dim (int, optional): Input embedding dimension (Default: ``512``). encoder_n_convolution (int, optional): Number of encoder convolutions (Default: ``3``). encoder_kernel_size (int, optional): Encoder kernel size (Default: ``5``). encoder_embedding_dim (int, optional): Encoder embedding dimension (Default: ``512``). decoder_rnn_dim (int, optional): Number of units in decoder LSTM (Default: ``1024``). decoder_max_step (int, optional): Maximum number of output mel spectrograms (Default: ``2000``). decoder_dropout (float, optional): Dropout probability for decoder LSTM (Default: ``0.1``). decoder_early_stopping (bool, optional): Continue decoding after all samples are finished (Default: ``True``). attention_rnn_dim (int, optional): Number of units in attention LSTM (Default: ``1024``). attention_hidden_dim (int, optional): Dimension of attention hidden representation (Default: ``128``). attention_location_n_filter (int, optional): Number of filters for attention model (Default: ``32``). attention_location_kernel_size (int, optional): Kernel size for attention model (Default: ``31``). attention_dropout (float, optional): Dropout probability for attention LSTM (Default: ``0.1``). prenet_dim (int, optional): Number of ReLU units in prenet layers (Default: ``256``). postnet_n_convolution (int, optional): Number of postnet convolutions (Default: ``5``). postnet_kernel_size (int, optional): Postnet kernel size (Default: ``5``). postnet_embedding_dim (int, optional): Postnet embedding dimension (Default: ``512``). gate_threshold (float, optional): Probability threshold for stop token (Default: ``0.5``). FPr皙?T rkN) mask_paddingrqn_symbolrsymbol_embedding_dimrKrrrrrrrJr6rLrMrrrtrsrrrrcst||_||_||_t|||_tjj |jj t ||||_ t|||| | | | | |||||||_t|||||_dS)N)r9r:rrqrrZ Embedding embeddingrrrrrencoderrdecoderrppostnet)r=rrqrrrrKrrrrrrrJr6rLrMrrrtrsrrrr>rrr:s0 zTacotron2.__init__)tokens token_lengthsrrrc Cs||dd}|||}|j|||d\}}}||} || } |jrt|} | |j| d| d} | ddd} | | d| | d| | dddddfd|| ||fS)aPass the input through the Tacotron2 model. This is in teacher forcing mode, which is generally used for training. The input ``tokens`` should be padded with zeros to length max of ``token_lengths``. The input ``mel_specgram`` should be padded with zeros to length max of ``mel_specgram_lengths``. Args: tokens (Tensor): The input tokens to Tacotron2 with shape `(n_batch, max of token_lengths)`. token_lengths (Tensor): The valid length of each sample in ``tokens`` with shape `(n_batch, )`. mel_specgram (Tensor): The target mel spectrogram with shape `(n_batch, n_mels, max of mel_specgram_lengths)`. mel_specgram_lengths (Tensor): The length of each mel spectrogram with shape `(n_batch, )`. Returns: [Tensor, Tensor, Tensor, Tensor]: Tensor Mel spectrogram before Postnet with shape `(n_batch, n_mels, max of mel_specgram_lengths)`. Tensor Mel spectrogram after Postnet with shape `(n_batch, n_mels, max of mel_specgram_lengths)`. Tensor The output for stop token at each time step with shape `(n_batch, max of mel_specgram_lengths)`. Tensor Sequence of attention weights from the decoder with shape `(n_batch, max of mel_specgram_lengths, max of token_lengths)`. rr")rrgNg@@) rrArrrrr2expandrqrZpermuteZ masked_fill_) r=rrrrembedded_inputsencoder_outputsrrZmel_specgram_postnetr1rrrrBs!     zTacotron2.forward)rr'rc Cs|j\}}|dur0t|g||j|j}|dusrr esp"5 6)Tr )rrNrTr )rtypingrrrrrrrZtorch.nnrr]__all__r$rstrrrr%r&r2Moduler3rIr`rprrr rrrrsF  .1W=HC