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.

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.

Demos

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

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/.