singularity7/Muggier

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import random8 9import numpy as np10import torch11 12from audiocraft.models import EncodecModel13from audiocraft.modules import SEANetEncoder, SEANetDecoder14from audiocraft.quantization import DummyQuantizer15 16 17class TestEncodecModel:18 19    def _create_encodec_model(self,20                              sample_rate: int,21                              channels: int,22                              dim: int = 5,23                              n_filters: int = 3,24                              n_residual_layers: int = 1,25                              ratios: list = [5, 4, 3, 2],26                              **kwargs):27        frame_rate = np.prod(ratios)28        encoder = SEANetEncoder(channels=channels, dimension=dim, n_filters=n_filters,29                                n_residual_layers=n_residual_layers, ratios=ratios)30        decoder = SEANetDecoder(channels=channels, dimension=dim, n_filters=n_filters,31                                n_residual_layers=n_residual_layers, ratios=ratios)32        quantizer = DummyQuantizer()33        model = EncodecModel(encoder, decoder, quantizer, frame_rate=frame_rate,34                             sample_rate=sample_rate, channels=channels, **kwargs)35        return model36 37    def test_model(self):38        random.seed(1234)39        sample_rate = 24_00040        channels = 141        model = self._create_encodec_model(sample_rate, channels)42        for _ in range(10):43            length = random.randrange(1, 10_000)44            x = torch.randn(2, channels, length)45            res = model(x)46            assert res.x.shape == x.shape47 48    def test_model_renorm(self):49        random.seed(1234)50        sample_rate = 24_00051        channels = 152        model_nonorm = self._create_encodec_model(sample_rate, channels, renormalize=False)53        model_renorm = self._create_encodec_model(sample_rate, channels, renormalize=True)54 55        for _ in range(10):56            length = random.randrange(1, 10_000)57            x = torch.randn(2, channels, length)58            codes, scales = model_nonorm.encode(x)59            codes, scales = model_renorm.encode(x)60            assert scales is not None61