a hN@sdZddlZddlmZddlmZmZddlmZm Z ddl m Z ddl m Z mZmZmZdd lmZmZerzddlZerddlZGd d d e d d ZGddded d ZGdddeZdgZdS)zProcessor class for DiaN)Path)OptionalUnion) AudioInputmake_list_of_audio) BatchFeature) AudioKwargsProcessingKwargsProcessorMixinUnpack)is_soundfile_availableis_torch_availablec@s:eZdZUeed<eed<eed<eeed<eed<dS)DiaAudioKwargs bos_token_id eos_token_id pad_token_id delay_pattern generationN)__name__ __module__ __qualname__int__annotations__listboolrrb/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/dia/processing_dia.pyr"s  rF)totalc@s@eZdZUeed<dddddddgd dd d d d idZdS)DiaProcessorKwargs audio_kwargsTrightF)paddingZ padding_sideZadd_special_tokensiii) r iD)rrrrr sampling_ratereturn_tensorspt) text_kwargsr common_kwargsN)rrrrr _defaultsrrrrr*s rc s eZdZdZdZdZdZfddZdee e e fe e e e eed d d Zd d e eeee d d ddZd!d e eeed d ddZd eeedddZe ee ee ee effeedddZed"eeee ee eddddZed eeedd dddZZS)# DiaProcessora Constructs a Dia processor which wraps a [`DiaFeatureExtractor`], [`DiaTokenizer`], and a [`DacModel`] into a single processor. It inherits, the audio feature extraction, tokenizer, and audio encode/decode functio- nalities. See [`~DiaProcessor.__call__`], [`~DiaProcessor.encode`], and [`~DiaProcessor.decode`] for more information. Args: feature_extractor (`DiaFeatureExtractor`): An instance of [`DiaFeatureExtractor`]. The feature extractor is a required input. tokenizer (`DiaTokenizer`): An instance of [`DiaTokenizer`]. The tokenizer is a required input. audio_tokenizer (`DacModel`): An instance of [`DacModel`] used to encode/decode audio into/from codebooks. It is is a required input. ZDiaFeatureExtractorZ DiaTokenizerZDacModelcstj|||ddS)N)audio_tokenizer)super__init__)selffeature_extractor tokenizerr2 __class__rrr4RszDiaProcessor.__init__NF)textaudio output_labelskwargsc' Kststd|durtd|jtfi|}|d}|d}|d}|dd} | dkrnt|jjd i} t|tr|g}n(t|t t frt d d |Dstd |j |fi|} | | |d d} |dd} |dd}|dd}|dd}| dus,|dus,|dus,| dur4td|rV|rVtd|d|d| djd}t| }t| }|dur2t|}|j|fi|}t|jjj}|ddjd|}g}g}t|d|dD]8\}}|jj}t|jdd||}||}||}tB|ddd|f|jj}|j |j!"dd}Wdn1sf0Y|stj#j$j%|d d!|d"}tj#j$j%|dd|ddddfd!| d"}|d|}||rdnd7}tj&dg|dg|tj'd#dddf} |(||(| qtj)|dd}tj)|dd}n@|rjtj*|d|f| tj'd#}tj+|d|ftj'd$}ntd%||jdkrtd&|d'|jdd(|jd}!|!|}"|j,||!|| d)d*}#tj*||!|f|tj-d+}$||$ddd|"f<|j.|$|| |#d,}%| |%|d-|r| d./ddddf}&d/|&|&|k<d/|&|&| k<|&"dd0||d1'| d0<| d.ddddf| d.<| d1ddddf| d1<t2| | d2S)3a Main method to prepare text(s) and audio to be fed as input to the model. The `audio` argument is forwarded to the DiaFeatureExtractor's [`~DiaFeatureExtractor.__call__`] and subsequently to the DacModel's [`~DacModel.encode`]. The `text` argument to [`~DiaTokenizer.__call__`]. Please refer to the docstring of the above methods for more information. zThe `DiaProcessor` relies on the `audio_tokenizer` which requires `torch` but we couldn't find it in your environment. You can install torch via `pip install torch`.Nz0You need to specify the `text` input to process.r.r r/r,r-z% only supports `return_tensors='pt'`.css|]}t|tVqdSN) isinstancestr).0trrr }z(DiaProcessor.__call__..zAInvalid input text. Please provide a string, or a list of stringsrrrrrTzTo enable processing for Dia, we need the `bos_token_id`, `eos_token_id`, `pad_token_id`, and `delay_pattern`. You may have accidentally overwritten one of those.z9Labels with `generation` is incompatible, got generation=z, output_labels=. input_idsr padding_maskZ input_valuesdim.)rrrrKrrZconstant)padmodevaluedtype)sizerQz;If you try to train, you should provide audio data as well.zNNeed the same amount of samples for both text and audio, but got text samples=z and audio samples = z instead.Fbszseq_len num_channelsrrevert)Z fill_valuerQr;rrprecomputed_idx)decoder_input_idsdecoder_attention_maskrZilabelsr[)dataZ tensor_type)3r ValueError _merge_kwargsrpopr9rr?r@rtupleallr7updateshapelenmaxrr6mathprodr2configZdownsampling_ratioszipZ hop_lengthceilsumtorchno_gradtodeviceencode audio_codes transposennZ functionalrMtensorlongappendcatfullZones build_indicesrapply_audio_delayclonereshape contiguousr)'r5r:r;r<r= output_kwargsr.r r/r,r] encodingsraudio_bos_token_idZaudio_eos_token_idaudio_pad_token_idrZ batch_sizerVZ max_delayZ input_audiosZcompression_rateZmax_encoded_sequence_lenrZr[rGZ base_pad_lenZcurrent_audio_lenZencoded_sequence_lenZ padding_lenrFZnum_valid_inputsZattention_maskZ max_seq_lenZ max_audio_lenrYZprefillZdelayed_decoder_input_idsr\rrr__call__Us             6 ,     $zDiaProcessor.__call__ torch.Tensor)rZaudio_prompt_lenr=returncKs|jtfi|}|d}|dd}|dd}|dd}|dusV|dusV|dur^td|durtj||jtjd}|d|j d} n"|dddddf|kj d d } |j d |dddddf|kj d d d } |j \} } } |j | | | |d d }|j |d d |d d d}g}txt| j dD]X}||dd| || |fd}||jj}|jj|dj}||q8Wdn1s0Y|S)a Decodes a batch of audio codebook sequences into their respective audio waveforms via the `audio_tokenizer`. See [`~DacModel.decode`] for more information. Args: decoder_input_ids (`torch.Tensor`): The complete output sequence of the decoder. audio_prompt_len (`int`): The audio prefix length (e.g. when using voice cloning). r rNrrzTo enable decoding for Dia, we need the `bos_token_id`, `pad_token_id`, and `delay_pattern`. You may have accidentally overwritten one of those.)rprQrrHrIrKTrSrXrL)N.)rr)r_rr`r^rmrurprvexpandrdrlrzr{rsrnrangeror2decodeZ audio_valuescpuZsqueezerw)r5rZrr=rr rrrZstart_of_generation_idxZend_of_generation_idxrTrUrVrYZoutput_sequencesZaudiosiZoutput_iZaudio_irrr batch_decodesV   ".   ".zDiaProcessor.batch_decodecKs<|jddkr$td|jdd|j||fi|dS)z Decodes a single sequence of audio codebooks into the respective audio waveform via the `audio_tokenizer`. See [`~DacModel.decode`] and [`~DiaProcessor.batch_decode`] for more information. rrKz5Expecting a single output to be decoded but received z samples instead.)rdr^r)r5rZrr=rrrrIs zDiaProcessor.decode)r[r=rcKsH|jtfi|}|d}|dd}|dur6td|jdt|S)z0Utility function to get the audio prompt length.r rNzTo enable the utility of retrieving the prompt length for Dia, we need the `delay_pattern`. You may have accidentally overwritten this.rK)r_rr`r^rdrf)r5r[r=rr rrrrget_audio_prompt_lenZs z!DiaProcessor.get_audio_prompt_len)r; saving_pathr=c Kststdt|}t|ttfr,|g}n(t|ttfrLtdd|DsTt dt |t |krlt d|j t fi|}|d}|d}t ||D]2\}}t|tjr|}t|||qdS)Nz/Please install `soundfile` to save audio files.css|]}t|ttfVqdSr>)r?r@r)rAprrrrCrDz*DiaProcessor.save_audio..zAInvalid input path. Please provide a string, or a list of stringsz5The number of audio and saving paths must be the samer r+)r ImportErrorrr?r@rrrarbr^rer_rrjrmZTensorrfloatnumpysfwrite) r5r;rr=rr r+Z audio_valuerrrr save_audioos(  zDiaProcessor.save_audio)rr)rTrUrVrrWrc Cstj|tjd}tj|tjddddf||d}|sT||ddddf}n||ddddf}t|d|d}tj|tjdddddf|||}tj|tjdddddf|||} tj|d|d| dgdd} || fS)a Precompute (sequence_idx, all_idx) so that out[seq, channel] = in[seq - delay[channel], channel] or in[seq, channel] = out[seq + delay[channel], channel] if `revert`. Negative sequence_idx => BOS; sequence_idx >= seq_len => PAD. rPN).NrrKrHrI) rmruZint32Zarangerclampstackr}rv) rTrUrVrrWZ delay_array sequence_idxvalid_sequence_idx batch_idx channel_idxall_idxrrrrzs ((( zDiaProcessor.build_indices)r;rrrYrc Cs|j}|\}}||}||}tj|dd\}}} |||| f|} |dk} ||jdk} t| |t| || } | S)a Applies or reverts the delay pattern to batched audio tokens using precomputed indices, inserting BOS where sequence_idx < 0 and PAD where sequence_idx >= seq_len. Args: audio: audio tokens of shape [bsz, seq_len, num_channels] pad_token_id: the PAD token bos_token_id: the BOS token precomputed_idx: from `build_indices` Returns: final_audio: delayed or reverted audio tokens of shape [bsz, seq_len, num_channels] rHrIrrK)rprormZunbindviewrRrdwhere)r;rrrYrprrrrrZgathered_audioZmask_bosZmask_padZ final_audiorrrr{s  zDiaProcessor.apply_audio_delay)NF)N)N)F)rrr__doc__Zfeature_extractor_classZtokenizer_classZaudio_tokenizer_classr4rr@rrrrr rrrrrrrr staticmethodrarzr{ __classcell__rrr8rr1>sj  1 J   ""r1)rrgpathlibrtypingrrZ audio_utilsrrZfeature_extraction_utilsrZprocessing_utilsr r r r utilsr rrmZ soundfilerrrr1__all__rrrrs$