DreamVAE
A distilled decoder for the ACE-Step 1.5
Oobleck VAE. Drop-in replacement: same [B, 64, T] latent input, same
[B, 2, 1920*T] stereo output at 48 kHz. 51.7M parameters (61% of the
teacher), 0.24 dB SNR degradation vs. the teacher across 20 diverse music
tracks, and 37.2 ms to decode 60 seconds of audio on an RTX 5090 TensorRT
FP16 engine (8.66x faster than the teacher in PyTorch, 1.45x faster than
the teacher's own TRT engine).
Usage
import json
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
repo = "daydreamlive/DreamVAE"
weights = hf_hub_download(repo, "model.safetensors")
config = hf_hub_download(repo, "config.json")
mod_py = hf_hub_download(repo, "modeling.py")
import sys, os
sys.path.insert(0, os.path.dirname(mod_py))
from modeling import FastOobleckDecoder
cfg = json.load(open(config))
model = FastOobleckDecoder(
channels=cfg["channels"],
input_channels=cfg["input_channels"],
audio_channels=cfg["audio_channels"],
upsampling_ratios=cfg["upsampling_ratios"],
channel_multiples=cfg["channel_multiples"],
).eval().to("cuda")
model.load_state_dict(load_file(weights))
# latents: [B, 64, T] -> audio: [B, 2, 1920 * T]
with torch.no_grad():
audio = model(latents)
Pair with the stock ACE-Step encoder; only the decoder is distilled.
ONNX export is at onnx/model.onnx. Build a TRT engine with
scripts/export_trt.py.
Results
Quality (20 diverse FMA tracks, 10 s clips):
| vs. original audio | SNR (dB) | STFT dist | Mel dist |
|---|---|---|---|
| Teacher OobleckDecoder | 5.0 (std 2.5) | 0.698 | 0.733 |
| DreamVAE | 4.7 (std 2.4) | 0.695 | 0.742 |
| Degradation | 0.24 dB | better | +0.008 |
The VAE's own reconstruction error (~5 dB SNR) dominates; DreamVAE's additional error is well under a quarter of that. STFT distance is marginally better for the student; mel distance is 1.1% worse. Informal listening finds the difference inaudible.
Speed (RTX 5090, 60 s stereo audio at 48 kHz, TensorRT FP16):
| Backend | Time (60 s) | vs. teacher PyTorch |
|---|---|---|
| Teacher PyTorch | 321.9 ms | 1.00x |
| Teacher TRT FP16 | 54.0 ms | 5.97x |
| DreamVAE PyTorch | 241.1 ms | 1.34x |
| DreamVAE TRT FP16 | 37.2 ms | 8.66x |
Distillation alone is modest (1.45x over the teacher's own TRT engine). The compelling number is the composition with TRT: what a deployment sees moving from stock PyTorch VAE to DreamVAE on TRT is 8.66x.
Architecture
| Teacher | DreamVAE | |
|---|---|---|
| Parameters | 84.4M | 51.7M (61%) |
| Channel multiples | [1, 2, 4, 8, 16] | [1, 2, 4, 8, 8] |
| Residuals/block | 3 (dil 1, 3, 9) | 2 (dil 1, 3) |
| Upsampling ratios | [10, 6, 4, 4, 2] | [10, 6, 4, 4, 2] |
| Activation | Snake1d | Snake1d |
| Hop length | 1920 | 1920 |
Trained for 650K steps on 8,000 FMA Small tracks (two-phase: 500K reconstruction + 150K adversarial). Released checkpoint is step 635K, selected as the HF-energy-match peak from a six-point trajectory over the adversarial phase.
Files
.
├── README.md
├── config.json architecture + training hyperparameters
├── model.safetensors FP32 weights (207 MB, 51.7M params)
├── modeling.py self-contained FastOobleckDecoder class
├── onnx/model.onnx ONNX export
└── scripts/ training, eval, speed bench, TRT export/verify
Limitations
- Signal-level metrics only (no MUSHRA / PESQ / preference tests).
- Evaluated on VAE round-trip, not end-to-end through the full text-to-DiT-to-VAE pipeline. FAD on generated samples is future work.
- One point on the quality-speed frontier (
[1,2,4,8,8], 2 residual units). Width variants not evaluated. - Trained in FP32. INT8 / FP8 quantization would likely stack another 1.5-2x on top of FP16 and is not explored here.
Citation
Technical note (methodology, training recipe with source-paper attribution, adversarial-phase trajectory, and negative results on encoder distillation) is forthcoming on arXiv. Citation block will be added here on release.
License
Apache 2.0, matching the upstream ACE-Step 1.5 license.