STFT-VAE

Inference for a continuous (VAE) neural audio codec built on the STFT. Audio is transformed to a complex spectrogram, a transformer encoder downsamples it to a low-rate continuous latent, and a transformer decoder + ISTFT reconstruct the waveform — phase is preserved through the STFT/ISTFT rather than learned by a vocoder.

Pretrained model: 24 kHz, 5 downsampling stages → 3.125 Hz latent, 128-dim continuous latent (no quantization). ~116M parameters.

Install

pip install stftvae            # library only
pip install "stftvae[cli]"     # + command-line tool (audio IO / resampling)

or from source:

pip install -e ".[cli]"

Usage

import torch
from stftvae import STFTVAE

vae = STFTVAE.from_pretrained("fluxions/stftvae", device="cuda")

audio = torch.randn(1, 24000)          # (1, T) mono @ 24 kHz

# round-trip
recons = vae.reconstruct(audio)        # (1, 1, T)

# or encode / decode separately
latent, length = vae.encode(audio)     # (1, 128, T_latent) continuous latent
recons = vae.decode(latent, length=length)

from_pretrained accepts a HuggingFace repo id or a local directory containing config.json + model.safetensors.

Command line (reconstruct a file; decodes/resamples any format via torchcodec):

python -m stftvae input.wav output.wav --model fluxions/stftvae --device cuda

Demo

A Gradio app in demo/ — upload or record audio, listen to the round-trip through the codec, and see the latent shape/rate:

pip install "stftvae[cli]" gradio
python demo/app.py          # → http://localhost:7860

Inputs are resampled to 24 kHz mono automatically (any format torchcodec can decode) and truncated to 30 s to keep CPU inference snappy. To deploy as a HuggingFace Space, create a Gradio Space and upload demo/app.py as app.py together with demo/requirements.txt.

Comparisons

Speech reconstruction quality vs. Mimi, SNAC, and MioCodec over 8 LibriTTS-R clips. STFT-VAE runs at a 3.125 Hz continuous latent — 4–15× lower frame rate than the others — yet outscores SNAC on every metric.

▶ Listen to the side-by-side comparison — inline players + objective metrics (SI-SDR / PESQ / STOI).

Reproduce it under comparisons/ (fetch_speech.py → run.py → run_miocodec.py → metrics.py → build_page.py).

Notes

• The encoder returns the deterministic posterior mean (mu); no sampling noise is added at inference.
• config.json (architecture) is required alongside model.safetensors — the weights don't encode layer shapes, so the config builds the model graph before loading.
• Dependencies: torch, safetensors, huggingface_hub (core); soundfile, torchcodec (only for the CLI).

Citation

@software{stftvae_2026,
  author = {Coultas Blum, Harry},
  month = {07},
  title = {{stftvae}},
  url = {https://github.com/fluxions-ai/stftvae},
  version = {0.1.0},
  year = {2026}
}

License

MIT — see LICENSE.