Descript Audio Codec (DAC)

@inproceedings{kumar2023high,
  title={High-Fidelity Audio Compression with Improved RVQGAN},
  author={Kumar, Rithesh and Seetharaman, Prem and Luebs, Alejandro and Kumar, Ishaan and Kumar, Kundan},
  booktitle={NeurIPS},
  year={2023}
}

arXiv

Pre-trained model weights

DAC is a GAN-trained codec built on residual vector quantization (RVQGAN): the encoder produces one latent frame per hop, and a stack of codebooks quantizes each frame residually, so the bitrate is chosen by how many codebook rows you keep. This notebook uses the 16 kHz mono variant.

import io
import math
import numpy as np
import torch
import torchaudio
import torchaudio.functional as taf
import matplotlib.pyplot as plt
from datasets import load_dataset
from IPython.display import Audio
from torchvision.transforms import ToPILImage

from compressors.audio_eval import normalize_audio, align_length, psnr_db
from compressors.dac.evaluate_rate_distortion import load_model

torch.set_num_threads(8)

Load the codec

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float32

model = load_model(device, torch_dtype)

model_rate = model.config.sampling_rate
codebook_size = model.config.codebook_size
bits_per_code = int(math.log2(codebook_size))
hop = int(np.prod(model.config.downsampling_ratios))

print(f"sampling_rate = {model_rate} Hz")
print(f"codebook_size = {codebook_size}  ⇒  {bits_per_code} bits/code")
print(f"hop           = {hop} samples  ⇒  {model_rate / hop:.0f} frames/s")

Skipping import of cpp extensions due to incompatible torch version 2.11.0+cu130 for torchao version 0.15.0             Please see https://github.com/pytorch/ao/issues/2919 for more info

sampling_rate = 16000 Hz
codebook_size = 1024  ⇒  10 bits/code
hop           = 320 samples  ⇒  50 frames/s

Load an example clip

A 5 s stereo 44.1 kHz excerpt from danjacobellis/music (24 full-length music files, train split; each row’s audio is a torchcodec AudioDecoder), normalized to zero-mean [-0.5, 0.5]. The source files are themselves MP3 (~258 kbps), so the reference has already been through one lossy codec.

CLIP_INDEX = 16        # which of the 24 example files
START_SECONDS = 150.0  # offset into the file, rounded to the nearest sample
DURATION_SECONDS = 5.0 # excerpt length in seconds

ds = load_dataset("danjacobellis/music", split="train")
samples = ds[CLIP_INDEX]["audio"].get_all_samples()
fs = samples.sample_rate
i0 = round(START_SECONDS * fs)
x = samples.data[:, i0 : i0 + round(DURATION_SECONDS * fs)]
x = normalize_audio(x)

n_samples = x.shape[-1]
duration = n_samples / fs
print(f"x.shape = {tuple(x.shape)}  ({duration:.1f} s at {fs} Hz)  range = [{x.min().item():.3f}, {x.max().item():.3f}]")

x.shape = (2, 220500)  (5.0 s at 44100 Hz)  range = [-0.500, 0.477]

Each channel is visualized as a log-mel spectrogram rendered directly as a PIL image — low frequencies at the bottom. The same helper is reused later to compare reconstructions against the original.

mel = torchaudio.transforms.MelSpectrogram(sample_rate=fs, n_fft=4096, hop_length=512)

def make_spectrogram(x_1ch):
    S = mel(x_1ch).log()
    S = ((S - S.mean()) / (3 * S.std()) + 0.5).clamp(0, 1)
    return ToPILImage()(S.flip(0))  # low frequencies at the bottom

for ch in range(x.shape[0]):
    print(f"channel {ch}")
    display(make_spectrogram(x[ch]))

channel 0

channel 1

Audio(x.numpy(), rate=fs, normalize=False)

Mono-on-stereo

The 16 kHz model is mono, but the reference is stereo and is never downmixed: resample 44.1 kHz → 16 kHz, stack the two channels as a batch of 2 with shape (2, 1, L), encode both, and sum the bits across channels.

x_16k = taf.resample(x, fs, model_rate)
x_in = x_16k.unsqueeze(1).to(device).to(torch_dtype)
print(f"x_in.shape = {tuple(x_in.shape)}")

x_in.shape = (2, 1, 80000)

Encode at n_quantizers=8

One encoder pass yields a (n_codebooks, n_frames) grid of codebook indices per channel — here 8 of the 12 codebooks the model trains with.

with torch.inference_mode():
    codes_8 = model.encode(x_in, n_quantizers=8).audio_codes

print(f"codes_8.shape = {tuple(codes_8.shape)}  (channels, codebooks, frames)")
print(f"index range   = [{int(codes_8.min())}, {int(codes_8.max())}]  of [0, {codebook_size})")

plt.figure(figsize=(10, 2))
plt.imshow(codes_8[0].cpu().numpy(), aspect="auto", interpolation="nearest", cmap="viridis")
plt.xlabel("frame")
plt.ylabel("codebook")
plt.title("RVQ code grid, left channel, n_quantizers=8")
plt.colorbar(fraction=0.02)
plt.show()

codes_8.shape = (2, 8, 250)  (channels, codebooks, frames)
index range   = [0, 1023]  of [0, 1024)

Bitrate scalability by truncation

RVQ quantizes greedily, stage by stage, so n_quantizers only truncates the loop: encoding the same input at n_quantizers=2 yields literally the first two rows of the deeper grid. Choosing the rate is choosing how many rows to keep.

with torch.inference_mode():
    codes_2 = model.encode(x_in, n_quantizers=2).audio_codes

print(f"codes_2.shape = {tuple(codes_2.shape)}")
print(f"codes_2 == codes_8[:, :2, :] : {torch.equal(codes_2, codes_8[:, :2, :])}")

plt.figure(figsize=(10, 0.8))
plt.imshow(codes_2[0].cpu().numpy(), aspect="auto", interpolation="nearest", cmap="viridis")
plt.xlabel("frame")
plt.ylabel("codebook")
plt.title("RVQ code grid, left channel, n_quantizers=2")
plt.show()

codes_2.shape = (2, 2, 250)
codes_2 == codes_8[:, :2, :] : True

Pack codes to bytes

The buffer round trip uses the n_quantizers=8 codes. Analytically each index is worth log₂(1024) = 10 bits; packing as uint16 spends 16, so the buffer is 1.6× the analytic size — an entropy coder over the indices would close most of that gap.

codes_np = codes_8.cpu().numpy().astype(np.uint16)
codes_shape = codes_np.shape                     # side info
buff = io.BytesIO(codes_np.tobytes())

bits_analytic = bits_per_code * codes_8.numel()  # summed over both channels
kbps = bits_analytic / duration / 1000
cr = 16 * 2 * n_samples / bits_analytic          # vs int16 stereo PCM
n_bytes = len(buff.getbuffer())

print(f"analytic     : {bits_analytic} bits  =  {kbps:.1f} kbps   CR = {cr:.0f}× vs int16 stereo PCM")
print(f"packed uint16: {n_bytes} bytes  =  {8 * n_bytes / duration / 1000:.1f} kbps")

analytic     : 40000 bits  =  8.0 kbps   CR = 176× vs int16 stereo PCM
packed uint16: 8000 bytes  =  12.8 kbps

Decode from the buffer

The decoder sees only buff (plus the code-grid shape as side info) — the codes are recovered from the bytes alone before any model call.

codes_rec = torch.from_numpy(
    np.frombuffer(buff.getvalue(), dtype=np.uint16).astype(np.int64).reshape(codes_shape)
).to(device)
print(f"codes bit-exact through the buffer: {torch.equal(codes_rec, codes_8)}")

with torch.inference_mode():
    y = model.decode(audio_codes=codes_rec).audio_values
    if y.dim() == 3:
        y = y.squeeze(1)

x_hat = align_length(taf.resample(y.float().cpu(), model_rate, fs), n_samples)
print(f"x_hat.shape = {tuple(x_hat.shape)}")
print(f"PSNR (n_quantizers=8) = {psnr_db(x, x_hat):.2f} dB")

codes bit-exact through the buffer: True

x_hat.shape = (2, 220500)
PSNR (n_quantizers=8) = 25.33 dB

# Side-by-side contrast (decoded from the in-memory n=2 codes, not the round trip)
with torch.inference_mode():
    y2 = model.decode(audio_codes=codes_2).audio_values
    if y2.dim() == 3:
        y2 = y2.squeeze(1)

x_hat_2 = align_length(taf.resample(y2.float().cpu(), model_rate, fs), n_samples)
print(f"PSNR (n_quantizers=2) = {psnr_db(x, x_hat_2):.2f} dB    at {bits_per_code * codes_2.numel() / duration / 1000:.1f} kbps")

PSNR (n_quantizers=2) = 26.11 dB    at 2.0 kbps

Reconstruction quality

A 16 kHz speech-domain codec evaluated on 44.1 kHz stereo music is out of its comfort zone: everything above 8 kHz is discarded before encoding even starts. On the mel images below that band limit is the dark, empty region spanning roughly the top quarter of both reconstructions — on a mel axis with the original’s 22.05 kHz ceiling, 8 kHz sits about 72% of the way up — while the original keeps harmonic structure all the way to the top. Below the cutoff the codec’s harmonics are visibly smeared relative to the original, slightly more so at n_quantizers=2. Waveform PSNR is low by construction and does not track the rate knob at all: on this clip the 2 kbps reconstruction measures 26.1 dB and the 8 kbps one 25.3 dB — four times the bits for a slightly lower PSNR — because the GAN-trained decoder synthesizes plausible texture rather than matching the waveform sample-for-sample, so the extra codebooks buy detail the metric barely sees (compare the two players below). This is exactly why the cross-codec comparison uses rate-matched sweeps over many clips and metrics rather than single numbers.

print("Original")
display(make_spectrogram(x[0]))
print("DAC, n_quantizers=8")
display(make_spectrogram(x_hat[0].float().cpu()))
print("DAC, n_quantizers=2")
display(make_spectrogram(x_hat_2[0].float().cpu()))

Original

DAC, n_quantizers=8

DAC, n_quantizers=2

Audio(x_hat.numpy(), rate=fs, normalize=False)

Audio(x_hat_2.numpy(), rate=fs, normalize=False)