o i@sdZddlZddlmZmZmZmZddlZddlm m Z ddl m Z ddlmZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZdd lm Z ddl!m"Z"m#Z#ddl$m%Z%ddl&m'Z'm(Z(dd l)mZ*GdddeZ+dS)z(Fastspeech2 related modules for ESPnet2.N)DictOptionalSequenceTuple)check_argument_types)force_gatherable) initialize)AbsTTS)FastSpeech2Loss)VariancePredictor) StyleEncoder)Encoder)DurationPredictor)LengthRegulator)make_non_pad_mask make_pad_mask)Postnet)PositionalEncodingScaledPositionalEncodingcseZdZdZ                                                     dded ed!ed"ed#ed$ed%ed&ed'ed(ed)ed*ed+ed,ed-ed.ed/ed0ed1ed2ed3ed4ed5ed6ed7ed8ed9ed:ed;eded?ed@edAedBedCedDedEedFedGedHedIedJedKedLedMedNedOedPedQedRedSedTedUedVedWeedXeedYeedZed[ed\ed]ed^ed_e ed`edaedbedceddedeedfedgedheffdidj Z    ddke j dle j dme j dne j doe j dpe j dqe j dre j dse j dte j duee j dvee j dwee j dxedye e j eee j fe j ffdzd{Z         dd|e j d}e j d~ee j dee j dee j dee j dee j duee j dvee j dwee j dededye e j fddZ         ddke j dmee j doee j due j dvee j dwee j dqee j dsee j dededyeee j ffddZde j due j dye j fddZd}e j dye j fddZddedeedfefddZZS) FastSpeech2aFastSpeech2 module. This is a module of FastSpeech2 described in `FastSpeech 2: Fast and High-Quality End-to-End Text to Speech`_. Instead of quantized pitch and energy, we use token-averaged value introduced in `FastPitch: Parallel Text-to-speech with Pitch Prediction`_. .. _`FastSpeech 2: Fast and High-Quality End-to-End Text to Speech`: https://arxiv.org/abs/2006.04558 .. _`FastPitch: Parallel Text-to-speech with Pitch Prediction`: https://arxiv.org/abs/2006.06873 ?conv1dTF transformer皙?legacyrel_pos rel_selfattnswish Nadd  r-@r.r/r/xavier_uniform?idimodimadimaheadselayerseunitsdlayersdunitspostnet_layers postnet_chans postnet_filtspostnet_dropout_ratepositionwise_layer_typepositionwise_conv_kernel_sizeuse_scaled_pos_encuse_batch_normencoder_normalize_beforedecoder_normalize_beforeencoder_concat_afterdecoder_concat_afterreduction_factor encoder_type decoder_typetransformer_enc_dropout_rate'transformer_enc_positional_dropout_rate!transformer_enc_attn_dropout_ratetransformer_dec_dropout_rate'transformer_dec_positional_dropout_rate!transformer_dec_attn_dropout_rateconformer_rel_pos_typeconformer_pos_enc_layer_typeconformer_self_attn_layer_typeconformer_activation_typeuse_macaron_style_in_conformeruse_cnn_in_conformer zero_triuconformer_enc_kernel_sizeconformer_dec_kernel_sizeduration_predictor_layersduration_predictor_chansduration_predictor_kernel_sizeduration_predictor_dropout_rateenergy_predictor_layersenergy_predictor_chansenergy_predictor_kernel_sizeenergy_predictor_dropoutenergy_embed_kernel_sizeenergy_embed_dropout#stop_gradient_from_energy_predictorpitch_predictor_layerspitch_predictor_chanspitch_predictor_kernel_sizepitch_predictor_dropoutpitch_embed_kernel_sizepitch_embed_dropout"stop_gradient_from_pitch_predictorspkslangs spk_embed_dimspk_embed_integration_typeuse_gst gst_tokens gst_headsgst_conv_layersgst_conv_chans_listgst_conv_kernel_sizegst_conv_stridegst_gru_layers gst_gru_units init_typeinit_enc_alphainit_dec_alpha use_maskinguse_weighted_maskingcKMsptsJt||_||_|d|_||_||_||_|8|_ |1|_ ||_ |=|_ d|_ |j r2tnt}Kd||fvrm|dkrU|dkrId}td| dkrTd } td n|d krf|dks_J| d kseJntd |tjj|||j d }L|dkrt||||||L||||K||| |d|_nN|dkrtd.id|d|d|d|d|d|Ld|d|d|d|d|d| d|d|"d|d| d |!d!|#d"|%d#|$|_nt|d$|j rt||>||?|@|A|B|C|D|Ed% |_d&|_|9d&ur|9dkr|9|_tj|9||_d&|_|:d&ur|:dkr|:|_tj|:||_d&|_|;d&ur2|;dkr2|;|_|<|_ |jd&urS|j d'krHtj!|j||_"n tj!||j||_"t#||'|(|)|*d(|_$t%||2|3|4|5d(|_&tj'tjj(d||6|6dd)d*tj)|7|_*t%||+|,|-|.d(|_+tj'tjj(d||/|/dd)d*tj)|0|_,t-|_.|dkrtd||||d&||||K||| |d|_/nL|dkrtd.iddd|d|d|d|dd&d|d|d|d|d|d| d|d|"d|d| d |!d!|#d"|&|_/nt|d$tj!||||_0| dkrd&n t1||| | | || d+|_2|j3|F|G|Hd,t4|I|Jd-|_5d&S)/aaInitialize FastSpeech2 module. Args: idim (int): Dimension of the inputs. odim (int): Dimension of the outputs. elayers (int): Number of encoder layers. eunits (int): Number of encoder hidden units. dlayers (int): Number of decoder layers. dunits (int): Number of decoder hidden units. postnet_layers (int): Number of postnet layers. postnet_chans (int): Number of postnet channels. postnet_filts (int): Kernel size of postnet. postnet_dropout_rate (float): Dropout rate in postnet. use_scaled_pos_enc (bool): Whether to use trainable scaled pos encoding. use_batch_norm (bool): Whether to use batch normalization in encoder prenet. encoder_normalize_before (bool): Whether to apply layernorm layer before encoder block. decoder_normalize_before (bool): Whether to apply layernorm layer before decoder block. encoder_concat_after (bool): Whether to concatenate attention layer's input and output in encoder. decoder_concat_after (bool): Whether to concatenate attention layer's input and output in decoder. reduction_factor (int): Reduction factor. encoder_type (str): Encoder type ("transformer" or "conformer"). decoder_type (str): Decoder type ("transformer" or "conformer"). transformer_enc_dropout_rate (float): Dropout rate in encoder except attention and positional encoding. transformer_enc_positional_dropout_rate (float): Dropout rate after encoder positional encoding. transformer_enc_attn_dropout_rate (float): Dropout rate in encoder self-attention module. transformer_dec_dropout_rate (float): Dropout rate in decoder except attention & positional encoding. transformer_dec_positional_dropout_rate (float): Dropout rate after decoder positional encoding. transformer_dec_attn_dropout_rate (float): Dropout rate in decoder self-attention module. conformer_rel_pos_type (str): Relative pos encoding type in conformer. conformer_pos_enc_layer_type (str): Pos encoding layer type in conformer. conformer_self_attn_layer_type (str): Self-attention layer type in conformer conformer_activation_type (str): Activation function type in conformer. use_macaron_style_in_conformer: Whether to use macaron style FFN. use_cnn_in_conformer: Whether to use CNN in conformer. zero_triu: Whether to use zero triu in relative self-attention module. conformer_enc_kernel_size: Kernel size of encoder conformer. conformer_dec_kernel_size: Kernel size of decoder conformer. duration_predictor_layers (int): Number of duration predictor layers. duration_predictor_chans (int): Number of duration predictor channels. duration_predictor_kernel_size (int): Kernel size of duration predictor. duration_predictor_dropout_rate (float): Dropout rate in duration predictor. pitch_predictor_layers (int): Number of pitch predictor layers. pitch_predictor_chans (int): Number of pitch predictor channels. pitch_predictor_kernel_size (int): Kernel size of pitch predictor. pitch_predictor_dropout_rate (float): Dropout rate in pitch predictor. pitch_embed_kernel_size (float): Kernel size of pitch embedding. pitch_embed_dropout_rate (float): Dropout rate for pitch embedding. stop_gradient_from_pitch_predictor: Whether to stop gradient from pitch predictor to encoder. energy_predictor_layers (int): Number of energy predictor layers. energy_predictor_chans (int): Number of energy predictor channels. energy_predictor_kernel_size (int): Kernel size of energy predictor. energy_predictor_dropout_rate (float): Dropout rate in energy predictor. energy_embed_kernel_size (float): Kernel size of energy embedding. energy_embed_dropout_rate (float): Dropout rate for energy embedding. stop_gradient_from_energy_predictor: Whether to stop gradient from energy predictor to encoder. spks (Optional[int]): Number of speakers. If set to > 1, assume that the sids will be provided as the input and use sid embedding layer. langs (Optional[int]): Number of languages. If set to > 1, assume that the lids will be provided as the input and use sid embedding layer. spk_embed_dim (Optional[int]): Speaker embedding dimension. If set to > 0, assume that spembs will be provided as the input. spk_embed_integration_type: How to integrate speaker embedding. use_gst (str): Whether to use global style token. gst_tokens (int): The number of GST embeddings. gst_heads (int): The number of heads in GST multihead attention. gst_conv_layers (int): The number of conv layers in GST. gst_conv_chans_list: (Sequence[int]): List of the number of channels of conv layers in GST. gst_conv_kernel_size (int): Kernel size of conv layers in GST. gst_conv_stride (int): Stride size of conv layers in GST. gst_gru_layers (int): The number of GRU layers in GST. gst_gru_units (int): The number of GRU units in GST. init_type (str): How to initialize transformer parameters. init_enc_alpha (float): Initial value of alpha in scaled pos encoding of the encoder. init_dec_alpha (float): Initial value of alpha in scaled pos encoding of the decoder. use_masking (bool): Whether to apply masking for padded part in loss calculation. use_weighted_masking (bool): Whether to apply weighted masking in loss calculation. rr conformerr!r"legacy_rel_poszFallback to conformer_pos_enc_layer_type = 'legacy_rel_pos' due to the compatibility. If you want to use the new one, please use conformer_pos_enc_layer_type = 'latest'.r#legacy_rel_selfattnzFallback to conformer_self_attn_layer_type = 'legacy_rel_selfattn' due to the compatibility. If you want to use the new one, please use conformer_pos_enc_layer_type = 'latest'.latestzUnknown rel_pos_type: )num_embeddings embedding_dim padding_idxr)r2 attention_dimattention_heads linear_units num_blocks input_layer dropout_ratepositional_dropout_rateattention_dropout_rate pos_enc_classnormalize_before concat_afterr>r?r2rrrrrrrrrrr>r? macaron_stylepos_enc_layer_typeselfattention_layer_typeactivation_typeuse_cnn_modulecnn_module_kernelrUz is not supported.) r2ro gst_token_dimrp conv_layersconv_chans_listconv_kernel_size conv_stride gru_layers gru_unitsNr*)r2n_layersn_chans kernel_sizerr') in_channels out_channelsrpadding)r2r3rrn_filtsrAr)rwrxry)rzr{)6rsuper__init__r2r3eosrFrGrHrirbr@rnrrrloggingwarning ValueErrortorchnn EmbeddingTransformerEncoderencoderConformerEncoderr gstrjsid_embrklid_embrlrmLinear projectionrduration_predictorr pitch_predictor SequentialConv1dDropout pitch_embedenergy_predictor energy_embedrlength_regulatordecoderfeat_outrpostnet_reset_parametersr criterion)Mselfr2r3r4r5r6r7r8r9r:r;r<r=r>r?r@rArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrUrVrWrXrYrZr[r\r]r^r_r`rarbrcrdrerfrgrhrirjrkrlrmrnrorprqrrrsrtrurvrwrxryrzr{rencoder_input_layer __class__rW/home/ubuntu/.local/lib/python3.10/site-packages/espnet2/tts/fastspeech2/fastspeech2.pyr0s 4                            zFastSpeech2.__init__text text_lengthsfeats feats_lengths durationsdurations_lengthspitch pitch_lengthsenergyenergy_lengthsspembssidslidsjoint_trainingreturnc& sP|ddd|f}|ddd|f}|ddd|f}|ddd|f}| ddd| f} |d}t|ddgdj}t|D] \}}j|||f<qP|d}|||| f\}}}}|}j|||||||| | | dd \}}}}}jdkr| fdd|D}t|}|ddd|f}j durd}j |||||||||||d \}} }!}"|| |!|"}#t | | |! |" d }$jd krjr|$jjjd jj d jd krjr|$jjjd jj d|s|$j|# dt|#|$|f|#j\}#}$}%|#|$|%fS|#|$|dur%|fS|fS)a@Calculate forward propagation. Args: text (LongTensor): Batch of padded token ids (B, T_text). text_lengths (LongTensor): Batch of lengths of each input (B,). feats (Tensor): Batch of padded target features (B, T_feats, odim). feats_lengths (LongTensor): Batch of the lengths of each target (B,). durations (LongTensor): Batch of padded durations (B, T_text + 1). durations_lengths (LongTensor): Batch of duration lengths (B, T_text + 1). pitch (Tensor): Batch of padded token-averaged pitch (B, T_text + 1, 1). pitch_lengths (LongTensor): Batch of pitch lengths (B, T_text + 1). energy (Tensor): Batch of padded token-averaged energy (B, T_text + 1, 1). energy_lengths (LongTensor): Batch of energy lengths (B, T_text + 1). spembs (Optional[Tensor]): Batch of speaker embeddings (B, spk_embed_dim). sids (Optional[Tensor]): Batch of speaker IDs (B, 1). lids (Optional[Tensor]): Batch of language IDs (B, 1). joint_training (bool): Whether to perform joint training with vocoder. Returns: Tensor: Loss scalar value. Dict: Statistics to be monitored. Tensor: Weight value if not joint training else model outputs. NrrconstantF)rrr is_inferencecsg|] }||jqSrrF.0olenrrr -sz'FastSpeech2.forward..) after_outs before_outsd_outsp_outse_outsysdspsesilensolens)l1_loss duration_loss pitch_loss energy_lossr) encoder_alpha) decoder_alpha)loss)maxsizeFpadr enumerater_forwardrFnewrrdictitemrGr@updaterembedalphadatarHrrdevice)&rrrrrrrrrrrrrrr batch_sizexsilrrrrrrrrrrrmax_olenrrrrrstatsweightrrrforwards)      zFastSpeech2.forwardrrrrrrrrrc sz|} || \}}jr|}||d}jdur0| d}||d}jdurD | d}||d}j durO ||}t | |j}jrf||d}n ||d}jr~||d}n ||d}| rj||}|dddd}|dddd}|||}||| }n,||}|dddd}|dddd}|||}||}|dur| sjdkr|fdd|D}n|}|}nd}||\}}||ddj}j dur'|}n| |dddd}|||||fS)Nrrr'csg|]}|jqSrrrrrrrsz(FastSpeech2._forward..r)! _source_maskrrnr unsqueezerjrviewrkrrl_integrate_with_spk_embedrtorrirdetachrbrr inferencer transposerrrFrrrrr3r)rrrrrrrrrrrrrx_maskshs_ style_embssid_embslid_embsd_masksrrrp_embse_embsolens_inh_maskszsrrrrrr_sb               zFastSpeech2._forwarduse_teacher_forcingc  Cs||} } ||||f\} }}}t| ddgd|j} tj| jdgtj| jd}| dd}}| dur9| d}| durB| d}| ri|d|d|d}}}|j |||||||||d \}}}}}n|j ||||||d| d\}}}}}t |d|d|d|dd S) aGenerate the sequence of features given the sequences of characters. Args: text (LongTensor): Input sequence of characters (T_text,). feats (Optional[Tensor): Feature sequence to extract style (N, idim). durations (Optional[Tensor): Groundtruth of duration (T_text + 1,). spembs (Optional[Tensor): Speaker embedding vector (spk_embed_dim,). sids (Optional[Tensor]): Speaker ID (1,). lids (Optional[Tensor]): Language ID (1,). pitch (Optional[Tensor]): Groundtruth of token-avg pitch (T_text + 1, 1). energy (Optional[Tensor]): Groundtruth of token-avg energy (T_text + 1, 1). alpha (float): Alpha to control the speed. use_teacher_forcing (bool): Whether to use teacher forcing. If true, groundtruth of duration, pitch and energy will be used. Returns: Dict[str, Tensor]: Output dict including the following items: * feat_gen (Tensor): Output sequence of features (T_feats, odim). * duration (Tensor): Duration sequence (T_text + 1,). * pitch (Tensor): Pitch sequence (T_text + 1,). * energy (Tensor): Energy sequence (T_text + 1,). rrr)dtyperN)rrrrrrT)rrrrr)feat_gendurationrr) rrrrtensorshapelongrrrr)rrrrrrrrrrrxyspembdperrrrrrroutsrrrrrrr sL $  "  zFastSpeech2.inferencercCsz|jdkr|t|}||d}|S|jdkr9t|dd|dd}|tj||gdd}|St d)aEIntegrate speaker embedding with hidden states. Args: hs (Tensor): Batch of hidden state sequences (B, T_text, adim). spembs (Tensor): Batch of speaker embeddings (B, spk_embed_dim). Returns: Tensor: Batch of integrated hidden state sequences (B, T_text, adim). r*rconcatr)dimzsupport only add or concat.) rmrr normalizerexpandrrcatNotImplementedError)rrrrrrr  s  z%FastSpeech2._integrate_with_spk_embedcCs"t|t|j}|dS)aMake masks for self-attention. Args: ilens (LongTensor): Batch of lengths (B,). Returns: Tensor: Mask tensor for self-attention. dtype=torch.uint8 in PyTorch 1.2- dtype=torch.bool in PyTorch 1.2+ (including 1.2) Examples: >>> ilens = [5, 3] >>> self._source_mask(ilens) tensor([[[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]]], dtype=torch.uint8) )rr next parametersrr)rrrrrrr#s zFastSpeech2._source_maskcCsj|dkr t|||jdkr|jrt||jjdj_|j dkr1|jr3t||j jdj_dSdSdS)Npytorchrr) rrGr@rrrrrrrHr)rrwrxryrrrr8s zFastSpeech2._reset_parameters)HrrrrrrrrrrrrTTTTFFrrrr r r r r r r!r"r#r$TTFr%r&r'rr(r r'rr(rr)rFr'rr(rr)rFNNNr*Fr+rrr,r(r'rr/r0r1r1FF)NNNF) NNNNNNNNFr1) NNNNNNNr1F)__name__ __module__ __qualname____doc__intfloatstrboolrrrrTensorrrrrr r rr __classcell__rrrrr!s     !"#$%&'()+,-.012345689:;<=>@ABCDEFGHIJKLNOPQRA            Z    S r),r6rtypingrrrrrtorch.nn.functionalr functionalr typeguardr espnet2.torch_utils.device_funcsrespnet2.torch_utils.initializerespnet2.tts.abs_ttsr espnet2.tts.fastspeech2.lossr *espnet2.tts.fastspeech2.variance_predictorr espnet2.tts.gst.style_encoderr -espnet.nets.pytorch_backend.conformer.encoderr r9espnet.nets.pytorch_backend.fastspeech.duration_predictorr7espnet.nets.pytorch_backend.fastspeech.length_regulatorr&espnet.nets.pytorch_backend.nets_utilsrr-espnet.nets.pytorch_backend.tacotron2.decoderr1espnet.nets.pytorch_backend.transformer.embeddingrr/espnet.nets.pytorch_backend.transformer.encoderrrrrrrs(