Spaces:
Sleeping
Sleeping
File size: 2,275 Bytes
541c6b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
"""
Tests for CLI.
"""
import subprocess
from pathlib import Path
import argbind
import numpy as np
import pytest
import torch
from audiotools import AudioSignal
from dac.__main__ import run
def setup_module(module):
data_dir = Path(__file__).parent / "assets"
data_dir.mkdir(exist_ok=True, parents=True)
input_dir = data_dir / "input"
input_dir.mkdir(exist_ok=True, parents=True)
for i in range(5):
signal = AudioSignal(np.random.randn(1000), 44_100)
signal.write(input_dir / f"sample_{i}.wav")
return input_dir
def teardown_module(module):
repo_root = Path(__file__).parent.parent
subprocess.check_output(["rm", "-rf", f"{repo_root}/tests/assets"])
@pytest.mark.parametrize("model_type", ["44khz", "24khz", "16khz"])
def test_reconstruction(model_type):
# Test encoding
input_dir = Path(__file__).parent / "assets" / "input"
output_dir = input_dir.parent / model_type / "encoded_output"
args = {
"input": str(input_dir),
"output": str(output_dir),
"device": "cuda" if torch.cuda.is_available() else "cpu",
"model_type": model_type,
}
with argbind.scope(args):
run("encode")
# Test decoding
input_dir = output_dir
output_dir = input_dir.parent / model_type / "decoded_output"
args = {
"input": str(input_dir),
"output": str(output_dir),
"model_type": model_type,
}
with argbind.scope(args):
run("decode")
def test_compression():
# Test encoding
input_dir = Path(__file__).parent / "assets" / "input"
output_dir = input_dir.parent / "encoded_output_quantizers"
args = {
"input": str(input_dir),
"output": str(output_dir),
"n_quantizers": 3,
"device": "cuda" if torch.cuda.is_available() else "cpu",
}
with argbind.scope(args):
run("encode")
# Open .dac file
dac_file = output_dir / "sample_0.dac"
artifacts = np.load(dac_file, allow_pickle=True)[()]
codes = artifacts["codes"]
# Ensure that the number of quantizers is correct
assert codes.shape[1] == 3
# Ensure that dtype of compression is uint16
assert codes.dtype == np.uint16
# CUDA_VISIBLE_DEVICES=0 python -m pytest tests/test_cli.py -s
|