KVAE-Audio

KVAE-Audio is a continuous, full-band (48 kHz) audio autoencoder. It compresses raw waveforms into compact continuous latents and reconstructs them with high fidelity across speech, music, and general sound. The model is designed not only for faithful reconstruction, but as a latent space for generative models — in our internal text-to-audio pipeline, swapping the autoencoder for KVAE-Audio improves generation quality under a fixed generator.

Inference instruction

Setup

First install PyTorch (adjust for your CUDA setup):

pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu128

Then clone the repository and install KVAE-Audio from the local package:

git clone https://github.com/kandinskylab/kvae-audio.git
pip install -e ./kvae-audio

Tested with Python 3.14.

Download weights

Checkpoint kvae-audio.pt is hosted on Hugging Face: kandinskylab/KVAE-Audio.

pip install huggingface_hub
hf download kandinskylab/KVAE-Audio kvae-audio.pt --local-dir .

Or from Python:

from huggingface_hub import hf_hub_download

weights_path = hf_hub_download("kandinskylab/KVAE-Audio", "kvae-audio.pt")

Run inference

Place input wavs in input_samples. Reconstructions are written to output_samples.

python scripts/infer.py \
  --weights_path kvae-audio.pt \
  --input_path input_samples \
  --output_path output_samples

Evaluate reconstructions

Compare reference wavs in input_samples against reconstructions in output_samples (same filenames). Writes metrics.csv to output_samples.

python scripts/eval.py \
  --input_path input_samples \
  --output_path output_samples

Python API

By default encode returns the mean latent (mu); pass sample=True for stochastic sampling. For end-to-end reconstruction, model(waveform, sample_rate)["audio"] is equivalent.

import torch
import soundfile as sf
from kvae_audio import KVAEAudio

model = KVAEAudio.load("kvae-audio.pt", map_location="cpu")
model.eval()

data, sr = sf.read("audio.wav")
waveform = torch.from_numpy(data).float().unsqueeze(0).unsqueeze(0)
length = waveform.shape[-1]

with torch.no_grad():
    latents, _, _, _ = model.encode(waveform, sample_rate=sr)
    audio = model.decode(latents)[..., :length]

sf.write("output.wav", audio.squeeze(0).T.numpy(), sr)

Evaluation results

Generative quality is established under a fixed generator — same DiT architecture, training data, and number of steps — varying only the autoencoder. We report objective generation metrics and blind human side-by-side below.

Evaluation of latent space qualities for generation

Audio Generation Examples

Below are qualitative examples generated from the same text prompts using four different models.

Example 1 (speech)

Prompt:

Low and gravelly, with a southern Russian accent, a man says slowly, <S>Я говорил тебе, что так и будет<E>. A sharp inhale precedes the line. The space is acoustically dry, with minimal room tone and a slight hiss.

---

Example 2 (background)

Prompt:

In a home kitchen, oil sizzles in a pan, a knife chops on a board, and a woman hums softly in Russian, a refrigerator humming beneath it all. Small kitchen, light reverberation. High-fidelity.

---

Example 3 (music)

Prompt:

In a large reverberant hall, a brass band launches into a lively march. Trumpets and cornets carry the melody while trombones drive the harmony, brisk and precise. Clean recording, natural reverb, faint equipment hiss.

---

AudioCaps test set

Song Describer

LibriSpeech test-clean

Reconstructions

Reconstruction is evaluated on open datasets across domains (the released weights directly substantiate these numbers). Baselines: MMAudio 44.1 kHz VAE, DACVAE from MovieGen Audio, SAME-L (Stable Audio 3 VAE).

AudioSet eval

MUSDB18-HQ

EARS

Citation

@misc{kvae_audio_2026,
    author = {Ivan Kirillov, Denis Parkhomenko, Alexander Ivanov, Azat Saginbaev, Egor Silvestrov, Denis Dimitrov},
    title = {KVAE-Audio: a full-band continuous audio tokenizer for generative models},
    howpublished = {\url{https://github.com/kandinskylab/kvae-audio}},
    year = 2026
}

Package layout

kvae-audio/
├── kvae_audio/
│   ├── model/
│   │   ├── kvae_audio.py
│   │   └── base.py
│   ├── metrics/loss.py
│   └── nn/layers.py
└── scripts/
    ├── infer.py
    └── eval.py