KVAE 1.0 on Habr | KVAE 2.0 on Habr | KVAE-Audio on Habr | Project page | Technical report

Hugging Face: Audio | Image 1.0 | Image 2.0 | Video 1.0 | Video 2.0 t4s8 | Video 2.0 t4s16

KVAE: audio, image, and video tokenizers

KVAE provides pretrained variational autoencoders for converting audio, images and videos into compact latent representations for diffusion generative models. All model classes use the same high-level workflow: from_pretrained -> encode -> latent distribution -> decode.

Available models

Modality	Python class	Hugging Face model	Inference
Audio	KVAEAudio	KVAE-Audio	inference_1d_kvae.py
Image	KVAEImage	KVAE-2D-1.0	inference_2d_kvae.py
Image	KVAEImage	KVAE-2D-2.0	inference_2d_kvae.py
Video	KVAEVideo	KVAE-3D-1.0	inference_3d_kvae.py
Video	KVAEVideo	KVAE-3D-2.0-t4s8	inference_3d_kvae.py
Video	KVAEVideo	KVAE-3D-2.0-t4s16	inference_3d_kvae.py

Highlights

KVAE-Audio is a continuous full-band 48 kHz tokenizer with 166.9M parameters and 64 latent channels. Under a fixed text-to-audio generator, it achieves the best CLAP, CE, PQ, and all reported FAD scores on AudioCaps among the compared autoencoders. On MUSDB18-HQ reconstruction, it leads all reported MEL, STFT, waveform, SI-SDR, SDR, and SNR metrics.

KVAE-Image 2.0 is the image tokenizer from KVAE 2.0, a family of image and video tokenizers designed to provide latent representations for diffusion models. It compresses RGB images by 8 x 8 into continuous 32-channel latents and reconstructs them with high fidelity.

KVAE-Video 2.0 is available with temporal compression 4 and spatial compression 8 x 8 or 16 x 16. The t4s8 model keeps 16 latent channels, while the t4s16 variant provides a more compact representation for higher spatial compression.

Show selected evaluation figures

KVAE-Audio latent-space qualities for generation

KVAE-Video 2.0 reconstruction

t4s8

t4s16

KVAE-Video 2.0 latent-space qualities for generation

Previous versions: KVAE 1.0

KVAE-2D-1.0 uses 8 x 8 spatial compression with 16 latent channels and provides the original image tokenizer released with KVAE 1.0.

KVAE-3D-1.0 uses 4 x 8 x 8 compression with 16 latent channels. It was evaluated at 540p, while the newer KVAE-Video 2.0 models target 720p evaluation and improved high-resolution processing.

Full metric tables, reconstruction comparisons, human evaluations, and audio generation examples are available in assets/docs/EVALUATION.md.

Quick start

Create an environment with Python 3.11 and the PyTorch 2.8.0 CUDA 12.8 build, then install this repository in editable mode:

conda create -n kvae_inference python=3.11
conda activate kvae_inference

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt
pip install --editable .

Inference

Run commands from the repository root. The first run downloads weights from Hugging Face Hub unless they are already cached.

Audio

python -m scripts.inference_1d_kvae \
    --dataset_folder ./assets/audio_test \
    --saving_folder ./outputs/audio

The released audio model expects mono 48 kHz input and runs in float32. Loading preserves the original channels and sample rate; it does not silently resample, downmix, crop, or normalize the input. Saved reconstructions are loudness-matched to their input using the same procedure as the original audio implementation.

Image

python -m scripts.inference_2d_kvae \
    --dataset_folder ./assets/image_test \
    --model KVAE_2.0 \
    --saving_folder ./outputs/images

Image inference uses bfloat16 and expects PNG inputs normalized to the model range. Use --img_size H,W to resize samples to a common shape and enable batching; without an explicit size, the inference script uses batch size one. Reconstructions can be saved as PNG files, and the script reports PSNR and LPIPS.

Video

python -m scripts.inference_3d_kvae \
    --dataset_folder ./assets/video_test \
    --model KVAE_2.0-t4s8 \
    --seg_len 16 \
    --saving_folder ./outputs/video

Video inference uses bfloat16 and expects one directory of PNG frames per video. --seg_len controls temporal chunking, while --input_norm selects the input normalization convention.

Temporal context is cached inside the causal convolution, residual, and sampling blocks. These caches are mutable, so one KVAEVideo instance cannot safely process independent samples concurrently or interleave their chunks. Use a separate model instance per concurrent stream. Sequential calls are supported because encode and decode reset their caches before and after each call.

Detailed Python API examples, input layouts, metrics, and script options are in assets/docs/INFERENCE.md.

A runnable example for all three modalities is available in scripts/inference_examples.ipynb.

Citation

@misc{kvae2026,
  title         = {KVAE: Family of Tokenizers for Multimodal Generative Models},
  author        = {Andrey Shutkin, Denis Parkhomenko, Ivan Kirillov,
                   Kirill Chernyshev, Kirill Malakhov, Ilia Vasiliev,
                   Ilia Trushkin, Valeriya Kobenko, David Chikovani,
                   Alexander Ivanov, Azat Saginbaev, Egor Silvestrov,
                   Ivan Mikheev, Konstantin Zakharov},
  year          = {2026},
  eprint        = {2608.05798},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  url           = {https://arxiv.org/abs/2608.05798}
}

License

The project is distributed under the terms of LICENSE.