MuVAE

MuVAE is a continuous music autoencoder (VAE) for 48 kHz stereo audio. It was trained from scratch, on the Oobleck / YuE2 VAE architecture. It compresses a waveform 1920× into 64-dim latents at 25 Hz, so it can serve as a latent space for music generation models.

Audio	48 kHz, stereo
Latents	64 channels × 25 frames/s (1920× downsampling)
Parameters	132.6M (encoder 66.2M, decoder 66.4M), FP32 EMA weights
Architecture	Oobleck conv encoder/decoder with SnakeBeta activations, channels 64 × (1, 2, 4, 8, 16, 32), strides (2, 2, 4, 4, 5, 6)

Usage

import soundfile as sf, torch
from transformers import AutoModel

vae = AutoModel.from_pretrained("mrfakename/MuVAE", trust_remote_code=True).eval()

audio, sr = sf.read("song.wav", dtype="float32", always_2d=True)   # 48 kHz stereo
x = torch.from_numpy(audio.T.copy())[None]                          # (1, 2, samples)

z = vae.encode(x)            # (1, 64, frames): posterior mean at 25 Hz
y = vae.decode_audio(z)      # (1, 2, 1920 * frames - 64), tiled decoding for long audio

• encode(x, sample=True) samples from the posterior instead of returning its mean.
• decode_audio decodes in bounded tiles, with enough context that the output matches a full decode, so whole songs fit in memory. Pass chunked=False to decode in one pass.
• The reconstruction is aligned to the input shifted by 32 samples, i.e. x[..., 32:32 + y.shape[-1]].
• from_pretrained(..., decoder_only=True) loads only the decoder.

Training

• Data: ~10k hours of music audio seen during training, as random 2.56 s stereo crops at 48 kHz.
• Steps: 116,228 steps, 128 crops per step (~5.5 minutes of audio).
• Losses:
  • multi-resolution STFT loss (FFT sizes 2048…64, log magnitude + spectral convergence, stereo sum/difference)
  • KL (weight 1e-4)
  • multi-resolution STFT discriminator: adversarial loss (0.1) + feature matching (5.0), starting at step 5k
  • gradient balancing across the STFT, adversarial and feature-matching terms
• Optimisation: bf16, AdamW (β = 0.8, 0.99), lr 1.5e-4 for the generator and 3e-4 for the discriminator, 1k warmup steps, then constant. Gradient clipping at 10. EMA of the weights (decay 0.999) is what's released.

Training curves

These are training metrics at selected steps, averaged over 50 steps.

step	multi-res STFT loss	feature matching	disc. loss
6,400	2.15	0.80	1.54
20,000	1.95	0.85	1.02
50,000	1.84	0.82	0.93
80,000	1.80	0.81	0.84
116,200	1.79	0.79	0.74

Demos

Source audio and its MuVAE reconstruction (48 kHz stereo).

	Source	Reconstruction
Sample 1
Sample 2
Sample 3
Sample 4

Credits

• Architecture and inference code follow the YuE VAE. MuVAE is an independent model: the weights were trained from scratch, not fine-tuned from YuE.
• The Oobleck autoencoder comes from stable-audio-tools (MIT, © 2023 Stability AI), and SnakeBeta from BigVGAN (MIT, © 2022 NVIDIA). See THIRD_PARTY_NOTICES.md and licenses/.