U ,:%eZ³ã @s.ddlZddlmZmZmZmZddlZddlmZmZddl m Z dgZ de e eeejjdœd d „Zd e e e e eeee ee fe eeejjd œ d d„Zeedœdd„ZGdd„dejƒZGdd„dejƒZGdd„dejƒZGdd„dejƒZGdd„dejƒZGdd„dejƒZGdd„dejƒZdS)!éN)ÚListÚOptionalÚTupleÚUnion)ÚnnÚTensor)Ú functionalÚ Tacotron2TÚlinear)Úin_dimÚout_dimÚbiasÚ w_init_gainÚreturncCs4tjj|||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)ÚtorchrÚLinearÚinitÚxavier_uniform_ÚweightÚcalculate_gain)r r r rr ©rúZ/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/torchaudio/models/tacotron2.pyÚ_get_linear_layer)s ré) Ú in_channelsÚ out_channelsÚ kernel_sizeÚstrideÚpaddingÚdilationr rrc Csl|dkr0|ddkrtdƒ‚t||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. Nérzkernel_size must be odd)rrr r!r r) Ú ValueErrorÚintrrÚConv1drrrr) rrrrr r!r rZconv1drrrÚ_get_conv1d_layer;s ù r&)ÚlengthsrcCsFt |¡ ¡}tjd||j|jd}|| d¡k ¡}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)ÚdeviceÚdtyper) rÚmaxÚitemZaranger(r)Ú unsqueezeÚbyteÚle)r'Úmax_lenZidsÚmaskrrrÚ_get_mask_from_lengthsis  r1cs:eZdZdZeeedœ‡fdd„ Zeedœdd„Z‡ZS)Ú_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_filterÚattention_kernel_sizeÚattention_hidden_dimc sHtƒ ¡t|ddƒ}td|||dddd|_t||ddd|_dS)Nrr"F)rr r rr!Útanh©r r)ÚsuperÚ__init__r$r&Ú location_convrÚlocation_dense)Úselfr3r4r5r ©Ú __class__rrr9ƒ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_attentionrrrÚforward˜s   z_LocationLayer.forward© Ú__name__Ú __module__Ú __qualname__Ú__doc__r$r9rrAÚ __classcell__rrr=rr2zs  ür2csdeZdZdZeeeeeddœ‡fdd„ Zeeeedœdd„Zeeeeeeeefd œd d „Z ‡Z S) Ú _AttentionaºLocally 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_dimÚencoder_embedding_dimr5Úattention_location_n_filterÚattention_location_kernel_sizercs\tƒ ¡t||ddd|_t||ddd|_t|ddd|_t|||ƒ|_tdƒ |_ dS)NFr6r7rrÚinf) r8r9rÚ query_layerÚ memory_layerÚvr2Úlocation_layerÚfloatÚscore_mask_value)r<rIrJr5rKrLr=rrr9¶s ÿýz_Attention.__init__)ÚqueryÚprocessed_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")rNr,rQrPrr6Úsqueeze)r<rTrUr?Zprocessed_queryZprocessed_attention_weightsZenergiesÚ alignmentrrrÚ_get_alignment_energiesËs   z"_Attention._get_alignment_energies)Úattention_hidden_stateÚmemoryrUr?r0rc CsN| |||¡}| ||j¡}tj|dd}t | d¡|¡}| d¡}||fS)a¹Pass 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``). r©Údim) rXZ masked_fillrSÚFZsoftmaxrZbmmr,rV) r<rYrZrUr?r0rWÚattention_weightsÚattention_contextrrrrAàs  z_Attention.forward) rCrDrErFr$r9rrXrrArGrrr=rrH«s  ù ùrHcs>eZdZdZeeeddœ‡fdd„ Zeedœdd„Z‡Z S) Ú_PrenetzÒPrenet 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…}t dd„t||ƒDƒ¡|_dS)NéÿÿÿÿcSsg|]\}}t||dd‘qS)Fr)r)Ú.0Zin_sizeZout_sizerrrÚ sz$_Prenet.__init__..)r8r9rÚ ModuleListÚzipÚlayers)r<r raZin_sizesr=rrr9 s  ÿz_Prenet.__init__©ÚxrcCs*|jD]}tjt ||ƒ¡ddd}q|S)zÙPass 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)ÚpÚtraining)rgr]ÚdropoutÚrelu)r<rir rrrrAs z_Prenet.forward) rCrDrErFr$rr9rrArGrrr=rr`sr`cs<eZdZdZeeeedœ‡fdd„ Zeedœdd„Z‡ZS)Ú_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_melsÚpostnet_embedding_dimÚpostnet_kernel_sizeÚpostnet_n_convolutionc s¶tƒ ¡t ¡|_t|ƒD]ˆ}|dkr,|n|}||dkr@|n|}||dkrTdnd}||dkrh|n|} |j t t|||dt |ddƒd|dt  | ¡¡¡qt |jƒ|_ dS)Nrrr r6r"©rrr r!r) r8r9rreÚ convolutionsÚrangeÚappendÚ Sequentialr&r$Ú BatchNorm1dÚlenÚn_convs) r<rprqrrrsÚirrZ init_gainZ num_featuresr=rrr9*s,   ù öÿz_Postnet.__init__rhcCsZt|jƒD]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``). rrj)rl)Ú enumeraterur{r]rmrr6rl)r<rir|ÚconvrrrrAJs z_Postnet.forwardrBrrr=rro s û rocs>eZdZdZeeeddœ‡fdd„ Zeeedœdd„Z‡ZS) Ú_Encodera§Encoder 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)rJÚencoder_n_convolutionÚencoder_kernel_sizerc sŠtƒ ¡t ¡|_t|ƒD]@}t t|||dt|ddƒdddt  |¡¡}|j  |¡qtj |t|dƒdddd|_ |j   ¡dS)Nrr"rnrtT)Ú batch_firstÚ bidirectional)r8r9rrerurvrxr&r$ryrwZLSTMÚlstmZflatten_parameters)r<rJr€rÚ_Z conv_layerr=rrr9ks0   ù ö  ûz_Encoder.__init__)riÚ input_lengthsrcCsv|jD]}t t ||ƒ¡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). rjrr"T)r‚) rur]rmrnrlr@ÚcpurÚutilsZrnnZpack_padded_sequencer„Zpad_packed_sequence)r<rir†r~Úoutputsr…rrrrAŒs  z_Encoder.forwardrBrrr=rr]sû!rcs2eZdZdZeeeeeeeeeeeeeeddœ‡fdd„ Zeedœdd„Z ee eeeeeeeefdœd d „Z eed œd d „Z eeee eeefdœdd„Z eeeeeeeeeeee eeeeeeeeef dœ dd„Zeeee eeefdœdd„Zeedœdd„Zejjeee eeeefdœdd„ƒZ‡ZS)Ú_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)rpÚn_frames_per_steprJÚdecoder_rnn_dimÚdecoder_max_stepÚdecoder_dropoutÚdecoder_early_stoppingrIr5rKrLÚattention_dropoutÚ prenet_dimÚgate_thresholdrcsÆtƒ ¡||_||_||_||_||_| |_||_||_ | |_ ||_ ||_ t ||| | gƒ|_t | ||¡|_t||| | | ƒ|_t |||d¡|_t||||ƒ|_t||dddd|_dS)NTrÚsigmoidr7)r8r9rpr‹rJrIrŒr‘rr’rrŽrr`ÚprenetrZLSTMCellÚ attention_rnnrHÚattention_layerÚ decoder_rnnrÚlinear_projectionÚ gate_layer)r<rpr‹rJrŒrrŽrrIr5rKrLrr‘r’r=rrr9¹s: ûÿz_Decoder.__init__)rZrcCs4| 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``). r©r)r(©Úsizer)r(rÚzerosrpr‹©r<rZÚn_batchr)r(Ú decoder_inputrrrÚ_get_initial_frameìs 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š) rœr)r(rrrIrŒrJr–rO)r<rZrŸZmax_timer)r(Úattention_hiddenÚattention_cellÚdecoder_hiddenÚ decoder_cellr^Úattention_weights_cumr_rUrrrÚ_initialize_decoder_statesýs*   øz#_Decoder._initialize_decoder_states)Údecoder_inputsrcCs@| dd¡}| | d¡t| d¡|jƒd¡}| 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)r@Úviewrœr$r‹)r<r¨rrrÚ_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")r@Ú contiguousÚshaperpr©)r<r«r¬r­r¯rrrÚ_parse_decoder_outputsCs  z_Decoder._parse_decoder_outputs) r r¢r£r¤r¥r^r¦r_rZrUr0rc  Csät ||fd¡} | | ||f¡\}}t ||j|j¡}tj| d¡| d¡fdd} | || | | | ¡\}}||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[) rÚcatr•r]rmrrlr,r–r—rŽr˜r™)r<r r¢r£r¤r¥r^r¦r_rZrUr0Z cell_inputr?Z decoder_hidden_attention_contextZdecoder_outputZgate_predictionrrrÚdecodecs8)ÿ  ÷z_Decoder.decode)rZÚmel_specgram_truthÚmemory_lengthsrc Cs| |¡ d¡}| |¡}tj||fdd}| |¡}t|ƒ}| |¡\}}} } } } } }ggg}}}t|ƒ|  d¡dkræ|t|ƒ}|  |||| | | | | |||¡ \ }}}}} } } } } ||  d¡g7}||  d¡g7}|| g7}qh|  t  |¡t  |¡t  |¡¡\}}}|||fS)aîDecoder 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)r¡r,rªrr±r”r1r§rzrœr²rVr°Ústack)r<rZr³r´r r¨r0r¢r£r¤r¥r^r¦r_rUZ mel_outputsr¬r­Z mel_outputÚ gate_outputr«rrrrA­sd   ÷ õö ÿ 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``). rršr›ržrrrÚ _get_go_frameùs z_Decoder._get_go_frame)rZr´rcCs–| 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 krLt d¡tj|dd}tj|dd}tj|dd}| |||¡\}}}||||fS)a’Decoder 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``). rršrzZReached max decoder steps. The generated spectrogram might not cover the whole transcript.r[)rœr(r·r1r§rrZint32Úboolrvrr”r²rwr,r@r“rVr’rÚallrzÚwarningsÚwarnr±r°)r<rZr´Z batch_sizer(r r0r¢r£r¤r¥r^r¦r_rUÚmel_specgram_lengthsÚfinishedZ mel_specgramsr¬r­r…r«r¶rrrÚinfer sx  ÷ õö ÿz_Decoder.infer)rCrDrErFr$rRr¸r9rr¡rr§rªr°r²rAr·rÚjitÚexportr¾rGrrr=rrŠ¥s`ð3þ 1 þ "ó K þ LrŠcs eZdZdZdeeeeeeeeeeeeeeeeeeeeeeddœ‡fdd„ Zeeeee eeeefdœdd„Z e j j deeee eeefdœdd„ƒZ‡ZS)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``). FéPé”réééééÐçš™™™™™¹?Té€é éérjN)Ú mask_paddingrpÚn_symbolr‹Úsymbol_embedding_dimrJr€rrŒrrŽrrIr5rKrLrr‘rsrrrqr’rcs‚tƒ ¡||_||_||_t ||¡|_tjj   |jj ¡t |||ƒ|_ t|||| | | | | ||||||ƒ|_t||||ƒ|_dS)N)r8r9rÍrpr‹rZ EmbeddingÚ embeddingrrrrrÚencoderrŠÚdecoderroÚpostnet)r<rÍrprÎr‹rÏrJr€rrŒrrŽrrIr5rKrLrr‘rsrrrqr’r=rrr9†s0 òzTacotron2.__init__)ÚtokensÚ token_lengthsr«r¼rc CsÀ| |¡ dd¡}| ||¡}|j|||d\}}}| |¡} || } |jr´t|ƒ} |  |j|   d¡|   d¡¡} |   ddd¡} |  | d¡|   | d¡|  | dd…ddd…fd¡|| ||fS)a´Pass 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")r´rgNg@@) rÐr@rÑrÒrÓrÍr1ÚexpandrprœZpermuteZ masked_fill_) r<rÔrÕr«r¼Úembedded_inputsÚencoder_outputsr¬r­Zmel_specgram_postnetr0rrrrA¹s"! ÿ    zTacotron2.forward)rÔr'rc Cs¢|j\}}|dkr0t |g¡ |¡ |j|j¡}|dk ssF ÿø÷ .1W=HC