Stable Audio Open VAE

@inproceedings{evans2025stable,
  title={Stable Audio Open},
  author={Evans, Zach and Parker, Julian D. and Carr, CJ and Zukowski, Zack and Taylor, Josiah and Pons, Jordi},
  booktitle={ICASSP},
  year={2025},
  organization={IEEE}
}

arXiv

Pre-trained model weights

The “Oobleck” VAE from Stable Audio Open is the latent space of a generative model, not a codec proper. Unlike the neural codecs in this collection, it is a continuous VAE: no quantizer, no entropy model, and no rate knob. The “codec” is the latent cast to fp16 at a fixed 64× dimensionality reduction — exactly one operating point. This tutorial walks that single round trip on one example music clip: encode, cast to fp16, pack the raw bytes to a buffer, and decode reading only from the buffer.

The repository is gated on the Hugging Face hub — accept the license once and load_model works.

import io
import numpy as np
import torch
import torchaudio
import matplotlib.pyplot as plt
from datasets import load_dataset
from torchvision.transforms import ToPILImage
from IPython.display import Audio
from compressors.audio_eval import normalize_audio, pad_audio, align_length, psnr_db
from compressors.oobleck.evaluate_rate_distortion import load_model

torch.set_num_threads(8)
device = 'cuda:1'
torch_dtype = torch.float16

Load the codec

fp16 is the canonical recipe (and very slow on CPU — use a GPU). The stride and latent width come straight from the model config.

vae = load_model(device, torch_dtype)

fs = vae.config.sampling_rate
latent_channels = vae.config.decoder_input_channels
stride = int(np.prod(vae.config.downsampling_ratios))
print(f'sample_rate         = {fs}')
print(f'latent_channels     = {latent_channels}')
print(f'downsampling_ratios = {list(vae.config.downsampling_ratios)}  -> temporal stride {stride}')
print(f'parameters          = {sum(p.numel() for p in vae.parameters()) / 1e6:.1f} M')

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

Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: 
```
pip install accelerate
```
.

/home/dgj335/g/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py:205: UserWarning: The `local_dir_use_symlinks` argument is deprecated and ignored in `hf_hub_download`. Downloading to a local directory does not use symlinks anymore.
  warnings.warn(
/home/dgj335/g/lib/python3.12/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)

sample_rate         = 44100
latent_channels     = 64
downsampling_ratios = [2, 4, 4, 8, 8]  -> temporal stride 2048
parameters          = 156.1 M

Load an example clip

The example clips come from danjacobellis/music — 24 full-length music files whose audio column decodes through a torchcodec AudioDecoder (the source files are themselves MP3 at ~258 kbps). The parameters below pick one file and a short excerpt; clip 16 is stereo 44.1 kHz, and this VAE is the only baseline in the audio registry whose native format matches it — 44.1 kHz stereo in, 44.1 kHz stereo out — so there is no resampling and no mono-on-stereo dance. Normalization follows the shared recipe: zero-mean, scaled to \([-0.5, 0.5]\).

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)
duration = x.shape[-1] / fs
print(f'x shape  = {tuple(x.shape)}')
print(f'duration = {duration:.2f} s')
print(f'range    = [{x.min():.3f}, {x.max():.3f}]')

x shape  = (2, 220500)
duration = 5.00 s
range    = [-0.500, 0.477]

To visualize the clip, compute a log-mel spectrogram per channel and display it directly as a PIL image — low frequencies at the bottom. The same make_spectrogram helper is reused later to compare the reconstruction 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, name in enumerate(['left', 'right']):
    print(name)
    display(make_spectrogram(x[ch]))

left

right

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

Encode

The general recipe right-pads the time axis to a multiple of \(2^{16}\) before encoding. Here the pad is not a no-op: 5 s at 44.1 kHz is 220,500 samples, padded up to \(262{,}144 = 4 \cdot 2^{16}\), and the latent covers the padded length. Following the harness convention, the rate below counts the bits of the padded latent while distortion is measured on the trimmed 220,500 samples. The encoder returns a distribution, not a tensor; the codec takes its mode. The posterior is tight — its std is small relative to the spread of the latent itself — so mode and sample are nearly the same thing.

x_padded = pad_audio(x.unsqueeze(0), 2**16).to(device).to(torch_dtype)
print(f'x_padded shape = {tuple(x_padded.shape)}')

with torch.inference_mode():
    dist = vae.encode(x_padded).latent_dist
    latent = dist.mode().to(torch.float16)

print(f'latent shape = {tuple(latent.shape)}  dtype = {latent.dtype}')
print(f'latent std   = {latent.float().std():.3f}')
print(f'posterior std (median) = {dist.std.float().median():.4f}')

x_padded shape = (1, 2, 262144)

latent shape = (1, 64, 128)  dtype = torch.float16
latent std   = 0.918
posterior std (median) = 0.0495

Visualize the latent

64 channels by 128 frames; each latent frame summarizes 2048 stereo samples (≈46 ms). With no codebook and no integer grid, this channel-by-time matrix is the closest thing this codec has to a “what does the code look like” picture.

z = latent[0].float().cpu().numpy()
v = np.abs(z).max()
plt.figure(figsize=(10, 3), dpi=120)
plt.imshow(z, aspect='auto', cmap='RdBu', vmin=-v, vmax=v, interpolation='nearest')
plt.colorbar(label='latent value')
plt.xlabel('latent frame')
plt.ylabel('channel')
plt.title(f'all {z.shape[1]} latent frames (~{z.shape[1] * stride / fs:.1f} s, padding included)')
plt.show()

fp16 as the quantizer

There is no quantizer in this model; the only lossy step between latent and bytes is the cast to fp16. The canonical recipe runs the whole model in fp16, so the latent above was born fp16 and the cast cost is hidden. To see what 16-bit storage actually gives up, encode once in float32 and cast.

vae32 = load_model(device, torch.float32)
with torch.inference_mode():
    z32 = vae32.encode(x_padded.float()).latent_dist.mode()
z16 = z32.half()

err = z32 - z16.float()
snr = 10 * np.log10(z32.pow(2).mean().item() / err.pow(2).mean().item())
print(f'latent range (fp32)  = [{z32.min():.3f}, {z32.max():.3f}]')
print(f'fp16 cast max |err|  = {err.abs().max():.2e}')
print(f'fp16 cast latent SNR = {snr:.1f} dB')

del vae32, z32, z16
torch.cuda.empty_cache()

Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: 
```
pip install accelerate
```
.

latent range (fp32)  = [-4.710, 4.011]
fp16 cast max |err|  = 1.66e-03
fp16 cast latent SNR = 73.6 dB

Pack to bytes

From the printed shape: the latent is \(1 \times 64 \times 128 = 8{,}192\) fp16 values covering the padded \(2 \times 262{,}144\) input samples — 64 padded input samples per latent value (equivalently, DR \(= 2 \cdot \text{stride} / \text{channels} = 2 \cdot 2048 / 64 = 64\), stereo counted). With fp16 latents and int16 PCM both 16 bits per value, the compression ratio would equal the dimensionality reduction — but only against the padded length. The harness convention counts the bits of the padded latent while the PCM denominator uses the true 220,500 samples, so the measured CR is \(64 \times 220{,}500 / 262{,}144 \approx 53.8\); CR \(=\) DR \(= 64\) is recovered exactly only when the clip length is a multiple of \(2^{16}\).

There is no entropy coding at all: the buffer is exactly 2 * latent.numel() bytes, so 16 * latent.numel() bits equals 8 * len(buff.getbuffer()) identically. The rate is fixed regardless of content and sits well above where the conventional codecs operate — this VAE compresses by dimensionality reduction for a latent-diffusion backbone, not by approaching the entropy of the source.

latent_shape = tuple(latent.shape)   # side info
buff = io.BytesIO(latent.cpu().numpy().astype(np.float16).tobytes())

bits = 16 * latent.numel()                 # padded latent — the harness convention
assert bits == 8 * len(buff.getbuffer())   # no entropy coding

n_samples = x.shape[-1]
pcm_bits = 16 * 2 * n_samples              # int16 stereo PCM of the unpadded clip
kbps = bits / duration / 1000
CR = pcm_bits / bits
pad_factor = x_padded.shape[-1] / n_samples
print(f'len(buff)  = {len(buff.getbuffer())} bytes')
print(f'bits       = {bits}')
print(f'kbps       = {kbps:.2f}')
print(f'CR         = {CR:.2f}  (vs int16 stereo PCM)')
print(f'pad factor = {pad_factor:.4f}  ->  CR × pad factor = {CR * pad_factor:.2f} = DR')

len(buff)  = 16384 bytes
bits       = 131072
kbps       = 26.21
CR         = 53.83  (vs int16 stereo PCM)
pad factor = 1.1889  ->  CR × pad factor = 64.00 = DR

Decode from the buffer

The decoder reads only from buff — the in-memory latent is not reused. The only side information is the latent shape and the fp16 dtype. Because the fp16 cast already happened before packing, the byte round trip is exact by construction; the torch.equal check below confirms it, so decoding from the buffer is bit-identical to decoding the in-memory latent.

latent_rec = torch.from_numpy(
    np.frombuffer(buff.getvalue(), dtype=np.float16).copy()
).reshape(latent_shape).to(device).to(torch_dtype)
print(f'bit-exact round trip: {torch.equal(latent_rec.cpu(), latent.cpu())}')

with torch.inference_mode():
    x_hat = vae.decode(latent_rec).sample
x_hat = align_length(x_hat[0].clamp(-0.5, 0.5).float().cpu(), x.shape[-1])
print(f'x_hat shape = {tuple(x_hat.shape)}')
print(f'PSNR = {psnr_db(x, x_hat):.2f} dB')

bit-exact round trip: True
x_hat shape = (2, 220500)
PSNR = 32.26 dB

Reconstruction

In the time domain, a zoomed 30 ms window shows the waveform tracking closely; the deviations that remain are what a 64× dimensionality reduction costs.

t0 = int(2.0 * fs)
win = slice(t0, t0 + int(0.03 * fs))
tt = np.arange(win.start, win.stop) / fs
plt.figure(figsize=(10, 2.5), dpi=120)
plt.plot(tt, x[0, win].numpy(), label='original')
plt.plot(tt, x_hat[0, win].numpy(), label='reconstruction', alpha=0.8)
plt.xlabel('time (s)')
plt.ylabel('amplitude (left)')
plt.legend(loc='upper right')
plt.grid(alpha=0.3)
plt.show()

The log-mel images (channel 0, same make_spectrogram helper as the input) tell the frequency-domain story. The clip’s structure survives: the dense stack of horizontal harmonic lines through the low and mid bands, the quiet top band, and the onsets where energy briefly reaches into the upper bands are all in place. The difference is in texture — the original’s harmonic lines are crisp and continuous, while the reconstruction renders the upper ones grainier and slightly smeared, the visible cost of summarizing 64 input samples per latent value.

for name, sig in [('original', x), ('reconstruction', x_hat)]:
    print(name)
    display(make_spectrogram(sig[0].float().cpu()))

original

reconstruction

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