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