o i A@sdZddlZddlmZmZddlZddlZddlmZddl m Z m Z m Z ddl mZddlmZmZmZdd lmZdd lmZdd lmZdd lmZmZdd lmZddlmZm Z m!Z!m"Z"m#Z#ddl$m%Z%m&Z&ddl'm(Z(m)Z)ddl*m+Z+ddl,m-Z-m.Z.e)/e0Z1dZ2dej3de4de4fddZ5 ddej3de4deej3fddZ6  dde7e4e4fde8d e4deej9d!e4d"ej:f d#d$Z;Gd%d&d&eZGd+d,d,ej?Z@Gd-d.d.ej?ZAGd/d0d0ej?ZBGd1d2d2ejj?ZCGd3d4d4ej?ZDGd5d6d6ej?ZEGd7d8d8ej?ZFGd9d:d:ej?ZGGd;d<dd>ej?ZIGd?d@d@ej?ZJGdAdBdBej?e%ZKGdCdDdDej?e%ZLGdEdFdFej?e%ZMGdGdHdHej?ZNGdIdJdJej?ZOGdKdLdLeZPGdMdNdNeZQe(GdOdPdPe&ZRGdQdRdReRZSGdSdTdTeRZTGdUdVdVeRZUGdWdXdXeRZVGdYdZdZeRZWGd[d\d\eRZXGd]d^d^eRZYGd_d`d`eRZZGdadbdbej?Z[Gdcddddej?Z\e(dedfGdgdhdheRZ]e(didfGdjdkdkeReZ^   l m n  o oddpeRdej_dqeej_deej9dre8dse8dte8dueej?dve`dwe`d"eej_e7ej_ej_fffdxdyZae(dzdfGd{d|d|eRZbe(d}dfGd~ddeRZcGdddej?Zde(ddfGddde&ZegdZfdS)zPyTorch SpeechT5 model.N)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossL1Loss)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)is_deepspeed_zero3_enabled)is_fsdp_managed_module)_prepare_4d_attention_mask!_prepare_4d_causal_attention_mask)GradientCheckpointingLayer)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsSeq2SeqLMOutputSeq2SeqModelOutputSeq2SeqSpectrogramOutput)EmbeddingAccessMixinPreTrainedModel)auto_docstringlogging)deprecate_kwarg)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_)r r!r"shifted_input_idsr*b/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/speecht5/modeling_speecht5.pyshift_tokens_right6s (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. rNr#Y)r$r%r&r()r-r.r/shifted_input_valuesr*r*r+shift_spectrograms_rightFs (r3r% 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)an Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for ASR](https://huggingface.co/papers/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_spanepsilonr5r4r6sequence_lengthr*r+compute_num_masked_spans  z6_compute_mask_indices..compute_num_masked_spanNr#csg|]}qSr*r*.0_)r?r*r+ sz)_compute_mask_indices..dtyperF)replace)r'nprandomranditemdetachsumtolistrangezerosboolchoicearangelen concatenateonesint32appendarray broadcast_toreshaper:put_along_axis)r%r4r5r/r6 batch_sizer@ input_lengths spec_aug_maskspec_aug_mask_idxsmax_num_masked_spanr;r<spec_aug_mask_idxdummy_mask_idxoffsetsr*r=r+_compute_mask_indices\s\    rec&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+rms  z%SpeechT5NoLayerNormConvLayer.__init__cCs||}||}|SN)rurwry hidden_statesr*r*r+forward  z$SpeechT5NoLayerNormConvLayer.forwardr__name__ __module__ __qualname__rmr __classcell__r*r*r|r+rgsrgcrf)SpeechT5LayerNormConvLayerrcs|t|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrhT)elementwise_affine)rlrmrnrorprrqrrrsrtru LayerNorm layer_normr rvrwrxr|r*r+rms  z#SpeechT5LayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)Nr#)ru transposerrwrr*r*r+rs     z"SpeechT5LayerNormConvLayer.forwardrrr*r*r|r+rsrcrf)SpeechT5GroupNormConvLayerrcst|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrhT) num_groups num_channelsaffine)rlrmrnrorprrqrrrsrtrur rvrw GroupNormrrxr|r*r+rms   z#SpeechT5GroupNormConvLayer.__init__cCs"||}||}||}|Sr~)rurrwrr*r*r+rs   z"SpeechT5GroupNormConvLayer.forwardrrr*r*r|r+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)rlrmoffsetrr make_weights)ryrrrr|r*r+rm"s z.SpeechT5SinusoidalPositionalEmbedding.__init__num_embeddingscCsB||||}t|dr|j|jj|jjd}|jd|dddS)NweightsrFdeviceF persistent) get_embeddinghasattrtorrFrregister_buffer)ryrrr 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'rrErdimr#N)mathlogtorchexprSint64float unsqueezecatsincosviewrPrget_default_dtype)rrrhalf_dimembr*r*r+r1s $& z3SpeechT5SinusoidalPositionalEmbedding.get_embeddingrr past_key_values_lengthcCs||\}}|||j||j}|jd|}||jdkr-|||j|j|j|j d| d ||d S)Nrrr#) size"create_position_ids_from_input_idsrrrrrrr index_selectrrL)ryr rbszseq_len position_idsmax_posr*r*r+rCs "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)ner9rcumsumtype_aslong)ryr rrmaskincremental_indicesr*r*r+rRs  zHSpeechT5SinusoidalPositionalEmbedding.create_position_ids_from_input_idsr~r)rrr__doc__r9rrmr staticmethodrrno_gradTensorrrrr*r*r|r+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)ripaddinggroups weight_normr) modifier_rankweight)namerparametrizations)rlrmrrq hidden_sizenum_conv_pos_embeddingsnum_conv_pos_embedding_groupsruutilsrrrr deepspeedzeroGatheredParametersr original0 original1weight_gweight_vregister_external_parameterSpeechT5SamePadLayerrr rvrw)ryrzrrrrr|r*r+rmes4     z(SpeechT5PositionalConvEmbedding.__init__cCs:|dd}||}||}||}|dd}|SNrr)rrurrwrr*r*r+rs     z'SpeechT5PositionalConvEmbedding.forwardrr*r*r|r+rds !rcs*eZdZdZdfdd ZddZZS) SpeechT5ScaledPositionalEncodingu[ Scaled positional encoding, see §3.2 in https://huggingface.co/papers/1809.08895 cst||}td|d}ttjd|dtjdtd| }t |||dddddf<t |||dddddf<|d}t |j d|ddtj|d |_||_ttd |_dS) NrrrrEg@peFrp?)rrPrSrrrrrrrrrlrmrrDropoutdropoutr Parametertensoralpha)ryrrmax_lenrpositiondiv_termr|r*r+rms .$$  z)SpeechT5ScaledPositionalEncoding.__init__cCs4||j|jddd|df}||}|S)Nr)rrrr)ryrr*r*r+rs& z(SpeechT5ScaledPositionalEncoding.forward)r)rrrrrmrrr*r*r|r+rs rcrf)"SpeechT5RelativePositionalEncodingcs.t||_||_tjd|||_dSr)rlrmr max_lengthrr Embeddingpe_k)ryrrr|r*r+rms z+SpeechT5RelativePositionalEncoding.__init__cCs|jd}td|j|jtjd}|dddf|dddf}|j |||j k<|jd|||jk<||j}||S)NrrrrF)r%rrSrrrrr)ryrrpos_seqr*r*r+rs    z*SpeechT5RelativePositionalEncoding.forward)rrr*r*r|r+rsrcr)rcs*t|ddkrd|_dSd|_dS)Nrrr)rlrmnum_pad_remove)ryrr|r*r+rms  zSpeechT5SamePadLayer.__init__cCs,|jdkr|ddddd|j f}|SNr)rrr*r*r+rs zSpeechT5SamePadLayer.forwardrr*r*r|r+rs rcs0eZdZdZfddZddZddZZS)SpeechT5FeatureEncoderz.Construct the features from raw audio waveformcstjdkr tddgfddtjdD}njdkr2fddtjD}n td jd t||_ d |_ d |_ dS) Ngrouprr{csg|] }t|ddqS)rr)rgrBirzr*r+rDsz3SpeechT5FeatureEncoder.__init__..rlayercg|]}t|dqS)r)rrrr*r+rDs z`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) rlrmfeat_extract_normrrOnum_feat_extract_layersr'r ModuleList conv_layersgradient_checkpointing_requires_grad)ryrzrr|rr+rms        zSpeechT5FeatureEncoder.__init__cCs|D]}d|_qd|_dS)NF) parameters requires_gradr)ryparamr*r*r+_freeze_parameterss  z)SpeechT5FeatureEncoder._freeze_parameterscCs:|dddf}|jr|jrd|_|jD]}||}q|SNT)rtrainingrr)ryr-r conv_layerr*r*r+rs    zSpeechT5FeatureEncoder.forward)rrrrrmrrrr*r*r|r+rs  rcr)SpeechT5FeatureProjectioncsJttj|jd|jd|_t|jd|j|_ t |j |_ dS)Nr#eps) rlrmrrrnlayer_norm_epsrLinearr projectionrfeat_proj_dropoutrryrzr|r*r+rms z"SpeechT5FeatureProjection.__init__cCs&||}||}||}||fSr~)rrr)ryrnorm_hidden_statesr*r*r+rs   z!SpeechT5FeatureProjection.forwardrr*r*r|r+r s r cseZdZfddZddZ  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)Nr1r)rlrmrzrfeature_encoderr feature_projectionmask_time_probmask_feature_probrrrrruniform_masked_spec_embedrpos_conv_embedrmax_speech_positionsr!pos_sinusoidal_embedrr|r*r+rms     z$SpeechT5SpeechEncoderPrenet.__init__cCs|jdSr~)rrryr*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)r!r/r) rr"_get_feature_vector_attention_maskr%r_mask_hidden_statesrrrrrPrr) ryr-r/r!extract_featuresrpositional_conv_embedding padding_mask positional_sinusoidal_embeddingsr*r*r+rs&    z#SpeechT5SpeechEncoderPrenet.forwardfeature_vector_lengthcCs|jdddddf}||tj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dgd dg }|S)Nr#rrrrr) r _get_feat_extract_output_lengthsrrrr%rPrFrrSfliprQ)ryr(r/non_padded_lengthsoutput_lengthsr]r*r*r+r"9s "z>SpeechT5SpeechEncoderPrenet._get_feature_vector_attention_maskr^cCs4dd}t|jj|jjD] \}}||||}q |S)zH Computes the output length of the convolutional layers cSstj|||dddS)Nfloor) rounding_moder)rdiv)r;rirjr*r*r+_conv_out_lengthNszVSpeechT5SpeechEncoderPrenet._get_feat_extract_output_lengths.._conv_out_length)ziprzrrrs)ryr^r1rirjr*r*r+r*Isz.)rlrmrzrrrOspeech_decoder_prenet_layerslayersrr@r final_layerrpositional_dropoutrencode_positionsspeaker_embedding_dimspeaker_embeds_layerrr|rr+rms   z$SpeechT5SpeechDecoderPrenet.__init__cCsJtj|d|d}|d|ddd}t|dk|ddd|S)Nrrr)r bernoullirrepeatrwhere)ry inputs_embedsrr all_masksr*r*r+_consistent_dropoutsz/SpeechT5SpeechDecoderPrenet._consistent_dropoutNr-speaker_embeddingscCs|}|jD]}tj||}|||jj}q||}||}|durKtj |}| d d| dd}t j||gdd}tj||}|S)Nrr#r)rBr functionalrelurMrzspeech_decoder_prenet_dropoutrCrE normalizerr9rrrrG)ryr-rNrKrr*r*r+rs    z#SpeechT5SpeechDecoderPrenet.forwardr~) rrrrmrMrrrrrr*r*r|r+r>s r>crf)SpeechT5BatchNormConvLayerrcst|dkr |j}n|j}||jdkr|j}n|j}tj|||jd|jdddd|_t ||_ ||jdkrCt |_ nd|_ t |j|_dS)NrrrF)rirjrrk)rlrmr?speech_decoder_postnet_unitsspeech_decoder_postnet_layersrrqspeech_decoder_postnet_kernelru BatchNorm1d batch_normTanhrwrspeech_decoder_postnet_dropoutr)ryrzr{rorpr|r*r+rms(    z#SpeechT5BatchNormConvLayer.__init__cCs6||}||}|jdur||}||}|Sr~)rurXrwrrr*r*r+rs     z"SpeechT5BatchNormConvLayer.forwardrrr*r*r|r+rSsrScs<eZdZfddZdejfddZdejfddZZS)SpeechT5SpeechDecoderPostnetcs^t|_tjjj|_tjj|_ t fddt j D|_ dS)Ncsg|]}t|qSr*)rSrrr*r+rDsz9SpeechT5SpeechDecoderPostnet.__init__..)rlrmrzrrrr?r.feat_outprob_outrrOrUrBrr|rr+rms  z%SpeechT5SpeechDecoderPostnet.__init__rcCsJ|||dd|jj}||}|||dd}|||fS)Nrr#)r\rrrzr?postnetr])ryroutputs_before_postnetoutputs_after_postnetlogitsr*r*r+rs  z$SpeechT5SpeechDecoderPostnet.forwardcCs0|dd}|jD]}||}q ||ddSr)rrB)ryr layer_outputrr*r*r+r^s   z$SpeechT5SpeechDecoderPostnet.postnet) rrrrmrrrr^rr*r*r|r+r[s  r[cs,eZdZfddZdejfddZZS)SpeechT5TextEncoderPrenetcs>t||_t|j|j|j|_t |j |j|j |_ dSr~) rlrmrzrr vocab_sizerr! embed_tokensrrDmax_text_positionsrErr|r*r+rms  z"SpeechT5TextEncoderPrenet.__init__r cCs||}||}|Sr~)rerE)ryr rKr*r*r+rrz!SpeechT5TextEncoderPrenet.forward)rrrrmrrrrr*r*r|r+rcs  rccsDeZdZfddZ  d dejdeejdeefddZ Z S) SpeechT5TextDecoderPrenetcsnt||_t|j|_|jrt |j nd|_ t |j |j |j|_t|j|jd|j |j|_dS)Nrr)rlrmrzrrrDrscale_embeddingrsqrtr embed_scalerrdr!rerrfembed_positionsrr|r*r+rms  z"SpeechT5TextDecoderPrenet.__init__Nr r/past_key_valuescCs|dur|}|d|d}ntdd}|dur-t|ts)|ddjdn|}|||}|||j }||7}| |}||fS)Nr#z'You have to specify `decoder_input_ids`rr) rrr' isinstancer r%get_seq_lengthrkrerjr)ryr r/rl input_shaper positionsrKr*r*r+rs  z!SpeechT5TextDecoderPrenet.forwardr;) rrrrmrrrr<r rrr*r*r|r+rgs rgcs<eZdZfddZdejfddZddZdd ZZ S) SpeechT5TextDecoderPostnetcs*t||_tj|j|jdd|_dS)NFrk)rlrmrzrrrrdlm_headrr|r*r+rm;s z#SpeechT5TextDecoderPostnet.__init__rcCs ||Sr~rsrr*r*r+r@ z"SpeechT5TextDecoderPostnet.forwardcC|jSr~rtrr*r*r+get_output_embeddingsCsz0SpeechT5TextDecoderPostnet.get_output_embeddingscCs ||_dSr~rtrynew_embeddingsr*r*r+set_output_embeddingsHruz0SpeechT5TextDecoderPostnet.set_output_embeddings) rrrrmrrrrwrzrr*r*r|r+rq:s  rqcseZdZdZ    ddededeed eed eed eef fd d Ze dddd       dde j dee j dee dee j dee j dee j dedee j de e j ee j ee ffddZZS)SpeechT5Attentionz Multi-headed attention from 'Attention Is All You Need' paper with relative position bias (see https://aclanthology.org/N18-2074.pdf) r1FTN embed_dim num_headsr is_decoderrk layer_idxcst||_||_||_|||_|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).grr)rlrmr|r}rhead_dimr'scalingr~rrrk_projv_projq_projout_proj)ryr|r}rr~rkrr|r*r+rmRs$    zSpeechT5Attention.__init__past_key_valuerl4.58new_nameversionrkey_value_statesr/layer_head_mask position_biasoutput_attentionscache_positionr7c Cs|du} |\} } } |||j} d}|dur1t|tr/|j|j}| r+|j}n|j }n|}| r5|n|}| rN|durN|rN|j |jj }|j |jj }nJ| |}||}|| d|j|jdd}|| d|j|jdd}|dur| s||nd}||||jd|i\}}| rt|trd|j|j<| |jd|jf}| | | |j|jdd} | j|} |j|}|j|}|d}t| |dd}|| |j| |fkrtd| |j| |fd ||dur#| | |jd|jd d}t||d d}|d d| |j|d |d}||7}|durX|| d| |fkrCtd | d| |fd ||| |j| ||}|| |j| |}tjj|dd }|dur||jfkr|td|jfd ||dddd|| |j| |}|| |j| |}|r|| |j| |}|| |j| |}nd}tjj||j|jd}t||}|| |j| |jfkrtd| |j| |jfd ||| |j| |j}|dd}|| | |j}| |}||fS)z#Input shape: Batch x Time x ChannelNFr#rrrTz$Attention weights should be of size z , but is rrz!Attention mask should be of size rz/Head mask for a single layer should be of size )rr z `attn_output` should be of size )!rrrrmr is_updatedgetrcross_attention_cacheself_attention_cacherBkeysvaluesrrrr}rrupdater[rbmmr' contiguousmatmulrrOsoftmaxrr r|r)ryrrrlr/rrrris_cross_attentionrtgt_lenrC query_statesrcurr_past_key_valuecurrent_states key_states value_states proj_shapesrc_len attn_weights reshape_q rel_pos_biasattn_weights_reshaped attn_probs attn_outputr*r*r+ros           "    "   zSpeechT5Attention.forward)r1FTN)NNNNNFN)rrrrr9rrrQrmrrrr tuplerrr*r*r|r+r{Ls\    r{cr)SpeechT5FeedForwardcsltt|j|_t|j||_t |j t r!t |j |_ n|j |_ t||j|_t|j|_dSr~)rlrmrractivation_dropoutintermediate_dropoutrrintermediate_denserm hidden_actstrr intermediate_act_fn output_densehidden_dropoutoutput_dropout)ryrzintermediate_sizer|r*r+rms  zSpeechT5FeedForward.__init__cCs6||}||}||}||}||}|Sr~)rrrrrrr*r*r+rs     zSpeechT5FeedForward.forwardrr*r*r|r+rs  rc s^eZdZdeffdd Z    d dejdeejdeejd eejd ef d d Z Z S)SpeechT5EncoderLayerrzcsjtt|j|j|jdd|_t|j |_ tj |j|j d|_ t||j|_tj |j|j d|_dS)NF)r|r}rr~r )rlrmr{rencoder_attention_headsattention_dropout attentionrrrrrrrrencoder_ffn_dim feed_forwardfinal_layer_normrr|r*r+rm s zSpeechT5EncoderLayer.__init__NFrr/rrrc Csh|}|j|||||d\}}||}||}||}|||}||}|f}|r2||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. )rr/rrr)rrrrr) ryrr/rrrresidualroutputsr*r*r+rs"    zSpeechT5EncoderLayer.forward)NNNF) rrrrrmrrrrQrrr*r*r|r+r s"rcseZdZddeffdd Zedddd       dd ejd eejd eejdeejdeejdeejdee dee dee deejfddZ Z S)SpeechT5DecoderLayerNrzcstt|j|j|jd|d|_t|j |_ tj |j|j d|_ t|j|j|jd|d|_tj |j|j d|_t||j|_tj |j|j d|_dS)NT)r|r}rr~rr )rr~r)rlrmr{rdecoder_attention_headsr self_attnrrrrrrself_attn_layer_norm encoder_attnencoder_attn_layer_normrdecoder_ffn_dimrr)ryrzrr|r*r+rmHs( zSpeechT5DecoderLayer.__init__rrlrrFTrr/encoder_hidden_statesencoder_attention_maskrcross_attn_layer_head_maskr use_cacherc  Cs|} |j|||||| d\}} ||}| |}||}d} |durA|} |j||||||| d\}} ||}| |}||}|||}||}|f}|rX|| | 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_values (`Cache`): 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. )rrlr/rrrN)rrr/rrlrr)rrrrrrr)ryrr/rrrrrlrrrrself_attn_weightscross_attn_weightsrr*r*r+r`sB       zSpeechT5DecoderLayer.forwardr~) NNNNNNFTN) rrrrrmrrrrr rQrrr*r*r|r+rGsB    rc@s2eZdZUeed<dZdZdZdej fddZ dS) SpeechT5PreTrainedModelrzspeecht5r-Tmodulec Cs|jj}t|tr-tjj|jjddt d|jj d|jj dtj |jjdnt|tr:|jjdnt|trat d|jj}tjj|jj| |dtjj|jj| |dn}t|tjr||jjjd|d|jdur{|jjnbt|tjtjtjfr|jj|jjdnIt|tjrtj|j|jdurt |j|j |j d}tjj|j| |dnt|tjr|jjjd|d|j dur|jj|j t!|d rtj|j"dSdS) zInitialize the weightsrrrmeanstdr)abr1Nr)#rzinitializer_rangermrrinitnormal_rurrriri in_channels constant_rkrrdatafill_r r in_featuresrrzero_rrrWrqkaiming_normal_rrrrr)ryrrkr*r*r+ _init_weightssF             z%SpeechT5PreTrainedModel._init_weightsN) rrrr__annotations__base_model_prefixmain_input_namesupports_gradient_checkpointingrModulerr*r*r*r+rs rczeZdZdZdeffdd 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`]. rzcs~ttjjjd|_tj|_ j |_ t fddt jD|_tjjj|_d|_|dS)Nr csg|]}tqSr*)rrArr*r+rDz,SpeechT5Encoder.__init__..F)rlrmrrrrrrrrencoder_layerdrop layerdroprrOencoder_layersrBrrencoder_max_relative_positionrkr post_initrr|rr+rms   zSpeechT5Encoder.__init__Nrr/ head_maskroutput_hidden_states return_dictr7cCs|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]@\} } |r{| |f} d} |jrtg}||jk} | r|r| ||||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/rrrr;rcs|] }|dur|VqdSr~r*rBvr*r*r+ Osz*SpeechT5Encoder.forward..last_hidden_stater attentions)rzrruse_return_dictrrFrrrkrrrrTrBr' enumerater rrJrrr)ryrr/rrrrr synced_gpusall_hidden_statesall_self_attentionsidx encoder_layerskip_the_layerdropout_probability layer_outputsr*r*r+rsb$           zSpeechT5Encoder.forwardNNNNNrrrrrrmrr=rrrQrrrrrr*r*r|r+rs. rcr)SpeechT5EncoderWithSpeechPrenetz Wrapper around SpeechT5Encoder that applies SpeechT5SpeechEncoderPrenet to convert the audio waveform data to hidden features. rzc,t|t||_t||_|dSr~)rlrmrprenetrwrapped_encoderrrr|r*r+rm^    z(SpeechT5EncoderWithSpeechPrenet.__init__Nr-r/rrrrr7c Cs*|||\}}|j||||||d}|SN)rr/rrrrrr ryr-r/rrrrrrr*r*r+rfs  z'SpeechT5EncoderWithSpeechPrenet.forwardrrr*r*r|r+rXs.  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. rzcrr~)rlrmrcrrrrrr|r*r+rmr z&SpeechT5EncoderWithTextPrenet.__init__cC |jSr~rget_input_embeddingsrr*r*r+rruz2SpeechT5EncoderWithTextPrenet.get_input_embeddingscC|j|dSr~rset_input_embeddingsryvaluer*r*r+rz2SpeechT5EncoderWithTextPrenet.set_input_embeddingsNr-r/rrrrr7c Cs$||}|j||||||d}|Sr r r r*r*r+rs  z%SpeechT5EncoderWithTextPrenet.forwardr)rrrrrrmrrrr=rrrQrrrrrr*r*r|r+r }s2 r cr)SpeechT5EncoderWithoutPrenet This wrapper class is a helper class to correctly load pretrained checkpoints when used in combination with [`SpeechT5Model`]. rzc"t|t||_|dSr~)rlrmrrrrr|r*r+rm   z%SpeechT5EncoderWithoutPrenet.__init__Nr-r/rrrrr7cCs|j||||||dSr )r)ryr-r/rrrrr*r*r+rs z$SpeechT5EncoderWithoutPrenet.forwardrrr*r*r|r+rs.  rceZdZdZdeffdd Z            ddeejdeej deejd eej d eej d eej d ee d ee dee dee dee deej de eeffddZZS)SpeechT5Decoderzt Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`SpeechT5DecoderLayer`] rzcsFtj|_tfddtjD|_d|_ | dS)Ncr))r)rrrr*r+rDz,SpeechT5Decoder.__init__..F) rlrmdecoder_layerdroprrrrOdecoder_layersrBrrrr|rr+rms   zSpeechT5Decoder.__init__Nrr/rrrcross_attn_head_maskrlrrrrrr7c  Cs| dur| n|jj} | dur| n|jj} |dur|n|jj}| dur$| n|jj} |dd} |jr?|jr?|r?t dd}|rR|durRt t |jdt |jd}|rct |t rct dt |}|durk|nd}t|| ||}|dur|durt||j| dd}tpt|}| rd nd}| rd nd}| r|durd nd}t||gd d gD](\}}|dur|dt|jkrtd |d t|jd|ddqt|jD]Y\}}| r||f}d}|jrtg}||jk}|r|sq||||||dur||nd|dur||nd|| || d }|d}| r3||df}|dur3||df}q| r<||f}| sMt dd|||||fDSt|||||dS)aA 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). 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. Nr#zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...FrzPassing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. You should pass an instance of `EncoderDecoderCache` instead, e.g. `past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`.r)rr*rr zThe `z` should be specified for rr)rrrrlrrrrrcsrr~r*rr*r*r+rysz*SpeechT5Decoder.forward..)rrlrrcross_attentions)rzrrrrrrr logger warning_oncer r rmrfrom_legacy_cachernrrrFrrr2rTrBr'rrrJrr)ryrr/rrrr rlrrrrrrorrrrall_cross_attentions attn_mask mask_namer decoder_layerrrrr*r*r+rsH              zSpeechT5Decoder.forward NNNNNNNNNNNNrrrrrrmrrr=r<rr rQrrrrrr*r*r|r+rsT       rcseZdZdZdeffdd Z             ddeejdeej deejd eej d eej d eej d eej d ee dee dee dee dee deej de eeffddZZS)SpeechT5DecoderWithSpeechPrenetz Wrapper around SpeechT5Decoder that applies SpeechT5SpeechDecoderPrenet to convert log-mel filterbanks to hidden features. rzcrr~)rlrmr>rrwrapped_decoderrrr|r*r+rmr z(SpeechT5DecoderWithSpeechPrenet.__init__Nr-r/rrrNrr rlrrrrrr7cCs2|||}|j|||||||| | | | | d }|SN) rr/rrrr rlrrrrrrr,)ryr-r/rrrNrr rlrrrrrdecoder_hidden_statesrr*r*r+rs z'SpeechT5DecoderWithSpeechPrenet.forward) NNNNNNNNNNNNNr*r*r*r|r+r+sZ       r+cseZdZdZdeffdd ZddZddZ dd ee j d ee j d ee j d ee j dee j dee j dee deedeedeedeedee j deeeffddZZS)SpeechT5DecoderWithTextPrenetz{ Wrapper around SpeechT5Decoder that applies SpeechT5TextDecoderPrenet to convert input tokens to hidden features. rzcrr~)rlrmrgrrr,rrr|r*r+rmr z&SpeechT5DecoderWithTextPrenet.__init__cCrr~rrr*r*r+rruz2SpeechT5DecoderWithTextPrenet.get_input_embeddingscCrr~rrr*r*r+rrz2SpeechT5DecoderWithTextPrenet.set_input_embeddingsNr-r/rrrr rlrrrrrr7c Cs8||||\} }|j| |||||||| | | | d }|Sr-r.)ryr-r/rrrr rlrrrrrr/rr*r*r+rs z%SpeechT5DecoderWithTextPrenet.forwardr))rrrrrrmrrrrr=r<rr rQrrrrrr*r*r|r+r0sX      r0cr)SpeechT5DecoderWithoutPrenetrrzcrr~)rlrmrr,rrr|r*r+rmrz%SpeechT5DecoderWithoutPrenet.__init__Nr-r/rrrr rlrrrrrr7c Cs&|j||||||||| | | | d } | Sr-)r,)ryr-r/rrrr rlrrrrrrr*r*r+rsz$SpeechT5DecoderWithoutPrenet.forwardr)r*r*r*r|r+r1sT       r1cs\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://huggingface.co/papers/1710.08969), adapted for multi-head attention. rzcst|j|_|j|_dSr~)rlrmguided_attention_loss_sigmasigmaguided_attention_loss_scalescalerr|r*r+rm#s  z-SpeechT5GuidedMultiheadAttentionLoss.__init__r input_masks output_masksr7cCsX||||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 r#rr)_make_guided_attention_masksrrrrr masked_selectr6)ryrr7r8guided_attn_masksmaskslosseslossr*r*r+r(s  z,SpeechT5GuidedMultiheadAttentionLoss.forwardc Cs|d}|d}tjt||jd|jdf|d}tt||D]\}\}} ||| |j|||d| d|f<q#| dS)Nr#rr)) rMrrPrTr%rr2_make_guided_attention_maskr4r) ryr7r8rr^r-r;rilenolenr*r*r+r9As  $& zASpeechT5GuidedMultiheadAttentionLoss._make_guided_attention_maskscCsdtjtj||dtj||ddd\}}||}||}dt||d d|dS)Nr)xy)indexingrr)rmeshgridrSrr)r; output_lengthr4rgrid_ygrid_xr*r*r+r?Ls    $z@SpeechT5GuidedMultiheadAttentionLoss._make_guided_attention_mask)rrrrrrmrr= BoolTensorrrr9rr?rr*r*r|r+r2s  r2csbeZdZdZdeffdd Z ddejdejdejd ejd ejd e ejd ej fd dZ Z S)SpeechT5SpectrogramLossz; Loss computation used by SpeechT5ForTextToSpeech. rzcsTt|j|_|j|_|j|_t|_tt dd|_ |jr(t ||_ dSdS)Ng@) pos_weight) rlrmuse_guided_attention_lossguided_attention_loss_num_headsr.r l1_criterionrrr bce_criterionr2attn_criterionrr|r*r+rm]s z SpeechT5SpectrogramLoss.__init__Nr/r_r`ralabelsr!r7cs<|dk}||}||}||}||||}|dddddf} tj| dt| dd| jgdd} | ddddf| } || }|| } || } j rtjfdd|Ddd} |dk}|dddddf}j dkr|ddj ddj f} | ||}| |7} | S)Nr0rrrrcs"g|] }|dddjfqSr~)rL)rBxrr*r+rDs"z3SpeechT5SpectrogramLoss.forward..) r:rMrrrVrrrrNrKr.rO)ryr/r_r`rarPr!r&l1_lossr< stop_labelsbce_lossr>attnr7r8 attn_lossr*rr+ris(   .   zSpeechT5SpectrogramLoss.forwardr~) rrrrrrmrr<r=rrrrr*r*r|r+rIXs&rIzv The bare SpeechT5 Encoder-Decoder Model outputting raw hidden-states without any specific pre- or post-nets.  custom_introc$s0eZdZ  d!dedeejdeejffdd ZddZd d Z d d Z d dZ e               d"dee jdee jdee jdee jdee jdee jdee jdeeee jdeedeedee jdeedeedeedee jdeee jeff dd ZZS)# SpeechT5ModelNrzencoderdecodercsJt|||_|durt|n||_|durt|n||_|dS)z encoder (`PreTrainedModel`, *optional*): The encoder model to use. decoder (`PreTrainedModel`, *optional*): The decoder model to use. N)rlrmrzrrZr1r[r)ryrzrZr[r|r*r+rms  zSpeechT5Model.__init__cCs0t|jtr |jSt|jtr|jStr~)rmrZr rr[r0NotImplementedErrorrr*r*r+rs    z"SpeechT5Model.get_input_embeddingscCs8t|jtr |j|t|jtr|j|dSdSr~)rmrZr rr[r0rr*r*r+rs   z"SpeechT5Model.set_input_embeddingscCrvr~)rZrr*r*r+ get_encoderszSpeechT5Model.get_encodercCs 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)rmrZrrr rr*r*r+r s z$SpeechT5Model.freeze_feature_encoderr-r/decoder_input_valuesdecoder_attention_maskrdecoder_head_maskr encoder_outputsrlrrNrrrrr7cCsh| 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://huggingface.co/papers/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/rrrrrrrrrN) r-r/rrrr rlrrrrr)rrlr/decoder_attentionsr!encoder_last_hidden_staterencoder_attentionsr*)rzrrrrrZrmrrTrrr"r%r[r+rrrlrrr!)ryr-r/r_r`rrar rbrlrrNrrrrr decoder_argsdecoder_outputsr*r*r+rsp*    zSpeechT5Model.forwardr;NNNNNNNNNNNNNNN)rrrrrrrrmrrr]r rrrr<r=rr rQrrrrr*r*r|r+rYs~      rYzB SpeechT5 Model with a speech encoder and a text decoder. c$seZdZdgZdeffdd 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 ede ede ede ede ede e jde e jdeeeff d d!ZZS)#SpeechT5ForSpeechToTextz#text_decoder_postnet.lm_head.weightrzc\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.) rlrmrdr'r}rr0rYrrqtext_decoder_postnetr)ryrzspeech_encoder text_decoderr|r*r+rm>     z SpeechT5ForSpeechToText.__init__cCrr~rr]rr*r*r+r]Rruz#SpeechT5ForSpeechToText.get_encodercCrr~r get_decoderrr*r*r+rrUruz#SpeechT5ForSpeechToText.get_decodercC|jdSr^r]rr rr*r*r+r Xz.SpeechT5ForSpeechToText.freeze_feature_encodercCrr~)rlrwrr*r*r+rw_ruz-SpeechT5ForSpeechToText.get_output_embeddingscCrr~)rlrzrxr*r*r+rzbrz-SpeechT5ForSpeechToText.set_output_embeddingsNr-r/decoder_input_idsr`rrar rbrlrrrrrPrr7cCs| dur| n|jj} |dur|durt||jj|jj}|j||||||||| | | | d|d}||d}d}|durMt}||d|jj |d}| sc|f|dd}|dura|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]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or 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://huggingface.co/papers/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" ... ) # 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/r_r`rrar rbrlrrrrrrr#r) r>rarlr/rcr!rdrre)rzrr,r!r"rrlrrrdrrlr/rcr!rdrre)ryr-r/rvr`rrar rbrlrrrrrPrrrar>loss_fctoutputr*r*r+resR[zSpeechT5ForSpeechToText.forwardrh)rrr_tied_weights_keysrrmr]rrr rwrzrrrr=r<rrr rQrrrrr*r*r|r+ri6sr      ri?r14@FmodelrN 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#)rr/rrrlrrrrcsg|]}|vr|qSr*r*rresult_spectrogramr*r+rDR rz$_generate_speech..rcsg|]}|qSr*r*rrr*r+rD[ r) batch_firstcs&g|]}tdt|qSr)r9rr:r)spectrogram_lengths waveformsr*r+rDu s&)*r'rzr!r9rrrZrrmrrr"r%r.r$r?r[r,rXrrr!squeezerlspeech_decoder_postnetr\rsigmoidr]rMrJrNrOrTstackrflattenr^rrrnn pad_sequence)"r|r-rNr/r}r~rrrrrr encoder_outrdmaxlenminlenoutput_sequence spectrogramr!rlrr/ decoder_outlast_decoder_outputspectrumnew_spectrogramprobmeet_thresholds meet_indexes spectrograms meet_indexrrwaveform_lengthsr*)rrrr+_generate_speechs       $   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 ede ede ede ede ede e jde 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)/SpeechT5ForTextToSpeechr rzcrj)Nrka 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.) rlrmrdr'r}r r+rYrr[rr)ryrz text_encoderspeech_decoderr|r*r+rm roz SpeechT5ForTextToSpeech.__init__r7cCsdSr r*)clsr*r*r+ can_generate sz$SpeechT5ForTextToSpeech.can_generatecCrr~rprr*r*r+r] ruz#SpeechT5ForTextToSpeech.get_encodercCrr~rqrr*r*r+rr ruz#SpeechT5ForTextToSpeech.get_decoderNr/r_r`rrar rbrlrrrrrNrPrSrcCs| dur| n|jj} |dur"|durt||jj|\}}|jjr"d} |j||||||||| | || | d|d}||d\}}}d}|durUt|j}|||||||j}| sk|f|dd}|duri|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://huggingface.co/papers/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/r_r`rrar rbrlrrNrrrrrr r>rrlr/rcr!rdrre)rzrr3r.rKrrrIr!rrlr/rcrdrre)ryr r/r_r`rrar rbrlrrrrrNrPrSrrr_r`rar> criterionrxr*r*r+r sfM   zSpeechT5ForTextToSpeech.forwardrzr1r{Fr}r~rrrrc 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.rrIr'r) ryr r/rNr}r~rrrrkwargsr]r*r*r+generate0 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) ryr rNr/r}r~rrrrr]r*r*r+generate_speech s(I z'SpeechT5ForTextToSpeech.generate_speechNNNNNNNNNNNNNNNNNNNrzr1r{NFF)rrrrrrm classmethodrQrr]rrrrrr<r=rrr rrrrrrrrrrr*r*r|r+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 ede ede ede ede ede e j de 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)-SpeechT5ForSpeechToSpeechrzcs@t|t|}t|}t||||_t||_|dSr~) rlrmrr+rYrr[rr)ryrzrmrr|r*r+rm s   z"SpeechT5ForSpeechToSpeech.__init__cCrr~rprr*r*r+r] ruz%SpeechT5ForSpeechToSpeech.get_encodercCrr~rqrr*r*r+rr ruz%SpeechT5ForSpeechToSpeech.get_decodercCrsr^rtrr*r*r+r  ruz0SpeechT5ForSpeechToSpeech.freeze_feature_encoderNr-r/r_r`rrar rbrlrrrrrNrPrSrr7cCs| dur| n|jj} |dur|durt||jj|\}}|j||||||||| | || | d|d}||d\}}}d}| sR|f|dd}|durP|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]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or 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://huggingface.co/papers/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" ... ) # 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]) ``` NTrrrr)rzrr3r.rrrrlr/rcr!rdrre)ryr-r/r_r`rrar rbrlrrrrrNrPrSrrrCrrar>rxr*r*r+r sNT z!SpeechT5ForSpeechToSpeech.forwardrzr1r{Fr}r~rrrrc 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]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or 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))rrPrr) ryr-rNr/r}r~rrrrr*r*r+r sJz)SpeechT5ForSpeechToSpeech.generate_speechrr)rrrrrmr]rrr rrrr=r<rrr rQrrrrrrrrrr*r*r|r+r s           rcs@eZdZdfdd ZdddZd d Zd d Zd dZZS)HifiGanResidualBlockrrr皙?csbt|_tfddttD_tfddttD_dS)Nc s2g|]}tjd||dqSr)rjdilationrrrq get_paddingrchannelsrriryr*r+rD s z1HifiGanResidualBlock.__init__..c s*g|]}tjddddqSrrrA)rriryr*r+rD s  ) rlrmleaky_relu_sloperrrOrTconvs1convs2)ryrrirrr|rr+rm s     zHifiGanResidualBlock.__init__rcCs|||dSrr*)ryrirr*r*r+r rz HifiGanResidualBlock.get_paddingcCsLtjj}ttjjdrtjjj}|jD]}||q|jD]}||qdSNr)rrrrrrrryrrr*r*r+apply_weight_norm s     z&HifiGanResidualBlock.apply_weight_normcCs4|jD]}tj|q|jD]}tj|qdSr~)rrrremove_weight_normrryrr*r*r+r s  z'HifiGanResidualBlock.remove_weight_normcCsXt|j|jD]"\}}|}tj||j}||}tj||j}||}||}q|Sr~)r2rrrrO leaky_relur)ryrconv1conv2rr*r*r+r s zHifiGanResidualBlock.forward)rrrr) rrrrmrrrrrr*r*r|r+r s   rz HiFi-GAN vocoder. cspeZdZUeed<dZdeffdd ZdejfddZ dd Z d d Z e d d de jde jfddZZS)SpeechT5HifiGanrzrc 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)rirjrrrr6)rlrmrTresblock_kernel_sizes num_kernelsupsample_rates num_upsamplesrrq model_in_dimupsample_initial_channelconv_prer upsamplerrr2upsample_kernel_sizesrXConvTranspose1d resblocksrOresblock_dilation_sizesrr conv_postrrrPrVr)ryrzr upsample_raterirrr|r*r+rm0 s>        zSpeechT5HifiGan.__init__rcCsJt|tjtjfr!|jjjd|jjd|j dur#|j j dSdSdS)zInitialize the weights.r1rN) rmrrqrrrrrzrrkr)ryrr*r*r+rV s  zSpeechT5HifiGan._init_weightscCs`tjj}ttjjdrtjjj}||j|jD]}||q|jD]}|q"||j dSr) rrrrrrrrrrrr*r*r+r] s      z!SpeechT5HifiGan.apply_weight_normcCsLtj|j|jD]}tj|q |jD]}|qtj|jdSr~)rrrrrrrrr*r*r+ri 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. rWr7cCs|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,)`. rrrrr#)rznormalize_beforerr6rrrrrOrrrOrrrrrrrtanhrr)ryr is_batchedrr res_statejwaveformr*r*r+rq s,         zSpeechT5HifiGan.forward)rrrrrrrmrrrrrrrr=rrr*r*r|r+r' s &  r)rirrrYrr)rNrr)grrtypingrrnumpyrHrrtorch.nnrrr activationsr cache_utilsr r r generationr integrations.deepspeedrintegrations.fsdprmodeling_attn_mask_utilsrrmodeling_layersrmodeling_outputsrrrrrmodeling_utilsrrrrrutils.deprecationrconfiguration_speecht5rr get_loggerrr"_HIDDEN_STATES_START_POSITIONrr9r,r3rrr<ndarrayrergrrrrrrrrrr rr>rSr[rcrgrqr{rrrrrrr rrr+r0r1r2rIrYrir=rQrrrrr__all__r*r*r*r+s           xE- )4(,'=f+%* B26-;=:   jx> w