o ZhB@sdZddlZddlmZmZmZmZddlZddl Z ddl Z ddl m Z ddl m Z mZmZddlmZddlmZdd lmZdd lmZdd lmZmZdd lmZmZmZmZm Z dd l!m"Z"ddl#m$Z$m%Z%ddl&m'Z'm(Z(e%)e*Z+dZ,de j-de.de.fddZ/ dde j-de.dee j-fddZ0  ddee.e.fde1de.dee j2de.dej3f d d!Z4Gd"d#d#e j5Z6Gd$d%d%e j5Z7Gd&d'd'e j5Z8Gd(d)d)e j5Z9Gd*d+d+e j5Z:Gd,d-d-e j5Z;Gd.d/d/e j j5ZGd4d5d5e j5Z?Gd6d7d7e j5Z@Gd8d9d9e j5ZAGd:d;d;e j5ZBGdd?d?e j5ZDGd@dAdAe j5ZEGdBdCdCe j5ZFGdDdEdEe j5ZGGdFdGdGe j5ZHGdHdIdIe j5ZIGdJdKdKe j5ZJe$GdLdMdMe"ZKGdNdOdOeKZLGdPdQdQeKZMGdRdSdSeKZNGdTdUdUeKZOGdVdWdWeKZPGdXdYdYeKZQGdZd[d[eKZRGd\d]d]eKZSGd^d_d_e j5ZTGd`dadae j5ZUe$dbdcGdddedeeKZVe$dfdcGdgdhdheKeZW   i j k  l lddmeKde jXdnee jXdee j2doe1dpe1dqe1dree j5dseYdteYdee jXee jXe jXfffdudvZZe$dwdcGdxdydyeKZ[e$dzdcGd{d|d|eKZ\Gd}d~d~e j5Z]e$ddcGddde"Z^gdZ_dS)zPyTorch SpeechT5 model.N)ListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossL1Loss)ACT2FN)GenerationMixin)is_deepspeed_zero3_enabled)is_fsdp_managed_module)_prepare_4d_attention_mask!_prepare_4d_causal_attention_mask)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsSeq2SeqLMOutputSeq2SeqModelOutputSeq2SeqSpectrogramOutput)PreTrainedModel)auto_docstringlogging)SpeechT5ConfigSpeechT5HifiGanConfig input_ids pad_token_iddecoder_start_token_idcCsh||j}|ddddf|ddddf<||dddf<|dur*td||dk||S)z1 Shift input ids one token to the right. Nrrz1self.model.config.pad_token_id has to be defined.i) new_zerosshapeclone ValueError masked_fill_)rrrZshifted_input_idsr%]/var/www/auris/lib/python3.10/site-packages/transformers/models/speecht5/modeling_speecht5.pyshift_tokens_right4s (r' input_valuesreduction_factorattention_maskcCs|dkr"|dd|dd|f}|dur"|dd|dd|f}||j}|ddddf|ddddf<||dkd||fS)zw Shift input spectrograms one timestep to the right. Also applies the reduction factor to the sequence length. rNrY)r r!r"r$)r(r)r*Zshifted_input_valuesr%r%r&shift_spectrograms_rightDs (r-r! mask_prob mask_length min_masksreturncs|\}dkr tdkrtdddtjdfdd}|dur:|d n fd d t|D}tj |ft d }g} |} | d krZ|S|D];} || } tjj t | d| dd} t | d kr}d}n| d }t| tj| | tjd |g} | | q\t| } t| dddddf|| f} | || } t ddddf}t||| f|| }| |} | dkrd| | dk<t|| dd |S)af Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for ASR](https://arxiv.org/abs/1904.08779). Note that this method is not optimized to run on TPU and should be run on CPU as part of the preprocessing during training. Args: shape: The shape for which to compute masks. This should be of a tuple of size 2 where the first element is the batch size and the second element is the length of the axis to span. mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of independently generated mask spans of length `mask_length` is computed by `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the actual percentage will be smaller. mask_length: size of the mask min_masks: minimum number of masked spans attention_mask: A (right-padded) attention mask which independently shortens the feature axis of each batch dimension. rz&`mask_length` has to be bigger than 0.zO`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: z and `sequence_length`: `csXt|}t|}|kr}|d|kr*t|dd}|S)z;Given input length, compute how many spans should be maskedrr)intmax) input_lengthnum_masked_spanepsilonr/r.r0sequence_lengthr%r&compute_num_masked_spans  z6_compute_mask_indices..compute_num_masked_spanNrcsg|]}qSr%r%.0_)r9r%r& sz)_compute_mask_indices..dtyperF)replace)r#nprandomranditemdetachsumtolistrangezerosboolchoicearangelenZ concatenateonesZint32appendarrayZ broadcast_toreshaper4Zput_along_axis)r!r.r/r*r0 batch_sizer: input_lengthsZ spec_aug_maskZspec_aug_mask_idxsZmax_num_masked_spanr5r6Zspec_aug_mask_idxZdummy_mask_idxoffsetsr%r7r&_compute_mask_indicesZs\    rVc&eZdZdfdd ZddZZS)SpeechT5NoLayerNormConvLayerrcsjt|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ dS)Nrr kernel_sizestridebias)super__init__conv_dim in_conv_dim out_conv_dimrConv1d conv_kernel conv_stride conv_biasconvr feat_extract_activation activationselfconfiglayer_id __class__r%r&r^s  z%SpeechT5NoLayerNormConvLayer.__init__cCs||}||}|SN)rfrhrj hidden_statesr%r%r&forward  z$SpeechT5NoLayerNormConvLayer.forwardr__name__ __module__ __qualname__r^rr __classcell__r%r%rmr&rXsrXcrW)SpeechT5LayerNormConvLayerrcs|t|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrYT)Zelementwise_affine)r]r^r_r`rarrbrcrdrerf LayerNorm layer_normr rgrhrirmr%r&r^s  z#SpeechT5LayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)Nr)rf transposer|rhrpr%r%r&rrs     z"SpeechT5LayerNormConvLayer.forwardrtrur%r%rmr&rzsrzcrW)SpeechT5GroupNormConvLayerrcst|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrYT)Z num_groupsZ num_channelsZaffine)r]r^r_r`rarrbrcrdrerfr rgrh GroupNormr|rirmr%r&r^s   z#SpeechT5GroupNormConvLayer.__init__cCs"||}||}||}|Sro)rfr|rhrpr%r%r&rrs   z"SpeechT5GroupNormConvLayer.forwardrtrur%r%rmr&rsrc seZdZdZddededeeffdd Zddededeefd d Zeddededeefd d Z e dde j defddZ dde j dedeefddZZS)%SpeechT5SinusoidalPositionalEmbeddingzDThis module produces sinusoidal positional embeddings of any length.N num_positions embedding_dim padding_idxcs4td|_||_||_|||j||dSN)r]r^offsetrr make_weights)rjrrrrmr%r&r^ s z.SpeechT5SinusoidalPositionalEmbedding.__init__num_embeddingscCsB||||}t|dr|j|jj|jjd}|jd|dddS)Nweightsr@deviceF persistent) get_embeddinghasattrtorr@rregister_buffer)rjrrrZ emb_weightsr%r%r&r's z2SpeechT5SinusoidalPositionalEmbedding.make_weightscCs|d}td|d}ttj|tjd| }tj|tjdd|d}tjt |t |gdd |d}|ddkrUtj|t |dgdd}|durad||ddf<| tS) z Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of "Attention Is All You Need". ri'rr?rdimrN)mathlogtorchexprMint64float unsqueezecatsincosviewrJrZget_default_dtype)rrrZhalf_dimembr%r%r&r/s $& z3SpeechT5SinusoidalPositionalEmbedding.get_embeddingrrpast_key_values_lengthcCs||\}}|||j||j}|jd|}||jdkr-|||j|j|j|j d| d ||d S)Nrrr) size"create_position_ids_from_input_idsrrrrrrr index_selectrrF)rjrrbszseq_lenZ position_idsZmax_posr%r%r&rrAs "z-SpeechT5SinusoidalPositionalEmbedding.forwardcCs6||}tj|dd|||}||S)a Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols are ignored. This is modified from fairseq's `utils.make_positions`. Args: x: torch.Tensor x: Returns: torch.Tensor rr)ner3rcumsumZtype_aslong)rjrrrmaskZincremental_indicesr%r%r&rPs  zHSpeechT5SinusoidalPositionalEmbedding.create_position_ids_from_input_idsrort)rvrwrx__doc__r3rr^r staticmethodrrno_gradTensorrrrryr%r%rmr&rs  rc$eZdZfddZddZZS)SpeechT5PositionalConvEmbeddingcs$ttj|j|j|j|jd|jd|_tjj }t tjj dr'tjj j }t r{ddl }|jj|jjdd||jddd|_Wdn1sLwYt |jdrd|jj jj}|jj jj}n|jj}|jj}|j|||j||n ||jddd|_t|j|_t|j|_dS) Nr)rZpaddinggroups weight_normr)Z modifier_rankweight)namerparametrizations)r]r^rrb hidden_sizenum_conv_pos_embeddingsZnum_conv_pos_embedding_groupsrfutilsrrrr deepspeedzeroZGatheredParametersrZ original0Z original1weight_gweight_vZregister_external_parameterSpeechT5SamePadLayerrr rgrh)rjrkrrrrrmr%r&r^cs4     z(SpeechT5PositionalConvEmbedding.__init__cCs:|dd}||}||}||}|dd}|SNrr)r~rfrrhrpr%r%r&rrs     z'SpeechT5PositionalConvEmbedding.forwardrur%r%rmr&rbs !rcs*eZdZdZdfdd ZddZZS) SpeechT5ScaledPositionalEncodinguS Scaled positional encoding, see §3.2 in https://arxiv.org/abs/1809.08895 cst||}td|d}ttjd|dtjdtd| }t |||dddddf<t |||dddddf<|d}t |j d|ddtj|d |_||_tjtd |_dS) Nrrrr?g@peFrp?)rrJrMrrrrrrrrr]r^rrDropoutdropoutr Parametertensoralpha)rjrrmax_lenrpositionZdiv_termrmr%r&r^s .$$  z)SpeechT5ScaledPositionalEncoding.__init__cCs4||j|jddd|df}||}|S)Nr)rrrr)rjrr%r%r&rrs& z(SpeechT5ScaledPositionalEncoding.forward)r)rvrwrxrr^rrryr%r%rmr&rs rcrW)"SpeechT5RelativePositionalEncodingcs.t||_||_tjd|||_dSr)r]r^r max_lengthrr Embeddingpe_k)rjrrrmr%r&r^s z+SpeechT5RelativePositionalEncoding.__init__cCs|jd}td|j|jtjd}|dddf|dddf}|j |||j k<|jd|||jk<||j}||S)Nrrrr@)r!rrMrrrrr)rjrqrZpos_seqr%r%r&rrs    z*SpeechT5RelativePositionalEncoding.forward)rrur%r%rmr&rsrcr)rcs*t|ddkrd|_dSd|_dS)Nrrr)r]r^num_pad_remove)rjrrmr%r&r^s  zSpeechT5SamePadLayer.__init__cCs,|jdkr|ddddd|j f}|SNr)rrpr%r%r&rrs zSpeechT5SamePadLayer.forwardrur%r%rmr&rs rcs0eZdZdZfddZddZddZZS)SpeechT5FeatureEncoderz.Construct the features from raw audio waveformcstjdkr tddgfddtjdD}njdkr2fddtjD}n td jd t||_ d |_ d |_ dS) Ngrouprrlcsg|] }t|ddqS)rr)rXr<irkr%r&r>sz3SpeechT5FeatureEncoder.__init__..rlayercsg|]}t|dqS)r)rzrrr%r&r>s z`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) r]r^Zfeat_extract_normrrIZnum_feat_extract_layersr#r ModuleList conv_layersgradient_checkpointing_requires_grad)rjrkrrmrr&r^s        zSpeechT5FeatureEncoder.__init__cCs|D]}d|_qd|_dS)NF) parameters requires_gradr)rjparamr%r%r&_freeze_parameterss  z)SpeechT5FeatureEncoder._freeze_parameterscCs\|dddf}|jr|jrd|_|jD]}|jr'|jr'|jr'||j|}q||}q|SNT)rtrainingrrr_gradient_checkpointing_func__call__)rjr(rqZ conv_layerr%r%r&rrs   zSpeechT5FeatureEncoder.forward)rvrwrxrr^rrrryr%r%rmr&rs  rcr)SpeechT5FeatureProjectioncsJttj|jd|jd|_t|jd|j|_ t |j |_ dS)NrZeps) r]r^rr{r_layer_norm_epsr|Linearr projectionrZfeat_proj_dropoutrrjrkrmr%r&r^s z"SpeechT5FeatureProjection.__init__cCs&||}||}||}||fSro)r|rr)rjrqZnorm_hidden_statesr%r%r&rrs   z!SpeechT5FeatureProjection.forwardrur%r%rmr&rs rcseZdZfddZddZ  ddejdeejdeej fd d Z d e dejfd d Z de eje ffddZ  ddej deej deejfddZZS)SpeechT5SpeechEncoderPrenetcs|t||_t||_t||_|jdks|jdkr(t t |j |_t||_t|j|jd|j |j|_dS)Nr,r)r]r^rkrfeature_encoderrfeature_projectionmask_time_probmask_feature_probrrrrruniform_masked_spec_embedrpos_conv_embedrmax_speech_positionsrpos_sinusoidal_embedrrmr%r&r^s     z$SpeechT5SpeechEncoderPrenet.__init__cCs|jdSro)rrrjr%r%r&freeze_feature_encodersz2SpeechT5SpeechEncoderPrenet.freeze_feature_encoderNr(r*mask_time_indicesc Cs||}|dd}|dur||jd|}||\}}|j|||d}||}||}|dur<|d}nt j |jddt j|j d}| |}||}||fS)Nrr)rr*r) rr~"_get_feature_vector_attention_maskr!r_mask_hidden_statesrrrrrJrr) rjr(r*rZextract_featuresrqZpositional_conv_embedding padding_maskZ positional_sinusoidal_embeddingsr%r%r&rrs&    z#SpeechT5SpeechEncoderPrenet.forwardfeature_vector_lengthcCs|jdddddf}||tj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dgd dg }|S)Nrrrrrr) r _get_feat_extract_output_lengthsrrrr!rJr@rrMfliprK)rjrr*Znon_padded_lengthsoutput_lengthsrSr%r%r&r=s "z>SpeechT5SpeechEncoderPrenet._get_feature_vector_attention_maskrTcCs4dd}t|jj|jjD] \}}||||}q |S)zH Computes the output length of the convolutional layers cSstj|||dddS)Nfloor)Z rounding_moder)rdiv)r5rZr[r%r%r&_conv_out_lengthRszVSpeechT5SpeechEncoderPrenet._get_feat_extract_output_lengths.._conv_out_length)ziprkrcrd)rjrTrrZr[r%r%r&rMszsz8SpeechT5SpeechDecoderPrenet.__init__..)r]r^rkrrrIZspeech_decoder_prenet_layerslayersrrr final_layerrpositional_dropoutrencode_positionsZspeaker_embedding_dimspeaker_embeds_layerrrmrr&r^s   z$SpeechT5SpeechDecoderPrenet.__init__cCsJtj|d|d}|d|ddd}t|dk|ddd|S)Nrrr)rZ bernoullirrepeatrwhere)rj inputs_embedsrrZ all_masksr%r%r&_consistent_dropoutsz/SpeechT5SpeechDecoderPrenet._consistent_dropoutNr(speaker_embeddingscCs|}|jD]}tj||}|||jj}q||}||}|durKtj |}| d d| dd}t j||gdd}tj||}|S)Nrrr)rr functionalZrelurrkZspeech_decoder_prenet_dropoutrr normalizerrrrrr)rjr(rrrr%r%r&rrs    z#SpeechT5SpeechDecoderPrenet.forwardro) rvrwrxr^rrrrrrryr%r%rmr&r s r crW)SpeechT5BatchNormConvLayerrcst|dkr |j}n|j}||jdkr|j}n|j}tj|||jd|jdddd|_t ||_ ||jdkrCt |_ nd|_ t |j|_dS)NrrrF)rZr[rr\)r]r^r Zspeech_decoder_postnet_unitsspeech_decoder_postnet_layersrrbZspeech_decoder_postnet_kernelrfZ BatchNorm1d batch_normZTanhrhrZspeech_decoder_postnet_dropoutr)rjrkrlr`rarmr%r&r^s(    z#SpeechT5BatchNormConvLayer.__init__cCs6||}||}|jdur||}||}|Sro)rfrrhrrpr%r%r&rrs     z"SpeechT5BatchNormConvLayer.forwardrtrur%r%rmr&rsrcs<eZdZfddZdejfddZdejfddZZS)SpeechT5SpeechDecoderPostnetcs^t|_tjjj|_tjj|_ t fddt j D|_ dS)Ncsg|]}t|qSr%)rrrr%r&r>sz9SpeechT5SpeechDecoderPostnet.__init__..)r]r^rkrrrr r)feat_outprob_outrrIrrrrmrr&r^s  z%SpeechT5SpeechDecoderPostnet.__init__rqcCsJ|||dd|jj}||}|||dd}|||fS)Nrr)rrrrkr postnetr )rjrqoutputs_before_postnetoutputs_after_postnetlogitsr%r%r&rrs  z$SpeechT5SpeechDecoderPostnet.forwardcCs0|dd}|jD]}||}q ||ddSr)r~r)rjrqZ layer_outputrr%r%r&r!s   z$SpeechT5SpeechDecoderPostnet.postnet) rvrwrxr^rrrrr!ryr%r%rmr&rs  rcs<eZdZfddZddZddZdejfdd ZZ S) SpeechT5TextEncoderPrenetcs>t||_t|j|j|j|_t |j |j|j |_ dSro) r]r^rkrr vocab_sizerr embed_tokensrrmax_text_positionsrrrmr%r&r^s  z"SpeechT5TextEncoderPrenet.__init__cC|jSror'rr%r%r&get_input_embeddings z.SpeechT5TextEncoderPrenet.get_input_embeddingscC ||_dSror*rjvaluer%r%r&set_input_embeddings z.SpeechT5TextEncoderPrenet.set_input_embeddingsrcCs||}||}|Sro)r'r)rjrrr%r%r&rrrsz!SpeechT5TextEncoderPrenet.forward) rvrwrxr^r+r0rrrrryr%r%rmr&r%s   r%c sZeZdZfddZddZddZ  d dejd eej d ee ej fd d Z Z S)SpeechT5TextDecoderPrenetcsnt||_t|j|_|jrt |j nd|_ t |j |j |j|_t|j|jd|j |j|_dS)Nrr)r]r^rkrrrrZscale_embeddingrsqrtr embed_scalerr&rr'rr(embed_positionsrrmr%r&r^s  z"SpeechT5TextDecoderPrenet.__init__cCr)ror*rr%r%r&r+'r,z.SpeechT5TextDecoderPrenet.get_input_embeddingscCr-ror*r.r%r%r&r0*r1z.SpeechT5TextDecoderPrenet.set_input_embeddingsNrr*past_key_valuescCs~|dur|}|d|d}ntd|dur"|ddjdnd}|||}|||j}||7}||}||fS)Nrz'You have to specify `decoder_input_ids`rr)rrr#r!r5r'r4r)rjrr*r6 input_shaperZ positionsrr%r%r&rr-s  z!SpeechT5TextDecoderPrenet.forwardr )rvrwrxr^r+r0rrrr rr rrryr%r%rmr&r2s  r2cs<eZdZfddZdejfddZddZdd ZZ S) SpeechT5TextDecoderPostnetcs*t||_tj|j|jdd|_dS)NFr\)r]r^rkrrrr&lm_headrrmr%r&r^Ds z#SpeechT5TextDecoderPostnet.__init__rqcCs ||Sror:rpr%r%r&rrIr1z"SpeechT5TextDecoderPostnet.forwardcCr)ror;rr%r%r&get_output_embeddingsLr,z0SpeechT5TextDecoderPostnet.get_output_embeddingscCr-ror;rjZnew_embeddingsr%r%r&set_output_embeddingsOr1z0SpeechT5TextDecoderPostnet.set_output_embeddings) rvrwrxr^rrrrr<r>ryr%r%rmr&r8Cs  r8cseZdZdZ   ddedededed ef fd d Zd ej d edefddZ      ddej de ej de e ej de ej de ej de ej dede ej e ej e e ej ffddZ ZS)SpeechT5Attentionz Multi-headed attention from 'Attention Is All You Need' paper with relative position bias (see https://aclanthology.org/N18-2074.pdf) r,FT embed_dim num_headsr is_decoderr\cst||_||_||_|||_|j||jkr'td|jd|d|jd|_||_t j |||d|_ t j |||d|_ t j |||d|_ t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).gr9)r]r^r@rArhead_dimr#scalingrBrrk_projv_projq_projout_proj)rjr@rArrBr\rmr%r&r^Ys"    zSpeechT5Attention.__init__rrrcCs ||||j|jddSr)rrArCr~ contiguous)rjrrrr%r%r&_shapets zSpeechT5Attention._shapeNrqkey_value_statespast_key_valuer*layer_head_mask position_biasoutput_attentionsr1cCs|du}|\} } } |||j} |r"|dur"|d} |d}nZ|r9|||d| } |||d| }nC|durh|||d| } |||d| }tj|d| gdd} tj|d|gdd}n|||d| } |||d| }|jr| |f}| |j d|j f}|| | | j |} | j |} |j |}| d}t | | dd}|| |j | |fkrtd| |j | |fd||dur|  | |j d|j  dd}t|| d d}| dd | |j |d|d}||7}|dur:|| d| |fkr%td | d| |fd|| | |j | ||}| | |j | |}tjj|dd}|dury||j fkr^td |j fd|| dddd| | |j | |}| | |j | |}|r| | |j | |}| | |j | |}nd}tjj||j|jd }t ||}|| |j | |j fkrtd | |j | |j fd|| | |j | |j }| dd}|| | |j}||}|||fS)z#Input shape: Batch x Time x ChannelNrrrrrz$Attention weights should be of size z , but is r}z!Attention mask should be of size z/Head mask for a single layer should be of size )rrz `attn_output` should be of size )rrGrDrJrErFrrrBrArCrZbmmr~r#rImatmulrrZsoftmaxrrrRr@rH)rjrqrKrLr*rMrNrOZis_cross_attentionrtgt_lenr=Z query_statesZ key_statesZ value_statesZ proj_shapeZsrc_len attn_weightsZ reshape_qZ rel_pos_biasZattn_weights_reshapedZ attn_probsZ attn_outputr%r%r&rrws      "    "    zSpeechT5Attention.forward)r,FT)NNNNNF)rvrwrxrr3rrKr^rrrJrrrrryr%r%rmr&r?SsP   r?cr)SpeechT5FeedForwardcsltt|j|_t|j||_t |j t r!t |j |_ n|j |_ t||j|_t|j|_dSro)r]r^rrZactivation_dropoutintermediate_dropoutrrintermediate_dense isinstanceZ hidden_actstrr intermediate_act_fn output_densehidden_dropoutoutput_dropout)rjrkZintermediate_sizermr%r&r^s  zSpeechT5FeedForward.__init__cCs6||}||}||}||}||}|Sro)rUrXrTrYr[rpr%r%r&rrs     zSpeechT5FeedForward.forwardrur%r%rmr&rSs  rSc s^eZdZdeffdd Z    d dejdeejdeejd eejd ef d d Z Z S)SpeechT5EncoderLayerrkcsjtt|j|j|jdd|_t|j |_ tj |j|j d|_ t||j|_tj |j|j d|_dS)NFr@rArrBr)r]r^r?rencoder_attention_headsattention_dropout attentionrrrZrr{rr|rSZencoder_ffn_dim feed_forwardfinal_layer_normrrmr%r&r^ s zSpeechT5EncoderLayer.__init__NFrqr*rMrNrOc Csj|}|j|||||d\}}}||}||}||}|||}||}|f} |r3| |f7} | S)as Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(config.encoder_attention_heads,)`. position_bias (`torch.FloatTensor`): relative position embeddings of size `(seq_len, seq_len, hidden_size // encoder_attention_heads)` output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. )rqr*rMrNrO)r`rr|rarb) rjrqr*rMrNrOresidualrRr=outputsr%r%r&rrs"    zSpeechT5EncoderLayer.forward)NNNF) rvrwrxrr^rrrrKrrryr%r%rmr&r\ s"r\cseZdZdeffdd Z        ddejdeejd eejd eejd eejd eejd eeejdee dee fddZ Z S)SpeechT5DecoderLayerrkcstt|j|j|jdd|_t|j |_ tj |j|j d|_ t|j|j|jdd|_tj |j|j d|_t||j|_tj |j|j d|_dS)NTr]r)rrB)r]r^r?rZdecoder_attention_headsr_ self_attnrrrZrr{rself_attn_layer_norm encoder_attnencoder_attn_layer_normrSZdecoder_ffn_dimrarbrrmr%r&r^Is$ zSpeechT5DecoderLayer.__init__NFTrqr*encoder_hidden_statesencoder_attention_maskrMcross_attn_layer_head_maskrLrO use_cachec Cs|} |dur |ddnd} |j|| |||d\}} } ||}| |}||}d}d}|dur_|} |dur=|ddnd}|j||||||d\}}}||}| |}||}| |} |||}||}|f}|rv|| |f7}| r}|| f7}|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(batch, seq_len, hidden_size)` encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of size `(decoder_attention_heads,)`. past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. Nr)rqrLr*rMrOr})rqrKr*rMrLrO)rfrrgrhrirarb)rjrqr*rjrkrMrlrLrOrmrcZself_attn_past_key_valueZself_attn_weightsZpresent_key_valueZcross_attn_present_key_valueZcross_attn_weightsZcross_attn_past_key_valuerdr%r%r&rr_sJ       zSpeechT5DecoderLayer.forward)NNNNNNFT) rvrwrxrr^rrrrrKrrryr%r%rmr&reHs:   rec@s$eZdZeZdZdZdZddZdS)SpeechT5PreTrainedModelspeecht5r(Tc Cst|tr*tjj|jjddtd|jj d|jj dtj |jj ddSt|t rRtd|jj}tjj|jj| |dtjj|jj | |ddSt|tjrr|jjjd|jjd|j durp|j jdSdSt|tjtjfr|j j|jjddSt|tjrtj|j|j durt|j|j |j d}tjj|j | |ddSdSt|tjr|jjjd|jjd|jdur|jj|jdSdSdS) zInitialize the weightsrrrmeanZstd)abr,Nr)rVrrinitnormal_rfrrr3rZZ in_channelsZ constant_r\rrZ in_featuresrrdatarkinitializer_rangezero_r{rZfill_rbZkaiming_normal_rrr)rjmodulekr%r%r& _init_weightss>          z%SpeechT5PreTrainedModel._init_weightsN) rvrwrxr config_classZbase_model_prefixmain_input_nameZsupports_gradient_checkpointingr{r%r%r%r&rns  rnczeZdZdZdeffdd Z     ddejdeej deej d ee d ee d ee d e e e ffd dZZS)SpeechT5Encoderzu Transformer encoder consisting of *config.encoder_layers* layers. Each layer is a [`SpeechT5EncoderLayer`]. rkcs~ttjjjd|_tj|_ j |_ t fddt jD|_tjjj|_d|_|dS)Nrcg|]}tqSr%)r\r;rr%r&r>z,SpeechT5Encoder.__init__..F)r]r^rr{rrr|rrZrZencoder_layerdrop layerdroprrIZencoder_layersrrr^Zencoder_max_relative_positionr5r post_initrrmrr&r^s   zSpeechT5Encoder.__init__Nrqr* head_maskrOoutput_hidden_states return_dictr1c Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur(t||j}||}||}||}t p=t |}|rBdnd} |rHdnd} |durk| dt |j krktdt |j d| ddt|j D]Z\} } |r{| |f} d} |jrtg}||jk} | r|r|jr|jr|| j|||dur|| nd||}n| ||||dur|| nd|d}|d}| rd }|r| |d f} qp|r| |f} |std d || | fDSt|| | d S)a Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, feature_size)`): Features extracted from the speech or text input by the encoder prenet. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Nr%rz&The head_mask should be specified for  layers, but it is for .F)r*rNrMrOr rcs|] }|dur|VqdSror%r<vr%r%r& Zsz*SpeechT5Encoder.forward..last_hidden_staterq attentions)rkrOruse_return_dictrr@r|rr5r rrrNrr# enumeraterrrDrrrrtupler)rjrqr*rrOrrrN synced_gpusall_hidden_statesall_self_attentionsidxZ encoder_layerskip_the_layerdropout_probability layer_outputsr%r%r&rrst$             zSpeechT5Encoder.forwardNNNNNrvrwrxrrr^rr rrrKrrrrrryr%r%rmr&rs. rcr~)SpeechT5EncoderWithSpeechPrenetz Wrapper around SpeechT5Encoder that applies SpeechT5SpeechEncoderPrenet to convert the audio waveform data to hidden features. rkc,t|t||_t||_|dSro)r]r^rprenetrwrapped_encoderrrrmr%r&r^i    z(SpeechT5EncoderWithSpeechPrenet.__init__Nr(r*rrOrrr1c Cs*|||\}}|j||||||d}|SN)rqr*rrOrrrr rjr(r*rrOrrrqrdr%r%r&rrqs  z'SpeechT5EncoderWithSpeechPrenet.forwardrrr%r%rmr&rcs.  rcseZdZdZdeffdd ZddZddZ dd ej d e ej d e ej d e e de e de e de eeffddZZS)SpeechT5EncoderWithTextPrenetz| Wrapper around SpeechT5Encoder that applies SpeechT5TextEncoderPrenet to convert the input_ids to hidden features. rkcrro)r]r^r%rrrrrrmr%r&r^rz&SpeechT5EncoderWithTextPrenet.__init__cC |jSrorr+rr%r%r&r+r1z2SpeechT5EncoderWithTextPrenet.get_input_embeddingscC|j|dSrorr0r.r%r%r&r0z2SpeechT5EncoderWithTextPrenet.set_input_embeddingsNr(r*rrOrrr1c Cs$||}|j||||||d}|Srrrr%r%r&rrs  z%SpeechT5EncoderWithTextPrenet.forwardr)rvrwrxrrr^r+r0rr rrrKrrrrrryr%r%rmr&rs2 rcr~)SpeechT5EncoderWithoutPrenet This wrapper class is a helper class to correctly load pretrained checkpoints when used in combination with [`SpeechT5Model`]. rkc"t|t||_|dSro)r]r^rrrrrmr%r&r^   z%SpeechT5EncoderWithoutPrenet.__init__Nr(r*rrOrrr1cCs|j||||||dSr)r)rjr(r*rrOrrr%r%r&rrs z$SpeechT5EncoderWithoutPrenet.forwardrrr%r%rmr&rs.  rceZdZdZdeffdd Z           ddeejdeej deejd eej d eej d eej d ee ejd ee dee dee dee de eeffddZZS)SpeechT5Decoderzt Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`SpeechT5DecoderLayer`] rkcsFtj|_tfddtjD|_d|_ | dS)Ncrr%)rer;rr%r&r>rz,SpeechT5Decoder.__init__..F) r]r^Zdecoder_layerdroprrrrIZdecoder_layersrrrrrmrr&r^s   zSpeechT5Decoder.__init__Nrqr*rjrkrcross_attn_head_maskr6rmrOrrr1c  Cs | dur| n|jj} | dur| n|jj} |dur|n|jj}| dur$| n|jj} |dd} |dur=|ddjdnd} t|| || }|durX|durXt||j | dd}t p^t |}|j rn|j rn|rntdd}| rrdnd}| rxdnd}| r|durdnd}|rdnd}t||gd d gD](\}}|dur|dt|jkrtd |d t|jd |ddqt|jD]\}}| r||f}d}|j rtg}||jk}|r|sq|dur||nd}|j r|j r||j|||||dur||nd|dur ||ndd| | }n||||||dur"||nd|dur,||nd|| |d }|d}|rF||| rAdndf7}| r\||df}|dur\||df}q| re||f}|rj|nd}| s}tdd|||||fDSt|||||dS)a Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, feature_size)`): Features extracted from the speech or text input by the decoder prenet. attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing cross-attention on hidden heads. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Nrrr)rQzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fr%rrzThe `z` should be specified for rr)r*rjrkrMrlrLrOrmr rcsrror%rr%r%r&rsz*SpeechT5Decoder.forward..)rr6rqrcross_attentions)rkrOrrmrrr!rrr@r rrrloggerZ warning_oncerrNrr#rrrDrrrrr)rjrqr*rjrkrrr6rmrOrrr7rrrrZall_cross_attentionsZnext_decoder_cacheZ attn_maskZ mask_namerZ decoder_layerrrrLrZ next_cacher%r%r&rrsI              zSpeechT5Decoder.forward NNNNNNNNNNNrvrwrxrrr^rrr r rrrKrrrrrryr%r%rmr&rsN       rcseZdZdZdeffdd Z            ddeejdeej deejd eej d eej d eej d eej d ee ejdee dee dee dee de eeffddZZS)SpeechT5DecoderWithSpeechPrenetz Wrapper around SpeechT5Decoder that applies SpeechT5SpeechDecoderPrenet to convert log-mel filterbanks to hidden features. rkcrro)r]r^r rrwrapped_decoderrrrmr%r&r^rz(SpeechT5DecoderWithSpeechPrenet.__init__Nr(r*rjrkrrrr6rmrOrrr1c  Cs0|||} |j| ||||||| | | | d }|SN) rqr*rjrkrrr6rmrOrrrr)rjr(r*rjrkrrrr6rmrOrrdecoder_hidden_statesrdr%r%r&rrs z'SpeechT5DecoderWithSpeechPrenet.forward) NNNNNNNNNNNNrr%r%rmr&rsT       rcseZdZdZdeffdd ZddZddZ dd ee j d ee j d ee j d ee j dee j dee j dee e j deedeedeedeedeeeffddZZS)SpeechT5DecoderWithTextPrenetz{ Wrapper around SpeechT5Decoder that applies SpeechT5TextDecoderPrenet to convert input tokens to hidden features. rkcrro)r]r^r2rrrrrrmr%r&r^rz&SpeechT5DecoderWithTextPrenet.__init__cCrrorrr%r%r&r+r1z2SpeechT5DecoderWithTextPrenet.get_input_embeddingscCrrorr.r%r%r&r0rz2SpeechT5DecoderWithTextPrenet.set_input_embeddingsNr(r*rjrkrrr6rmrOrrr1c  Cs6||||\} }|j| |||||||| | | d } | Srr)rjr(r*rjrkrrr6rmrOrrrrdr%r%r&rrsz%SpeechT5DecoderWithTextPrenet.forwardr)rvrwrxrrr^r+r0rrr r rrrKrrrrrryr%r%rmr&rsR      rcr)SpeechT5DecoderWithoutPrenetrrkcrro)r]r^rrrrrmr%r&r^ rz%SpeechT5DecoderWithoutPrenet.__init__Nr(r*rjrkrrr6rmrOrrr1c Cs$|j||||||||| | | d } | Sr)r) rjr(r*rjrkrrr6rmrOrrrdr%r%r&rrs z$SpeechT5DecoderWithoutPrenet.forwardrrr%r%rmr&rsN       rcs\eZdZdZdeffdd Zdejdejdejdej fd d Z d d Z e d dZ ZS)$SpeechT5GuidedMultiheadAttentionLossz Guided attention loss from the paper [Efficiently Trainable Text-to-Speech System Based on Deep Convolutional Networks with Guided Attention](https://arxiv.org/abs/1710.08969), adapted for multi-head attention. rkcst|j|_|j|_dSro)r]r^Zguided_attention_loss_sigmasigmaZguided_attention_loss_scalescalerrmr%r&r^6s  z-SpeechT5GuidedMultiheadAttentionLoss.__init__r input_masks output_masksr1cCsX||||j}|d|d@}||jd}||}t||}|j|S)aY Compute the attention loss. Args: attentions (`torch.FloatTensor` of shape `(batch_size, layers * heads, output_sequence_length, input_sequence_length)`): Batch of multi-head attention weights input_masks (`torch.BoolTensor` of shape `(batch_size, input_sequence_length)`): Input attention mask as booleans. output_masks (`torch.BoolTensor` of shape `(batch_size, output_sequence_length)`): Target attention mask as booleans. Returns: `torch.Tensor` with the loss value rr}r)_make_guided_attention_masksrrrrrq masked_selectr)rjrrrguided_attn_masksmasksZlosseslossr%r%r&rr;s  z,SpeechT5GuidedMultiheadAttentionLoss.forwardc Cs|d}|d}tjt||jd|jdf|d}tt||D]\}\}} ||| |j|||d| d|f<q#| dS)Nrrr) rGrrJrNr!rr_make_guided_attention_maskrr) rjrrrrTrrrZilenZolenr%r%r&rTs  $& zASpeechT5GuidedMultiheadAttentionLoss._make_guided_attention_maskscCsdtjtj||dtj||ddd\}}||}||}dt||d d|dS)NrZxy)Zindexingrr)rZmeshgridrMrr)r5Z output_lengthrrZgrid_yZgrid_xr%r%r&r_s    $z@SpeechT5GuidedMultiheadAttentionLoss._make_guided_attention_mask)rvrwrxrrr^rr Z BoolTensorrrrrrrryr%r%rmr&r0s  rcsbeZdZdZdeffdd Z ddejdejdejd ejd ejd e ejd ej fd dZ Z S)SpeechT5SpectrogramLossz; Loss computation used by SpeechT5ForTextToSpeech. rkcsTt|j|_|j|_|j|_t|_tt dd|_ |jr(t ||_ dSdS)Ng@)Z pos_weight) r]r^use_guided_attention_lossguided_attention_loss_num_headsr)r l1_criterionrrr bce_criterionrattn_criterionrrmr%r&r^ps z SpeechT5SpectrogramLoss.__init__Nr*r"r#r$labelsrr1cs<|dk}||}||}||}||||}|dddddf} tj| dt| dd| jgdd} | ddddf| } || }|| } || } j rtjfdd|Ddd} |dk}|dddddf}j dkr|ddj ddj f} | ||}| |7} | S)Nr+rrrrcs"g|] }|dddjfqSro)r)r<xrr%r&r>s"z3SpeechT5SpectrogramLoss.forward..) rrrrrOrrrrrr)r)rjr*r"r#r$rrrZl1_lossr stop_labelsZbce_lossrZattnrrZ attn_lossr%rr&rr|s(   .   zSpeechT5SpectrogramLoss.forwardro) rvrwrxrrr^rr r rrrrryr%r%rmr&rks&rzv The bare SpeechT5 Encoder-Decoder Model outputting raw hidden-states without any specific pre- or post-nets. Z custom_introc"s6eZdZ  d"dedeejdeejffdd ZddZd d Z d d Z d dZ ddZ e              d#deejdeejdeejdeejdeejdeejdeejdeeeejdeeeejdeedeejdeedeedeedeeejeffd d!ZZS)$ SpeechT5ModelNrkencoderdecodercsJt|||_|durt|n||_|durt|n||_|dS)z encoder (`PreTrainedModel`, *optional*): The encoder model to use. decoder (`PreTrainedModel`, *optional*): The decoder model to use. N)r]r^rkrrrrr)rjrkrrrmr%r&r^s  zSpeechT5Model.__init__cCs0t|jtr |jSt|jtr|jStro)rVrrr+rrNotImplementedErrorrr%r%r&r+s    z"SpeechT5Model.get_input_embeddingscCs8t|jtr |j|t|jtr|j|dSdSro)rVrrr0rrr.r%r%r&r0s   z"SpeechT5Model.set_input_embeddingscCr)ro)rrr%r%r& get_encoderr,zSpeechT5Model.get_encodercCr)ro)rrr%r%r& get_decoderr,zSpeechT5Model.get_decodercCs t|jtr|jjdSdSz Calling this function will disable the gradient computation for the feature encoder so that its parameter will not be updated during training. N)rVrrrrrr%r%r&rs z$SpeechT5Model.freeze_feature_encoderr(r*decoder_input_valuesdecoder_attention_maskrdecoder_head_maskrencoder_outputsr6rmrrOrrr1cCsf| dur| n|jj} | dur| n|jj} | dur| n|jj} |dur$|n|jj}|dur8|j|||| | |d}n$|r\t|ts\t|dt|dkrM|dndt|dkrX|dndd}|durtt|jt rt|jj |dj d|}n|}t|j trd| i}ni}|j d |||d|||| | | | |d |}|s||St|j|j|j|j|j|j|j|jd S) a input_values (`torch.Tensor` of shape `(batch_size, sequence_length)`): Depending on which encoder is being used, the `input_values` are either: float values of the input raw speech waveform, or indices of input sequence tokens in the vocabulary, or hidden states. decoder_input_values (`torch.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*): Depending on which decoder is being used, the `decoder_input_values` are either: float values of log-mel filterbank features extracted from the raw speech waveform, or indices of decoder input sequence tokens in the vocabulary, or hidden states. decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_values`. Causal mask will also be used by default. If you want to change padding behavior, you should read [`SpeechT5Decoder._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. N)r(r*rrOrrrrrrr) r(r*rjrkrrr6rmrOrr)rr6rdecoder_attentionsrencoder_last_hidden_staterjencoder_attentionsr%)rkrOrrmrrrVrrNrrrr!rrrrr6rqrr)rjr(r*rrrrrrr6rmrrOrrrkZ decoder_argsZdecoder_outputsr%r%r&rrsn)    zSpeechT5Model.forwardr NNNNNNNNNNNNNN)rvrwrxrrrModuler^r+r0rrrrrrr r rrKrrrrryr%r%rmr&rsz      rzB SpeechT5 Model with a speech encoder and a text decoder. c"s(eZdZdgZdeffdd ZddZddZd d Zd d Z d dZ e              d#de e jde e jde e jde e jde e jde e jde e jde eee jde eee jde ede ede ede ede e jdeeeffdd Zed!d"ZZS)$SpeechT5ForSpeechToTextz#text_decoder_postnet.lm_head.weightrkc\t||jdurtd|jdt|}t|}t||||_t ||_ | dS)NYou are trying to instantiate a with a configuration that does not define the vocabulary size of the language model head. Please instantiate the model as follows: `SpeechT5ForSpeechToText.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration.) r]r^r&r#rnrrrror8text_decoder_postnetr)rjrkspeech_encoderZ text_decoderrmr%r&r^R     z SpeechT5ForSpeechToText.__init__cCrrororrr%r%r&rfr1z#SpeechT5ForSpeechToText.get_encodercCrrororrr%r%r&rir1z#SpeechT5ForSpeechToText.get_decodercC|jdSrrrrrr%r%r&rlz.SpeechT5ForSpeechToText.freeze_feature_encodercCrro)rr<rr%r%r&r<sr1z-SpeechT5ForSpeechToText.get_output_embeddingscCrro)rr>r=r%r%r&r>vrz-SpeechT5ForSpeechToText.set_output_embeddingsNr(r*decoder_input_idsrrrrrr6rmrOrrrr1cCs| dur| n|jj} |dur|durt||jj|jj}|j||||||||| | | | dd }||d}d}|durLt}||d|jj |d}| sb|f|dd}|dur`|f|S|St |||j |j |j |j|j|j|jd S)a input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Float values of input raw speech waveform. Values can be obtained by loading a *.flac* or *.wav* audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via the soundfile library (*pip install soundfile*). To prepare the array into `input_values`, the [`SpeechT5Processor`] should be used for padding and conversion into a tensor of type `torch.FloatTensor`. See [`SpeechT5Processor.__call__`] for details. decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`SpeechT5Tokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) SpeechT5 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_values`. Causal mask will also be used by default. If you want to change padding behavior, you should read [`SpeechT5Decoder._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Label indices can be obtained using [`SpeechT5Tokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. Example: ```python >>> from transformers import SpeechT5Processor, SpeechT5ForSpeechToText >>> from datasets import load_dataset >>> dataset = load_dataset( ... "hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True ... ) # doctest: +IGNORE_RESULT >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_asr") >>> model = SpeechT5ForSpeechToText.from_pretrained("microsoft/speecht5_asr") >>> # audio file is decoded on the fly >>> inputs = processor(audio=dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") >>> predicted_ids = model.generate(**inputs, max_length=100) >>> # transcribe speech >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) >>> transcription[0] 'mister quilter is the apostle of the middle classes and we are glad to welcome his gospel' ``` ```python >>> inputs["labels"] = processor(text_target=dataset[0]["text"], return_tensors="pt").input_ids >>> # compute loss >>> loss = model(**inputs).loss >>> round(loss.item(), 2) 19.68 ``` NT) r(r*rrrrrrr6rmrOrrrrr) rr$r6rrrrrjr)rkrr'rrrorrrr&rr6rrrrrjr)rjr(r*rrrrrrr6rmrOrrrrdr$rZloss_fctoutputr%r%r&rrysPYzSpeechT5ForSpeechToText.forwardcs.d}|D]}|tfdd|Df7}q|S)Nr%c3s$|] }|d|jVqdS)rN)rrr)r<Z past_statebeam_idxr%r&r s"z9SpeechT5ForSpeechToText._reorder_cache..)r)r6rZreordered_pastZ layer_pastr%rr&_reorder_cache s z&SpeechT5ForSpeechToText._reorder_cacher)rvrwrxZ_tied_weights_keysrr^rrrr<r>rrrr r rrrKrrrrrrryr%r%rmr&rJsr       r?r,4@Fmodelr threshold minlenratio maxlenratiovocoderoutput_cross_attentionsreturn_output_lengthsc " s|durtd|durd||jjk} n|} |d} |jj|| dd} | j} t|jjt r?|jjj | dj d| } t| d||jj }t| d||jj }| | d|jj}g}g}d}d}i |d7}|jj ||}|jjj|ddddfd| | |d|dd}|r|tj|jdd|jd}|j}|j|}|| |jj |jj}|||dddddf| d|jj}tj||fdd}t|j|}||krql||krtj|dd|k}t|d}nt t!|}fd d |D}t!|dkr3t"|}|#dd$dd }|j%|}|D] }|||<q)t!| kr;nqmfd d t t!D}| s| dkrU|dn tj&j'j(j)|dd }|duri||}n|}|rtj|d d}| dkr|j| t|d| g|ddR}||f}|Sgt | D] } || dq|durtj&j'j(j)|dd }|f}ngtj&j'j(j)|dd }||fdd D}!|!f}|rtj|d d}|j| t|d| g|ddR}g||R}|S)Na`speaker_embeddings` must be specified. For example, you can use a speaker embeddings by following the code snippet provided in this link: https://huggingface.co/datasets/Matthijs/cmu-arctic-xvectors rrT)r(r*rr)rqr*rjrkr6rmrOrrcsg|]}|vr|qSr%r%rresult_spectrogramr%r&r>l sz$_generate_speech..rcsg|]}|qSr%r%rrr%r&r>u r)Z batch_firstcs&g|]}tdt|qSr)r3rr4r)spectrogram_lengths waveformsr%r&r> s&)*r#rkrr3rrorrrVrrrr!r)r r rrrPrrrsqueezer6speech_decoder_postnetrrZsigmoidr rGrrHrIrNstackr~flattenr!rrZrnnZ pad_sequence)"rr(rr*rrrrrrrkrZ encoder_outrmaxlenZminlenZoutput_sequence spectrogramrr6rrZ decoder_outZlast_decoder_outputZspectrumZnew_spectrogramZprobZmeet_thresholdsZ meet_indexesZ spectrogramsZ meet_indexrdrZwaveform_lengthsr%)rrrr&_generate_speech s       $   5&      rzB SpeechT5 Model with a text encoder and a speech decoder. c&seZdZdZdeffdd ZedefddZdd Z d d Z e d,de e jd e e jde e jde e jde e jde e jde e jde eee jde eee jde ede ede ede ede e jde e jde e jdeeeff"ddZe    ! !d-de jd e e jde e jd"ed#ed$ed%e ejd&ed'edee jee je jfffd(d)Ze    ! !d-de jde e jd e e jd"ed#ed$ed%e ejd&ed'edee jee je jfffd*d+ZZS).SpeechT5ForTextToSpeechrrkcr)Nra with a configuration that does not define the vocabulary size of the language model head. Please instantiate the model as follows: `SpeechT5ForTextToSpeech.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration.) r]r^r&r#rnrrrrorrr)rjrkZ text_encoderspeech_decoderrmr%r&r^ rz SpeechT5ForTextToSpeech.__init__r1cCsdSrr%)clsr%r%r& can_generate sz$SpeechT5ForTextToSpeech.can_generatecCrrorrr%r%r&r r1z#SpeechT5ForTextToSpeech.get_encodercCrrorrr%r%r&r r1z#SpeechT5ForTextToSpeech.get_decoderNr*rrrrrrr6rmrOrrrrrcCs| dur| n|jj} |dur"|durt||jj|\}}|jjr"d} |j||||||||| | || | dd}||d\}}}d}|durTt|j}|||||||j}| sj|f|dd}|durh|f|S|St |||j |j |j |j|j |j|jd S)a input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`SpeechT5Tokenizer`]. See [`~PreTrainedTokenizer.encode`] and [`~PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) decoder_input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`): Float values of input mel spectrogram. SpeechT5 uses an all-zero spectrum as the starting token for `decoder_input_values` generation. If `past_key_values` is used, optionally only the last `decoder_input_values` have to be input (see `past_key_values`). decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_values`. Causal mask will also be used by default. If you want to change padding behavior, you should read [`SpeechT5Decoder._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. labels (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`, *optional*): Float values of target mel spectrogram. Timesteps set to `-100.0` are ignored (masked) for the loss computation. Spectrograms can be obtained using [`SpeechT5Processor`]. See [`SpeechT5Processor.__call__`] for details. stop_labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Binary tensor indicating the position of the stop token in the sequence. Example: ```python >>> from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, set_seed >>> import torch >>> processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts") >>> model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts") >>> vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan") >>> inputs = processor(text="Hello, my dog is cute", return_tensors="pt") >>> speaker_embeddings = torch.zeros((1, 512)) # or load xvectors from a file >>> set_seed(555) # make deterministic >>> # generate speech >>> speech = model.generate(inputs["input_ids"], speaker_embeddings=speaker_embeddings, vocoder=vocoder) >>> speech.shape torch.Size([15872]) ``` NTr(r*rrrrrrr6rmrrOrrrr rrr6rrrrrjr)rkrr-r)rrorrrrr6rrrrjr)rjrr*rrrrrrr6rmrOrrrrrrdr"r#r$r criterionrr%r%r&rr sdL   zSpeechT5ForTextToSpeech.forwardrr,rFrrrrrrc Ks^|dur"|d} |d| kr"|ddkr|| d}ntdt|||||||||| S)aE Converts a sequence of input tokens into a sequence of mel spectrograms, which are subsequently turned into a speech waveform using a vocoder. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`SpeechT5Tokenizer`]. See [`~PreTrainedTokenizer.encode`] and [`~PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Attention mask from the tokenizer, required for batched inference to signal to the model where to ignore padded tokens from the input_ids. speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. threshold (`float`, *optional*, defaults to 0.5): The generated sequence ends when the predicted stop token probability exceeds this value. minlenratio (`float`, *optional*, defaults to 0.0): Used to calculate the minimum required length for the output sequence. maxlenratio (`float`, *optional*, defaults to 20.0): Used to calculate the maximum allowed length for the output sequence. vocoder (`nn.Module`, *optional*): The vocoder that converts the mel spectrogram into a speech waveform. If `None`, the output is the mel spectrogram. output_cross_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of the decoder's cross-attention layers. return_output_lengths (`bool`, *optional*, defaults to `False`): Whether or not to return the concrete spectrogram/waveform lengths. Returns: `tuple(torch.FloatTensor)` comprising various elements depending on the inputs: - when `return_output_lengths` is False - **spectrogram** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrogram. - **waveform** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(num_frames,)` -- The predicted speech waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. - when `return_output_lengths` is True - **spectrograms** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrograms that are padded to the maximum length. - **spectrogram_lengths** (*optional*, returned when no `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each spectrogram. - **waveforms** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, num_frames)` -- The predicted speech waveforms that are padded to the maximum length. - **waveform_lengths** (*optional*, returned when a `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(batch_size, config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. NrrzUThe first dimension of speaker_embeddings must be either 1 or the same as batch_size.rrr#r) rjrr*rrrrrrrkwargsrSr%r%r&generateH s(E z SpeechT5ForTextToSpeech.generatec Cs^|dur"|d} |d| kr"|ddkr|| d}ntdt|||||||||| S)a Converts a sequence of input tokens into a sequence of mel spectrograms, which are subsequently turned into a speech waveform using a vocoder. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`SpeechT5Tokenizer`]. See [`~PreTrainedTokenizer.encode`] and [`~PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) threshold (`float`, *optional*, defaults to 0.5): The generated sequence ends when the predicted stop token probability exceeds this value. minlenratio (`float`, *optional*, defaults to 0.0): Used to calculate the minimum required length for the output sequence. maxlenratio (`float`, *optional*, defaults to 20.0): Used to calculate the maximum allowed length for the output sequence. vocoder (`nn.Module`, *optional*, defaults to `None`): The vocoder that converts the mel spectrogram into a speech waveform. If `None`, the output is the mel spectrogram. output_cross_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of the decoder's cross-attention layers. return_output_lengths (`bool`, *optional*, defaults to `False`): Whether or not to return the concrete spectrogram/waveform lengths. Returns: `tuple(torch.FloatTensor)` comprising various elements depending on the inputs: - when `return_output_lengths` is False - **spectrogram** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrogram. - **waveform** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(num_frames,)` -- The predicted speech waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. - when `return_output_lengths` is True - **spectrograms** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrograms that are padded to the maximum length. - **spectrogram_lengths** (*optional*, returned when no `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each spectrogram. - **waveforms** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, num_frames)` -- The predicted speech waveforms that are padded to the maximum length. - **waveform_lengths** (*optional*, returned when a `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(batch_size, config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. NrrzUThe first dimension of speaker_embeddings must be either 1 or the same as batch size.r) rjrrr*rrrrrrrSr%r%r&generate_speech s(I z'SpeechT5ForTextToSpeech.generate_speechNNNNNNNNNNNNNNNNNNrr,rNFF)rvrwrxr}rr^ classmethodrKr rrrrrr r rrrrrrrrrrrrryr%r%rmr&r s          [   rzD SpeechT5 Model with a speech encoder and a speech decoder. c&seZdZdeffdd ZddZddZdd Ze d*d e e j d e e j d e e j de e j de e j de e j de e j de eee j de eee j de ede ede ede ede e j de e j de e j deeeff"ddZe    ! !d+d e j de e j d e e j d"ed#ed$ed%e ejd&ed'ede j fd(d)ZZS),SpeechT5ForSpeechToSpeechrkcs@t|t|}t|}t||||_t||_|dSro) r]r^rrrrorrr)rjrkrrrmr%r&r^ s   z"SpeechT5ForSpeechToSpeech.__init__cCrrorrr%r%r&r r1z%SpeechT5ForSpeechToSpeech.get_encodercCrrorrr%r%r&r r1z%SpeechT5ForSpeechToSpeech.get_decodercCrrrrr%r%r&r rz0SpeechT5ForSpeechToSpeech.freeze_feature_encoderNr(r*rrrrrrr6rmrOrrrrrr1cCs| dur| n|jj} |dur|durt||jj|\}}|j||||||||| | || | dd}||d\}}}d}| sQ|f|dd}|durO|f|S|St|||j|j|j |j |j |j |j d S)a input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Float values of input raw speech waveform. Values can be obtained by loading a *.flac* or *.wav* audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via the soundfile library (*pip install soundfile*). To prepare the array into `input_values`, the [`SpeechT5Processor`] should be used for padding and conversion into a tensor of type `torch.FloatTensor`. See [`SpeechT5Processor.__call__`] for details. decoder_input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`): Float values of input mel spectrogram. SpeechT5 uses an all-zero spectrum as the starting token for `decoder_input_values` generation. If `past_key_values` is used, optionally only the last `decoder_input_values` have to be input (see `past_key_values`). decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_values`. Causal mask will also be used by default. If you want to change padding behavior, you should read [`SpeechT5Decoder._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. labels (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`, *optional*): Float values of target mel spectrogram. Spectrograms can be obtained using [`SpeechT5Processor`]. See [`SpeechT5Processor.__call__`] for details. stop_labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Binary tensor indicating the position of the stop token in the sequence. Example: ```python >>> from transformers import SpeechT5Processor, SpeechT5ForSpeechToSpeech, SpeechT5HifiGan, set_seed >>> from datasets import load_dataset >>> import torch >>> dataset = load_dataset( ... "hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True ... ) # doctest: +IGNORE_RESULT >>> dataset = dataset.sort("id") >>> sampling_rate = dataset.features["audio"].sampling_rate >>> processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_vc") >>> model = SpeechT5ForSpeechToSpeech.from_pretrained("microsoft/speecht5_vc") >>> vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan") >>> # audio file is decoded on the fly >>> inputs = processor(audio=dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") >>> speaker_embeddings = torch.zeros((1, 512)) # or load xvectors from a file >>> set_seed(555) # make deterministic >>> # generate speech >>> speech = model.generate_speech(inputs["input_values"], speaker_embeddings, vocoder=vocoder) >>> speech.shape torch.Size([77824]) ``` NTr rrr )rkrr-r)rorrr6rrrrrjr)rjr(r*rrrrrrr6rmrOrrrrrrdr=rr$rrr%r%r&rr$ sLR z!SpeechT5ForSpeechToSpeech.forwardrr,rFrrrrrrc Cs2|dur tjd|jd}t|||||||||| S)a Converts a raw speech waveform into a sequence of mel spectrograms, which are subsequently turned back into a speech waveform using a vocoder. Args: input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Float values of input raw speech waveform. Values can be obtained by loading a *.flac* or *.wav* audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via the soundfile library (*pip install soundfile*). To prepare the array into `input_values`, the [`SpeechT5Processor`] should be used for padding and conversion into a tensor of type `torch.FloatTensor`. See [`SpeechT5Processor.__call__`] for details. speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*): Tensor containing the speaker embeddings. attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) threshold (`float`, *optional*, defaults to 0.5): The generated sequence ends when the predicted stop token probability exceeds this value. minlenratio (`float`, *optional*, defaults to 0.0): Used to calculate the minimum required length for the output sequence. maxlenratio (`float`, *optional*, defaults to 20.0): Used to calculate the maximum allowed length for the output sequence. vocoder (`nn.Module`, *optional*, defaults to `None`): The vocoder that converts the mel spectrogram into a speech waveform. If `None`, the output is the mel spectrogram. output_cross_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of the decoder's cross-attention layers. return_output_lengths (`bool`, *optional*, defaults to `False`): Whether or not to return the concrete spectrogram/waveform lengths. Returns: `tuple(torch.FloatTensor)` comprising various elements depending on the inputs: - when `return_output_lengths` is False - **spectrogram** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrogram. - **waveform** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(num_frames,)` -- The predicted speech waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. - when `return_output_lengths` is True - **spectrograms** (*optional*, returned when no `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, output_sequence_length, config.num_mel_bins)` -- The predicted log-mel spectrograms that are padded to the maximum length. - **spectrogram_lengths** (*optional*, returned when no `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each spectrogram. - **waveforms** (*optional*, returned when a `vocoder` is provided) `torch.FloatTensor` of shape `(batch_size, num_frames)` -- The predicted speech waveforms that are padded to the maximum length. - **waveform_lengths** (*optional*, returned when a `vocoder` is provided) `List[Int]` -- A list of all the concrete lengths for each waveform. - **cross_attentions** (*optional*, returned when `output_cross_attentions` is `True`) `torch.FloatTensor` of shape `(batch_size, config.decoder_layers, config.decoder_attention_heads, output_sequence_length, input_sequence_length)` -- The outputs of the decoder's cross-attention layers. N)rir)rrJrr) rjr(rr*rrrrrrr%r%r&r sIz)SpeechT5ForSpeechToSpeech.generate_speechrr)rvrwrxrr^rrrrrrr r rrrKrrrrrrrrrryr%r%rmr&r s        ~   rcs@eZdZdfdd ZdddZd d Zd d Zd dZZS)HifiGanResidualBlockr rr 皙?csbt|_tfddttD_tfddttD_dS)Nc s2g|]}tjd||dqSr)r[dilationrrrb get_paddingrchannelsrrZrjr%r&r> s z1HifiGanResidualBlock.__init__..c s*g|]}tjddddqSrrr;)rrZrjr%r&r> s  ) r]r^leaky_relu_sloperrrIrNconvs1convs2)rjrrZrr rmrr&r^ s     zHifiGanResidualBlock.__init__rcCs|||dSrr%)rjrZrr%r%r&r rz HifiGanResidualBlock.get_paddingcCsLtjj}ttjjdrtjjj}|jD]}||q|jD]}||qdSNr)rrrrrr!r"rjrrr%r%r&apply_weight_norm s     z&HifiGanResidualBlock.apply_weight_normcCs4|jD]}tj|q|jD]}tj|qdSro)r!rrremove_weight_normr"rjrr%r%r&r&* s  z'HifiGanResidualBlock.remove_weight_normcCsXt|j|jD]"\}}|}tj||j}||}tj||j}||}||}q|Sro)rr!r"rr leaky_relur )rjrqZconv1Zconv2rcr%r%r&rr0 s zHifiGanResidualBlock.forward)r rrr) rvrwrxr^rr%r&rrryr%r%rmr&r s   rz HiFi-GAN vocoder. csbeZdZeZdZdeffdd ZddZddZd d Z e d d de j d e j fddZ ZS)SpeechT5HifiGanrrkc sNt|t|j|_t|j|_tj|j |j dddd|_ t |_ tt|j|jD]$\}\}}|j tj|j d||j d|d||||ddq-t |_tt|j D]#}|j d|d}t|j|jD]\}}|jt||||jqpq^tj|ddddd|_|dt|j |dt|j |dS)Nrr )rZr[rrrqr)r]r^rNZresblock_kernel_sizes num_kernelsZupsample_rates num_upsamplesrrbZ model_in_dimZupsample_initial_channelconv_prer upsamplerrrZupsample_kernel_sizesrPZConvTranspose1d resblocksrIZresblock_dilation_sizesrr  conv_postrrrJrOr)rjrkrZ upsample_raterZrrrmr%r&r^D s>        zSpeechT5HifiGan.__init__cCsJt|tjtjfr!|jjjd|jjd|j dur#|j j dSdSdS)zInitialize the weights.r,rpN) rVrrrbrrvrurkrwr\rx)rjryr%r%r&r{j s  zSpeechT5HifiGan._init_weightscCs`tjj}ttjjdrtjjj}||j|jD]}||q|jD]}|q"||j dSr#) rrrrrr-r.r/r%r0r$r%r%r&r%q s      z!SpeechT5HifiGan.apply_weight_normcCsLtj|j|jD]}tj|q |jD]}|qtj|jdSro)rrr&r-r.r/r0r'r%r%r&r&} s    z"SpeechT5HifiGan.remove_weight_norma Converts a log-mel spectrogram into a speech waveform. Passing a batch of log-mel spectrograms returns a batch of speech waveforms. Passing a single, un-batched log-mel spectrogram returns a single, un-batched speech waveform. rr1cCs|jjr ||j|j}|dk}|s|d}|dd}||}t|j D]8}t j ||jj }|j||}|j||j|}td|jD]}||j||j||7}qK||j}q)t j |}||}t|}|s|dddd}|S|d}|S)a spectrogram (`torch.FloatTensor`): Tensor containing the log-mel spectrograms. Can be batched and of shape `(batch_size, sequence_length, config.model_in_dim)`, or un-batched and of shape `(sequence_length, config.model_in_dim)`. Returns: `torch.FloatTensor`: Tensor containing the speech waveform. If the input spectrogram is batched, will be of shape `(batch_size, num_frames,)`. If un-batched, will be of shape `(num_frames,)`. r rrrr)rkZnormalize_beforerqrrrr~r-rIr,rrr(r r.r/r+r0rtanhrr)rjrZ is_batchedrqrZ res_statejZwaveformr%r%r&rr s,         zSpeechT5HifiGan.forward)rvrwrxrr|r}r^r{r%r&rrr rrryr%r%rmr&r); s&  r))rrrrrnr))rNrr)`rrtypingrrrrnumpyrBrZtorch.utils.checkpointrZtorch.nnrrr Z activationsr Z generationr Zintegrations.deepspeedr Zintegrations.fsdprZmodeling_attn_mask_utilsrrZmodeling_outputsrrrrrZmodeling_utilsrrrrZconfiguration_speecht5rrZ get_loggerrvrZ_HIDDEN_STATES_START_POSITIONrr3r'r-rr ZndarrayrVrrXrzrrrrrrrrrr rrr%r2r8r?rSr\rernrrrrrrrrrrrrr rKrrrrr)__all__r%r%r%r&s          xE- /4(+!=l% %* P04+;=@   ht> w