o j_ @srddlZddlZddlmZddlmmZddlmZm Z m Z ddl m Z ddl mZmZddZddZGd d d ejZGd d d e ZGd ddejZGdddejZddZGdddejZGdddejZedkreddddddZeeddddejd d!d"ed#d gejd$d%d"ed&d'gZe edddejd(d)d"ed#dgdSdS)*N) GPT2ConfigGPT2PreTrainedModelLogitsProcessorList)!CausalLMOutputWithCrossAttentions)AttentionBlockTypicalLogitsWarpercCs"tj|jd|jd|f|jdS)Nrdevice)torchzerosshaper )rangedimrY/home/kuhnn/.local/lib/python3.10/site-packages/TTS/tts/layers/tortoise/autoregressive.pynull_position_embeddings "rcCs$|ot|t|d|ddjfS)Nr)lenr )trrr_ps$rcs(eZdZdZfddZddZZS)ResBlockzA Basic residual convolutional block that uses GroupNorm. c sZtttj||dddt|d|ttj||dddt|d||_dS)Nr kernel_sizepadding)super__init__nn SequentialConv1d GroupNormReLUnet)selfchan __class__rrrs  zResBlock.__init__cCst|||SN)Frelur$)r%xrrrforward$szResBlock.forward)__name__ __module__ __qualname____doc__rr- __classcell__rrr'rrs  rcs`eZdZfddZddZd ddZ              d dd Zed d ZZ S)GPT2InferenceModelcs6t|||_||_||_t|||_||_dSr)) rr transformertext_pos_embedding embeddingsrr lm_headkv_cache)r%configgpt text_pos_embr6normlinearr8r'rrr)s  zGPT2InferenceModel.__init__cCs ||_dSr))cached_mel_emb)r%mel_embrrr store_mel_emb1s z GPT2InferenceModel.store_mel_embNcKs|dd}|js d}|r'|dddfd}|dur'|dddfd}|dd}|dd}|durZ|durZ|dd}||dkd|rY|dddfd}nd}|||d|||dS) Ntoken_type_idsattention_mask position_idsrr use_cache) input_idspast_key_valuesrErDrCrA)getr8 unsqueezelongcumsum masked_fill_)r%rFrGkwargsrArCrDrrrprepare_inputs_for_generation4s.   z0GPT2InferenceModel.prepare_inputs_for_generationcCs`|jdusJ|dus J| dusJ|dur|n|jj}|jjd}|jddkri|dd|df}||}|||}|jjd|jdkr\|j|jd|jjdd}n|j}tj||gdd}n||}||j |jd||j }|j |||||||| | | | |d }|d}| |}|s|f|ddSt d||j|j|j|jdS)Nrrr) inputs_embedsrGrCrArD head_maskencoder_hidden_statesencoder_attention_maskrEoutput_attentionsoutput_hidden_states return_dict)losslogitsrG hidden_states attentionscross_attentions)r>r9use_return_dictr r6r5repeat_interleaver catget_fixed_embeddingr r4r7rrGrYrZr[)r%rFrGrCrArDrQrPrRrSlabelsrErTrUrVmel_len text_inputstext_embr?embtransformer_outputsrY lm_logitsrrrr-RsV    "  zGPT2InferenceModel.forwardcstfdd|DS)a> This function is used to re-order the :obj:`past_key_values` cache if :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step. c3s&|]}tfdd|DVqdS)c3s$|] }|d|jVqdS)rN) index_selecttor ).0 past_statebeam_idxrr s"z>GPT2InferenceModel._reorder_cache...Ntuple)ri layer_pastrkrrrms  z4GPT2InferenceModel._reorder_cache..rn)pastrlrrkr_reorder_caches z!GPT2InferenceModel._reorder_cacher))NNNNNNNNNNNNNN) r.r/r0rr@rNr- staticmethodrrr2rrr'rr3(s*    Dr3cs.eZdZ    dfdd ZddZZS) ConditioningEncoderFc s`tg}tj||dd|_t|D] }|t||qtj||_ ||_ ||_ ||_ dS)Nr)r) rrrr!initrappendrr attnrdo_checkpointingmean) r%spec_dim embedding_dim attn_blocksnum_attn_headsrzr{ryar'rrrs    zConditioningEncoder.__init__cCs<||}||}|jr|jddS|dddddfS)NrOr)rwryr{)r%r,hrrrr-s   zConditioningEncoder.forward)rurvFFr.r/r0rr-r2rrr'rrtsrtcs.eZdZdfdd ZddZddZZS) LearnedPositionEmbeddings{Gz?cs0tt|||_|jjjjd|ddS)Nr{std)rrr Embeddingrdweightdatanormal_)r%seq_len model_dimrwr'rrrs z"LearnedPositionEmbeddings.__init__cCs"|jd}|tjd||jdS)Nrrr )r rdr aranger )r%r,slrrrr-s z!LearnedPositionEmbeddings.forwardcCs"|tjd||d|d|S)Nrr r)rdr r)r%inddevrrrr_rz-LearnedPositionEmbeddings.get_fixed_embedding)r)r.r/r0rr-r_r2rrr'rrsrc Cslddlm}m}|d||||||||| d}||} | `tjt|d| _| `| t||t||ddfS)z7 GPT-2 implemented by the HuggingFace library. r)r GPT2Model vocab_size n_positionsn_ctxn_embdn_layern_headgradient_checkpointingrErON) transformersrrwpe functoolspartialrwter) layersrheadsmax_mel_seq_lenmax_text_seq_len checkpointingrr gpt_configr:rrrbuild_hf_gpt_transformers* rcs&eZdZdfdd ZddZZS) MelEncoderPrcst|_ttj|ddddtjfddt|Dtjddddddtd dttjfd dt|Dtjdddddtd ttjfd dt|D |_ d|_ dS) Nrvrrrcg|]}tdqS)rvrri_channelsrr z'MelEncoder.__init__..r)rstridercr)rrrrrrrrrcsg|]}tqSrrrrrrrs) rrrrr r!rr"r#encoder reduction)r%r mel_channelsresblocks_per_reductionr'rrrs  zMelEncoder.__init__cCs"|jD]}||}q|dddS)Nrrr)rpermute)r%r,errrr-s  zMelEncoder.forward)rrrrrr'rrsrcseZdZ          dfdd Zd ddZddZddZ   d!ddZddZ   d"ddZ    d#ddZ Z S)$ UnifiedVoicerxrrN  FTcst||_| dur||n| |_d|_| |_| |_| |_||_||_ ||_ ||_ ||_ ||_ ||_td||d|_t|j|d||_|rSt|j||_nt|dd|_t||||j d|j |j d|\|_|_|_|_|_| rtjtdd|dd d |_tjtdd|dd d |_ nd|_d|_ t!||_"t#||j|d|_$t#||j|_%|jg}|r|&|j|D] }|j'j(j)d dd qdS) a Args: layers: Number of layers in transformer stack. model_dim: Operating dimensions of the transformer heads: Number of transformer heads. Must be divisible by model_dim. Recommend model_dim//64 max_text_tokens: Maximum number of text tokens that will be encountered by model. max_mel_tokens: Maximum number of MEL tokens that will be encountered by model. max_conditioning_inputs: Maximum number of conditioning inputs provided to the model. If (1), conditioning input can be of format (b,80,s), otherwise (b,n,80,s). mel_length_compression: The factor between and . Used to compute MEL code padding given wav input length. number_text_tokens: start_text_token: stop_text_token: number_mel_codes: start_mel_token: stop_mel_token: train_solo_embeddings: use_mel_codes_as_input: checkpointing: Nrr)rr)rrrT) requires_gradrr)*rrnumber_text_tokensstart_text_tokenstop_text_tokennumber_mel_codesstart_mel_tokenstop_mel_tokenrrmax_mel_tokensmax_text_tokensrmax_conditioning_inputsmel_length_compressionrtconditioning_encoderrrtext_embedding mel_embeddingrrr:mel_pos_embeddingr5mel_layer_pos_embeddingtext_layer_pos_embedding Parameterr randnmel_solo_embeddingtext_solo_embedding LayerNorm final_normLinear text_headmel_headrxrrr)r%rrrrrrrrrrrrtrain_solo_embeddingsuse_mel_codes_as_inputrtypesr6moduler'rrrs^ &    zUnifiedVoice.__init__c Cs`|j|jd}t|j|||j|j|jddd}t||j|j|j |j |j |d|_ |j |j_ dS)NrFTr)r8)rrrrrrr3r:rrrrinference_modelr)r%r8 seq_lengthrrrrpost_init_gpt2_configbs*  z"UnifiedVoice.post_init_gpt2_configcCs(tj|d|d}tj|d|d}||fS)N)rrvaluerr)r*pad)r%input start_token stop_tokeninptarrrr build_aligned_inputs_and_targets{sz-UnifiedVoice.build_aligned_inputs_and_targetscCsTtj||jdd}tt|D]}||d}||jdkr'|j|||df<q|S)a" Given mel tokens that are derived from a padded audio clip and the actual lengths of each batch element in that audio clip, reformats the tokens with STOP_MEL_TOKEN in place of the zero padding. This is required preformatting to create a working TTS model. trunc) rounding_moderrBN)r divrrrr r)r%mel_input_tokens wav_lengths mel_lengthsb actual_endrrrset_mel_paddings zUnifiedVoice.set_mel_paddingc Cs*|durtj|||gdd}n tj||gdd}|j|d|d} |r%| jS| jddddf} || } |rZ| dd|jd|jd|jdf| dd|jd dffS| ddd|jdf} || } | ddd} |dur| dd|jd df} || } | ddd} | | fS| S)NrrOT)rPrVrTrr)r r^r:rZlast_hidden_staterr r) r%speech_conditioning_inputs first_inputs first_head second_inputs second_head get_attns return_latentrdgpt_outenc first_logits second_logitsrrr get_logitss2  zUnifiedVoice.get_logitscCsnt|jdkr |dn|}g}t|jdD]}|||dd|fqtj|dd}|jdd}|S)NrrrO) rr rIrrxrr stackr{)r%speech_conditioning_inputcondsjrrrget_conditionings  zUnifiedVoice.get_conditioningc  Cs|dur |d|d}| rA|} |ddd| f}||j} |ddd| f}|durA|ddddd| df}|||}tj|d|jd}tj|d|jd}|d}|||j |j\}}| || |}|||j |j\}}|durt|d}n|}| |}|||}|r|j|||j||j| | d\}}| r|dddd fSn|j|||j||j| | d\}}| r|dddd fS| r|St||}t||}|||fS) a Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode (actuated by `text_first`). speech_conditioning_input: MEL float tensor, (b,1024) text_inputs: long tensor, (b,t) text_lengths: long tensor, (b,) mel_inputs: long tensor, (b,m) wav_lengths: long tensor, (b,) raw_mels: MEL float tensor (b,80,s) If return_attentions is specified, only logits are returned. If return_latent is specified, loss & logits are not computed or returned. Only the predicted latents are returned. If clip_inputs is True, the inputs will be clipped to the smallest input size across each input modality. NrrBrvrr)rr)rr)rImaxrrr*rrrrrrr5rrrrrr cross_entropyrJr{)r%speech_conditioning_latentrb text_lengths mel_codesrr text_firstraw_melsreturn_attentionsr clip_inputs max_text_len max_mel_lenr text_targetsrc mel_targetsmel_inpr? text_logits mel_logits loss_textloss_melrrrr-st         zUnifiedVoice.forward?c Kstj|d|jd}|||j|j\}} ||||} |d} tj | | gdd} |j | tj | j d| j d| j dfdtj|jd} |j| dddf<| j d}|dur_| }n'||j ddkslJd| |d} |||j dd}tj | |gdd}|rtt|d gnt}|dur||jdn||}|j j|f|j|j|j|||d |}|dd|dfS) NrrrrOr) fill_valuedtyper rBzQThe number of return sequences must be divisible by the number of input sequences)mass) bos_token_id pad_token_id eos_token_id max_lengthlogits_processornum_return_sequences)r*rrrrrr5rIr r^rr@fullr rJr rrepeatrrrgenerater)r%r rb input_tokensr%max_generate_lengthtypical_sampling typical_masshf_generate_kwargsrrcrrd fake_inputs trunc_indexinputsr$r#genrrrinference_speech'sV       zUnifiedVoice.inference_speech)rrrrrrrrNrrrFTTr)T)NNFF)NTNFFT)NrNFr) r.r/r0rrrrrrr-r2r2rrr'rrsP [ * dr__main__rrvT)rrrrrrrri r)rr)highsize r)rrii2)rr)!rr torch.nnrtorch.nn.functional functionalr*rrrrtransformers.modeling_outputsr"TTS.tts.layers.tortoise.arch_utilsrrrrModulerr3rtrrrrr.r:rrandinttensorl text_forwardrrrrsN  {`