maolan_generate/acestep/
vae.rs

1//! Burn port of the decoder side of the diffusers `AutoencoderOobleck` VAE
2//! used by ACE-Step 1.5 (see `vae/config.json` in the checkpoint).
3//!
4//! Maps DiT output latents of shape `[B, T, decoder_input_channels]` (25 Hz,
5//! time-major) to 48 kHz stereo audio of shape `[B, audio_channels,
6//! T * hop_length]`, where `hop_length` is the product of
7//! `downsampling_ratios` (1920 for the released model).
8//!
9//! All convolutions are weight-normalized, matching
10//! `torch.nn.utils.weight_norm` in the reference: each conv stores
11//! `weight_g` `[out, 1, 1]` and `weight_v` (the unnormalized kernel) and the
12//! effective weight is `g * v / (||v|| + 1e-12)` with the norm taken per
13//! output channel over the remaining dims.
14//!
15//! # Canonical burnpack tensor names
16//!
17//! Tensor names are the module paths of [`OobleckDecoder`]. The offline
18//! converter produces them from the official
19//! `vae/diffusion_pytorch_model.safetensors` names by stripping the leading
20//! `decoder.` prefix (e.g. `decoder.block.0.conv_t1.weight_v` becomes
21//! `block.0.conv_t1.weight_v`):
22//!
23//! - `conv1.{weight_g, weight_v, bias}` — input conv, k=7, p=3
24//! - `block.{i}.snake1.{alpha, beta}`
25//! - `block.{i}.conv_t1.{weight_g, weight_v, bias}` — upsampling transpose
26//!   conv, k = 2*stride, p = ceil(stride/2)
27//! - `block.{i}.res_unit{1,2,3}.snake{1,2}.{alpha, beta}`
28//! - `block.{i}.res_unit{1,2,3}.conv1.{weight_g, weight_v, bias}` — k=7,
29//!   dilation 1/3/9, p = 3*dilation
30//! - `block.{i}.res_unit{1,2,3}.conv2.{weight_g, weight_v, bias}` — k=1
31//! - `snake1.{alpha, beta}`
32//! - `conv2.{weight_g, weight_v}` — output conv, k=7, p=3, **no bias**
33//!
34//! `block.{i}` enumerates the decoder blocks in decode order, i.e. using
35//! `downsampling_ratios` reversed (`[10, 6, 4, 4, 2]` for the released
36//! model).
37
38use std::path::Path;
39
40use anyhow::{Context, Result};
41use burn::module::{Module, Param};
42use burn::prelude::Backend;
43use burn::tensor::Tensor;
44use burn::tensor::module::{conv_transpose1d, conv1d};
45use burn::tensor::ops::{ConvOptions, ConvTransposeOptions};
46use burn_store::{BurnpackStore, ModuleSnapshot};
47use serde::{Deserialize, Serialize};
48
49/// Configuration for the Oobleck VAE, parsed from `vae/config.json`.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct OobleckVaeConfig {
52    /// Latent dimension produced by the DiT (input channels of the decoder).
53    pub decoder_input_channels: usize,
54    /// Base channel count of the decoder.
55    pub decoder_channels: usize,
56    /// Per-stage channel multiples applied to `decoder_channels`.
57    pub channel_multiples: Vec<usize>,
58    /// Encoder downsampling ratios; the decoder upsamples by these reversed.
59    pub downsampling_ratios: Vec<usize>,
60    /// Number of audio channels in the decoded output (2 = stereo).
61    pub audio_channels: usize,
62    /// Sampling rate of the decoded audio in Hz.
63    pub sampling_rate: usize,
64}
65
66impl OobleckVaeConfig {
67    pub fn load(path: &Path) -> Result<Self> {
68        let text = std::fs::read_to_string(path)
69            .with_context(|| format!("failed to read VAE config from {}", path.display()))?;
70        serde_json::from_str(&text)
71            .with_context(|| format!("failed to parse VAE config from {}", path.display()))
72    }
73
74    /// Number of audio samples produced per latent frame.
75    pub fn hop_length(&self) -> usize {
76        self.downsampling_ratios.iter().product()
77    }
78}
79
80/// Snake activation with learned per-channel `alpha` and `beta` in log scale:
81/// `x + sin(exp(alpha) * x)^2 / (exp(beta) + 1e-9)`.
82#[derive(Module, Debug)]
83pub struct Snake1d<B: Backend> {
84    pub alpha: Param<Tensor<B, 3>>,
85    pub beta: Param<Tensor<B, 3>>,
86}
87
88impl<B: Backend> Snake1d<B> {
89    pub fn new(channels: usize, device: &B::Device) -> Self {
90        Self {
91            alpha: Param::from_tensor(Tensor::zeros([1, channels, 1], device)),
92            beta: Param::from_tensor(Tensor::zeros([1, channels, 1], device)),
93        }
94    }
95
96    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
97        let alpha = self.alpha.val().exp();
98        let beta = self.beta.val().exp();
99        x.clone() + (alpha * x).sin().powf_scalar(2.0) / (beta + 1e-9)
100    }
101}
102
103/// Weight-normalized 1D convolution (stride 1 "same" convs and the k=7
104/// input/output convs of the decoder).
105#[derive(Module, Debug)]
106pub struct WnConv1d<B: Backend> {
107    weight_g: Param<Tensor<B, 3>>,
108    weight_v: Param<Tensor<B, 3>>,
109    bias: Option<Param<Tensor<B, 1>>>,
110    padding: usize,
111    dilation: usize,
112}
113
114impl<B: Backend> WnConv1d<B> {
115    fn new(
116        in_channels: usize,
117        out_channels: usize,
118        kernel_size: usize,
119        padding: usize,
120        dilation: usize,
121        bias: bool,
122        device: &B::Device,
123    ) -> Self {
124        Self {
125            weight_g: Param::from_tensor(Tensor::ones([out_channels, 1, 1], device)),
126            weight_v: Param::from_tensor(Tensor::zeros(
127                [out_channels, in_channels, kernel_size],
128                device,
129            )),
130            bias: bias.then(|| Param::from_tensor(Tensor::zeros([out_channels], device))),
131            padding,
132            dilation,
133        }
134    }
135
136    /// Effective weight `g * v / (||v|| + 1e-12)`, norm per output channel.
137    fn weight(&self) -> Tensor<B, 3> {
138        let g = self.weight_g.val();
139        let v = self.weight_v.val();
140        let out_channels = v.dims()[0];
141        let v_norm = v
142            .clone()
143            .powf_scalar(2.0)
144            .sum_dim(2)
145            .sum_dim(1)
146            .sqrt()
147            .reshape([out_channels, 1, 1]);
148        g * v / (v_norm + 1e-12)
149    }
150
151    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
152        let bias = self.bias.as_ref().map(|bias| bias.val());
153        conv1d(
154            x,
155            self.weight(),
156            bias,
157            ConvOptions::new([1], [self.padding], [self.dilation], 1),
158        )
159    }
160}
161
162/// Weight-normalized 1D transpose convolution used for upsampling.
163#[derive(Module, Debug)]
164pub struct WnConvTranspose1d<B: Backend> {
165    weight_g: Param<Tensor<B, 3>>,
166    weight_v: Param<Tensor<B, 3>>,
167    bias: Param<Tensor<B, 1>>,
168    stride: usize,
169    padding: usize,
170}
171
172impl<B: Backend> WnConvTranspose1d<B> {
173    fn new(in_channels: usize, out_channels: usize, stride: usize, device: &B::Device) -> Self {
174        Self {
175            weight_g: Param::from_tensor(Tensor::ones([in_channels, 1, 1], device)),
176            weight_v: Param::from_tensor(Tensor::zeros(
177                [in_channels, out_channels, 2 * stride],
178                device,
179            )),
180            bias: Param::from_tensor(Tensor::zeros([out_channels], device)),
181            stride,
182            padding: stride.div_ceil(2),
183        }
184    }
185
186    /// Effective weight `g * v / (||v|| + 1e-12)`, norm per input channel
187    /// (dim 0 of the transpose-conv kernel, as in `weight_norm` with dim=0).
188    fn weight(&self) -> Tensor<B, 3> {
189        let g = self.weight_g.val();
190        let v = self.weight_v.val();
191        let in_channels = v.dims()[0];
192        let v_norm = v
193            .clone()
194            .powf_scalar(2.0)
195            .sum_dim(2)
196            .sum_dim(1)
197            .sqrt()
198            .reshape([in_channels, 1, 1]);
199        g * v / (v_norm + 1e-12)
200    }
201
202    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
203        conv_transpose1d(
204            x,
205            self.weight(),
206            Some(self.bias.val()),
207            ConvTransposeOptions::new([self.stride], [self.padding], [0], [1], 1),
208        )
209    }
210}
211
212/// Residual unit: `x + conv2(snake2(conv1(snake1(x))))` with a dilated k=7
213/// conv and a k=1 conv. With same-padding the sequence length is preserved,
214/// so the length-mismatch crop in the reference is a no-op and is omitted.
215#[derive(Module, Debug)]
216pub struct OobleckResidualUnit<B: Backend> {
217    pub snake1: Snake1d<B>,
218    pub conv1: WnConv1d<B>,
219    pub snake2: Snake1d<B>,
220    pub conv2: WnConv1d<B>,
221}
222
223impl<B: Backend> OobleckResidualUnit<B> {
224    pub fn new(channels: usize, dilation: usize, device: &B::Device) -> Self {
225        Self {
226            snake1: Snake1d::new(channels, device),
227            conv1: WnConv1d::new(channels, channels, 7, 3 * dilation, dilation, true, device),
228            snake2: Snake1d::new(channels, device),
229            conv2: WnConv1d::new(channels, channels, 1, 0, 1, true, device),
230        }
231    }
232
233    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
234        let residual = x.clone();
235        let out = self.conv1.forward(self.snake1.forward(x));
236        let out = self.conv2.forward(self.snake2.forward(out));
237        residual + out
238    }
239}
240
241/// Decoder block: Snake -> ConvTranspose1d upsampling -> 3 residual units
242/// with dilations 1, 3, 9.
243#[derive(Module, Debug)]
244pub struct OobleckDecoderBlock<B: Backend> {
245    pub snake1: Snake1d<B>,
246    pub conv_t1: WnConvTranspose1d<B>,
247    pub res_unit1: OobleckResidualUnit<B>,
248    pub res_unit2: OobleckResidualUnit<B>,
249    pub res_unit3: OobleckResidualUnit<B>,
250}
251
252impl<B: Backend> OobleckDecoderBlock<B> {
253    pub fn new(in_channels: usize, out_channels: usize, stride: usize, device: &B::Device) -> Self {
254        Self {
255            snake1: Snake1d::new(in_channels, device),
256            conv_t1: WnConvTranspose1d::new(in_channels, out_channels, stride, device),
257            res_unit1: OobleckResidualUnit::new(out_channels, 1, device),
258            res_unit2: OobleckResidualUnit::new(out_channels, 3, device),
259            res_unit3: OobleckResidualUnit::new(out_channels, 9, device),
260        }
261    }
262
263    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
264        let x = self.conv_t1.forward(self.snake1.forward(x));
265        let x = self.res_unit1.forward(x);
266        let x = self.res_unit2.forward(x);
267        self.res_unit3.forward(x)
268    }
269}
270
271/// Oobleck VAE decoder. See the module-level docs for the canonical burnpack
272/// tensor names.
273#[derive(Module, Debug)]
274pub struct OobleckDecoder<B: Backend> {
275    pub conv1: WnConv1d<B>,
276    pub block: Vec<OobleckDecoderBlock<B>>,
277    pub snake1: Snake1d<B>,
278    pub conv2: WnConv1d<B>,
279    /// Audio samples per latent frame (product of `downsampling_ratios`).
280    pub hop_length: usize,
281    latent_channels: usize,
282}
283
284impl<B: Backend> OobleckDecoder<B> {
285    pub fn new(config: &OobleckVaeConfig, device: &B::Device) -> Self {
286        // Channel multiples with the prepended base multiple, as in the
287        // reference: [1] + channel_multiples.
288        let mut multiples = Vec::with_capacity(config.channel_multiples.len() + 1);
289        multiples.push(1);
290        multiples.extend_from_slice(&config.channel_multiples);
291
292        let num_stages = config.downsampling_ratios.len();
293        let top_channels = config.decoder_channels * multiples[num_stages];
294
295        let conv1 = WnConv1d::new(
296            config.decoder_input_channels,
297            top_channels,
298            7,
299            3,
300            1,
301            true,
302            device,
303        );
304
305        // Decode order uses the downsampling ratios reversed.
306        let block = (0..num_stages)
307            .map(|i| {
308                let stride = config.downsampling_ratios[num_stages - 1 - i];
309                let in_channels = config.decoder_channels * multiples[num_stages - i];
310                let out_channels = config.decoder_channels * multiples[num_stages - i - 1];
311                OobleckDecoderBlock::new(in_channels, out_channels, stride, device)
312            })
313            .collect();
314
315        let snake1 = Snake1d::new(config.decoder_channels, device);
316        let conv2 = WnConv1d::new(
317            config.decoder_channels,
318            config.audio_channels,
319            7,
320            3,
321            1,
322            false,
323            device,
324        );
325
326        Self {
327            conv1,
328            block,
329            snake1,
330            conv2,
331            hop_length: config.hop_length(),
332            latent_channels: config.decoder_input_channels,
333        }
334    }
335
336    /// Load decoder weights from a burnpack file using the canonical tensor
337    /// names documented at the top of this module.
338    pub fn from_burnpack(
339        config: &OobleckVaeConfig,
340        path: &Path,
341        device: &B::Device,
342    ) -> Result<Self> {
343        let mut model = Self::new(config, device);
344        let mut store = BurnpackStore::from_file(path).zero_copy(true);
345        model.load_from(&mut store).map_err(|err| {
346            anyhow::anyhow!("failed to load VAE decoder from {}: {err}", path.display())
347        })?;
348        Ok(model)
349    }
350
351    /// Forward pass on channel-major latents `[B, decoder_input_channels, T]`
352    /// returning audio `[B, audio_channels, T * hop_length]`.
353    pub fn forward(&self, latents: Tensor<B, 3>) -> Tensor<B, 3> {
354        let x = self.conv1.forward(latents);
355        let x = self
356            .block
357            .iter()
358            .fold(x, |hidden, block| block.forward(hidden));
359        let x = self.snake1.forward(x);
360        self.conv2.forward(x)
361    }
362
363    /// Decode time-major DiT latents `[B, T, decoder_input_channels]` into
364    /// audio `[B, audio_channels, T * hop_length]`.
365    ///
366    /// Every released ratio is even, in which case each transpose conv
367    /// upsamples by exactly its stride and the output length is already
368    /// `T * hop_length`; the final trim/pad only matters for hypothetical
369    /// odd ratios, where ConvTranspose1d falls one sample short per stage.
370    pub fn decode(&self, latents: Tensor<B, 3>) -> Tensor<B, 3> {
371        let [batch, frames, channels] = latents.dims();
372        assert_eq!(
373            channels, self.latent_channels,
374            "expected {} latent channels, got {channels}",
375            self.latent_channels
376        );
377        let audio = self.forward(latents.swap_dims(1, 2));
378
379        let target = frames * self.hop_length;
380        let [_, _, length] = audio.dims();
381        if length > target {
382            audio.slice([0..batch, 0..self.conv2_channels(), 0..target])
383        } else if length < target {
384            let device = audio.device();
385            let padding = Tensor::zeros([batch, self.conv2_channels(), target - length], &device);
386            Tensor::cat(vec![audio, padding], 2)
387        } else {
388            audio
389        }
390    }
391
392    fn conv2_channels(&self) -> usize {
393        self.conv2.weight_v.dims()[0]
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use burn::backend::ndarray::{NdArray, NdArrayDevice};
401
402    type TestBackend = NdArray<f32>;
403
404    fn tiny_config() -> OobleckVaeConfig {
405        OobleckVaeConfig {
406            decoder_input_channels: 4,
407            decoder_channels: 8,
408            channel_multiples: vec![1, 2],
409            downsampling_ratios: vec![2, 3],
410            audio_channels: 2,
411            sampling_rate: 48_000,
412        }
413    }
414
415    #[test]
416    fn decode_produces_stereo_audio_of_expected_length() {
417        let device = NdArrayDevice::default();
418        let decoder = OobleckDecoder::<TestBackend>::new(&tiny_config(), &device);
419
420        let latents = Tensor::<TestBackend, 3>::random(
421            [1, 5, 4],
422            burn::tensor::Distribution::Normal(0.0, 1.0),
423            &device,
424        );
425        let audio = decoder.decode(latents);
426
427        assert_eq!(audio.dims(), [1, 2, 5 * 6]);
428        let values: Vec<f32> = audio.into_data().to_vec().unwrap();
429        assert!(values.iter().all(|v| v.is_finite()));
430    }
431
432    #[test]
433    fn decode_upsamples_each_stage() {
434        let device = NdArrayDevice::default();
435        let config = tiny_config();
436        let decoder = OobleckDecoder::<TestBackend>::new(&config, &device);
437
438        // Stage strides are the downsampling ratios reversed: [3, 2].
439        let x = Tensor::<TestBackend, 3>::zeros([1, 4, 5], &device);
440        let x = decoder.conv1.forward(x);
441        assert_eq!(x.dims(), [1, 16, 5]);
442        let x = decoder.block[0].forward(x);
443        assert_eq!(x.dims(), [1, 8, 14]); // odd stride 3: (5-1)*3 - 4 + 6
444        let x = decoder.block[1].forward(x);
445        assert_eq!(x.dims(), [1, 8, 28]); // even stride 2: 14 * 2
446    }
447
448    #[test]
449    fn residual_unit_preserves_shape() {
450        let device = NdArrayDevice::default();
451        let unit = OobleckResidualUnit::<TestBackend>::new(8, 3, &device);
452        let x = Tensor::<TestBackend, 3>::random(
453            [2, 8, 11],
454            burn::tensor::Distribution::Normal(0.0, 1.0),
455            &device,
456        );
457        assert_eq!(unit.forward(x).dims(), [2, 8, 11]);
458    }
459
460    #[test]
461    fn loads_real_vae_config() {
462        let config = OobleckVaeConfig::load(
463            &Path::new(env!("CARGO_MANIFEST_DIR"))
464                .join("src/acestep/testdata/oobleck_vae_config.json"),
465        )
466        .unwrap();
467        assert_eq!(config.decoder_input_channels, 64);
468        assert_eq!(config.decoder_channels, 128);
469        assert_eq!(config.channel_multiples, vec![1, 2, 4, 8, 16]);
470        assert_eq!(config.downsampling_ratios, vec![2, 4, 4, 6, 10]);
471        assert_eq!(config.audio_channels, 2);
472        assert_eq!(config.sampling_rate, 48_000);
473        assert_eq!(config.hop_length(), 1920);
474    }
475}
maolan_generate/acestep/vae.rs

maolan_generate/acestep/
vae.rs
