| 57 | |
| 58 | |
| 59 | def test_compression(): |
| 60 | # Test encoding |
| 61 | input_dir = Path(__file__).parent / "assets" / "input" |
| 62 | output_dir = input_dir.parent / "encoded_output_quantizers" |
| 63 | args = { |
| 64 | "input": str(input_dir), |
| 65 | "output": str(output_dir), |
| 66 | "n_quantizers": 3, |
| 67 | "device": "cuda" if torch.cuda.is_available() else "cpu", |
| 68 | } |
| 69 | with argbind.scope(args): |
| 70 | run("encode") |
| 71 | |
| 72 | # Open .dac file |
| 73 | dac_file = output_dir / "sample_0.dac" |
| 74 | artifacts = np.load(dac_file, allow_pickle=True)[()] |
| 75 | codes = artifacts["codes"] |
| 76 | |
| 77 | # Ensure that the number of quantizers is correct |
| 78 | assert codes.shape[1] == 3 |
| 79 | |
| 80 | # Ensure that dtype of compression is uint16 |
| 81 | assert codes.dtype == np.uint16 |
| 82 | |
| 83 | |
| 84 | # CUDA_VISIBLE_DEVICES=0 python -m pytest tests/test_cli.py -s |