0;ji;)ddlmZmZmZddlZddlmZddlZddl Z ddl m cm Z ddl m Z mZddlmZddlmZddefd Zd ed ed ed efdZdeded efdZ ddejd edefdZGdde jZdS))ListUnionOptionalN) as_strided)nnTensor)VoiceEncConfig)melspectrogramseq_lenc<|td|D}n|td|DksJt|dtr d|D}d}t|dtjr|}|dj}n d|D}t ||g|djddR}tj|||dj |}t|D]!\}}|||d| df<"|S) a Given a list of length B of array-like objects of shapes (Ti, ...), packs them in a single tensor of shape (B, T, ...) by padding each individual array on the right. :param arrays: a list of array-like objects of matching shapes except for the first axis. :param seq_len: the value of T. It must be the maximum of the lengths Ti of the arrays at minimum. Will default to that value if None. :param pad_value: the value to pad the arrays with. :return: a (B, T, ...) tensor Nc34K|]}t|VdSNlen.0arrays g/root/voice-cloning/.venv/lib/python3.11/site-packages/chatterbox/models/voice_encoder/voice_encoder.py zpack..s(55Uc%jj555555c34K|]}t|VdSrrrs rrzpack..s(==Uc%jj======rrc6g|]}tj|S)nprrs r zpack.."s 666e"(5//666rc6g|]}tj|Sr)torch as_tensorrs rrzpack..*s">>>e5?5))>>>rr )dtypedevice) max isinstancelistrrr!rshapefullr enumeratesize) arraysr pad_valuer!tensors packed_shape packed_tensoritensors rpackr0sJ55f55555#==f=========&)T""766v666F&)U\**?">>v>>>LL'AGAJ,NW]^^^Mw''33 6,2 a&++a..()) rn_framesstep min_coveragehpc|dksJ|j}tt||z |zd|\}}|dks|||z z|z |kr|dz }|||dz zz}||fSNrr )ve_partial_framesdivmodr")r1r2r3r4win_sizen_wins remaindertarget_ns r get_num_winsr=6s a<<<<#Hs8h#6#=qAA4HHFI {{yHtO4@LPP! $&1*--H 8 roverlapratecd|cxkrdksnJ|-ttj|jd|z z}n1ttj|j|z |jz }d|cxkr |jksnJ|Sr6)intrroundr7 sample_rate)r>r?r4 frame_steps rget_frame_steprEEs     !       |""6!g+"FGGHH 2>D#8Br?r3rD n_partials target_lenr%rPpartialss rstride_as_partialsrTTs& | q r22J*#c((J bQQJ CHHnc27JS,A2;+OQR#S#STUU c#hh  +:+ **RZs* + +C -r{ ;E{1~ *CKNCKNKG#ug..H OrcDeZdZeffd ZedZdejfdZ ddej de fd Z e d ejfd Ze d ejd ejfdZ ddee eejffdZ ddeejdee fdZxZS) VoiceEncoderc t||_tj|jj|jjdd|_|jr|j tj |jj|jj |_ tj tjdgd|_tj tjdgd|_dS)NT) num_layers batch_firstg$@) requires_gradg)super__init__r4rLSTMrMve_hidden_sizelstmflatten_lstm_paramsflatten_parametersLinearspeaker_embed_sizeproj Parameterrr/similarity_weightsimilarity_bias)selfr4 __class__s rr]zVoiceEncoder.__init__xs GDG,dg.DQR`deee ! + I ( ( * * *Idg4dg6PQQ "$elC5.A.AQU!V!V!V!|EL#,?,?tTTTrcNt|jSr)next parametersr!)ris rr!zVoiceEncoder.devicesDOO%%&&--rmelsc|jjri|dks|dkr9t d|d|||\}\}}||d}|jjrtj |}|tj |ddz S)a Computes the embeddings of a batch of partial utterances. :param mels: a batch of unscaled mel spectrograms of same duration as a float32 tensor of shape (B, T, M) where T is hp.ve_partial_frames :return: the embeddings as a float32 tensor of shape (B, E) where E is hp.speaker_embed_size. Embeddings are L2-normed and thus lay in the range [-1, 1]. rr zMels outside [0, 1]. Min=z, Max=Tdimkeepdim) r4normalized_melsminr" Exceptionr`re ve_final_reluFrelurlinalgnorm)rirn_hidden raw_embedss rforwardzVoiceEncoder.forwards 7 " X Q$((**q..V VV$((**VVWW W4;FAYYvbz** 7  , ++JEL--ja-NNNNrrFNrGr?ctj|r|n|}t||jt fd|D\}}t ||dz } | dkrqtj|d| jj fdtj } tj || |j fd}fdt ||DtfdDsJtjt!t#jt'|pt'z } tj fd| Ddt#jdgt#j|f} fd t | d d | dd D} tj| } | tj| dd z }|S)z Computes the embeddings of a batch of full utterances with gradients. :param mels: (B, T, M) unscaled mels :return: (B, E) embeddings on CPU c3FK|]}t|jVdSr)r=r4)rlrDr3ris rrz)VoiceEncoder.inference..s6'm'm_` Q LRVRY(Z(Z'm'm'm'm'm'mrr r)r rrctg|]4\}}t|D]}||z|zjjz 5Sr)ranger4r7)rrH n_partialr.rDris rrz*VoiceEncoder.inference..se   YyAQAQ  <= JJ1J JJ K    rc3DK|]}dj|jkVdS)rNr%)rpartialrSs rrz)VoiceEncoder.inference..s1NN'8A;$ 5NNNNNNrc&g|] }|Srr)rbatchris rrz*VoiceEncoder.inference..s!#V#V#VEDDKK#V#V#VrcRg|]#\}}tj||d$S)rr)rmean)rstartendpartial_embedss rrz*VoiceEncoder.inference..s5rrrzucejc !:BBBrrrrNrpTrq)r is_tensortolistrEr4zipr"r(r&rMrOcattor!allstackrArceilrchunkcpurLcumsumrzr{)rirnmel_lensr>r?r3 batch_sizerQ target_lenslen_diffpadn_chunksslicesr~embedsrDrrSs` ` @@@r inferencezVoiceEncoder.inferencesa).(A(AO8??$$$x$GT47;; "%'m'm'm'm'm'mdl'm'm'm"n K{##diill2 a<<*diillHdg6FGRWR_```C9dCFF4;$7$78a@@@D     "%dJ"7"7   NNNNXNNNNNNNN;x((rws8}} 0Kc(mmLMMNN#V#V#V#VX^^H=U=U#V#V#V\]^^^bbdd!bi &;&; <==rrrrUXY_`cac`cYdflmnmomofpUqUqrrr [,, el// 4/PPP r utt_embedsc|jdksJtj|d}|tj|dz S)z Takes an array of L2-normalized utterance embeddings, computes the mean embedding and L2-normalize it to get a speaker embedding. r)axis)ndimrrrzr{)rs rutt_to_spk_embedzVoiceEncoder.utt_to_spk_embedsF !####WZa000 BINN:q9999rembeds_xembeds_yc|jdkr|nt|}|jdkr|nt|}||zS)z` Cosine similarity for L2-normalized utterance embeddings or speaker embeddings r )rrVr)rrs rvoice_similarityzVoiceEncoder.voice_similaritysW (}1188|7T7TU]7^7^'}1188|7T7TU]7^7^(""rF c ttrLdDtfdDs JddD}tt j5|j|j|fd|i| }dddn #1swxYwY|r| |n|S)a Convenience function for deriving utterance or speaker embeddings from mel spectrograms. :param mels: unscaled mels strictly within [0, 1] as either a (B, T, M) tensor or a list of (Ti, M) arrays. :param mel_lens: if passing mels as a tensor, individual mel lengths :param as_spk: whether to return utterance embeddings or a single speaker embedding :param kwargs: args for inference() :returns: embeds as a (B, E) float32 numpy array if is False, else as a (E,) array c6g|]}tj|Sr)rasarrayrrHs rrz1VoiceEncoder.embeds_from_mels..s 444BJsOO444rc3\K|]&}|jddjdkV'dS)r rNr)rmrns rrz0VoiceEncoder.embeds_from_mels..s9DD!qwqzT!W]1%55DDDDDDrzMels aren't in (B, T, M) formatc(g|]}|jdS)rrrs rrz1VoiceEncoder.embeds_from_mels..s555 ! 555rrN) r#rrr0rinference_moderrr!numpyr)rirnras_spkrkwargsrs ` rembeds_from_melszVoiceEncoder.embeds_from_melssM dD ! ! 44t444DDDDDtDDDDD g gFg g gD55555H::D ! # # q q' (<(>CCwavs trim_top_dbc jjkrfd|D}rfd|D}d|vrd|d<fd|D}j|f||d|S)z Wrapper around embeds_from_mels :param trim_top_db: this argument was only added for the sake of compatibility with metavoice's implementation cTg|]$}tj|jjd%S) kaiser_fast)orig_sr target_srres_type)librosaresampler4rC)rwavrCris rrz1VoiceEncoder.embeds_from_wavs..sC kTWEXcpqqqrc^g|])}tj|d*S))top_dbr)reffectstrim)rrrs rrz1VoiceEncoder.embeds_from_wavs.. s3UUUGO(([(AA!DUUUrr?g?cDg|]}t|jjSr)r r4T)rwris rrz1VoiceEncoder.embeds_from_wavs..s(;;;q$'**,;;;r)rr)r4rCr)rirrCrrrrrns` ` ` rembeds_from_wavszVoiceEncoder.embeds_from_wavss $'- - -D  VUUUUPTUUUD    F6N;;;;d;;;$t$TZ&ZZZSYZZZr)rFNrGN)NFr)Frr)__name__ __module__ __qualname__r r]propertyr!r FloatTensorrrfloatr staticmethodrndarrayrrrrrrr __classcell__)rjs@rrVrVws(.** U U U U U U..X.OE-OOOO.%%el%%%%%N:RZ:::\:#2:####\#^`KK&$rz"223KKKK<%' [[2:[ e_ [[[[[[[[rrV)Nr)rFNrG)typingrrrrrnumpy.lib.stride_tricksrrrtorch.nn.functionalr functionalrxrconfigr melspecr rAr0rr=rErrTModulerVrrrrs)(((((((((...... """"""#############L                    $           F[[[[[[[[[[29[[[[[[[[[[r