a h Z@sdZddlZddlmZddlmZmZddlZddlm Z ddl m m Z ddl mZddlmZmZdd lmZd d lmZeGd d d eZeGdddeZeGdddeZGddde jZGddde jZGddde jZGddde jZGddde jZGddde jZGddde jZ Gd d!d!e jZ!eGd"d#d#eZ"ed$d%Gd&d'd'e"Z#d'd#gZ$dS)(zTransformers Xcodec model.N) dataclass)OptionalUnion)PreTrainedAudioTokenizerBase) ModelOutputauto_docstring) AutoModel) XcodecConfigc@s6eZdZUdZdZeejed<dZ eej ed<dS) XcodecOutputao Args: audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`, *optional*): Discrete code indices computed using `model.encode`. audio_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`, *optional*) Decoded audio values obtained using the decoder part of Xcodec. N audio_codes audio_values) __name__ __module__ __qualname____doc__rrtorch LongTensor__annotations__r FloatTensorrrf/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/xcodec/modeling_xcodec.pyr s r c@s$eZdZUdZdZeejed<dS)XcodecEncoderOutputz Args: audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`, *optional*): Discrete code indices computed using `model.encode`. Nr) rrrrrrrrrrrrrr-s rc@s$eZdZUdZdZeejed<dS)XcodecDecoderOutputz Args: audio_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`, *optional*): Decoded audio values obtained using the decoder part of Xcodec. Nr) rrrrrrrrrrrrrr8s rcs@eZdZdZeeeedfdd ZejejdddZ Z S) ResidualUnitzFResidual block for SemanticEncoder and SemanticDecoder used in Xcodec.)config in_channels out_channelsdilationc s\tt|_|jdd|}tj|||jd||ddd|_tj||ddd|_dS)Nr r F)stridepaddingr groupsbias)rr kernel_sizer$) super__init__nnZELU activationZunit_kernel_sizeConv1dconv1conv2)selfrrrr r" __class__rrr'Fs   zResidualUnit.__init__ hidden_statereturncCs0||}||}||}||}||SN)r)r+r,)r-r1Z output_tensorrrrforwardVs     zResidualUnit.forward) rrrrr intr'rTensorr4 __classcell__rrr.rrCsrcs<eZdZeeeedfdd ZejejdddZZ S)SemanticEncoderBlockrrrr!csdttfddjD|_|dkr4dnd|}|dd}tj||||dd|_dS)Ncsg|]}t|qSrr.0r rrrr bz1SemanticEncoderBlock.__init__..r rr Tr%r!r"r$)r&r'r( ModuleListblock_dilations res_unitsr*conv)r-rrrr!Zkernelr"r.r=rr'_s  zSemanticEncoderBlock.__init__r0cCs"|jD] }||}q||}|Sr3)rCrDr-r1unitrrrr4js   zSemanticEncoderBlock.forward rrrr r5r'rr6r4r7rrr.rr8^s r8cs0eZdZfddZejejdddZZS)SemanticEncodercstt|jt|jkr&tdtj|j|j|j d|j ddd|_ |j}g}t |jD]4\}}t |j|j|}|t ||||g7}|}q^t||_dS)Nz:Number of strides must match the number of channel_ratios.r r Fr$)r&r'lenstrideschannel_ratios ValueErrorr(r*semantic_hidden_sizer%rD enumerater5r8rA conv_blocks)r-rrrPir!rr.rrr'rs$  zSemanticEncoder.__init__r0cCs"||}|jD] }||}q|Sr3)rDrPr-r1blockrrrr4s   zSemanticEncoder.forwardrrrr'rr6r4r7rrr.rrHqs rHcs<eZdZeeeedfdd ZejejdddZZ S)SemanticDecoderBlockr9c st|dkr,tj|ddddd|_nBd|}|dd}|ddkrPdnd}tj|||||dd|_tfd d jD|_dS) Nr rTr@r rFrIcsg|]}t|qSrr:r;rrrrr>r?z1SemanticDecoderBlock.__init__..) r&r'r(r*rDConvTranspose1drArBrC)r-rrrr!r%r"output_paddingr.rVrr's&   zSemanticDecoderBlock.__init__r0cCs"||}|jD] }||}q|Sr3)rDrCrErrrr4s   zSemanticDecoderBlock.forwardrGrrr.rrUsrUcs0eZdZfddZejejdddZZS)SemanticDecodercsttj|jt|j|jd|jd|jddd|_g}t |j D]b\}}t|j|j|}|t |jdkrt|j|j|d}n|j}|t ||||g7}qJt ||_tj|j|j|jd|jddd|_dS)Nrr r F)rrr%r!r"r$)r!r"r$)r&r'r(r*rNr5rLr%r+rOrKrJrUrArPr,)r-rrPrQr!rrr.rrr's2  zSemanticDecoder.__init__r0cCs,||}|jD] }||}q||}|Sr3)r+rPr,rRrrrr4s     zSemanticDecoder.forwardrTrrr.rrYs rYcs8eZdZdZfddZddZddZdd ZZS) XcodecEuclideanCodebookz!Codebook with Euclidean distance.csjtt|j|j}|j|_|dtdg|dt|j|d||d|dS)NinitedTZ cluster_sizeembedZ embed_avg) r&r'rZzeros codebook_sizeZ codebook_dimZregister_bufferr6clone)r-rr\r.rrr's  z XcodecEuclideanCodebook.__init__cCsV|j}|djddd}|d|||djddd }|jddj}|S)Nr r T)Zkeepdimrdim)r\tpowsummaxindices)r- hidden_statesr\Z scaled_statesdist embed_indrrrquantizes  &z XcodecEuclideanCodebook.quantizecCs8|j}|d|df}||}|j|dd}|S)Nr_)shapeZreshaperjview)r-rgrkrirrrencodes  zXcodecEuclideanCodebook.encodecCst||j}|Sr3)FZ embeddingr\)r-ri quantizedrrrdecodeszXcodecEuclideanCodebook.decode) rrrrr'rjrmrpr7rrr.rrZs  rZcs6eZdZdZedfdd ZddZddZZS) XcodecVectorQuantizationzY Vector quantization implementation. Currently supports only euclidean distance. rcstt||_dSr3)r&r'rZcodebookr-rr.rrr's z!XcodecVectorQuantization.__init__cCs|ddd}|j|}|SNrr r )permutersrm)r-rgZembed_inrrrrms zXcodecVectorQuantization.encodecCs|j|}|ddd}|Sru)rsrprv)r-rirjrrrrps zXcodecVectorQuantization.decode) rrrrr r'rmrpr7rrr.rrqsrqcsheZdZdZedfdd ZddZdedd d Zde j e j d d d Z e j e j dddZ Z S) XcodecResidualVectorQuantizationzv Residual vector quantization implementation. Follows Algorithm 1 in https://huggingface.co/papers/2107.03312 rrcsFttfddtjD|_j|_j|_j|_dS)Ncsg|] }tqSr)rq)r<_rrrrr>r?z=XcodecResidualVectorQuantization.__init__..) r&r'r(rArangenum_quantizers quantizers frame_rater]rtr.rrrr's   z)XcodecResidualVectorQuantization.__init__cCst|j|jdS)zReturn bandwidth per quantizer.i)mathlog2r]r|)r-rrrget_bandwidth_per_quantizerszsrz$The Xcodec neural audio codec model.)Z custom_introc seZdZfddZeejdddZej ej dddZ e dej e ee eeej efd d d Ze dej e eeej efd ddZe dej e ej e ee eeeej ej fefdddZZS) XcodecModelcst|||_|jd|_t|j}|j|_ |j |_ | |j t ||_t||_t|j|_t|j|j|_t|j|jj|_t|j|jj|_t||_dS)Nr )r&r'rZ hop_lengthpadr from_configZacoustic_model_configencoderrdecoderr_adjust_dac_decoderrHencoder_semanticrYZdecoder_semanticZsemantic_model_configevalsemantic_modelr(rZ hidden_sizefcZfc1fc2rwr)r-rZacoustic_modelr.rrr's      zXcodecModel.__init__)rcCsh|D]8}t|tjrt|jtr.|jdn|j}|df|_qt|drdt|jtj rdt |_dS)z DAC implemented in Xcodec is slightly different from the HF version. DAC in Xcodec adjusts the output padding in every ConvTranspose1d in the decoder and removes the final `nn.Tanh` activation function. rr tanhN) rrr(rWr!tuplerXrrZTanhZIdentity)rrr!rrrrs   zXcodecModel._adjust_dac_decoder)rr2cCs|dddddf}t||j|jf}t$|j|dd}|j}Wdn1s\0Ytj|dd}|jddS)NrT)Zoutput_hidden_statesr r`)rnrrZno_gradrrgrr)r-routputsrgZstackedrrr_extract_semantic_featuress $z&XcodecModel._extract_semantic_featuresN)rr return_dictr2c Cs8|dur |n|jj}|jd}|dkr4td||durJ|jjd}n&||jjvrptd|d|jjd||}||dd}| |}|jd|jdkr| t |ddd ddf|j |j f d}t j||gdd }||dddd}|j||} | d d} |s0| St| S) ac input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`): Float values of the input audio waveform. bandwidth (`float`, *optional*): The target bandwidth in (kbps) supports only values in `config.target_bandwidths`. Defaults to the highest available bandwidth `4.0` kbps. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`]. Returns: `torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)` containing the discrete encoded audio codes. Nr zAudio must be mono, but got r_z)This model doesn't support the bandwidth z. Select one of .r rr`)rrrkrMZtarget_bandwidthsrdetachr transposerrnrZ unsqueezercatrrrmr) r-rrrZchannelsZe_semantic_inputZ e_semanticZ e_acousticrrrrrrms,   2 zXcodecModel.encode)rrr2cCs`|dur |n|jj}|dd}|j|}||dddd}||}|sX|St|S)a audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`): Discrete code indices computed using `model.encode`. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] Returns: Decoded audio values of shape `(batch_size, channels, num_samples)` obtained using the decoder part of Xcodec. Nrr r )rrrrrprrr)r-rrroZquantized_acousticrrrrrps   zXcodecModel.decode)rrrrr2cCsl|dur |n|jj}|jd}|dur6|j||dd}|j||dddd|f}|s`||fSt||dS)a input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`): The raw float values of the input audio waveform. audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`: Discrete code indices computed using `model.encode`. bandwidth (`float`, *optional*): Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth. bandwidth (`float`, *optional*): Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth. return_dict (`bool`, *optional*): Whether to return a [`XcodecOutput`] instead of a plain tuple. Returns: `XcodecOutput` or tuple `(audio_codes, audio_values)`: - `audio_codes` of shape `(batch_size, num_quantizers, codes_length)`: the quantized discrete codes. - `audio_values` of shape `(batch_size, channels, num_samples)`: the reconstructed audio waveform given the codes. Example: ```python >>> from datasets import load_dataset >>> from transformers import AutoFeatureExtractor, XcodecModel >>> model_id = "hf-audio/xcodec-hubert-librispeech" >>> model = XcodecModel.from_pretrained(model_id) >>> feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) >>> audio_sample = dataset[0]['audio']['array'] >>> inputs = feature_extractor(raw_audio=audio_sample, return_tensors="pt") >>> outputs = model(**inputs) >>> audio_codes = outputs.audio_codes >>> audio_values = outputs.audio_values ``` Nr_F)rr.)rr)rrrkrmrpr )r-rrrrlengthrrrrr4s. zXcodecModel.forward)NN)N)NNN)rrrr' staticmethodr(Modulerrrrrr6rfloatboolrrrmrrprr r4r7rrr.rr}s>    1 r)%rr} dataclassesrtypingrrrZtorch.nnr(Ztorch.nn.functionalZ functionalrnZmodeling_utilsrrrrautor Zconfiguration_xcodecr r rrrrr8rHrUrYrZrqrwrr__all__rrrrs<        ( 2>4