o ei @sdZddlZddlmZddlmZmZddlZddl Z ddl m Z ddl m m ZddlmZzddlmZWneyHddlmZYnweeZdgZdd d dd Zd d dddZddZddZddZ    d4dedededefddZe jj d d!Z!Gd"d#d#e j"Z#Gd$d%d%e j"Z$Gd&d'd'e j"Z%Gd(d)d)e j"Z&Gd*d+d+e j"Z'Gd,d-d-e j"Z(Gd.d/d/e j"Z)Gd0d1d1e j"Z*Gd2d3d3e j"Z+dS)5a, This lobe enables the integration of pretrained discrete DAC model. Reference: http://arxiv.org/abs/2306.06546 Reference: https://descript.notion.site/Descript-Audio-Codec-11389fce0ce2419891d6591a68f814d5 Reference: https://github.com/descriptinc/descript-audio-codec Author * Shubham Gupta 2023 N)Path)ListUnion) get_logger) weight_norm1.0.00.0.10.0.40.0.5))44khz8kbps)24khzr )16khzr )r 16kbpszWhttps://github.com/descriptinc/descript-audio-codec/releases/download/0.0.1/weights.pthz]https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.4/weights_24khz.pthz]https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.5/weights_16khz.pthzdhttps://github.com/descriptinc/descript-audio-codec/releases/download/1.0.0/weights_44khz_16kbps.pth))r rr )r r r )rr r )r rrcOttj|i|S)aH Apply weight normalization to a 1D convolutional layer. Arguments --------- *args : tuple Variable length argument list for nn.Conv1d. **kwargs : dict Arbitrary keyword arguments for nn.Conv1d. Returns ------- torch.nn.Module The weight-normalized nn.Conv1d layer. )rnnConv1dargskwargsrc/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/speechbrain/lobes/models/discrete/dac.pyWNConv1dHrcOr)an Apply weight normalization to a 1D transposed convolutional layer. Arguments --------- *args : tuple Variable length argument list for nn.ConvTranspose1d. **kwargs : dict Arbitrary keyword arguments for nn.ConvTranspose1d. Returns ------- torch.nn.Module The weight-normalized nn.ConvTranspose1d layer. )rrConvTranspose1drrrrWNConvTranspose1d[rrcCs6t|tjrtjj|jddtj|jddSdS)z= Initialize the weights of a 1D convolutional layer. g{Gz?)stdrN) isinstancerrinit trunc_normal_weight constant_bias)mrrr init_weightsns r$r r latest model_type model_bitratetag local_pathcCs|}|}|dvsJd|dvsJd|dkr"t||f}t|||fd}td||durAtd|d ||durUtd |d |d |d }| s}|j j d d dddl }||}|j dkrwtd|j ||j|S)a? Downloads a specified model file based on model type, bitrate, and tag, saving it to a local path. Arguments --------- model_type : str, optional The type of model to download. Can be '44khz', '24khz', or '16khz'. Default is '44khz'. model_bitrate : str, optional The bitrate of the model. Can be '8kbps' or '16kbps'. Default is '8kbps'. tag : str, optional A specific version tag for the model. Default is 'latest'. local_path : Path, optional The local file path where the model will be saved. If not provided, a default path will be used. Returns ------- Path The local path where the model is saved. Raises ------ ValueError If the model type or bitrate is not supported, or if the model cannot be found or downloaded. )r r rz6model_type must be one of '44khz', '24khz', or '16khz')r rz1model_bitrate must be one of '8kbps', or '16kbps'r%NzDownload link: zCould not find model with tag z and model type z.cache/descript/dac/weights__z.pthT)parentsexist_okrz1Could not download model. Received response code )lower__MODEL_LATEST_TAGS____MODEL_URLS__getloggerinfo ValueErrorrhomeexistsparentmkdirrequests status_code write_bytescontent)r&r'r(r) download_linkr9responserrrdownloadws>       r?cCsN|j}||d|dd}||dt||d}||}|S)a1 Applies the 'snake' activation function on the input tensor. This function reshapes the input tensor, applies a modified sine function to it, and then reshapes it back to its original shape. Arguments --------- x : torch.Tensor The input tensor to which the snake activation function will be applied. alpha : float A scalar value that modifies the sine function within the snake activation. Returns ------- torch.Tensor The transformed tensor after applying the snake activation function. rg& .>)shapereshape reciprocaltorchsinpow)xalpharCrrrsnakes $ rKcsneZdZdZdededeffdd Zdejfdd Zd ejfd d Z d ejfd dZ dejfddZ Z S)VectorQuantizea An implementation for Vector Quantization Implementation of VQ similar to Karpathy's repo: https://github.com/karpathy/deep-vector-quantization Additionally uses following tricks from Improved VQGAN (https://arxiv.org/pdf/2110.04627.pdf): 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space for improved codebook usage 2. l2-normalized codes: Converts euclidean distance to cosine similarity which improves training stability Arguments --------- input_dim : int Dimensionality of input codebook_size : int Size of codebook codebook_dim : int Dimensionality of codebook input_dim codebook_size codebook_dimcsHt||_||_t||dd|_t||dd|_t|||_ dS)Nr@ kernel_size) super__init__rNrOrin_projout_projr Embeddingcodebook)selfrMrNrO __class__rrrSs zVectorQuantize.__init__zcCs|||}||\}}tj||ddddg}tj||ddddg}|||}||}|||||fS)aQuantized the input tensor using a fixed codebook and returns the corresponding codebook vectors Arguments --------- z : torch.Tensor[B x D x T] Returns ------- torch.Tensor[B x D x T] Quantized continuous representation of input torch.Tensor[1] Commitment loss to train encoder to predict vectors closer to codebook entries torch.Tensor[1] Codebook loss to update the codebook torch.Tensor[B x T] Codebook indices (quantized discrete representation of input) torch.Tensor[B x D x T] Projected latents (continuous representation of input before quantization) none) reductionr@rB)rTdecode_latentsFmse_lossdetachmeanrU)rXr[z_ez_qindicescommitment_loss codebook_lossrrrforwards  zVectorQuantize.forwardembed_idcCst||jjS)a Embeds an ID using the codebook weights. This method utilizes the codebook weights to embed the given ID. Arguments --------- embed_id : torch.Tensor The tensor containing IDs that need to be embedded. Returns ------- torch.Tensor The embedded output tensor after applying the codebook weights. )r_ embeddingrWr rXrirrr embed_code-szVectorQuantize.embed_codecCs||ddS)a Decodes the embedded ID by transposing the dimensions. This method decodes the embedded ID by applying a transpose operation to the dimensions of the output tensor from the `embed_code` method. Arguments --------- embed_id : torch.Tensor The tensor containing embedded IDs. Returns ------- torch.Tensor The decoded tensor r@rB)rl transposerkrrr decode_code?szVectorQuantize.decode_codelatentsc Cs|dddd|d}|jj}t|}t|}|djdddd|| |djddd }| j ddd}|d}| |}| ||}| |} | |fS)a Decodes latent representations into discrete codes by comparing with the codebook. Arguments --------- latents : torch.Tensor The latent tensor representations to be decoded. Returns ------- Tuple[torch.Tensor, torch.Tensor] A tuple containing the decoded latent tensor (`z_q`) and the indices of the codes. rrBr@rAT)keepdimdim)permuterDsizerWr r_ normalizerHsumtmaxnumelviewrn) rXro encodingsrWdist max_indicesbrwrerdrrrr^Rs       zVectorQuantize.decode_latents) __name__ __module__ __qualname____doc__intrSrFTensorrhrlrnr^ __classcell__rrrYrrLs *rLc szeZdZdZ     ddeded ed eeefd ef fd d ZddefddZ de j fddZ de j fddZ ZS)ResidualVectorQuantizea Introduced in SoundStream: An end2end neural audio codec https://arxiv.org/abs/2107.03312 Arguments --------- input_dim : int, optional, by default 512 n_codebooks : int, optional, by default 9 codebook_size : int, optional, by default 1024 codebook_dim : Union[int, list], optional, by default 8 quantizer_dropout : float, optional, by default 0.0 Example ------- Using a pretrained RVQ unit. >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest") >>> quantizer = dac.quantizer >>> continuous_embeddings = torch.randn(1, 1024, 100) # Example shape: [Batch, Channels, Time] >>> discrete_embeddings, codes, _, _, _ = quantizer(continuous_embeddings)  rM n_codebooksrNrOquantizer_dropoutcshtttrfddt|D||_|_|_t fddt|D|_ ||_ dS)Ncsg|]}qSrr).0r*rOrr sz3ResidualVectorQuantize.__init__..csg|] }t|qSr)rL)rirOrNrMrrrs) rRrSrrrangerrOrNr ModuleList quantizersr)rXrMrrNrOrrYrrrSs   zResidualVectorQuantize.__init__N n_quantizerscCsrd}|}d}d}g}g}|dur|j}|jrLt|jdf|jd}td|jd|jdf} t|jd|j} | d| |d| <||j }t |j D]R\} } |jdur`| |kr`nD| |\} }}}}tj |jdf| |j d|k}|| |ddddf}|| }||| 7}||| 7}||||qQtj|dd}tj|dd}|||||fS)aQuantized the input tensor using a fixed set of `n` codebooks and returns the corresponding codebook vectors Arguments --------- z : torch.Tensor Shape [B x D x T] n_quantizers : int, optional No. of quantizers to use (n_quantizers < self.n_codebooks ex: for quantizer dropout) Note: if `self.quantizer_dropout` is True, this argument is ignored when in training mode, and a random number of quantizers is used. Returns ------- z : torch.Tensor[B x D x T] Quantized continuous representation of input codes : torch.Tensor[B x N x T] Codebook indices for each codebook (quantized discrete representation of input) latents : torch.Tensor[B x N*D x T] Projected latents (continuous representation of input before quantization) vq/commitment_loss : torch.Tensor[1] Commitment loss to train encoder to predict vectors closer to codebook entries vq/codebook_loss : torch.Tensor[1] Codebook loss to update the codebook rNr@F) fill_valuedevicerq)rtrainingrFonesrCrandintrrtor enumeraterfullrbappendstackcat)rXr[rrdresidualrfrgcodebook_indicesrodropout n_dropoutr quantizerz_q_icommitment_loss_icodebook_loss_i indices_iz_e_imaskcodesrrrrhsJ   zResidualVectorQuantize.forwardrcCsxd}g}|jd}t|D]$}|j||dd|ddf}|||j||}||}q |tj|dd|fS)aKGiven the quantized codes, reconstruct the continuous representation Arguments --------- codes : torch.Tensor[B x N x T] Quantized discrete representation of input Returns ------- torch.Tensor[B x D x T] Quantized continuous representation of input rr@Nrq)rCrrrnrrUrFr)rXrrdz_prrz_p_irrrr from_codess   "  z!ResidualVectorQuantize.from_codesroc Csd}g}g}tdgdd|jD}t||jdkdjdddd}t|D]8}||||d}} |j||dd|| ddf\} } || || |j| | } || }q+|t j |ddt j |ddfS) aGiven the unquantized latents, reconstruct the continuous representation after quantization. Arguments --------- latents : torch.Tensor[B x N x T] Continuous representation of input after projection Returns ------- torch.Tensor[B x D x T] Quantized representation of full-projected space torch.Tensor[B x D x T] Quantized representation of latent space rcSsg|]}|jqSrr)rqrrrrsz7ResidualVectorQuantize.from_latents..r@T)axiskeepdimsNrq) npcumsumrwhererCrxrr^rrUrFrr) rXrordrrdimsrrjkrcodes_irrrr from_latents s&     z#ResidualVectorQuantize.from_latents)rrrrrN)rrrrrrlistfloatrSrhrFrrrrrrrYrrxs* Mrcs(eZdZdZfddZddZZS)Snake1dz A PyTorch module implementing the Snake activation function in 1D. Arguments --------- channels : int The number of channels in the input tensor. cs$tttd|d|_dS)Nr@)rRrSr ParameterrFrrJ)rXchannelsrYrrrS<s zSnake1d.__init__cCs t||jSz} Arguments --------- x : torch.Tensor Returns ------- torch.Tensor )rKrJrXrIrrrrh@s zSnake1d.forward)rrrrrSrhrrrrYrr2s rcsBeZdZdZd dedeffdd Zdejd ejfd d Z Z S) ResidualUnita A residual unit module for convolutional neural networks. Arguments --------- dim : int, optional The number of channels in the input tensor. Default is 16. dilation : int, optional The dilation rate for the convolutional layers. Default is 1. r@rrdilationc sLtd|d}tt|t||d||dt|t||dd|_dS)NrB)rQrpaddingr@rP)rRrSr Sequentialrrblock)rXrrrpadrYrrrS[s    zResidualUnit.__init__rIreturncCsD||}|jd|jdd}|dkr|d|| f}||S)| Arguments --------- x : torch.Tensor Returns ------- torch.Tensor rArBr.)rrC)rXrIyrrrrrhes zResidualUnit.forwardrr@) rrrrrrSrFrtensorrhrrrrYrrNs  rcs<eZdZdZd dedeffdd Zdejfd d ZZ S) EncoderBlocka  An encoder block module for convolutional neural networks. This module constructs an encoder block consisting of a series of ResidualUnits and a final Snake1d activation followed by a weighted normalized 1D convolution. This block can be used as part of an encoder in architectures like autoencoders. Arguments --------- dim : int, optional The number of output channels. Default is 16. stride : int, optional The stride for the final convolutional layer. Default is 1. rr@rrstridecsnttt|dddt|dddt|dddt|dt|d|d||t|dd|_ dS)NrBr@rrrQrr) rRrSrrrrrmathceilr)rXrrrrYrrrSs    zEncoderBlock.__init__rIcC ||Srrrrrrrh zEncoderBlock.forwardr) rrrrrrSrFrrhrrrrYrrvsrcsBeZdZdZdgddfdededeffdd Zd d ZZS) Encodera A PyTorch module for the Encoder part of DAC. Arguments --------- d_model : int, optional The initial dimensionality of the model. Default is 64. strides : list, optional A list of stride values for downsampling in each EncoderBlock. Default is [2, 4, 8, 8]. d_latent : int, optional The dimensionality of the output latent space. Default is 64. Example ------- Creating an Encoder instance >>> encoder = Encoder() >>> audio_input = torch.randn(1, 1, 44100) # Example shape: [Batch, Channels, Time] >>> continuous_embedding = encoder(audio_input) Using a pretrained encoder. >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest") >>> encoder = dac.encoder >>> audio_input = torch.randn(1, 1, 44100) # Example shape: [Batch, Channels, Time] >>> continuous_embeddings = encoder(audio_input) @rBrrd_modelstridesd_latentc sttd|dddg|_|D]}|d9}|jt||dg7_q|jt|t||dddg7_tj|j|_||_dS)Nr@rrrQrrB)r) rRrSrrrrrrenc_dim)rXrrrrrYrrrSs  zEncoder.__init__cCrrrrrrrrhrzEncoder.forward) rrrrrrrSrhrrrrYrrsrcs:eZdZdZ d dededeffdd Zd d ZZS) DecoderBlocka A PyTorch module representing a block within the Decoder architecture. Arguments --------- input_dim : int, optional The number of input channels. Default is 16. output_dim : int, optional The number of output channels. Default is 8. stride : int, optional The stride for the transposed convolution, controlling the upsampling. Default is 1. rrr@rM output_dimrc sZttt|t||d||t|ddt|ddt|ddt|dd|_ dS)NrBrr@rrr) rRrSrrrrrrrr)rXrMrrrYrrrSs      zDecoderBlock.__init__cCrrrrrrrrh zDecoderBlock.forward)rrr@)rrrrrrSrhrrrrYrrsrc sBeZdZdZ d dededeedeffdd Zd d ZZS) DecoderaP A PyTorch module for the Decoder part of DAC. Arguments --------- input_channel : int The number of channels in the input tensor. channels : int The base number of channels for the convolutional layers. rates : list A list of stride rates for each decoder block d_out: int The out dimension of the final conv layer, Default is 1. Example ------- Creating a Decoder instance >>> decoder = Decoder(256, 1536, [8, 8, 4, 2]) >>> discrete_embeddings = torch.randn(2, 256, 200) # Example shape: [Batch, Channels, Time] >>> recovered_audio = decoder(discrete_embeddings) Using a pretrained decoder. Note that the actual input should be proper discrete representation. Using randomly generated input here for illustration of use. >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest") >>> decoder = dac.decoder >>> discrete_embeddings = torch.randn(1, 1024, 500) # Example shape: [Batch, Channels, Time] >>> recovered_audio = decoder(discrete_embeddings) r@ input_channelrratesd_outc stt||dddg}t|D]\}}|d|}|d|d} |t|| |g7}q|t| t| |dddtg7}tj||_ dS)NrrrrBr@) rRrSrrrrrTanhrmodel) rXrrrrlayersrrrMrrYrrrS3s  zDecoder.__init__cCrr)rrrrrrhNrzDecoder.forward)r@) rrrrrrrSrhrrrrYrrs$rc!seZdZdZdgdddgdddd d d d d ddd d fdedeedededeedededeeefdedede de de de dedef fdd Z d*d!e j d"efd#d$Z d%e j fd&d'Z  d+d!e j ded"efd(d)ZZS),DACa Discrete Autoencoder Codec (DAC) for audio data encoding and decoding. This class implements an autoencoder architecture with quantization for efficient audio processing. It includes an encoder, quantizer, and decoder for transforming audio data into a compressed latent representation and reconstructing it back into audio. This implementation supports both initializing a new model and loading a pretrained model. Arguments --------- encoder_dim : int Dimensionality of the encoder. encoder_rates : List[int] Downsampling rates for each encoder layer. latent_dim : int, optional Dimensionality of the latent space, automatically calculated if None. decoder_dim : int Dimensionality of the decoder. decoder_rates : List[int] Upsampling rates for each decoder layer. n_codebooks : int Number of codebooks for vector quantization. codebook_size : int Size of each codebook. codebook_dim : Union[int, list] Dimensionality of each codebook entry. quantizer_dropout : bool Whether to use dropout in the quantizer. sample_rate : int Sample rate of the audio data. model_type : str Type of the model to load (if pretrained). model_bitrate : str Bitrate of the model to load (if pretrained). tag : str Specific tag of the model to load (if pretrained). load_path : str, optional Path to load the pretrained model from, automatically downloaded if None. strict : bool Whether to strictly enforce the state dictionary match. load_pretrained : bool Whether to load a pretrained model. Example ------- Creating a new DAC instance: >>> dac = DAC() >>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time] >>> tokens, embeddings = dac(audio_data) Loading a pretrained DAC instance: >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest") >>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time] >>> tokens, embeddings = dac(audio_data) The tokens and the discrete embeddings obtained above or from other sources can be decoded: >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest") >>> audio_data = torch.randn(1, 1, 16000) # Example shape: [Batch, Channels, Time] >>> tokens, embeddings = dac(audio_data) >>> decoded_audio = dac.decode(embeddings) rrNi)rrrrBrrrFiDr r r% encoder_dim encoder_rates latent_dim decoder_dim decoder_ratesrrNrOr sample_rater&r'r( load_pathstrictload_pretrainedcsDt||_||_||_||_| |_||_||_||_ ||_ | |_ |rQ|s6t | | | d}t d|t|d}|d}|dD] \}}t|||qFt|j|_|j durh|jdt|j|_ t|j|j|j |_t|j |j|j|j |j d|_t|j |j|j|_|t|r|j|d|d ||_ dSdS) N)r&r'r(zObtained load path as: cpumetadatarrB)rMrrNrOr state_dict)r)!rRrSrrrrrrrNrOrrr?r2r3rFloaditemssetattrrprod hop_lengthlenrencoderrrrdecoderapplyr$load_state_dictr)rXrrrrrrrNrOrrr&r'r(rrr model_dictrkeyvaluerYrrrSsX      z DAC.__init__ audio_datarcCs.||}|||\}}}}}|||||fS)aEncode given audio data and return quantized latent codes Arguments --------- audio_data : torch.Tensor[B x 1 x T] Audio data to encode n_quantizers : int, optional Number of quantizers to use, by default None If None, all quantizers are used. Returns ------- "z" : torch.Tensor[B x D x T] Quantized continuous representation of input "codes" : torch.Tensor[B x N x T] Codebook indices for each codebook (quantized discrete representation of input) "latents" : torch.Tensor[B x N*D x T] Projected latents (continuous representation of input before quantization) "vq/commitment_loss" : torch.Tensor[1] Commitment loss to train encoder to predict vectors closer to codebook entries "vq/codebook_loss" : torch.Tensor[1] Codebook loss to update the codebook "length" : int Number of samples in input audio )rr)rXr rr[rrorfrgrrrencodes z DAC.encoder[cCr)a9Decode given latent codes and return audio data Arguments --------- z : torch.Tensor Shape [B x D x T] Quantized continuous representation of input Returns ------- torch.Tensor: shape B x 1 x length Decoded audio data. )r)rXr[rrrdecodes z DAC.decodec CsT|jd}t||j|j|}tj|d|f}|||\}}}}}||fS)aModel forward pass Arguments --------- audio_data : torch.Tensor[B x 1 x T] Audio data to encode sample_rate : int, optional Sample rate of audio data in Hz, by default None If None, defaults to `self.sample_rate` n_quantizers : int, optional Number of quantizers to use, by default None. If None, all quantizers are used. Returns ------- "tokens" : torch.Tensor[B x N x T] Codebook indices for each codebook (quantized discrete representation of input) "embeddings" : torch.Tensor[B x D x T] Quantized continuous representation of input rAr)rCrrrr functionalrr ) rXr rrlength right_padr[rr*rrrrhs z DAC.forwardr)NN)rrrrrrrrboolstrrSrFrr r rhrrrrYrr\sB     F &r)r r r%N),rrpathlibrtypingrrnumpyrrFtorch.nnrtorch.nn.functionalr r_speechbrain.utils.loggerrtorch.nn.utils.parametrizationsr ImportErrortorch.nn.utilsrr2SUPPORTED_VERSIONSr/r0rrr$rr?jitscriptrKModulerLrrrrrrrrrrrrsn       Q ;(-B.I