codingutf8 2023 meta platforms inc and affiliates and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch encodec model import math from dataclasses import dataclass from typing import list optional tuple union import torch import torch utils checkpoint from torch import nn from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationencodec import encodecconfig logger logging getloggername general docstring configfordoc encodecconfig encodecpretrainedmodelarchivelist facebookencodec24khz facebookencodec48khz see all encodec models at https huggingface comodels filterencodec dataclass class encodecoutputmodeloutput audiocodes torch floattensor none audiovalues torch floattensor none dataclass class encodecencoderoutputmodeloutput audiocodes torch floattensor none audioscales torch floattensor none dataclass class encodecdecoderoutputmodeloutput audiovalues torch floattensor none class encodecconv1dnn module warn user on unusual setup between dilation and stride see padforconv1d length hiddenstates shape1 nframes length kernelsize paddingtotal stride 1 ideallength math ceilnframes 1 stride kernelsize paddingtotal return ideallength length staticmethod def pad1dhiddenstates torch tensor paddings tupleint int mode str zero value float 0 0 length hiddenstates shape1 paddingleft paddingright paddings if not mode reflect return nn functional padhiddenstates paddings mode value maxpad maxpaddingleft paddingright extrapad 0 if length maxpad extrapad maxpad length 1 hiddenstates nn functional padhiddenstates 0 extrapad padded nn functional padhiddenstates paddings mode value end padded shape1 extrapad return padded end def forwardself hiddenstates kernelsize self conv kernelsize0 stride self conv stride0 dilation self conv dilation0 kernelsize kernelsize 1 dilation 1 effective kernel size with dilations paddingtotal kernelsize stride extrapadding self getextrapaddingforconv1dhiddenstates kernelsize stride paddingtotal if self causal left padding for causal hiddenstates self pad1dhiddenstates paddingtotal extrapadding modeself padmode else asymmetric padding required for odd strides paddingright paddingtotal 2 paddingleft paddingtotal paddingright hiddenstates self pad1d hiddenstates paddingleft paddingright extrapadding modeself padmode hiddenstates self convhiddenstates if self normtype timegroupnorm hiddenstates self normhiddenstates return hiddenstates class encodecconvtranspose1dnn module we will only trim fixed padding extra padding from padforconv1d would be removed at the very end when keeping only the right length for the output as removing it here would require also passing the length at the matching layer in the encoder trim the padding on the right according to the specified ratio if trimrightratio 1 0 trim everything from right asymmetric padding required for odd strides unpad lstm without worrying about the hidden state nor the layout of the data expects input as convolutional layout residual block from seanet model as used by encodec seanet encoder as used by encodec def initself config encodecconfig super init model encodecconv1dconfig config audiochannels config numfilters config kernelsize scaling 1 downsample to raw audio scale for ratio in reversedconfig upsamplingratios currentscale scaling config numfilters add residual layers for j in rangeconfig numresiduallayers model encodecresnetblockconfig currentscale config dilationgrowthratej 1 add downsampling layers model nn elu model encodecconv1dconfig currentscale currentscale 2 kernelsizeratio 2 strideratio scaling 2 model encodeclstmconfig scaling config numfilters model nn elu model encodecconv1dconfig scaling config numfilters config hiddensize config lastkernelsize self layers nn modulelistmodel def forwardself hiddenstates for layer in self layers hiddenstates layerhiddenstates return hiddenstates class encodecdecodernn module upsample to raw audio scale add upsampling layers add residual layers add final layers codebook with euclidean distance def initself config encodecconfig super init embed torch zerosconfig codebooksize config codebookdim self codebooksize config codebooksize self registerbufferinited torch tensortrue self registerbufferclustersize torch zerosconfig codebooksize self registerbufferembed embed self registerbufferembedavg embed clone def quantizeself hiddenstates embed self embed t scaledstates hiddenstates pow2 sum1 keepdimtrue dist scaledstates 2 hiddenstates embed embed pow2 sum0 keepdimtrue embedind dist maxdim1 indices return embedind def encodeself hiddenstates shape hiddenstates shape preprocess hiddenstates hiddenstates reshape1 shape1 quantize embedind self quantizehiddenstates postprocess embedind embedind viewshape 1 return embedind def decodeself embedind quantize nn functional embeddingembedind self embed return quantize class encodecvectorquantizationnn module def initself config encodecconfig super init self codebook encodeceuclideancodebookconfig def encodeself hiddenstates hiddenstates hiddenstates permute0 2 1 embedin self codebook encodehiddenstates return embedin def decodeself embedind quantize self codebook decodeembedind quantize quantize permute0 2 1 return quantize class encodecresidualvectorquantizernn module return numquantizers based on specified target bandwidth bwperq math log2self codebooksize self framerate numquantizers self numquantizers if bandwidth is not none and bandwidth 0 0 numquantizers intmax1 math floorbandwidth 1000 bwperq return numquantizers def encodeself embeddings torch tensor bandwidth optionalfloat none torch tensor numquantizers self getnumquantizersforbandwidthbandwidth residual embeddings allindices for layer in self layers numquantizers indices layer encoderesidual quantized layer decodeindices residual residual quantized allindices appendindices outindices torch stackallindices return outindices def decodeself codes torch tensor torch tensor an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights if isinstancemodule nn linear module weight data normalmean0 0 stdself config initializerrange if module bias is not none module bias data zero elif isinstancemodule nn layernorm nn groupnorm module bias data zero module weight data fill1 0 elif isinstancemodule nn conv1d nn init kaimingnormalmodule weight if module bias is not none k math sqrtmodule groups module inchannels module kernelsize0 nn init uniformmodule bias ak bk elif isinstancemodule nn embedding module weight data normalmean0 0 stdself config initializerrange if module paddingidx is not none module weight datamodule paddingidx zero elif isinstancemodule nn lstm for name param in module namedparameters if weight in name nn init xavieruniformparam elif bias in name nn init constantparam 0 0 encodecstartdocstring r this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config encodecconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights addstartdocstrings the encodec neural audio codec model encodecstartdocstring class encodecmodelencodecpretrainedmodel def initself config encodecconfig super initconfig self config config self encoder encodecencoderconfig self decoder encodecdecoderconfig self quantizer encodecresidualvectorquantizerconfig self bitspercodebook intmath log2self config codebooksize if 2self bitspercodebook self config codebooksize raise valueerrorthe codebooksize must be a power of 2 initialize weights and apply final processing self postinit def getencoderself return self encoder def getdecoderself return self decoder def encodeframe self inputvalues torch tensor bandwidth float paddingmask int tupletorch tensor optionaltorch tensor length inputvalues shape1 duration length self config samplingrate if self config chunklengths is not none and duration 1e5 self config chunklengths raise runtimeerrorfduration of frame duration is longer than chunk self config chunklengths scale none if self config normalize if the padding is non zero inputvalues inputvalues paddingmask mono torch suminputvalues 1 keepdimtrue inputvalues shape1 scale mono pow2 meandim1 keepdimtrue sqrt 1e8 inputvalues inputvalues scale embeddings self encoderinputvalues codes self quantizer encodeembeddings bandwidth codes codes transpose0 1 return codes scale def encode self inputvalues torch tensor paddingmask torch tensor none bandwidth optionalfloat none returndict optionalbool none uniontupletorch tensor optionaltorch tensor encodecencoderoutput returndict returndict if returndict is not none else self config returndict if bandwidth is none bandwidth self config targetbandwidths0 if bandwidth not in self config targetbandwidths raise valueerror fthis model doesn t support the bandwidth bandwidth fselect one of self config targetbandwidths channels inputlength inputvalues shape if channels 1 or channels 2 raise valueerrorfnumber of audio channels must be 1 or 2 but got channels chunklength self config chunklength if chunklength is none chunklength inputlength stride inputlength else stride self config chunkstride if paddingmask is none paddingmask torch oneslikeinputvalues bool encodedframes scales step chunklength stride if inputlength stride step 0 raise valueerror the input length is not properly padded for batched chunked decoding make sure to pad the input correctly for offset in range0 inputlength step stride mask paddingmask offset offset chunklength bool frame inputvalues offset offset chunklength encodedframe scale self encodeframeframe bandwidth mask encodedframes appendencodedframe scales appendscale encodedframes torch stackencodedframes if not returndict return encodedframes scales return encodecencoderoutputencodedframes scales staticmethod def linearoverlapaddframes listtorch tensor stride int generic overlap add with linear fadeinfadeout supporting complex scenario e g more than 2 frames per position the core idea is to use a weight function that is a triangle with a maximum value at the middle of the chunk we use this weighting when summing the frames and divide by the sum of weights for each positions at the end thus if a frame is the only one to cover a position the weighting is a noop if 2 frames cover a position s t i e s offset of second frame starts t end of first frame then the weight function for each one is t s t t with t a given offset after the final normalization the weight of the second frame at position t is t s t s t t t s t s which is exactly what we want if more than 2 frames overlap at a given point we hope that by induction something sensible happens if lenframes 0 raise valueerrorframes cannot be an empty list device frames0 device dtype frames0 dtype shape frames0 shape 1 totalsize stride lenframes 1 frames1 shape1 framelength frames0 shape1 timevec torch linspace0 1 framelength 2 devicedevice dtypedtype1 1 weight 0 5 timevec 0 5 abs sumweight torch zerostotalsize devicedevice dtypedtype out torch zerosshape totalsize devicedevice dtypedtype offset int 0 for frame in frames framelength frame shape1 out offset offset framelength weight framelength frame sumweightoffset offset framelength weight framelength offset stride if sumweight min 0 raise valueerrorfsumweight minimum element must be bigger than zero sumweight return out sumweight def decodeframeself codes torch tensor scale optionaltorch tensor none torch tensor codes codes transpose0 1 embeddings self quantizer decodecodes outputs self decoderembeddings if scale is not none outputs outputs scale view1 1 1 return outputs def decode self audiocodes torch tensor audioscales torch tensor paddingmask optionaltorch tensor none returndict optionalbool none uniontupletorch tensor torch tensor encodecdecoderoutput returndict returndict or self config returndict chunklength self config chunklength if chunklength is none if lenaudiocodes 1 raise valueerrorfexpected one frame got lenaudiocodes audiovalues self decodeframeaudiocodes0 audioscales0 else decodedframes for frame scale in zipaudiocodes audioscales frames self decodeframeframe scale decodedframes appendframes audiovalues self linearoverlapadddecodedframes self config chunkstride or 1 truncate based on padding mask if paddingmask is not none and paddingmask shape1 audiovalues shape1 audiovalues audiovalues paddingmask shape1 if not returndict return audiovalues return encodecdecoderoutputaudiovalues addstartdocstringstomodelforwardencodecinputsdocstring replacereturndocstringsoutputtypeencodecoutput configclassconfigfordoc def forward self inputvalues torch tensor paddingmask optionaltorch tensor none bandwidth optionalfloat none audiocodes optionaltorch tensor none audioscales optionaltorch tensor none returndict optionalbool none uniontupletorch tensor torch tensor encodecoutput r returns examples python from datasets import loaddataset from transformers import autoprocessor encodecmodel dataset loaddatasetashraqesc50 audiosample datasettrainaudio0array modelid facebookencodec24khz model encodecmodel frompretrainedmodelid processor autoprocessor frompretrainedmodelid inputs processorrawaudioaudiosample returntensorspt outputs modelinputs audiocodes outputs audiocodes audiovalues outputs audiovalues returndict returndict or self config returndict if paddingmask is none paddingmask torch oneslikeinputvalues bool if audiocodes is not none and audioscales is none raise valueerroryou specified audiocodes but did not specify the audioscales if audioscales is not none and audiocodes is none raise valueerroryou specified audioscales but did not specify the audiocodes if audioscales is none and audiocodes is none audiocodes audioscales self encodeinputvalues paddingmask bandwidth false audiovalues self decodeaudiocodes audioscales paddingmask returndictreturndict0 if not returndict return audiocodes audiovalues return encodecoutputaudiocodesaudiocodes audiovaluesaudiovalues coding utf 8 2023 meta platforms inc and affiliates and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch encodec model general docstring see all encodec models at https huggingface co models filter encodec args audio_codes torch floattensor of shape batch_size nb_chunks chunk_length optional discret code embeddings computed using model encode audio_values torch flaottensor of shape batch_size sequence_length optional decoded audio values obtained using the decoder part of encodec args audio_codes torch floattensor of shape batch_size nb_chunks chunk_length optional discret code embeddings computed using model encode audio_scales torch tensor of shape batch_size nb_chunks optional scaling factor for each audio_codes input this is used to unscale each chunk of audio when decoding args audio_values torch floattensor of shape batch_size segment_length optional decoded audio values obtained using the decoder part of encodec conv1d with asymmetric or causal padding and normalization warn user on unusual setup between dilation and stride see pad_for_conv1d tiny wrapper around torch nn functional pad just to allow for reflect padding on small input if this is the case we insert extra 0 padding to the right before the reflection happens effective kernel size with dilations left padding for causal asymmetric padding required for odd strides convtranspose1d with asymmetric or causal padding and normalization we will only trim fixed padding extra padding from pad_for_conv1d would be removed at the very end when keeping only the right length for the output as removing it here would require also passing the length at the matching layer in the encoder trim the padding on the right according to the specified ratio if trim_right_ratio 1 0 trim everything from right asymmetric padding required for odd strides unpad lstm without worrying about the hidden state nor the layout of the data expects input as convolutional layout residual block from seanet model as used by encodec seanet encoder as used by encodec downsample to raw audio scale add residual layers add downsampling layers seanet decoder as used by encodec upsample to raw audio scale add upsampling layers add residual layers add final layers codebook with euclidean distance pre process quantize post process vector quantization implementation currently supports only euclidean distance residual vector quantizer return num_quantizers based on specified target bandwidth encode a given input tensor with the specified frame rate at the given bandwidth the rvq encode method sets the appropriate number of quantizers to use and returns indices for each quantizer decode the given codes to the quantized representation an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config encodecconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args input_values torch floattensor of shape batch_size channels sequence_length optional raw audio input converted to float and padded to the approriate length in order to be encoded using chunks of length self chunk_length and a stride of config chunk_stride padding_mask torch booltensor of shape batch_size channels sequence_length optional mask to avoid computing scaling factors on padding token indices can we avoid computing conv on these mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked tip warning true padding_mask should always be passed unless the input was truncated or not padded this is because in order to process tensors effectively the input audio should be padded so that input_length stride step with step chunk_length stride this ensures that all chunks are of the same shape tip bandwidth float optional the target bandwidth must be one of config target_bandwidths if none uses the smallest possible bandwidth bandwidth is represented as a thousandth of what it is e g 6kbps bandwidth is represented as bandwidth 6 0 audio_codes torch floattensor of shape batch_size nb_chunks chunk_length optional discret code embeddings computed using model encode audio_scales torch tensor of shape batch_size nb_chunks optional scaling factor for each audio_codes input return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple initialize weights and apply final processing encodes the given input using the underlying vqvae if config normalize is set to true the input is first normalized the padding mask is required to compute the correct scale if the padding is non zero encodes the input audio waveform into discrete codes args input_values torch tensor of shape batch_size channels sequence_length float values of the input audio waveform padding_mask torch tensor of shape batch_size channels sequence_length padding mask used to pad the input_values bandwidth float optional the target bandwidth must be one of config target_bandwidths if none uses the smallest possible bandwidth bandwidth is represented as a thousandth of what it is e g 6kbps bandwidth is represented as bandwidth 6 0 returns a list of frames containing the discrete encoded codes for the input audio waveform along with rescaling factors for each chunk when normalize is true each frames is a tuple codebook scale with codebook of shape batch_size num_codebooks frames generic overlap add with linear fade in fade out supporting complex scenario e g more than 2 frames per position the core idea is to use a weight function that is a triangle with a maximum value at the middle of the chunk we use this weighting when summing the frames and divide by the sum of weights for each positions at the end thus if a frame is the only one to cover a position the weighting is a no op if 2 frames cover a position s t i e s offset of second frame starts t end of first frame then the weight function for each one is t s t t with t a given offset after the final normalization the weight of the second frame at position t is t s t s t t t s t s which is exactly what we want if more than 2 frames overlap at a given point we hope that by induction something sensible happens decodes the given frames into an output audio waveform note that the output might be a bit bigger than the input in that case any extra steps at the end can be trimmed args audio_codes torch floattensor of shape batch_size nb_chunks chunk_length optional discret code embeddings computed using model encode audio_scales torch tensor of shape batch_size nb_chunks optional scaling factor for each audio_codes input padding_mask torch tensor of shape batch_size channels sequence_length padding mask used to pad the input_values return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple truncate based on padding mask returns examples python from datasets import load_dataset from transformers import autoprocessor encodecmodel dataset load_dataset ashraq esc50 audio_sample dataset train audio 0 array model_id facebook encodec_24khz model encodecmodel from_pretrained model_id processor autoprocessor from_pretrained model_id inputs processor raw_audio audio_sample return_tensors pt outputs model inputs audio_codes outputs audio_codes audio_values outputs audio_values
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from ...modeling_utils import PreTrainedModel
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_encodec import EncodecConfig


logger = logging.get_logger(__name__)



_CONFIG_FOR_DOC = "EncodecConfig"


ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/encodec_24khz",
    "facebook/encodec_48khz",
    
]


@dataclass
class EncodecOutput(ModelOutput):
    

    audio_codes: torch.FloatTensor = None
    audio_values: torch.FloatTensor = None


@dataclass
class EncodecEncoderOutput(ModelOutput):
    

    audio_codes: torch.FloatTensor = None
    audio_scales: torch.FloatTensor = None


@dataclass
class EncodecDecoderOutput(ModelOutput):
    

    audio_values: torch.FloatTensor = None


class EncodecConv1d(nn.Module):
    

    def __init__(
        self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1
    ):
        super().__init__()
        self.causal = config.use_causal_conv
        self.pad_mode = config.pad_mode
        self.norm_type = config.norm_type

        if self.norm_type not in ["weight_norm", "time_group_norm"]:
            raise ValueError(
                f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'
            )

        
        if stride > 1 and dilation > 1:
            logger.warning(
                "EncodecConv1d has been initialized with stride > 1 and dilation > 1"
                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
            )

        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation)
        if self.norm_type == "weight_norm":
            self.conv = nn.utils.weight_norm(self.conv)
        elif self.norm_type == "time_group_norm":
            self.norm = nn.GroupNorm(1, out_channels)

    @staticmethod
    def _get_extra_padding_for_conv1d(
        hidden_states: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
    ) -> int:
        
        length = hidden_states.shape[-1]
        n_frames = (length - kernel_size + padding_total) / stride + 1
        ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
        return ideal_length - length

    @staticmethod
    def _pad1d(hidden_states: torch.Tensor, paddings: Tuple[int, int], mode: str = "zero", value: float = 0.0):
        
        length = hidden_states.shape[-1]
        padding_left, padding_right = paddings
        if not mode == "reflect":
            return nn.functional.pad(hidden_states, paddings, mode, value)

        max_pad = max(padding_left, padding_right)
        extra_pad = 0
        if length <= max_pad:
            extra_pad = max_pad - length + 1
            hidden_states = nn.functional.pad(hidden_states, (0, extra_pad))
        padded = nn.functional.pad(hidden_states, paddings, mode, value)
        end = padded.shape[-1] - extra_pad
        return padded[..., :end]

    def forward(self, hidden_states):
        kernel_size = self.conv.kernel_size[0]
        stride = self.conv.stride[0]
        dilation = self.conv.dilation[0]
        kernel_size = (kernel_size - 1) * dilation + 1  
        padding_total = kernel_size - stride
        extra_padding = self._get_extra_padding_for_conv1d(hidden_states, kernel_size, stride, padding_total)

        if self.causal:
            
            hidden_states = self._pad1d(hidden_states, (padding_total, extra_padding), mode=self.pad_mode)
        else:
            
            padding_right = padding_total // 2
            padding_left = padding_total - padding_right
            hidden_states = self._pad1d(
                hidden_states, (padding_left, padding_right + extra_padding), mode=self.pad_mode
            )

        hidden_states = self.conv(hidden_states)

        if self.norm_type == "time_group_norm":
            hidden_states = self.norm(hidden_states)

        return hidden_states


class EncodecConvTranspose1d(nn.Module):
    

    def __init__(self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1):
        super().__init__()
        self.causal = config.use_causal_conv
        self.trim_right_ratio = config.trim_right_ratio
        self.norm_type = config.norm_type
        if self.norm_type not in ["weight_norm", "time_group_norm"]:
            raise ValueError(
                f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'
            )

        self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride)
        if config.norm_type == "weight_norm":
            self.conv = nn.utils.weight_norm(self.conv)
        elif config.norm_type == "time_group_norm":
            self.norm = nn.GroupNorm(1, out_channels)

        if not (self.causal or self.trim_right_ratio == 1.0):
            raise ValueError("`trim_right_ratio` != 1.0 only makes sense for causal convolutions")

    def forward(self, hidden_states):
        kernel_size = self.conv.kernel_size[0]
        stride = self.conv.stride[0]
        padding_total = kernel_size - stride

        hidden_states = self.conv(hidden_states)

        if self.norm_type == "time_group_norm":
            hidden_states = self.norm(hidden_states)

        
        
        
        
        if self.causal:
            
            
            padding_right = math.ceil(padding_total * self.trim_right_ratio)
        else:
            
            padding_right = padding_total // 2

        padding_left = padding_total - padding_right

        
        end = hidden_states.shape[-1] - padding_right
        hidden_states = hidden_states[..., padding_left:end]
        return hidden_states


class EncodecLSTM(nn.Module):
    

    def __init__(self, config, dimension):
        super().__init__()
        self.lstm = nn.LSTM(dimension, dimension, config.num_lstm_layers)

    def forward(self, hidden_states):
        hidden_states = hidden_states.permute(2, 0, 1)
        hidden_states = self.lstm(hidden_states)[0] + hidden_states
        hidden_states = hidden_states.permute(1, 2, 0)
        return hidden_states


class EncodecResnetBlock(nn.Module):
    

    def __init__(self, config: EncodecConfig, dim: int, dilations: List[int]):
        super().__init__()
        kernel_sizes = (config.residual_kernel_size, 1)
        if len(kernel_sizes) != len(dilations):
            raise ValueError("Number of kernel sizes should match number of dilations")

        hidden = dim // config.compress
        block = []
        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
            in_chs = dim if i == 0 else hidden
            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
            block += [nn.ELU()]
            block += [EncodecConv1d(config, in_chs, out_chs, kernel_size, dilation=dilation)]
        self.block = nn.ModuleList(block)

        if config.use_conv_shortcut:
            self.shortcut = EncodecConv1d(config, dim, dim, kernel_size=1)
        else:
            self.shortcut = nn.Identity()

    def forward(self, hidden_states):
        residual = hidden_states
        for layer in self.block:
            hidden_states = layer(hidden_states)

        return self.shortcut(residual) + hidden_states


class EncodecEncoder(nn.Module):
    

    def __init__(self, config: EncodecConfig):
        super().__init__()
        model = [EncodecConv1d(config, config.audio_channels, config.num_filters, config.kernel_size)]
        scaling = 1

        
        for ratio in reversed(config.upsampling_ratios):
            current_scale = scaling * config.num_filters
            
            for j in range(config.num_residual_layers):
                model += [EncodecResnetBlock(config, current_scale, [config.dilation_growth_rate**j, 1])]
            
            model += [nn.ELU()]
            model += [EncodecConv1d(config, current_scale, current_scale * 2, kernel_size=ratio * 2, stride=ratio)]
            scaling *= 2

        model += [EncodecLSTM(config, scaling * config.num_filters)]
        model += [nn.ELU()]
        model += [EncodecConv1d(config, scaling * config.num_filters, config.hidden_size, config.last_kernel_size)]

        self.layers = nn.ModuleList(model)

    def forward(self, hidden_states):
        for layer in self.layers:
            hidden_states = layer(hidden_states)
        return hidden_states


class EncodecDecoder(nn.Module):
    

    def __init__(self, config: EncodecConfig):
        super().__init__()
        scaling = int(2 ** len(config.upsampling_ratios))
        model = [EncodecConv1d(config, config.hidden_size, scaling * config.num_filters, config.kernel_size)]

        model += [EncodecLSTM(config, scaling * config.num_filters)]

        
        for ratio in config.upsampling_ratios:
            current_scale = scaling * config.num_filters
            
            model += [nn.ELU()]
            model += [
                EncodecConvTranspose1d(config, current_scale, current_scale // 2, kernel_size=ratio * 2, stride=ratio)
            ]
            
            for j in range(config.num_residual_layers):
                model += [EncodecResnetBlock(config, current_scale // 2, (config.dilation_growth_rate**j, 1))]
            scaling //= 2

        
        model += [nn.ELU()]
        model += [EncodecConv1d(config, config.num_filters, config.audio_channels, config.last_kernel_size)]
        self.layers = nn.ModuleList(model)

    def forward(self, hidden_states):
        for layer in self.layers:
            hidden_states = layer(hidden_states)
        return hidden_states


class EncodecEuclideanCodebook(nn.Module):
    

    def __init__(self, config: EncodecConfig):
        super().__init__()
        embed = torch.zeros(config.codebook_size, config.codebook_dim)

        self.codebook_size = config.codebook_size

        self.register_buffer("inited", torch.Tensor([True]))
        self.register_buffer("cluster_size", torch.zeros(config.codebook_size))
        self.register_buffer("embed", embed)
        self.register_buffer("embed_avg", embed.clone())

    def quantize(self, hidden_states):
        embed = self.embed.t()
        scaled_states = hidden_states.pow(2).sum(1, keepdim=True)
        dist = -(scaled_states - 2 * hidden_states @ embed + embed.pow(2).sum(0, keepdim=True))
        embed_ind = dist.max(dim=-1).indices
        return embed_ind

    def encode(self, hidden_states):
        shape = hidden_states.shape
        
        hidden_states = hidden_states.reshape((-1, shape[-1]))
        
        embed_ind = self.quantize(hidden_states)
        
        embed_ind = embed_ind.view(*shape[:-1])
        return embed_ind

    def decode(self, embed_ind):
        quantize = nn.functional.embedding(embed_ind, self.embed)
        return quantize


class EncodecVectorQuantization(nn.Module):
    

    def __init__(self, config: EncodecConfig):
        super().__init__()
        self.codebook = EncodecEuclideanCodebook(config)

    def encode(self, hidden_states):
        hidden_states = hidden_states.permute(0, 2, 1)
        embed_in = self.codebook.encode(hidden_states)
        return embed_in

    def decode(self, embed_ind):
        quantize = self.codebook.decode(embed_ind)
        quantize = quantize.permute(0, 2, 1)
        return quantize


class EncodecResidualVectorQuantizer(nn.Module):
    

    def __init__(self, config: EncodecConfig):
        super().__init__()
        self.codebook_size = config.codebook_size
        self.frame_rate = config.frame_rate
        self.num_quantizers = config.num_quantizers
        self.layers = nn.ModuleList([EncodecVectorQuantization(config) for _ in range(config.num_quantizers)])

    def get_num_quantizers_for_bandwidth(self, bandwidth: Optional[float] = None) -> int:
        
        bw_per_q = math.log2(self.codebook_size) * self.frame_rate
        num_quantizers = self.num_quantizers
        if bandwidth is not None and bandwidth > 0.0:
            num_quantizers = int(max(1, math.floor(bandwidth * 1000 / bw_per_q)))
        return num_quantizers

    def encode(self, embeddings: torch.Tensor, bandwidth: Optional[float] = None) -> torch.Tensor:
        
        num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth)
        residual = embeddings
        all_indices = []
        for layer in self.layers[:num_quantizers]:
            indices = layer.encode(residual)
            quantized = layer.decode(indices)
            residual = residual - quantized
            all_indices.append(indices)
        out_indices = torch.stack(all_indices)
        return out_indices

    def decode(self, codes: torch.Tensor) -> torch.Tensor:
        
        quantized_out = torch.tensor(0.0, device=codes.device)
        for i, indices in enumerate(codes):
            layer = self.layers[i]
            quantized = layer.decode(indices)
            quantized_out = quantized_out + quantized
        return quantized_out


class EncodecPreTrainedModel(PreTrainedModel):
    

    config_class = EncodecConfig
    base_model_prefix = "encodec"
    main_input_name = "input_values"

    def _init_weights(self, module):
        
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Conv1d):
            nn.init.kaiming_normal_(module.weight)
            if module.bias is not None:
                k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
                nn.init.uniform_(module.bias, a=-k, b=k)
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LSTM):
            for name, param in module.named_parameters():
                if "weight" in name:
                    nn.init.xavier_uniform_(param)
                elif "bias" in name:
                    nn.init.constant_(param, 0.0)


ENCODEC_START_DOCSTRING = r


ENCODEC_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The EnCodec neural audio codec model.",
    ENCODEC_START_DOCSTRING,
)
class EncodecModel(EncodecPreTrainedModel):
    def __init__(self, config: EncodecConfig):
        super().__init__(config)
        self.config = config

        self.encoder = EncodecEncoder(config)
        self.decoder = EncodecDecoder(config)

        self.quantizer = EncodecResidualVectorQuantizer(config)

        self.bits_per_codebook = int(math.log2(self.config.codebook_size))
        if 2**self.bits_per_codebook != self.config.codebook_size:
            raise ValueError("The codebook_size must be a power of 2.")

        
        self.post_init()

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def _encode_frame(
        self, input_values: torch.Tensor, bandwidth: float, padding_mask: int
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        
        length = input_values.shape[-1]
        duration = length / self.config.sampling_rate

        if self.config.chunk_length_s is not None and duration > 1e-5 + self.config.chunk_length_s:
            raise RuntimeError(f"Duration of frame ({duration}) is longer than chunk {self.config.chunk_length_s}")

        scale = None
        if self.config.normalize:
            
            input_values = input_values * padding_mask
            mono = torch.sum(input_values, 1, keepdim=True) / input_values.shape[1]
            scale = mono.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-8
            input_values = input_values / scale

        embeddings = self.encoder(input_values)
        codes = self.quantizer.encode(embeddings, bandwidth)
        codes = codes.transpose(0, 1)
        return codes, scale

    def encode(
        self,
        input_values: torch.Tensor,
        padding_mask: torch.Tensor = None,
        bandwidth: Optional[float] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], EncodecEncoderOutput]:
        
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if bandwidth is None:
            bandwidth = self.config.target_bandwidths[0]
        if bandwidth not in self.config.target_bandwidths:
            raise ValueError(
                f"This model doesn't support the bandwidth {bandwidth}. "
                f"Select one of {self.config.target_bandwidths}."
            )

        _, channels, input_length = input_values.shape

        if channels < 1 or channels > 2:
            raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}")

        chunk_length = self.config.chunk_length
        if chunk_length is None:
            chunk_length = input_length
            stride = input_length
        else:
            stride = self.config.chunk_stride

        if padding_mask is None:
            padding_mask = torch.ones_like(input_values).bool()

        encoded_frames = []
        scales = []

        step = chunk_length - stride
        if (input_length % stride) - step != 0:
            raise ValueError(
                "The input length is not properly padded for batched chunked decoding. Make sure to pad the input correctly."
            )

        for offset in range(0, input_length - step, stride):
            mask = padding_mask[..., offset : offset + chunk_length].bool()
            frame = input_values[:, :, offset : offset + chunk_length]
            encoded_frame, scale = self._encode_frame(frame, bandwidth, mask)
            encoded_frames.append(encoded_frame)
            scales.append(scale)

        encoded_frames = torch.stack(encoded_frames)

        if not return_dict:
            return (encoded_frames, scales)

        return EncodecEncoderOutput(encoded_frames, scales)

    @staticmethod
    def _linear_overlap_add(frames: List[torch.Tensor], stride: int):
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        if len(frames) == 0:
            raise ValueError("`frames` cannot be an empty list.")

        device = frames[0].device
        dtype = frames[0].dtype
        shape = frames[0].shape[:-1]
        total_size = stride * (len(frames) - 1) + frames[-1].shape[-1]

        frame_length = frames[0].shape[-1]
        time_vec = torch.linspace(0, 1, frame_length + 2, device=device, dtype=dtype)[1:-1]
        weight = 0.5 - (time_vec - 0.5).abs()

        sum_weight = torch.zeros(total_size, device=device, dtype=dtype)
        out = torch.zeros(*shape, total_size, device=device, dtype=dtype)
        offset: int = 0

        for frame in frames:
            frame_length = frame.shape[-1]
            out[..., offset : offset + frame_length] += weight[:frame_length] * frame
            sum_weight[offset : offset + frame_length] += weight[:frame_length]
            offset += stride

        if sum_weight.min() == 0:
            raise ValueError(f"`sum_weight` minimum element must be bigger than zero: {sum_weight}`")

        return out / sum_weight

    def _decode_frame(self, codes: torch.Tensor, scale: Optional[torch.Tensor] = None) -> torch.Tensor:
        codes = codes.transpose(0, 1)
        embeddings = self.quantizer.decode(codes)
        outputs = self.decoder(embeddings)
        if scale is not None:
            outputs = outputs * scale.view(-1, 1, 1)
        return outputs

    def decode(
        self,
        audio_codes: torch.Tensor,
        audio_scales: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecDecoderOutput]:
        
        return_dict = return_dict or self.config.return_dict

        chunk_length = self.config.chunk_length
        if chunk_length is None:
            if len(audio_codes) != 1:
                raise ValueError(f"Expected one frame, got {len(audio_codes)}")
            audio_values = self._decode_frame(audio_codes[0], audio_scales[0])
        else:
            decoded_frames = []

            for frame, scale in zip(audio_codes, audio_scales):
                frames = self._decode_frame(frame, scale)
                decoded_frames.append(frames)

            audio_values = self._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1)

        
        if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]:
            audio_values = audio_values[..., : padding_mask.shape[-1]]

        if not return_dict:
            return (audio_values,)
        return EncodecDecoderOutput(audio_values)

    @add_start_docstrings_to_model_forward(ENCODEC_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=EncodecOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_values: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        bandwidth: Optional[float] = None,
        audio_codes: Optional[torch.Tensor] = None,
        audio_scales: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecOutput]:
        r
        return_dict = return_dict or self.config.return_dict

        if padding_mask is None:
            padding_mask = torch.ones_like(input_values).bool()

        if audio_codes is not None and audio_scales is None:
            raise ValueError("You specified `audio_codes` but did not specify the `audio_scales`")

        if audio_scales is not None and audio_codes is None:
            raise ValueError("You specified `audio_scales` but did not specify the `audio_codes`")

        if audio_scales is None and audio_codes is None:
            audio_codes, audio_scales = self.encode(input_values, padding_mask, bandwidth, False)

        audio_values = self.decode(audio_codes, audio_scales, padding_mask, return_dict=return_dict)[0]
        if not return_dict:
            return (audio_codes, audio_values)

        return EncodecOutput(audio_codes=audio_codes, audio_values=audio_values)