Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import random | |
| import numpy as np | |
| import torch | |
| from audiocraft.models import EncodecModel | |
| from audiocraft.modules import SEANetEncoder, SEANetDecoder | |
| from audiocraft.quantization import DummyQuantizer | |
| class TestEncodecModel: | |
| def _create_encodec_model(self, | |
| sample_rate: int, | |
| channels: int, | |
| dim: int = 5, | |
| n_filters: int = 3, | |
| n_residual_layers: int = 1, | |
| ratios: list = [5, 4, 3, 2], | |
| **kwargs): | |
| frame_rate = np.prod(ratios) | |
| encoder = SEANetEncoder(channels=channels, dimension=dim, n_filters=n_filters, | |
| n_residual_layers=n_residual_layers, ratios=ratios) | |
| decoder = SEANetDecoder(channels=channels, dimension=dim, n_filters=n_filters, | |
| n_residual_layers=n_residual_layers, ratios=ratios) | |
| quantizer = DummyQuantizer() | |
| model = EncodecModel(encoder, decoder, quantizer, frame_rate=frame_rate, | |
| sample_rate=sample_rate, channels=channels, **kwargs) | |
| return model | |
| def test_model(self): | |
| random.seed(1234) | |
| sample_rate = 24_000 | |
| channels = 1 | |
| model = self._create_encodec_model(sample_rate, channels) | |
| for _ in range(10): | |
| length = random.randrange(1, 10_000) | |
| x = torch.randn(2, channels, length) | |
| res = model(x) | |
| assert res.x.shape == x.shape | |
| def test_model_renorm(self): | |
| random.seed(1234) | |
| sample_rate = 24_000 | |
| channels = 1 | |
| model_nonorm = self._create_encodec_model(sample_rate, channels, renormalize=False) | |
| model_renorm = self._create_encodec_model(sample_rate, channels, renormalize=True) | |
| for _ in range(10): | |
| length = random.randrange(1, 10_000) | |
| x = torch.randn(2, channels, length) | |
| codes, scales = model_nonorm.encode(x) | |
| codes, scales = model_renorm.encode(x) | |
| assert scales is not None | |