codingutf8 2023 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 testing suite for the pytorch encodec model import copy import inspect import os import tempfile import unittest from typing import dict list tuple import numpy as np from datasets import audio loaddataset from transformers import autoprocessor encodecconfig from transformers testingutils import istorchavailable requiretorch slow torchdevice from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin configzeroinit floatstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import encodecmodel def prepareinputsdict config inputidsnone inputvaluesnone decoderinputidsnone attentionmasknone decoderattentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone if inputids is not none encoderdict inputids inputids else encoderdict inputvalues inputvalues decoderdict decoderinputids decoderinputids if decoderinputids is not none else return encoderdict decoderdict requiretorch class encodecmodeltester def init self parent batchsize needs to be an even number if the model has some outputs with batch dim 0 batchsize12 numchannels2 istrainingfalse intermediatesize40 hiddensize32 numfilters8 numresiduallayers1 upsamplingratios8 4 numlstmlayers1 codebooksize64 self parent parent self batchsize batchsize self numchannels numchannels self istraining istraining self intermediatesize intermediatesize self hiddensize hiddensize self numfilters numfilters self numresiduallayers numresiduallayers self upsamplingratios upsamplingratios self numlstmlayers numlstmlayers self codebooksize codebooksize def prepareconfigandinputsself inputvalues floatstensorself batchsize self numchannels self intermediatesize scale1 0 config self getconfig inputsdict inputvalues inputvalues return config inputsdict def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict def getconfigself return encodecconfig audiochannelsself numchannels chunkinsecnone hiddensizeself hiddensize numfiltersself numfilters numresiduallayersself numresiduallayers upsamplingratiosself upsamplingratios numlstmlayersself numlstmlayers codebooksizeself codebooksize def createandcheckmodelforwardself config inputsdict model encodecmodelconfigconfig totorchdevice eval inputvalues inputsdictinputvalues result modelinputvalues self parent assertequal result audiovalues shape self batchsize self numchannels self intermediatesize requiretorch class encodecmodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses encodecmodel if istorchavailable else isencoderdecoder true testpruning false testheadmasking false testresizeembeddings false pipelinemodelmapping featureextraction encodecmodel if istorchavailable else inputname inputvalues def prepareforclassself inputsdict modelclass returnlabelsfalse model does not have attention and does not support returning hidden states inputsdict super prepareforclassinputsdict modelclass returnlabelsreturnlabels if outputattentions in inputsdict inputsdict popoutputattentions if outputhiddenstates in inputsdict inputsdict popoutputhiddenstates return inputsdict def setupself self modeltester encodecmodeltesterself self configtester configtester self configclassencodecconfig hiddensize37 commonproperties hastextmodalityfalse def testconfigself self configtester runcommontests def testmodelforwardself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelforwardconfigandinputs def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames inputvalues paddingmask bandwidth self assertlistequalargnames lenexpectedargnames expectedargnames unittest skipthe encodecmodel is not transformers based thus it does not have inputsembeds logics def testinputsembedsself pass unittest skipthe encodecmodel is not transformers based thus it does not have inputsembeds logics def testmodelcommonattributesself pass unittest skipthe encodecmodel is not transformers based thus it does not have the usual attention logic def testretaingradhiddenstatesattentionsself pass unittest skipthe encodecmodel is not transformers based thus it does not have the usual attention logic def testtorchscriptoutputattentionsself pass unittest skipthe encodecmodel is not transformers based thus it does not have the usual hiddenstates logic def testtorchscriptoutputhiddenstateself pass def createandchecktorchscriptself config inputsdict if not self testtorchscript return configsnoinit configzeroinitconfig to be sure we have no nan configsnoinit torchscript true configsnoinit returndict false for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit model totorchdevice model eval inputs self prepareforclassinputsdict modelclass maininputname modelclass maininputname try maininput inputsmaininputname modelmaininput tracedmodel torch jit tracemodel maininput except runtimeerror self failcouldn t trace module with tempfile temporarydirectory as tmpdirname ptfilename os path jointmpdirname tracedmodel pt try torch jit savetracedmodel ptfilename except exception self failcouldn t save module try loadedmodel torch jit loadptfilename except exception self failcouldn t load module model totorchdevice model eval loadedmodel totorchdevice loadedmodel eval modelstatedict model statedict loadedmodelstatedict loadedmodel statedict nonpersistentbuffers for key in loadedmodelstatedict keys if key not in modelstatedict keys nonpersistentbufferskey loadedmodelstatedictkey loadedmodelstatedict key value for key value in loadedmodelstatedict items if key not in nonpersistentbuffers self assertequalsetmodelstatedict keys setloadedmodelstatedict keys modelbuffers listmodel buffers for nonpersistentbuffer in nonpersistentbuffers values foundbuffer false for i modelbuffer in enumeratemodelbuffers if torch equalnonpersistentbuffer modelbuffer foundbuffer true break self asserttruefoundbuffer modelbuffers popi modelbuffers listmodel buffers for nonpersistentbuffer in nonpersistentbuffers values foundbuffer false for i modelbuffer in enumeratemodelbuffers if torch equalnonpersistentbuffer modelbuffer foundbuffer true break self asserttruefoundbuffer modelbuffers popi modelsequal true for layername p1 in modelstatedict items if layername in loadedmodelstatedict p2 loadedmodelstatedictlayername if p1 data nep2 data sum 0 modelsequal false self asserttruemodelsequal avoid memory leak without this each call increase ram usage by 20mb even with this call there are still memory leak by 0 04mb self cleartorchjitclassregistry unittest skipthe encodecmodel is not transformers based thus it does not have the usual attention logic def testattentionoutputsself pass def testfeedforwardchunkingself originalconfig inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses torch manualseed0 config copy deepcopyoriginalconfig config chunklengths none config overlap none config samplingrate 10 model modelclassconfig model totorchdevice model eval inputs self prepareforclassinputsdict modelclass inputsinputvalues inputsinputvalues repeat1 1 10 hiddenstatesnochunk modelinputs0 torch manualseed0 config chunklengths 1 config overlap 0 config samplingrate 10 model modelclassconfig model totorchdevice model eval hiddenstateswithchunk modelinputs0 self asserttruetorch allclosehiddenstatesnochunk hiddenstateswithchunk atol1e3 unittest skipthe encodecmodel is not transformers based thus it does not have the usual hiddenstates logic def testhiddenstatesoutputself pass def testdeterminismself config inputsdict self modeltester prepareconfigandinputsforcommon def checkdeterminismfirst second outputs are not tensors but list since each sequence don t have the same framelength out1 first cpu numpy out2 second cpu numpy out1 out1np isnanout1 out2 out2np isnanout2 maxdiff np amaxnp absout1 out2 self assertlessequalmaxdiff 1e5 for modelclass in self allmodelclasses model modelclassconfig model totorchdevice model eval with torch nograd first modelself prepareforclassinputsdict modelclass0 second modelself prepareforclassinputsdict modelclass0 if isinstancefirst tuple and isinstancesecond tuple for tensor1 tensor2 in zipfirst second checkdeterminismtensor1 tensor2 else checkdeterminismfirst second def testmodeloutputsequivalenceself config inputsdict self modeltester prepareconfigandinputsforcommon def setnantensortozerot tt t 0 return t def checkequivalencemodel tupleinputs dictinputs additionalkwargs with torch nograd tupleoutput modeltupleinputs returndictfalse additionalkwargs dictoutput modeldictinputs returndicttrue additionalkwargs def recursivechecktupleobject dictobject if isinstancetupleobject list tuple for tupleiterablevalue dictiterablevalue in ziptupleobject dictobject recursivechecktupleiterablevalue dictiterablevalue elif isinstancetupleobject dict for tupleiterablevalue dictiterablevalue in zip tupleobject values dictobject values recursivechecktupleiterablevalue dictiterablevalue elif tupleobject is none return else self asserttrue torch allclose setnantensortozerotupleobject setnantensortozerodictobject atol1e5 msg tuple and dict output are not equal difference f torch maxtorch abstupleobject dictobject tuple has nan f torch isnantupleobject any and inf torch isinftupleobject dict has f nan torch isnandictobject any and inf torch isinfdictobject recursivechecktupleoutput dictoutput for modelclass in self allmodelclasses model modelclassconfig model totorchdevice model eval tupleinputs self prepareforclassinputsdict modelclass dictinputs self prepareforclassinputsdict modelclass checkequivalencemodel tupleinputs dictinputs def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters uniforminitparms conv ignoreinit lstm if param requiresgrad if anyx in name for x in uniforminitparms self asserttrue 1 0 param data mean 1e9 round 1e9 item 1 0 msgfparameter name of model modelclass seems not properly initialized elif not anyx in name for x in ignoreinit self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized def testidentityshortcutself config inputsdict self modeltester prepareconfigandinputs config useconvshortcut false self modeltester createandcheckmodelforwardconfig inputsdict def normalizearr norm np linalg normarr normalizedarr arr norm return normalizedarr def computermsearr1 arr2 arr1normalized normalizearr1 arr2normalized normalizearr2 return np sqrtarr1normalized arr2normalized 2 mean slow requiretorch class encodecintegrationtestunittest testcase def testintegration24khzself expectedrmse 1 5 0 0025 24 0 0 0015 expectedcodesums 1 5 371955 24 0 6659962 librispeechdummy loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation modelid facebookencodec24khz model encodecmodel frompretrainedmodelid totorchdevice processor autoprocessor frompretrainedmodelid librispeechdummy librispeechdummy castcolumnaudio audiosamplingrateprocessor samplingrate audiosample librispeechdummy1audioarray inputs processor rawaudioaudiosample samplingrateprocessor samplingrate returntensorspt totorchdevice for bandwidth expectedrmse in expectedrmse items with torch nograd use max bandwith for best possible reconstruction encoderoutputs model encodeinputsinputvalues bandwidthfloatbandwidth audiocodesums a0 sum cpu item for a in encoderoutputs0 make sure audio encoded codes are correct self assertlistequalaudiocodesums expectedcodesumsbandwidth audiocodes scales encoderoutputs totuple inputvaluesdec model decodeaudiocodes scales inputspaddingmask0 inputvaluesencdec model inputsinputvalues inputspaddingmask bandwidthfloatbandwidth 1 make sure forward and decode gives same result self asserttruetorch allcloseinputvaluesdec inputvaluesencdec atol1e3 make sure shape matches self asserttrueinputsinputvalues shape inputvaluesencdec shape arr inputsinputvalues0 cpu numpy arrencdec inputvaluesencdec0 cpu numpy make sure audios are more or less equal the rmse of two random gaussian noise vectors with n0 1 is around 1 0 rmse computermsearr arrencdec self asserttruermse expectedrmse def testintegration48khzself expectedrmse 3 0 0 001 24 0 0 0005 expectedcodesums 3 0 144259 146765 156435 176871 161971 24 0 1568553 1294948 1306190 1464747 1663150 librispeechdummy loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation modelid facebookencodec48khz model encodecmodel frompretrainedmodelid totorchdevice model model eval processor autoprocessor frompretrainedmodelid librispeechdummy librispeechdummy castcolumnaudio audiosamplingrateprocessor samplingrate audiosample librispeechdummy1audioarray transform mono to stereo audiosample np arrayaudiosample audiosample inputs processorrawaudioaudiosample samplingrateprocessor samplingrate returntensorspt to torchdevice for bandwidth expectedrmse in expectedrmse items with torch nograd use max bandwith for best possible reconstruction encoderoutputs model encode inputsinputvalues inputspaddingmask bandwidthfloatbandwidth returndictfalse audiocodesums a0 sum cpu item for a in encoderoutputs0 make sure audio encoded codes are correct self assertlistequalaudiocodesums expectedcodesumsbandwidth audiocodes scales encoderoutputs inputvaluesdec model decodeaudiocodes scales inputspaddingmask0 inputvaluesencdec model inputsinputvalues inputspaddingmask bandwidthfloatbandwidth 1 make sure forward and decode gives same result self asserttruetorch allcloseinputvaluesdec inputvaluesencdec atol1e3 make sure shape matches self asserttrueinputsinputvalues shape inputvaluesencdec shape arr inputsinputvalues0 cpu numpy arrencdec inputvaluesencdec0 cpu numpy make sure audios are more or less equal the rmse of two random gaussian noise vectors with n0 1 is around 1 0 rmse computermsearr arrencdec self asserttruermse expectedrmse def testbatch48khzself expectedrmse 3 0 0 001 24 0 0 0005 expectedcodesums 3 0 72410 79137 76694 90854 73023 82980 72707 54842 85561 81870 76953 48967 79315 85442 81479 107241 24 0 72410 79137 76694 90854 73023 82980 72707 54842 85561 81870 76953 48967 79315 85442 81479 107241 librispeechdummy loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation modelid facebookencodec48khz model encodecmodel frompretrainedmodelid totorchdevice processor autoprocessor frompretrainedmodelid chunklengths1 overlap0 01 librispeechdummy librispeechdummy castcolumnaudio audiosamplingrateprocessor samplingrate audiosamples np arrayaudiosamplearray audiosamplearray for audiosample in librispeechdummy2 audio inputs processorrawaudioaudiosamples samplingrateprocessor samplingrate returntensorspt inputvalues inputsinputvalues totorchdevice for bandwidth expectedrmse in expectedrmse items with torch nograd use max bandwith for best possible reconstruction encoderoutputs model encodeinputvalues bandwidthfloatbandwidth returndictfalse audiocodesums0 a00 sum cpu item for a in encoderoutputs0 audiocodesums1 a01 sum cpu item for a in encoderoutputs0 make sure audio encoded codes are correct self assertlistequalaudiocodesums0 expectedcodesumsbandwidth0 self assertlistequalaudiocodesums1 expectedcodesumsbandwidth1 audiocodes scales encoderoutputs inputvaluesdec model decodeaudiocodes scales0 inputvaluesencdec modelinputvalues bandwidthfloatbandwidth1 make sure forward and decode gives same result self asserttruetorch allcloseinputvaluesdec inputvaluesencdec atol1e3 make sure shape matches self asserttrueinputvalues shape inputvaluesencdec shape arr inputvalues0 cpu numpy arrencdec inputvaluesencdec0 cpu numpy make sure audios are more or less equal the rmse of two random gaussian noise vectors with n0 1 is around 1 0 rmse computermsearr arrencdec self asserttruermse expectedrmse coding utf 8 2023 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 testing suite for the pytorch encodec model batch_size needs to be an even number if the model has some outputs with batch dim 0 model does not have attention and does not support returning hidden states signature parameters is an ordereddict so arg_names order is deterministic to be sure we have no nan avoid memory leak without this each call increase ram usage by 20mb even with this call there are still memory leak by 0 04mb outputs are not tensors but list since each sequence don t have the same frame_length use max bandwith for best possible reconstruction make sure audio encoded codes are correct make sure forward and decode gives same result make sure shape matches make sure audios are more or less equal the rmse of two random gaussian noise vectors with n 0 1 is around 1 0 transform mono to stereo use max bandwith for best possible reconstruction make sure audio encoded codes are correct make sure forward and decode gives same result make sure shape matches make sure audios are more or less equal the rmse of two random gaussian noise vectors with n 0 1 is around 1 0 use max bandwith for best possible reconstruction make sure audio encoded codes are correct make sure forward and decode gives same result make sure shape matches make sure audios are more or less equal the rmse of two random gaussian noise vectors with n 0 1 is around 1 0
import copy
import inspect
import os
import tempfile
import unittest
from typing import Dict, List, Tuple

import numpy as np
from datasets import Audio, load_dataset

from transformers import AutoProcessor, EncodecConfig
from transformers.testing_utils import (
    is_torch_available,
    require_torch,
    slow,
    torch_device,
)

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
    ModelTesterMixin,
    _config_zero_init,
    floats_tensor,
)
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import EncodecModel


def prepare_inputs_dict(
    config,
    input_ids=None,
    input_values=None,
    decoder_input_ids=None,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if input_ids is not None:
        encoder_dict = {"input_ids": input_ids}
    else:
        encoder_dict = {"input_values": input_values}

    decoder_dict = {"decoder_input_ids": decoder_input_ids} if decoder_input_ids is not None else {}

    return {**encoder_dict, **decoder_dict}


@require_torch
class EncodecModelTester:
    def __init__(
        self,
        parent,
        
        batch_size=12,
        num_channels=2,
        is_training=False,
        intermediate_size=40,
        hidden_size=32,
        num_filters=8,
        num_residual_layers=1,
        upsampling_ratios=[8, 4],
        num_lstm_layers=1,
        codebook_size=64,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.num_channels = num_channels
        self.is_training = is_training
        self.intermediate_size = intermediate_size
        self.hidden_size = hidden_size
        self.num_filters = num_filters
        self.num_residual_layers = num_residual_layers
        self.upsampling_ratios = upsampling_ratios
        self.num_lstm_layers = num_lstm_layers
        self.codebook_size = codebook_size

    def prepare_config_and_inputs(self):
        input_values = floats_tensor([self.batch_size, self.num_channels, self.intermediate_size], scale=1.0)
        config = self.get_config()
        inputs_dict = {"input_values": input_values}
        return config, inputs_dict

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict

    def get_config(self):
        return EncodecConfig(
            audio_channels=self.num_channels,
            chunk_in_sec=None,
            hidden_size=self.hidden_size,
            num_filters=self.num_filters,
            num_residual_layers=self.num_residual_layers,
            upsampling_ratios=self.upsampling_ratios,
            num_lstm_layers=self.num_lstm_layers,
            codebook_size=self.codebook_size,
        )

    def create_and_check_model_forward(self, config, inputs_dict):
        model = EncodecModel(config=config).to(torch_device).eval()

        input_values = inputs_dict["input_values"]
        result = model(input_values)
        self.parent.assertEqual(
            result.audio_values.shape, (self.batch_size, self.num_channels, self.intermediate_size)
        )


@require_torch
class EncodecModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (EncodecModel,) if is_torch_available() else ()
    is_encoder_decoder = True
    test_pruning = False
    test_headmasking = False
    test_resize_embeddings = False
    pipeline_model_mapping = {"feature-extraction": EncodecModel} if is_torch_available() else {}
    input_name = "input_values"

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        
        inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
        if "output_attentions" in inputs_dict:
            inputs_dict.pop("output_attentions")
        if "output_hidden_states" in inputs_dict:
            inputs_dict.pop("output_hidden_states")
        return inputs_dict

    def setUp(self):
        self.model_tester = EncodecModelTester(self)
        self.config_tester = ConfigTester(
            self, config_class=EncodecConfig, hidden_size=37, common_properties=[], has_text_modality=False
        )

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_model_forward(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_forward(*config_and_inputs)

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = ["input_values", "padding_mask", "bandwidth"]
            self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have `inputs_embeds` logics")
    def test_inputs_embeds(self):
        pass

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have `inputs_embeds` logics")
    def test_model_common_attributes(self):
        pass

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic")
    def test_retain_grad_hidden_states_attentions(self):
        pass

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic")
    def test_torchscript_output_attentions(self):
        pass

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `hidden_states` logic")
    def test_torchscript_output_hidden_state(self):
        pass

    def _create_and_check_torchscript(self, config, inputs_dict):
        if not self.test_torchscript:
            return

        configs_no_init = _config_zero_init(config)  
        configs_no_init.torchscript = True
        configs_no_init.return_dict = False
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            model.to(torch_device)
            model.eval()
            inputs = self._prepare_for_class(inputs_dict, model_class)

            main_input_name = model_class.main_input_name

            try:
                main_input = inputs[main_input_name]
                model(main_input)
                traced_model = torch.jit.trace(model, main_input)
            except RuntimeError:
                self.fail("Couldn't trace module.")

            with tempfile.TemporaryDirectory() as tmp_dir_name:
                pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt")

                try:
                    torch.jit.save(traced_model, pt_file_name)
                except Exception:
                    self.fail("Couldn't save module.")

                try:
                    loaded_model = torch.jit.load(pt_file_name)
                except Exception:
                    self.fail("Couldn't load module.")

            model.to(torch_device)
            model.eval()

            loaded_model.to(torch_device)
            loaded_model.eval()

            model_state_dict = model.state_dict()
            loaded_model_state_dict = loaded_model.state_dict()

            non_persistent_buffers = {}
            for key in loaded_model_state_dict.keys():
                if key not in model_state_dict.keys():
                    non_persistent_buffers[key] = loaded_model_state_dict[key]

            loaded_model_state_dict = {
                key: value for key, value in loaded_model_state_dict.items() if key not in non_persistent_buffers
            }

            self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys()))

            model_buffers = list(model.buffers())
            for non_persistent_buffer in non_persistent_buffers.values():
                found_buffer = False
                for i, model_buffer in enumerate(model_buffers):
                    if torch.equal(non_persistent_buffer, model_buffer):
                        found_buffer = True
                        break

                self.assertTrue(found_buffer)
                model_buffers.pop(i)

            model_buffers = list(model.buffers())
            for non_persistent_buffer in non_persistent_buffers.values():
                found_buffer = False
                for i, model_buffer in enumerate(model_buffers):
                    if torch.equal(non_persistent_buffer, model_buffer):
                        found_buffer = True
                        break

                self.assertTrue(found_buffer)
                model_buffers.pop(i)

            models_equal = True
            for layer_name, p1 in model_state_dict.items():
                if layer_name in loaded_model_state_dict:
                    p2 = loaded_model_state_dict[layer_name]
                    if p1.data.ne(p2.data).sum() > 0:
                        models_equal = False

            self.assertTrue(models_equal)

            
            
            self.clear_torch_jit_class_registry()

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic")
    def test_attention_outputs(self):
        pass

    def test_feed_forward_chunking(self):
        (original_config, inputs_dict) = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            torch.manual_seed(0)
            config = copy.deepcopy(original_config)
            config.chunk_length_s = None
            config.overlap = None
            config.sampling_rate = 10

            model = model_class(config)
            model.to(torch_device)
            model.eval()
            inputs = self._prepare_for_class(inputs_dict, model_class)
            inputs["input_values"] = inputs["input_values"].repeat(1, 1, 10)

            hidden_states_no_chunk = model(**inputs)[0]

            torch.manual_seed(0)
            config.chunk_length_s = 1
            config.overlap = 0
            config.sampling_rate = 10

            model = model_class(config)
            model.to(torch_device)
            model.eval()

            hidden_states_with_chunk = model(**inputs)[0]
            self.assertTrue(torch.allclose(hidden_states_no_chunk, hidden_states_with_chunk, atol=1e-3))

    @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `hidden_states` logic")
    def test_hidden_states_output(self):
        pass

    def test_determinism(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        def check_determinism(first, second):
            
            out_1 = first.cpu().numpy()
            out_2 = second.cpu().numpy()
            out_1 = out_1[~np.isnan(out_1)]
            out_2 = out_2[~np.isnan(out_2)]
            max_diff = np.amax(np.abs(out_1 - out_2))
            self.assertLessEqual(max_diff, 1e-5)

        for model_class in self.all_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                first = model(**self._prepare_for_class(inputs_dict, model_class))[0]
                second = model(**self._prepare_for_class(inputs_dict, model_class))[0]

            if isinstance(first, tuple) and isinstance(second, tuple):
                for tensor1, tensor2 in zip(first, second):
                    check_determinism(tensor1, tensor2)
            else:
                check_determinism(first, second)

    def test_model_outputs_equivalence(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        def set_nan_tensor_to_zero(t):
            t[t != t] = 0
            return t

        def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}):
            with torch.no_grad():
                tuple_output = model(**tuple_inputs, return_dict=False, **additional_kwargs)
                dict_output = model(**dict_inputs, return_dict=True, **additional_kwargs)

                def recursive_check(tuple_object, dict_object):
                    if isinstance(tuple_object, (List, Tuple)):
                        for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object):
                            recursive_check(tuple_iterable_value, dict_iterable_value)
                    elif isinstance(tuple_object, Dict):
                        for tuple_iterable_value, dict_iterable_value in zip(
                            tuple_object.values(), dict_object.values()
                        ):
                            recursive_check(tuple_iterable_value, dict_iterable_value)
                    elif tuple_object is None:
                        return
                    else:
                        self.assertTrue(
                            torch.allclose(
                                set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5
                            ),
                            msg=(
                                "Tuple and dict output are not equal. Difference:"
                                f" {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`:"
                                f" {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has"
                                f" `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}."
                            ),
                        )

                recursive_check(tuple_output, dict_output)

        for model_class in self.all_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
            dict_inputs = self._prepare_for_class(inputs_dict, model_class)
            check_equivalence(model, tuple_inputs, dict_inputs)

    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                uniform_init_parms = ["conv"]
                ignore_init = ["lstm"]
                if param.requires_grad:
                    if any(x in name for x in uniform_init_parms):
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )
                    elif not any(x in name for x in ignore_init):
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )

    def test_identity_shortcut(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
        config.use_conv_shortcut = False
        self.model_tester.create_and_check_model_forward(config, inputs_dict)


def normalize(arr):
    norm = np.linalg.norm(arr)
    normalized_arr = arr / norm
    return normalized_arr


def compute_rmse(arr1, arr2):
    arr1_normalized = normalize(arr1)
    arr2_normalized = normalize(arr2)
    return np.sqrt(((arr1_normalized - arr2_normalized) ** 2).mean())


@slow
@require_torch
class EncodecIntegrationTest(unittest.TestCase):
    def test_integration_24kHz(self):
        expected_rmse = {
            "1.5": 0.0025,
            "24.0": 0.0015,
        }
        expected_codesums = {
            "1.5": [371955],
            "24.0": [6659962],
        }
        librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
        model_id = "facebook/encodec_24khz"

        model = EncodecModel.from_pretrained(model_id).to(torch_device)
        processor = AutoProcessor.from_pretrained(model_id)

        librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
        audio_sample = librispeech_dummy[-1]["audio"]["array"]

        inputs = processor(
            raw_audio=audio_sample,
            sampling_rate=processor.sampling_rate,
            return_tensors="pt",
        ).to(torch_device)

        for bandwidth, expected_rmse in expected_rmse.items():
            with torch.no_grad():
                
                encoder_outputs = model.encode(inputs["input_values"], bandwidth=float(bandwidth))

                audio_code_sums = [a[0].sum().cpu().item() for a in encoder_outputs[0]]

                
                self.assertListEqual(audio_code_sums, expected_codesums[bandwidth])

                audio_codes, scales = encoder_outputs.to_tuple()
                input_values_dec = model.decode(audio_codes, scales, inputs["padding_mask"])[0]
                input_values_enc_dec = model(
                    inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth)
                )[-1]

            
            self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3))

            
            self.assertTrue(inputs["input_values"].shape == input_values_enc_dec.shape)

            arr = inputs["input_values"][0].cpu().numpy()
            arr_enc_dec = input_values_enc_dec[0].cpu().numpy()

            
            
            rmse = compute_rmse(arr, arr_enc_dec)
            self.assertTrue(rmse < expected_rmse)

    def test_integration_48kHz(self):
        expected_rmse = {
            "3.0": 0.001,
            "24.0": 0.0005,
        }
        expected_codesums = {
            "3.0": [144259, 146765, 156435, 176871, 161971],
            "24.0": [1568553, 1294948, 1306190, 1464747, 1663150],
        }
        librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
        model_id = "facebook/encodec_48khz"

        model = EncodecModel.from_pretrained(model_id).to(torch_device)
        model = model.eval()
        processor = AutoProcessor.from_pretrained(model_id)

        librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
        audio_sample = librispeech_dummy[-1]["audio"]["array"]

        
        audio_sample = np.array([audio_sample, audio_sample])

        inputs = processor(raw_audio=audio_sample, sampling_rate=processor.sampling_rate, return_tensors="pt").to(
            torch_device
        )

        for bandwidth, expected_rmse in expected_rmse.items():
            with torch.no_grad():
                
                encoder_outputs = model.encode(
                    inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth), return_dict=False
                )
                audio_code_sums = [a[0].sum().cpu().item() for a in encoder_outputs[0]]

                
                self.assertListEqual(audio_code_sums, expected_codesums[bandwidth])
                audio_codes, scales = encoder_outputs
                input_values_dec = model.decode(audio_codes, scales, inputs["padding_mask"])[0]
                input_values_enc_dec = model(
                    inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth)
                )[-1]

            
            self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3))

            
            self.assertTrue(inputs["input_values"].shape == input_values_enc_dec.shape)

            arr = inputs["input_values"][0].cpu().numpy()
            arr_enc_dec = input_values_enc_dec[0].cpu().numpy()

            
            
            rmse = compute_rmse(arr, arr_enc_dec)
            self.assertTrue(rmse < expected_rmse)

    def test_batch_48kHz(self):
        expected_rmse = {
            "3.0": 0.001,
            "24.0": 0.0005,
        }
        expected_codesums = {
            "3.0": [
                [72410, 79137, 76694, 90854, 73023, 82980, 72707, 54842],
                [85561, 81870, 76953, 48967, 79315, 85442, 81479, 107241],
            ],
            "24.0": [
                [72410, 79137, 76694, 90854, 73023, 82980, 72707, 54842],
                [85561, 81870, 76953, 48967, 79315, 85442, 81479, 107241],
            ],
        }
        librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
        model_id = "facebook/encodec_48khz"

        model = EncodecModel.from_pretrained(model_id).to(torch_device)
        processor = AutoProcessor.from_pretrained(model_id, chunk_length_s=1, overlap=0.01)

        librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))

        audio_samples = [
            np.array([audio_sample["array"], audio_sample["array"]])
            for audio_sample in librispeech_dummy[-2:]["audio"]
        ]

        inputs = processor(raw_audio=audio_samples, sampling_rate=processor.sampling_rate, return_tensors="pt")
        input_values = inputs["input_values"].to(torch_device)
        for bandwidth, expected_rmse in expected_rmse.items():
            with torch.no_grad():
                
                encoder_outputs = model.encode(input_values, bandwidth=float(bandwidth), return_dict=False)
                audio_code_sums_0 = [a[0][0].sum().cpu().item() for a in encoder_outputs[0]]
                audio_code_sums_1 = [a[0][1].sum().cpu().item() for a in encoder_outputs[0]]

                
                self.assertListEqual(audio_code_sums_0, expected_codesums[bandwidth][0])
                self.assertListEqual(audio_code_sums_1, expected_codesums[bandwidth][1])

                audio_codes, scales = encoder_outputs
                input_values_dec = model.decode(audio_codes, scales)[0]
                input_values_enc_dec = model(input_values, bandwidth=float(bandwidth))[-1]

            
            self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3))

            
            self.assertTrue(input_values.shape == input_values_enc_dec.shape)

            arr = input_values[0].cpu().numpy()
            arr_enc_dec = input_values_enc_dec[0].cpu().numpy()

            
            
            rmse = compute_rmse(arr, arr_enc_dec)
            self.assertTrue(rmse < expected_rmse)