3jr SSKJrJrJr SSKrSSKrSSKJr SSKJ r \(aSSK J r J r SSK J r SSKJr S r"S S \5r"S S \5r\\\-r"SS5r"SS\5r"SS5r"SS\5r"SS5r"SS\5r"SS\5r"SS\5r"SS\5r"S S!\5r"S"S#5r"S$S%5r \\\\\\\\\\ / r!g)&) TYPE_CHECKINGAny TypedDictN)get_torch_device)snapshot_downloadWhisperForConditionalGenerationWhisperProcessor)log)get_model_path>cB\rSrSr%Sr\\S'\R\S'Sr g) AudioTensorz%Comfy's representation of AUDIO data. sample_ratewaveformN) __name__ __module__ __qualname____firstlineno____doc__int__annotations__torchTensor__static_attributes__rB/home/wildlama/comfy/ComfyUI/custom_nodes/comfy_mtb/nodes/audio.pyrrs/llrrcH\rSrSr%Sr\\S'\\\\ 4\S'\\S'Sr g) WhisperDataz# UHnTRU5v M g7fN) is_stereo).0audioclss r %MtbAudio.is_stereo../s@s}}U++s!rr ) isinstancer(anyshape)r5r-s` r r2MtbAudio.is_stereo)s> fd # #@@@ @*%++A.!3 3rr4common_sample_ratecUSnX!:waG[R"SUSU35 [RR X!S9nUU"US5S.$U$)NrzResampling audio from z to ) orig_freqnew_freqrrr)r debug torchaudio transformsResample)r4r= current_rate resamplers r resampleMtbAudio.resample3ss]+  - II(d;M:NO #--66&7I 2%eJ&78  LrcxUSRSS:Xa#US[R"USUS/SS9S.$U$)Nrr8rdimrA)r;rcat)r4s r to_stereoMtbAudio.to_stereoDsS   " "1 % *$]3!II:&j(9: Lrc[UVs/sHo"SPM sn5nUVs/sHo RX#5PM nnURU5nU(a UVs/sHo RU5PM nnXU4$s snfs snfs snf)Nr)maxrHr2rN)r5r-r4max_sample_rateresampled_audiosr2s r preprocess_audiosMtbAudio.preprocess_audiosPsH]3HI?E >DULL 0f  MM&) 8HI8HummE*8HFI?33I JsA;BBrN)rrrrr classmethod AudioDataboolr2 staticmethodrrrHrNr(tuplerTrrrr r+r+&s*44 44       4+& 4 tK $+ , 4 4rr+c.\rSrSr%SrS\S'S\S'Srg) WhisperPipelineazWhisper model pipeline.r processorr modelrN)rrrrrrrrrr r\r\as!!! ,,rr\cB\rSrSrSr\S5rSrSrSr Sr S Sjr S r g ) MTB_LoadWhisperh!Load Whisper model and processor.c*S/SQSS040SSSSS .40S .$) N model_size) tinysmallmediumz medium.enbaselargezlarge-v2zlarge-v3zlarge-v3-turbodefaultrfdownload_missingBOOLEANFzSDownload missing models if missing,otherwise they must be in ComfyUI/models/whisper)rktooltiprequiredoptionalrr5s r INPUT_TYPESMTB_LoadWhisper.INPUT_TYPESksJ ' "##(O % #  rWHISPER_PIPELINE)pipeline mtb/audioloadcjSSKJnJn [S5nSU3nXV- nUR 5(deUR 5(dPU(d [ S5eUR SS9 UR SS9 [SU3S/S QUR5S S 9 [5n[R"S US USU35 UR"UR55n UR"UR55RU5n U R5 U RS 5 XS.4$)rcrrwhisperzwhisper-z+Models not found and download_missing=FalseT)exist_okzopenai/)z *.msgpackz*.binz*.h5F)repo_idresume_downloadignore_patterns local_dirlocal_dir_use_symlinkszLoading Whisper model z on z from )r^r_) transformersr r r exists RuntimeErrormkdirras_posixrr rBfrom_pretrainedtoevalrequires_grad_) selfrerlr r whisper_dirtag model_dirdevicer^r_s r ryMTB_LoadWhisper.loads2 %Y/  %% ""$$ (8(8(:(:#"A!!4!0.!%cUO$($B'002+0 "# $ZLVHF9+ N %44Y5G5G5IJ /??     "V*    U#'8::rrN)rfF) rrrrrrVrs RETURN_TYPES RETURN_NAMESCATEGORYFUNCTIONryrrrr rarahs3+  @)L LHH);rracN\rSrSrSr\S5rSrSrSr S S\ S\ 4S jjr S r g ) MTB_AudioToText'Transcribe audio to text using Whisper.cBSSSS/[/SQ5-SS04SSS04S .0$) NrpruAUDIOauto) enfresdeitptnlruzhjakorkrmT)rwr4r&return_timestamps)sortedrrs r rsMTB_AudioToText.INPUT_TYPESsQ 1#H  '#&'0)T1B%C-  r)STRINGWHISPER_OUTPUT transcriberxrwr4cv USnUSnURnURU[5nUSn[R"SUR 35 [ UR 5S:XaURS5n[ UR 5S:XaWUR SS:XaURSS9nO4UR S S:XaURS S9nOURS5nUS n S n X-n UR S n X- n [R"S U SS35 /n/n/nSnSn[SX5GHn[UU -U 5nUUUnUU - nURU5 [R"SUSSUU - SS35 URR=(d Sn[R"S U45nU"UU SS9R R#U5n[R$"5 UR'UUR#U5SUS:XaSOUUSSSUS9 nSSS5 UR(R+WS5n/nUHnUR-S5(aUR/S5(asUSSnUR1S S!5R35(a5[5U5nUU:aUU- nUU-n URSU SS35 UnMURU5 MURU5 M UR9U5 UR;US"S#9Sn!URU!5 GM S$n"US:Xa[R"S%5 [R$"5 U"USU U SS9R R#U5n#UR=U#5Sn$UR(R+U$R?S 5RA55n%U%R-S5(aU%SSOS$n"[R"S&U"35 SSS5 S(RGU5n'U'U"UUUS).n(U'U(4$!,(df  GN"=f![6a URU5 GM"f=f!,(df  N`=f![Ba#n&[RD"S'U&35 Sn&A&NSn&A&ff=f)*rr^r_rzProcessed waveform shape: rr rKr8rAudio duration: .2fszProcessing chunk z.1fzs - ir) sampling_ratereturn_tensorsrrN?)attention_masktaskr&rno_repeat_ngram_size num_beamslength_penalty max_length<||>.T)skip_special_tokensrzDetecting languagezDetected language: zLanguage detection failed:  )r$r&tokensr4 chunk_offsets)$rrHWHISPER_SAMPLE_RATEr rBr;lensqueezemeanrangeminappendconfigrronesinput_featuresrno_gradgenerate tokenizerconvert_ids_to_tokens startswithendswithreplaceisdigitfloat ValueErrorextend batch_decodedetect_languageargmaxitem Exceptionwarningjoin))rrwr4r&rr^r_rrrchunk_duration chunk_samples total_samplestotal_duration all_tokensall_textr last_timeaccumulated_offset chunk_start chunk_endchunk_waveform chunk_offsetrrr predicted_ids chunk_tokensadjusted_tokenstokentime_strtime_val adjusted_time chunk_textdetected_languagefirst_chunk_featurespredicted_probslanguage_tokenefull_transcriptionwhisper_outputs) r rMTB_AudioToText.transcribes[) ! e%89$ .x~~.>?@ x~~ ! #''*H x~~ ! #~~a A%#==Q=/"a'#==Q=/#++A.M* &4  r* &4 $^C$8:;     MAKK-7GI%k)>',XH ()3 2i ? 2,47I,IM+22R c7J"3MN(0I+2259$**51'&*   o ."//40J OOJ '}B@! v  ? ./]]_+4 -0&1'+,%nRRZ ) ',&;&;,''O&/%8%8%N%N'..r2779&N *44T::'q,!& II 34E3FGH%%.!XXh/') *  ">11e!D&6'..u556%_( ? 9!=>> ?s\;/QAQ#Q+R 9B'Q: R  Q Q76Q7: RR R R8R33R8rN)rT)rrrrrrVrsrrrr\rrrrrr rrsP1  80LHH  R2!R2R2R2rrc>\rSrSrSr\S5rSrSrSr S Sjr Sr g ) MTB_ProcessWhisperOutputiv/Process Whisper output into timestamped chunks.cSSSSSSSS.4S.0$) Nrp)rFLOATrg$@g?rkrrQstep)rmin_chunk_lengthrrrs r rs$MTB_ProcessWhisperOutput.INPUT_TYPESys- "5 #CcJ%  r)rWHISPER_CHUNKSprocessrxcUSnUSn/nUSRSUS- n[R"SUSS35 [U5HupxUR S 5(dMUR S 5(dM5US S n U R S S5R5(aN[U 5n SU s=::aU::a4O M|URXz45 [R"SUSU 35 MMM /n [U5S:a[[U5S- 5HunXWupXWS-upX- U:aMX.s3 )LL..1::d3CA)0? ?)r$ timestampc3# UH5nURS5(aURS5(aM1Uv M7 g7frrrs r r6rs3 4LL..1::d3CA4rNrr$r&)r$r%r&)r;r rB enumeraterrrrrrrrrrstripr)rrrrr4timestamp_tokensaudio_durationirrrr% start_pos start_timeend_posend_timerr$ prev_chunk prev_durationresults r r  MTB_ProcessWhisperOutput.processs)w'z*004u]7KK $^C$8:;!&)HA%%%..*>*>$Qr{H''R088::#(?:N:,33QMBIIqcH:&>?;; *  1 $3/0145(8(;% $4U$;!(+;;%!mg> xx )  ::<<MM$(JJL*4)?6, $4R$8 !I3v;?*xx #MO4  ::<<%+BZ &{3A6(5a89&$'&6$$2!-$15EE (, .8-C#6*&z2  f%v--C"sAI".I"" I10I1rN)r) rrrrrrVrsrrrr rrrr rrvs.9    0LHHS.rrcN\rSrSrSr\S5rSrSrSr Sr S\ S \ S \ 4S jr S rg ) MTB_AudioCutiz%Basic audio cutter, values are in ms.c .SSSSSSSS.4SSSSSS.4S .0$) Nrprr@@rg~.Ar8r)r4lengthoffsetrrrs r rsMTB_AudioCut.INPUT_TYPESsH ##)"' !  #C!L  rr) cut_audiorxcutr4r'r(cUSn[X4-S- 5n[U[X$-S- 5-USRS5nUSSS2SS2XV24nUUS.4$)NrrrrA)rrr;)rr4r'r(r start_idxend_idx cut_waveforms r r+MTB_AudioCut.cutsM* ,t34  F0478 8 *  # #B ' Z(Ay/@)@A  +(   rrN)rrrrrrVrsrrrrrrr+rrrr r$r$sG/  (L!LHH  e U rr$cP\rSrSrSr\S5rSrSrSr Sr S\ S \ \ 4S jr S rg ) MTB_AudioStacki zStack/Overlay audio inputs (dynamic inputs). - pad audios to the longest inputs. - resample audios to the highest sample rate in the inputs. - convert them all to stereo if one of the inputs is. c S00$)Nrprrrs r rsMTB_AudioStack.INPUT_TYPESs Brr) stacked_audiorxstackkwargsr.c UR[UR555up#n[UVs/sHoUSRSPM sn5n/nUHbn[ R "SU(aSOSXeSRS- 45n[ R"USU/SS9n URU 5 Md [ R"USS9RSS9n UU S.4$s snf)Nrrr8r rKrrA) rTr(valuesrQr;rzerosrMrr7sum) rr8r-r2max_rater4r padded_audiospadding padded_audiostacked_waveforms r r7MTB_AudioStack.stacks&*&<&<  !' #86J6% +11"56JK ,. Ekk"Az!2!8!8!<<G!99eJ&7%ArJL   .!;;}!<@@Q@G (,   !KsC)rN)rrrrrrVrsrrrrrrZr7rrrr r3r3 sI   L%LHH k eK.@ rr3cJ\rSrSrSr\S5rSrSrSr Sr S\ S \ 4S jr S rg ) MTB_AudioSequencei:a"Sequence audio inputs (dynamic inputs). - adding silence_duration between each segment can now also be negative to overlap the clips, safely bound to the the input length. - resample audios to the highest sample rate in the inputs. - convert them all to stereo if one of the inputs is. cSSSSSSSS.400$) Nrpsilence_durationrrg8ig{Gz?rrrrs r rsMTB_AudioSequence.INPUT_TYPESCs- " #F3M%  rr)sequenced_audiorxsequencerFr8c UR[UR555up4n/n[U5GHupxUS:aUS:a?[R "SU(aSOS[ X-545n URU 5 OUS:a[ [U5U-5n USn [U U RSUSRS5n U S:aXU SS2SS2U *S24USSS2SS2SU 24-n U SS2SS2SU *24US'URU 5 USSS2SS2U S24US'URUS5 GM [R"USS9n UU S.4$)Nrr8r rrrKrA) rTr(r:rrr;rrabsrr;rM)rrFr8r-r2r=rIrr4silenceoverlapprevious_audio overlap_partsequenced_waveforms r rIMTB_AudioSequence.sequenceSs&*&<&<  !' #8(*!&)HA1u#a'#kk!*A 0 ;<GOOG,%)!#&6"7("BCG%-b\N!&,,R0j)//3G {*1a'?;#J/1hwh?@%(6aIgXIo'F   5,1*,=aGHn,Mj) OOE*- .9*<#YYxR8 (.   rrN)rrrrrrVrsrrrrrrrIrrrr rDrD:sD  L'LHH* * +* rrDcT\rSrSrSr\S5rSrSrSr Sr S\ S \ S \ \ 4S jrS rg )MTB_AudioResampleiz*Resample audio to a different sample rate.c  SSSSSSSSS .4S .0$) NrprINTrr-idz1Target sample rate in Hz. Whisper requires 16000.rkrrQrrn)r4rrrrs r rsMTB_AudioResample.INPUT_TYPESs5 ##(#% ##V     rr)resampled_audiorxresample_audior4rr.c*URX5nU4$r1)rH)rr4r resampleds r rZ MTB_AudioResample.resample_audiosMM%5 |rrN)rrrrrrVrsrrrrrrrZrZrrrr rSrSsO4  "L'LHH /2 { rrSch\rSrSrSr\S5rSrSrSr Sr SS\ S \ S \ S \ S \S \\ 4 SjjrSrg)MTB_AudioIsolateSpeakeriz3Isolate or mute specific speakers using WhisperDatac :SSSSSS04SS/SS04S S S S S SS.4S.0$)Nrprr rrk SPEAKER_00isolatemuterY@rr& z-Fade duration in milliseconds to avoid clicksrW)r4 whisper_datatarget_speakermodefade_msrrrs r rs#MTB_AudioIsolateSpeaker.INPUT_TYPESsW # 3#+i-F"G#V,y).DE#("% "#R     rr)processed_audiorx process_audior4rgrhrirjr.c.[US- US-5nUS:Xa[R"US5O[R"US5nUSHxnUR S5(dMX8S;n US:XaU (dUS:XdM7U (aM@[USS US-5n [USS US-5n S USS2X24'Mz US :a[R "S S U5n [R "US S S24US SS 24:g5S S -n U HgnX:dM XRS U- ::dM!US U4S :XaUSS2XU-24==U -ss'MCUSS2X- U24==U RS 5-ss'Mi USU-nUSUS .4$)Nr&rrcrr%speakerrdrrr8rrrA) rr zeros_like ones_likegetlinspacewherer;flip)rr4rgrhrirj fade_samplesmaskchunkspeaker_present start_sample end_samplefade transitions trans_idxprocessed_waveforms r rm%MTB_AudioIsolateSpeaker.process_audiosGf,m0DDE y    U:. /z!23 "(+E99Y'',i0@@O !o"+&q)E-,@@  !{!3A!6}9M!MN 36Q //0, ! >>!Q 5D++d1ab5kT!SbS&\&AB1EIK( -!ZZ]\%AAAyL)Q.Q ,D DDEMEQ 89 DDE IIaLE)#:.5 %]3.   rrN)rcre)rrrrrrVrsrrrrrr"r'rrZrmrrrr r_r_s~=  (L'LHH 7 7 "7  7  7  7  { 7 7 rr_cP\rSrSrSr\S5rSrSrSr Sr Sr S S jr S r g ) MTB_ProcessWhisperDiarizationizNProcess Whisper chunks with speaker diarization using either pyannote or NeMo.c @SSSS/SS04SSSS SS .4S .S S S/SS 040S.$)NrarpyannotenemorkrUr r8rfr)whisper_chunksr4backend num_speakersrcudacpurorrrs r rs)MTB_ProcessWhisperDiarization.INPUT_TYPESsY#6#'09j2IJ !!BB! FE?Y,?@  rrar rxcSSKJn SSKJn UR SSS9nUR [R"U55 U"5nU"USSUS S .UUS 9nSSS5 /n WRS S 9H/upn U RU RU RU S.5 M1 U $![a [ S5ef=f!,(df  Nm=f)z%Process audio using pyannote backend.r)Pipeline) ProgressHookzBpyannote.audio not found. Install with: pip install pyannote.audioz pyannote/speaker-diarization-3.1N)use_auth_tokenrr)rr)rhookT) yield_labelstartendro) pyannote.audior#pyannote.audio.pipelines.utils.hookr ImportErrorrrrr itertracksrrr) rr4rrrrrwr diarizationspeaker_segmentsturn_ros r process_pyannote.MTB_ProcessWhisperDiarization.process_pyannotes  / H ++ .t,   ELL() ^t" %j 1! 4#(#7* K + 6 64 6 H DW  # #!ZZ88& !I ; T  ^s B- C-C Cc@SSKJs Jn URR R S5RU5nURUSSUSUS9n/nUH$nURUSUS S US 3S .5 M& U$![a [S5ef=f) z!Process audio using NeMo backend.rNz;NeMo not found. Install with: pip install nemo_toolkit[asr]z+nvidia/speakerverification_en_titanet_largerr)r4rrrrSPEAKER_ror) nemo.collections.asr collectionsasrrmodelsClusteringDiarizerrrdiarizer) rr4rrnemo_asrr_rrsegments r process_nemo*MTB_ProcessWhisperDiarization.process_nemo4s  3 3 22BB 9 "V* mm #A&m,%$ "G  # #$W-"5>!)')*<)=> # 3 M  s BBc<US:XaURX$U5nOURX$U5nUSHfnUSup[5n UH-n U SU ::dMU SU:dMU RU S5 M/ U (a[ U 5SUS'MaSUS'Mh U4$) Nrr%rrrrorunknown)rrsetaddr() rrr4rrrrrxrrchunk_speakersrs r r %MTB_ProcessWhisperDiarization.processSs j #44V   $00fM #H-E%*;%7 "K UN+G$ 1+5"&&wy'9: ,#'#7#:i #,i .  rrN)rr r)rrrrrrVrsrrrrrr rrrr rrsDX  'LHH" H F !rrc>\rSrSrSr\S5rSrSrSr Sr Sr S r g ) MTB_AudioDurationitz#Get audio duration in milliseconds.cSSS00$)Nrpr4rrrrs r rsMTB_AudioDuration.INPUT_TYPESws   r)rU) duration_ms get_durationrxcUSnUSn[URSU- S-5n[R"SUSUS- SS35 U4$) Nrrrr-rzms (rzs))rr;r rB)rr4rrrs r rMTB_AudioDuration.get_durationsc$M* 8>>"- ;tCD  {m4 d0B3/Gr J ~rrN) rrrrrrVrsrrrrrrrrr rrts2-  L#LHH rr)"typingrrrrrCcomfy.model_managementrhuggingface_hubrrr r r utilsr rrr"r(rWr+r\rarrr$r3rDrSr_rr __nodes__rrr rs00 3- ")) ${+ + 8484v-i-R;R;jv2hv2rf.f.R+ 8+ \* X* ZC C LBT hT ny!y!x:!  r