DAC-JAX: A JAX Implementation of the Descript Audio Codec

TL;DR

DAC-JAX is an open-source JAX implementation of the Descript Audio Codec that enables efficient audio compression and decompression, supporting model weight reuse and device parallelism, with performance benchmarks on different GPU types.

Contribution

This work introduces a JAX-based implementation of DAC that supports model reuse, device parallelism, and provides performance benchmarking across GPU types.

Findings

DAC-JAX outperforms original DAC on consumer GPUs for all chunk sizes.

On high-performance GPUs, DAC-JAX is better for small chunks but worse for large chunks.

The implementation supports processing long audio files via overlapping chunks.

Abstract

We present an open-source implementation of the Descript Audio Codec (DAC) using Google's JAX ecosystem of Flax, Optax, Orbax, AUX, and CLU. Our codebase enables the reuse of model weights from the original PyTorch DAC, and we confirm that the two implementations produce equivalent token sequences and decoded audio if given the same input. We provide a training and fine-tuning script which supports device parallelism, although we have only verified it using brief training runs with a small dataset. Even with limited GPU memory, the original DAC can compress or decompress a long audio file by processing it as a sequence of overlapping "chunks." We implement this feature in JAX and benchmark the performance on two types of GPUs. On a consumer-grade GPU, DAC-JAX outperforms the original DAC for compression and decompression at all chunk sizes. However, on a high-performance, cluster-based GPU, DAC-JAX outperforms the original DAC for small chunk sizes but performs worse for large chunks.