| |
| |
| |
| |
| |
|
|
| 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 |
|
|
