Skip to main content

OobleckDecoder

maolan_generate::acestep::vae

Struct OobleckDecoder 

Source 
pub struct OobleckDecoder<B: Backend> {
    pub conv1: WnConv1d<B>,
    pub block: Vec<OobleckDecoderBlock<B>>,
    pub snake1: Snake1d<B>,
    pub conv2: WnConv1d<B>,
    pub hop_length: usize,
    /* private fields */
}
Expand description

Oobleck VAE decoder. See the module-level docs for the canonical burnpack tensor names.

Fields§

§conv1: WnConv1d<B>§block: Vec<OobleckDecoderBlock<B>>§snake1: Snake1d<B>§conv2: WnConv1d<B>§hop_length: usize

Audio samples per latent frame (product of downsampling_ratios).

Implementations§

Source§

impl<B: Backend> OobleckDecoder<B>

Source

pub fn new(config: &OobleckVaeConfig, device: &B::Device) -> Self

Source

pub fn from_burnpack( config: &OobleckVaeConfig, path: &Path, device: &B::Device, ) -> Result<Self>

Load decoder weights from a burnpack file using the canonical tensor names documented at the top of this module.

Examples found in repository?
examples/decode_silence.rs (lines 21-25)
13fn main() -> anyhow::Result<()> {
14    let model_dir = std::env::args()
15        .nth(1)
16        .map(std::path::PathBuf::from)
17        .unwrap_or_else(|| std::path::PathBuf::from("/home/meka/repos/ace"));
18    let device = Default::default();
19
20    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
21    let vae = OobleckDecoder::<B>::from_burnpack(
22        &vae_config,
23        &model_dir.join("acestep-vae.bpk"),
24        &device,
25    )?;
26    let silence =
27        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
28    let [_, frames, _] = silence.silence_latent.dims();
29    println!("silence latent: {frames} frames");
30
31    let latents = silence.slice(750);
32    let audio = vae.decode(latents);
33    let [_, channels, samples] = audio.dims();
34    let values: Vec<f32> = audio
35        .into_data()
36        .convert::<f32>()
37        .to_vec()
38        .map_err(|e| anyhow::anyhow!("{e}"))?;
39    let peak = values.iter().fold(0.0_f32, |a, v| a.max(v.abs()));
40    let rms = (values.iter().map(|v| v * v).sum::<f32>() / values.len() as f32).sqrt();
41    println!("decoded: {channels} ch x {samples} samples");
42    println!("peak {peak:.6}  rms {rms:.6}");
43    println!(
44        "{}",
45        if peak < 1e-3 {
46            "OK: VAE decodes silence correctly"
47        } else {
48            "SUSPECT: VAE decoder output is far from silence"
49        }
50    );
51
52    // Also dump the latent stats so we can eyeball them.
53    let latents = SilenceLatent::<B>::from_burnpack(
54        Path::new(&model_dir).join("silence_latent.bpk").as_path(),
55        &device,
56    )?
57    .slice(4);
58    let v: Vec<f32> = latents
59        .into_data()
60        .convert::<f32>()
61        .to_vec()
62        .map_err(|e| anyhow::anyhow!("{e}"))?;
63    let lpeak = v.iter().fold(0.0_f32, |a, v| a.max(v.abs()));
64    let lrms = (v.iter().map(|v| v * v).sum::<f32>() / v.len() as f32).sqrt();
65    println!("latent peak {lpeak:.6} rms {lrms:.6}");
66    Ok(())
67}
More examples
Hide additional examples
examples/fsq_probe.rs (lines 75-79)
61fn main() -> anyhow::Result<()> {
62    let model_dir = std::env::args()
63        .nth(1)
64        .unwrap_or_else(|| "/home/meka/repos/ace".to_string());
65    let model_dir = Path::new(&model_dir);
66    let device = Default::default();
67
68    let dit_config = AceStepConfig::load(&model_dir.join("dit_config.json"))?;
69    let condition = AceStepCondition::<B>::from_burnpack(
70        &dit_config,
71        &model_dir.join("acestep-condition.bpk"),
72        &device,
73    )?;
74    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
75    let vae = OobleckDecoder::<B>::from_burnpack(
76        &vae_config,
77        &model_dir.join("acestep-vae.bpk"),
78        &device,
79    )?;
80    let silence =
81        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
82
83    // Splitmix64 random codes.
84    let mut state = 42_u64;
85    let mut rand_codes = Vec::new();
86    for _ in 0..20 {
87        state = state
88            .wrapping_mul(6364136223846793005)
89            .wrapping_add(1442695040888963407);
90        rand_codes.push(((state >> 33) % 64000) as u32);
91    }
92
93    let lm_codes: Vec<u32> = vec![
94        61890, 51649, 53753, 53753, 56314, 53754, 55802, 4538, 4538, 4538, 5050, 5050, 5050, 5050,
95        5050, 5050, 5050, 5050, 5050, 43513,
96    ];
97
98    // Structure check: for a few distinct codes, dump the detokenizer output
99    // frames and measure how much outputs differ between codes.
100    for code in [0u32, 12345, 53754] {
101        let tensor = Tensor::<B, 2, Int>::from_data(TensorData::new(vec![code], [1, 1]), &device);
102        let hints = condition.codes_to_hints(tensor); // [1, 5, 64]
103        let values: Vec<f32> = hints
104            .into_data()
105            .convert::<f32>()
106            .to_vec()
107            .map_err(|e| anyhow::anyhow!("{e}"))?;
108        let frame_means: Vec<f32> = values
109            .as_chunks::<64>()
110            .0
111            .iter()
112            .map(|f| f.iter().sum::<f32>() / 64.0)
113            .collect();
114        let frame_rms: Vec<f32> = values.as_chunks::<64>().0.iter().map(|f| rms(f)).collect();
115        println!(
116            "code {code:>6}: frame means {:?}",
117            frame_means
118                .iter()
119                .map(|v| format!("{v:.4}"))
120                .collect::<Vec<_>>()
121        );
122        println!(
123            "           frame rms   {:?}",
124            frame_rms
125                .iter()
126                .map(|v| format!("{v:.4}"))
127                .collect::<Vec<_>>()
128        );
129    }
130
131    // Time-variance vs channel-variance of hint latents for random codes.
132    let tensor =
133        Tensor::<B, 2, Int>::from_data(TensorData::new(rand_codes.clone(), [1, 20]), &device);
134    let hints = condition.codes_to_hints(tensor); // [1, 100, 64]
135    let values: Vec<f32> = hints
136        .into_data()
137        .convert::<f32>()
138        .to_vec()
139        .map_err(|e| anyhow::anyhow!("{e}"))?;
140    let (frames, _) = values.as_chunks::<64>();
141    // variance of per-frame means (time structure) vs mean per-frame variance (channel structure)
142    let frame_means: Vec<f32> = frames
143        .iter()
144        .map(|f| f.iter().sum::<f32>() / 64.0)
145        .collect();
146    let tm = frame_means.iter().sum::<f32>() / frame_means.len() as f32;
147    let time_var =
148        frame_means.iter().map(|v| (v - tm) * (v - tm)).sum::<f32>() / frame_means.len() as f32;
149    let chan_var = frames
150        .iter()
151        .map(|f| {
152            let m = f.iter().sum::<f32>() / 64.0;
153            f.iter().map(|v| (v - m) * (v - m)).sum::<f32>() / 64.0
154        })
155        .sum::<f32>()
156        / frames.len() as f32;
157    println!("random-code hints: time var {time_var:.6}  channel var {chan_var:.6}");
158
159    run_case(
160        "random codes ",
161        &rand_codes,
162        &condition,
163        &vae,
164        &silence,
165        &device,
166    )?;
167    run_case(
168        "LM codes     ",
169        &lm_codes,
170        &condition,
171        &vae,
172        &silence,
173        &device,
174    )?;
175    run_case(
176        "constant 0   ",
177        &[0u32; 20],
178        &condition,
179        &vae,
180        &silence,
181        &device,
182    )?;
183    Ok(())
184}
examples/dit_dump_cpu.rs (lines 88-92)
63fn main() -> anyhow::Result<()> {
64    let mut args = std::env::args().skip(1);
65    let model_dir = PathBuf::from(args.next().unwrap_or_else(|| "/home/meka/repos/ace".into()));
66    let out_dir = PathBuf::from(args.next().unwrap_or_else(|| "/var/tmp/dump_ours".into()));
67    let noise_path = args.next().map(PathBuf::from);
68    std::fs::create_dir_all(&out_dir)?;
69
70    let device = Default::default();
71
72    let dit_config = AceStepConfig::load(&model_dir.join("dit_config.json"))?;
73    let text_config = Qwen3Config::load(&model_dir.join("qwen3_config.json"))?;
74    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
75    eprintln!("loading components...");
76    let text_encoder = Qwen3Model::<B>::from_burnpack(
77        &text_config,
78        &model_dir.join("qwen3-encoder.bpk"),
79        &device,
80    )?;
81    let condition = AceStepCondition::<B>::from_burnpack(
82        &dit_config,
83        &model_dir.join("acestep-condition.bpk"),
84        &device,
85    )?;
86    let dit =
87        AceStepDiT::<B>::from_burnpack(&dit_config, &model_dir.join("acestep-dit.bpk"), &device)?;
88    let vae = OobleckDecoder::<B>::from_burnpack(
89        &vae_config,
90        &model_dir.join("acestep-vae.bpk"),
91        &device,
92    )?;
93    let silence =
94        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
95    let tokenizer = tokie::Tokenizer::from_json(model_dir.join("tokenizer.json"))
96        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
97
98    // ---- text + lyric encoding (causal, official prompts) ----
99    let metas = build_metas_block(Some(120.0), Some("A minor"), Some("4/4"), 4);
100    let text_prompt = build_dit_text_prompt("Metal guitar with a lot of distortion", &metas);
101    let ids = tokenizer.encode(&text_prompt, false).ids;
102    let mut ids: Vec<i64> = ids.into_iter().map(i64::from).collect();
103    ids.push(151643); // explicit EOS (official add_eos=true)
104    let n_text = ids.len();
105    let ids_t = Tensor::<B, 2, Int>::from_data(TensorData::new(ids, [1, n_text]), &device);
106    let text_hidden = text_encoder.forward(ids_t, true);
107    let text_hidden_vec = vec_of(text_hidden.clone());
108    dump_tensor(
109        &out_dir,
110        "text_hidden",
111        &text_hidden_vec,
112        &[n_text as i32, 1024],
113    );
114
115    let lyric_ids = tokenizer
116        .encode(
117            maolan_generate::acestep::pipeline::INSTRUMENTAL_LYRIC_PROMPT,
118            false,
119        )
120        .ids;
121    let mut lyric_ids: Vec<i64> = lyric_ids.into_iter().map(i64::from).collect();
122    lyric_ids.push(151643); // explicit EOS (official add_eos=true)
123    let n_lyric = lyric_ids.len();
124    let lyric_ids_t =
125        Tensor::<B, 2, Int>::from_data(TensorData::new(lyric_ids, [1, n_lyric]), &device);
126    // Official: lyric branch is a raw embed_tokens lookup (no transformer).
127    let lyric_hidden = text_encoder.embed_tokens.forward(lyric_ids_t);
128    let lyric_embed_vec = vec_of(lyric_hidden.clone());
129    dump_tensor(
130        &out_dir,
131        "lyric_embed",
132        &lyric_embed_vec,
133        &[n_lyric as i32, 1024],
134    );
135    let lyric_mask = Tensor::<B, 2, Int>::ones([1, n_lyric], &device);
136
137    // ---- conditioning: use the oracle's enc_hidden/context when provided ----
138    let mut oracle_context: Option<Tensor<B, 3>> = None;
139    let enc = if let Some(dir) = std::env::var_os("MAOLAN_ORACLE_DUMP_DIR") {
140        let dir = PathBuf::from(dir);
141        let (enc_shape, enc_data) = load_bin(&dir.join("enc_hidden.bin"));
142        eprintln!("oracle enc_hidden: {enc_shape:?}");
143        let (ctx_shape, ctx_data) = load_bin(&dir.join("context.bin"));
144        eprintln!("oracle context: {ctx_shape:?}");
145        let enc = Tensor::<B, 3>::from_data(
146            TensorData::new(enc_data, [1, enc_shape[0] as usize, enc_shape[1] as usize]),
147            &device,
148        );
149        let ctx = Tensor::<B, 3>::from_data(TensorData::new(ctx_data, [1, 96, 128]), &device);
150        oracle_context = Some(ctx);
151        enc
152    } else {
153        condition.encode(
154            text_hidden,
155            lyric_hidden,
156            lyric_mask,
157            silence.timbre_reference(),
158        )
159    };
160    let [_, enc_len, enc_dim] = enc.dims();
161    let enc_vec = vec_of(enc.clone());
162    dump_tensor(
163        &out_dir,
164        "enc_hidden",
165        &enc_vec,
166        &[enc_len as i32, enc_dim as i32],
167    );
168
169    // ---- FSQ hints from the given codes (env MAOLAN_ACESTEP_CODES or the
170    // built-in oracle sequence) ----
171    let codes: Vec<u32> = std::env::var_os("MAOLAN_ACESTEP_CODES")
172        .map(|raw| {
173            raw.to_string_lossy()
174                .split(',')
175                .filter_map(|part| part.trim().parse::<u32>().ok())
176                .collect()
177        })
178        .unwrap_or_else(|| ORACLE_CODES.to_vec());
179    let n_codes = codes.len();
180    let codes_tensor =
181        Tensor::<B, 2, Int>::from_data(TensorData::new(codes, [1, n_codes]), &device);
182    let hints = condition.codes_to_hints(codes_tensor);
183    let [_, hint_len, _] = hints.dims();
184    let hints_vec = vec_of(hints.clone());
185    dump_tensor(&out_dir, "detok_output", &hints_vec, &[hint_len as i32, 64]);
186
187    // ---- noise ----
188    let noise = if let Some(path) = noise_path {
189        let (shape, data) = load_bin(&path);
190        let frames = shape[0] as usize;
191        Tensor::<B, 3>::from_data(TensorData::new(data, [1, frames, 64]), &device)
192    } else {
193        maolan_generate::acestep::pipeline::seeded_latent_noise(0, hint_len, 64, &device)
194    };
195    let [_, frames, _] = noise.dims();
196
197    // src latents: hints cropped or silence-padded to the noise's frame count
198    let src = if hint_len >= frames {
199        hints.narrow(1, 0, frames)
200    } else {
201        let padding = silence.slice(frames - hint_len);
202        Tensor::cat(vec![hints, padding], 1)
203    };
204    let chunk_mask = Tensor::ones([1, frames, 64], &device);
205    let context = Tensor::cat(vec![src.clone(), chunk_mask], 2);
206    let context = oracle_context.unwrap_or(context);
207    let context_vec = vec_of(context.clone());
208    dump_tensor(&out_dir, "context", &context_vec, &[frames as i32, 128]);
209
210    let noise_vec = vec_of(noise.clone());
211    dump_tensor(&out_dir, "noise", &noise_vec, &[frames as i32, 64]);
212
213    // ---- DiT loop with per-step dumps ----
214    let kv = dit.prepare_cross_kv(enc);
215    let mut xt = noise;
216    let total = TURBO_TIMESTEPS.len();
217    for (index, &t_cur) in TURBO_TIMESTEPS.iter().enumerate() {
218        let xt_vec = vec_of(xt.clone());
219        dump_tensor(
220            &out_dir,
221            &format!("dit_step{index}_xt"),
222            &xt_vec,
223            &[frames as i32, 64],
224        );
225        let v = dit.forward_with_kv(xt.clone(), t_cur, context.clone(), &kv);
226        let v_vec = vec_of(v.clone());
227        dump_tensor(
228            &out_dir,
229            &format!("dit_step{index}_vt"),
230            &v_vec,
231            &[frames as i32, 64],
232        );
233        let dt = if index + 1 == total {
234            t_cur
235        } else {
236            t_cur - TURBO_TIMESTEPS[index + 1]
237        };
238        xt = xt - v * dt;
239    }
240    let x0_vec = vec_of(xt.clone());
241    dump_tensor(&out_dir, "dit_x0", &x0_vec, &[frames as i32, 64]);
242
243    // ---- VAE decode ----
244    let audio = vae.decode(xt);
245    let [_, channels, samples] = audio.dims();
246    let audio_vec = vec_of(audio);
247    dump_tensor(
248        &out_dir,
249        "vae_audio",
250        &audio_vec,
251        &[channels as i32, samples as i32],
252    );
253    println!("dumps written to {}", out_dir.display());
254    Ok(())
255}
examples/dit_dump.rs (lines 103-107)
63fn main() -> anyhow::Result<()> {
64    let mut args = std::env::args().skip(1);
65    let model_dir = PathBuf::from(args.next().unwrap_or_else(|| "/home/meka/repos/ace".into()));
66    let out_dir = PathBuf::from(args.next().unwrap_or_else(|| "/var/tmp/dump_ours".into()));
67    let noise_path = args.next().map(PathBuf::from);
68    std::fs::create_dir_all(&out_dir)?;
69
70    let device = burn::backend::wgpu::WgpuDevice::default();
71    burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
72        &device,
73        burn::backend::wgpu::RuntimeOptions {
74            memory_config: burn::backend::wgpu::MemoryConfiguration::ExclusivePages,
75            ..Default::default()
76        },
77    );
78
79    let variant_sft = std::env::var("MAOLAN_ACESTEP_VARIANT")
80        .map(|v| v == "sft")
81        .unwrap_or(false);
82    let prefix = if variant_sft { "sft-" } else { "" };
83
84    let dit_config = AceStepConfig::load(&model_dir.join(format!("{prefix}dit_config.json")))?;
85    let text_config = Qwen3Config::load(&model_dir.join("qwen3_config.json"))?;
86    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
87    eprintln!("loading components (prefix '{prefix}')...");
88    let text_encoder = Qwen3Model::<B>::from_burnpack(
89        &text_config,
90        &model_dir.join("qwen3-encoder.bpk"),
91        &device,
92    )?;
93    let condition = AceStepCondition::<B>::from_burnpack(
94        &dit_config,
95        &model_dir.join(format!("{prefix}acestep-condition.bpk")),
96        &device,
97    )?;
98    let dit = AceStepDiT::<B>::from_burnpack(
99        &dit_config,
100        &model_dir.join(format!("{prefix}acestep-dit.bpk")),
101        &device,
102    )?;
103    let vae = OobleckDecoder::<B>::from_burnpack(
104        &vae_config,
105        &model_dir.join("acestep-vae.bpk"),
106        &device,
107    )?;
108    let silence =
109        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
110    let tokenizer = tokie::Tokenizer::from_json(model_dir.join("tokenizer.json"))
111        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
112
113    // ---- text + lyric encoding (causal, official prompts) ----
114    let metas = build_metas_block(Some(120.0), Some("A minor"), Some("4/4"), 4);
115    let text_prompt = build_dit_text_prompt("Metal guitar with a lot of distortion", &metas);
116    let ids = tokenizer.encode(&text_prompt, false).ids;
117    let mut ids: Vec<i64> = ids.into_iter().map(i64::from).collect();
118    ids.push(151643); // explicit EOS (official add_eos=true)
119    let n_text = ids.len();
120    let ids_t = Tensor::<B, 2, Int>::from_data(TensorData::new(ids, [1, n_text]), &device);
121    let text_hidden = text_encoder.forward(ids_t, true);
122    let text_hidden_vec = vec_of(text_hidden.clone());
123    dump_tensor(
124        &out_dir,
125        "text_hidden",
126        &text_hidden_vec,
127        &[n_text as i32, 1024],
128    );
129
130    let lyric_ids = tokenizer
131        .encode(
132            maolan_generate::acestep::pipeline::INSTRUMENTAL_LYRIC_PROMPT,
133            false,
134        )
135        .ids;
136    let mut lyric_ids: Vec<i64> = lyric_ids.into_iter().map(i64::from).collect();
137    lyric_ids.push(151643); // explicit EOS (official add_eos=true)
138    let n_lyric = lyric_ids.len();
139    let lyric_ids_t =
140        Tensor::<B, 2, Int>::from_data(TensorData::new(lyric_ids, [1, n_lyric]), &device);
141    // Official: lyric branch is a raw embed_tokens lookup (no transformer).
142    let lyric_hidden = text_encoder.embed_tokens.forward(lyric_ids_t);
143    let lyric_embed_vec = vec_of(lyric_hidden.clone());
144    dump_tensor(
145        &out_dir,
146        "lyric_embed",
147        &lyric_embed_vec,
148        &[n_lyric as i32, 1024],
149    );
150    let lyric_mask = Tensor::<B, 2, Int>::ones([1, n_lyric], &device);
151
152    // ---- conditioning: use the oracle's enc_hidden/context when provided ----
153    let mut oracle_context: Option<Tensor<B, 3>> = None;
154    let enc = if let Some(dir) = std::env::var_os("MAOLAN_ORACLE_DUMP_DIR") {
155        let dir = PathBuf::from(dir);
156        let (enc_shape, enc_data) = load_bin(&dir.join("enc_hidden.bin"));
157        eprintln!("oracle enc_hidden: {enc_shape:?}");
158        let (ctx_shape, ctx_data) = load_bin(&dir.join("context.bin"));
159        eprintln!("oracle context: {ctx_shape:?}");
160        let enc = Tensor::<B, 3>::from_data(
161            TensorData::new(enc_data, [1, enc_shape[0] as usize, enc_shape[1] as usize]),
162            &device,
163        );
164        let ctx = Tensor::<B, 3>::from_data(TensorData::new(ctx_data, [1, 96, 128]), &device);
165        oracle_context = Some(ctx);
166        enc
167    } else {
168        condition.encode(
169            text_hidden,
170            lyric_hidden,
171            lyric_mask,
172            silence.timbre_reference(),
173        )
174    };
175    let [_, enc_len, enc_dim] = enc.dims();
176    let enc_vec = vec_of(enc.clone());
177    dump_tensor(
178        &out_dir,
179        "enc_hidden",
180        &enc_vec,
181        &[enc_len as i32, enc_dim as i32],
182    );
183
184    // ---- FSQ hints from the given codes (env MAOLAN_ACESTEP_CODES or the
185    // built-in oracle sequence) ----
186    let codes: Vec<u32> = std::env::var_os("MAOLAN_ACESTEP_CODES")
187        .map(|raw| {
188            raw.to_string_lossy()
189                .split(',')
190                .filter_map(|part| part.trim().parse::<u32>().ok())
191                .collect()
192        })
193        .unwrap_or_else(|| ORACLE_CODES.to_vec());
194    let n_codes = codes.len();
195    let codes_tensor =
196        Tensor::<B, 2, Int>::from_data(TensorData::new(codes, [1, n_codes]), &device);
197    let hints = condition.codes_to_hints(codes_tensor);
198    let [_, hint_len, _] = hints.dims();
199    let hints_vec = vec_of(hints.clone());
200    dump_tensor(&out_dir, "detok_output", &hints_vec, &[hint_len as i32, 64]);
201
202    // ---- noise ----
203    let noise = if let Some(path) = noise_path {
204        let (shape, data) = load_bin(&path);
205        let frames = shape[0] as usize;
206        Tensor::<B, 3>::from_data(TensorData::new(data, [1, frames, 64]), &device)
207    } else {
208        maolan_generate::acestep::pipeline::seeded_latent_noise(0, hint_len, 64, &device)
209    };
210    let [_, frames, _] = noise.dims();
211
212    // src latents: hints cropped or silence-padded to the noise's frame count
213    let src = if hint_len >= frames {
214        hints.narrow(1, 0, frames)
215    } else {
216        let padding = silence.slice(frames - hint_len);
217        Tensor::cat(vec![hints, padding], 1)
218    };
219    let chunk_mask = Tensor::ones([1, frames, 64], &device);
220    let context = Tensor::cat(vec![src.clone(), chunk_mask], 2);
221    let context = oracle_context.unwrap_or(context);
222    let context_vec = vec_of(context.clone());
223    dump_tensor(&out_dir, "context", &context_vec, &[frames as i32, 128]);
224
225    let noise_vec = vec_of(noise.clone());
226    dump_tensor(&out_dir, "noise", &noise_vec, &[frames as i32, 64]);
227
228    // ---- DiT loop with per-step dumps (schedule from is_turbo) ----
229    let timesteps: &[f32] = if dit_config.is_turbo {
230        &TURBO_TIMESTEPS
231    } else {
232        &SFT_TIMESTEPS
233    };
234    let kv = dit.prepare_cross_kv(enc);
235    let mut xt = noise;
236    let total = timesteps.len();
237    for (index, &t_cur) in timesteps.iter().enumerate() {
238        let xt_vec = vec_of(xt.clone());
239        dump_tensor(
240            &out_dir,
241            &format!("dit_step{index}_xt"),
242            &xt_vec,
243            &[frames as i32, 64],
244        );
245        let v = dit.forward_with_kv(xt.clone(), t_cur, context.clone(), &kv);
246        let v_vec = vec_of(v.clone());
247        dump_tensor(
248            &out_dir,
249            &format!("dit_step{index}_vt"),
250            &v_vec,
251            &[frames as i32, 64],
252        );
253        let dt = if index + 1 == total {
254            t_cur
255        } else {
256            t_cur - timesteps[index + 1]
257        };
258        xt = xt - v * dt;
259    }
260    let x0_vec = vec_of(xt.clone());
261    dump_tensor(&out_dir, "dit_x0", &x0_vec, &[frames as i32, 64]);
262
263    // ---- VAE decode ----
264    let audio = vae.decode(xt);
265    let [_, channels, samples] = audio.dims();
266    let audio_vec = vec_of(audio);
267    dump_tensor(
268        &out_dir,
269        "vae_audio",
270        &audio_vec,
271        &[channels as i32, samples as i32],
272    );
273    println!("dumps written to {}", out_dir.display());
274    Ok(())
275}
Source

pub fn forward(&self, latents: Tensor<B, 3>) -> Tensor<B, 3>

Forward pass on channel-major latents [B, decoder_input_channels, T] returning audio [B, audio_channels, T * hop_length].

Source

pub fn decode(&self, latents: Tensor<B, 3>) -> Tensor<B, 3>

Decode time-major DiT latents [B, T, decoder_input_channels] into audio [B, audio_channels, T * hop_length].

Every released ratio is even, in which case each transpose conv upsamples by exactly its stride and the output length is already T * hop_length; the final trim/pad only matters for hypothetical odd ratios, where ConvTranspose1d falls one sample short per stage.

Examples found in repository?
examples/fsq_probe.rs (line 43)
19fn run_case(
20    name: &str,
21    codes: &[u32],
22    condition: &AceStepCondition<B>,
23    vae: &OobleckDecoder<B>,
24    silence: &SilenceLatent<B>,
25    device: &burn::tensor::Device<B>,
26) -> anyhow::Result<()> {
27    let n = codes.len();
28    let tensor = Tensor::<B, 2, Int>::from_data(TensorData::new(codes.to_vec(), [1, n]), device);
29    let hints = condition.codes_to_hints(tensor);
30    let silence_frames = silence.slice(n * 5);
31    let diff = (hints.clone() - silence_frames)
32        .abs()
33        .into_data()
34        .convert::<f32>()
35        .to_vec::<f32>()
36        .map_err(|e| anyhow::anyhow!("{e}"))?;
37    let hint_values: Vec<f32> = hints
38        .clone()
39        .into_data()
40        .convert::<f32>()
41        .to_vec()
42        .map_err(|e| anyhow::anyhow!("{e}"))?;
43    let audio = vae.decode(hints);
44    let audio_values: Vec<f32> = audio
45        .into_data()
46        .convert::<f32>()
47        .to_vec()
48        .map_err(|e| anyhow::anyhow!("{e}"))?;
49    let peak = audio_values.iter().fold(0.0_f32, |a, v| a.max(v.abs()));
50    let mean_abs_diff = diff.iter().sum::<f32>() / diff.len() as f32;
51    println!(
52        "{name}: hints rms {:.4}  mean|hints-silence| {:.4}  audio rms {:.4}  audio peak {:.4}",
53        rms(&hint_values),
54        mean_abs_diff,
55        rms(&audio_values),
56        peak
57    );
58    Ok(())
59}
More examples
Hide additional examples
examples/decode_silence.rs (line 32)
13fn main() -> anyhow::Result<()> {
14    let model_dir = std::env::args()
15        .nth(1)
16        .map(std::path::PathBuf::from)
17        .unwrap_or_else(|| std::path::PathBuf::from("/home/meka/repos/ace"));
18    let device = Default::default();
19
20    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
21    let vae = OobleckDecoder::<B>::from_burnpack(
22        &vae_config,
23        &model_dir.join("acestep-vae.bpk"),
24        &device,
25    )?;
26    let silence =
27        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
28    let [_, frames, _] = silence.silence_latent.dims();
29    println!("silence latent: {frames} frames");
30
31    let latents = silence.slice(750);
32    let audio = vae.decode(latents);
33    let [_, channels, samples] = audio.dims();
34    let values: Vec<f32> = audio
35        .into_data()
36        .convert::<f32>()
37        .to_vec()
38        .map_err(|e| anyhow::anyhow!("{e}"))?;
39    let peak = values.iter().fold(0.0_f32, |a, v| a.max(v.abs()));
40    let rms = (values.iter().map(|v| v * v).sum::<f32>() / values.len() as f32).sqrt();
41    println!("decoded: {channels} ch x {samples} samples");
42    println!("peak {peak:.6}  rms {rms:.6}");
43    println!(
44        "{}",
45        if peak < 1e-3 {
46            "OK: VAE decodes silence correctly"
47        } else {
48            "SUSPECT: VAE decoder output is far from silence"
49        }
50    );
51
52    // Also dump the latent stats so we can eyeball them.
53    let latents = SilenceLatent::<B>::from_burnpack(
54        Path::new(&model_dir).join("silence_latent.bpk").as_path(),
55        &device,
56    )?
57    .slice(4);
58    let v: Vec<f32> = latents
59        .into_data()
60        .convert::<f32>()
61        .to_vec()
62        .map_err(|e| anyhow::anyhow!("{e}"))?;
63    let lpeak = v.iter().fold(0.0_f32, |a, v| a.max(v.abs()));
64    let lrms = (v.iter().map(|v| v * v).sum::<f32>() / v.len() as f32).sqrt();
65    println!("latent peak {lpeak:.6} rms {lrms:.6}");
66    Ok(())
67}
examples/acestep_phases.rs (line 140)
23fn main() -> anyhow::Result<()> {
24    let mut args = std::env::args().skip(1);
25    let model_dir = args
26        .next()
27        .unwrap_or_else(|| "/home/meka/repos/ace".to_string());
28    let out_prefix = args
29        .next()
30        .unwrap_or_else(|| "/tmp/acestep_phase".to_string());
31    let caption = args
32        .next()
33        .unwrap_or_else(|| "Metal guitar with a lot of distortion".to_string());
34    let bpm: Option<f32> = args.next().and_then(|v| v.parse().ok()).or(Some(120.0));
35    let key_scale = args.next().or_else(|| Some("A minor".to_string()));
36    let time_signature = args.next().or_else(|| Some("4/4".to_string()));
37    let length_ms: usize = args.next().and_then(|v| v.parse().ok()).unwrap_or(4000);
38
39    let device = burn::backend::wgpu::WgpuDevice::default();
40    burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
41        &device,
42        burn::backend::wgpu::RuntimeOptions {
43            memory_config: burn::backend::wgpu::MemoryConfiguration::ExclusivePages,
44            ..Default::default()
45        },
46    );
47
48    let variant = if std::env::var("MAOLAN_ACESTEP_VARIANT")
49        .map(|v| v == "sft")
50        .unwrap_or(false)
51    {
52        maolan_generate::acestep::AceStepVariant::Sft
53    } else {
54        maolan_generate::acestep::AceStepVariant::Turbo
55    };
56    let paths = AceStepModelPaths::resolve(Path::new(&model_dir), variant)?;
57    let mut progress = |phase: &str, p: f32, op: &str| {
58        eprintln!("[{phase}] {:.0}% {op}", p * 100.0);
59    };
60    let pipeline = AceStepPipeline::<B>::load(&paths, &device, &mut progress)?;
61
62    let metadata = GenerateMetadata {
63        bpm,
64        key_scale: key_scale.as_deref(),
65        time_signature: time_signature.as_deref(),
66    };
67    let mut trace = AceStepTrace::default();
68    let (audio, meta) =
69        pipeline.generate_traced(&caption, &metadata, length_ms, 0, &mut progress, &mut trace)?;
70
71    // ---- Phase dumps ----
72    println!(
73        "\n===== TEXT PROMPT (caption branch) =====\n{}",
74        trace.text_prompt
75    );
76    println!("===== LYRIC PROMPT =====\n{}", trace.lyric_prompt);
77    println!("===== LM CoT BLOCK =====\n{}", trace.cot_block);
78    println!("===== LM CODES ({} total) =====", trace.codes.len());
79    println!("all codes: {:?}", trace.codes);
80    let mut sorted = trace.codes.clone();
81    sorted.sort_unstable();
82    sorted.dedup();
83    println!(
84        "unique: {}, min: {}, max: {}",
85        sorted.len(),
86        sorted.first().unwrap_or(&0),
87        sorted.last().unwrap_or(&0)
88    );
89    println!(
90        "\n===== CONDITIONING =====\nenc mean {:.6}  enc std {:.6}",
91        trace.enc_mean, trace.enc_std
92    );
93    println!("hints latent rms: {:.6}", trace.hints_latent_rms);
94    println!("final latent rms: {:.6}", trace.final_latent_rms);
95    println!("\n===== DIT STEPS =====\nstep  t        xt_rms   v_rms");
96    for (i, step) in trace.dit_steps.iter().enumerate() {
97        println!(
98            "{i:>4}  {:.4}   {:.4}   {:.4}",
99            step.t, step.xt_rms, step.v_rms
100        );
101    }
102
103    // ---- WAVs ----
104    if let Some((interleaved, channels, frames)) = &trace.hints_audio {
105        let path = format!("{out_prefix}_hints.wav");
106        write_wav_from_f32_interleaved(
107            interleaved,
108            *channels,
109            *frames,
110            meta.sample_rate_hz,
111            Path::new(&path),
112        )?;
113        println!("\nwrote {path} (VAE decode of LM hints, DiT bypassed)");
114    }
115
116    let [_, channels, frames] = audio.dims();
117    let channel_major: Vec<f32> = audio
118        .into_data()
119        .convert::<f32>()
120        .to_vec()
121        .map_err(|e| anyhow::anyhow!("{e}"))?;
122    let mut interleaved = vec![0.0_f32; channel_major.len()];
123    for (ch, samples) in channel_major.chunks_exact(frames).enumerate() {
124        for (frame, sample) in samples.iter().enumerate() {
125            interleaved[frame * channels + ch] = *sample;
126        }
127    }
128    let path = format!("{out_prefix}_final.wav");
129    write_wav_from_f32_interleaved(
130        &interleaved,
131        channels,
132        frames,
133        meta.sample_rate_hz,
134        Path::new(&path),
135    )?;
136    println!("wrote {path} (full pipeline)");
137
138    // Silence decode reference.
139    let silence = SilenceLatent::<B>::from_burnpack(&paths.silence_latent_bpk, &device)?;
140    let silence_audio = pipeline.vae.decode(silence.slice(frames / 1920));
141    let [_, sch, sframes] = silence_audio.dims();
142    let smaj: Vec<f32> = silence_audio
143        .into_data()
144        .convert::<f32>()
145        .to_vec()
146        .map_err(|e| anyhow::anyhow!("{e}"))?;
147    let mut sinter = vec![0.0_f32; smaj.len()];
148    for (ch, samples) in smaj.chunks_exact(sframes).enumerate() {
149        for (frame, sample) in samples.iter().enumerate() {
150            sinter[frame * sch + ch] = *sample;
151        }
152    }
153    let path = format!("{out_prefix}_silence.wav");
154    write_wav_from_f32_interleaved(&sinter, sch, sframes, meta.sample_rate_hz, Path::new(&path))?;
155    println!("wrote {path} (VAE decode of silence latent)");
156
157    Ok(())
158}
examples/dit_dump_cpu.rs (line 244)
63fn main() -> anyhow::Result<()> {
64    let mut args = std::env::args().skip(1);
65    let model_dir = PathBuf::from(args.next().unwrap_or_else(|| "/home/meka/repos/ace".into()));
66    let out_dir = PathBuf::from(args.next().unwrap_or_else(|| "/var/tmp/dump_ours".into()));
67    let noise_path = args.next().map(PathBuf::from);
68    std::fs::create_dir_all(&out_dir)?;
69
70    let device = Default::default();
71
72    let dit_config = AceStepConfig::load(&model_dir.join("dit_config.json"))?;
73    let text_config = Qwen3Config::load(&model_dir.join("qwen3_config.json"))?;
74    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
75    eprintln!("loading components...");
76    let text_encoder = Qwen3Model::<B>::from_burnpack(
77        &text_config,
78        &model_dir.join("qwen3-encoder.bpk"),
79        &device,
80    )?;
81    let condition = AceStepCondition::<B>::from_burnpack(
82        &dit_config,
83        &model_dir.join("acestep-condition.bpk"),
84        &device,
85    )?;
86    let dit =
87        AceStepDiT::<B>::from_burnpack(&dit_config, &model_dir.join("acestep-dit.bpk"), &device)?;
88    let vae = OobleckDecoder::<B>::from_burnpack(
89        &vae_config,
90        &model_dir.join("acestep-vae.bpk"),
91        &device,
92    )?;
93    let silence =
94        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
95    let tokenizer = tokie::Tokenizer::from_json(model_dir.join("tokenizer.json"))
96        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
97
98    // ---- text + lyric encoding (causal, official prompts) ----
99    let metas = build_metas_block(Some(120.0), Some("A minor"), Some("4/4"), 4);
100    let text_prompt = build_dit_text_prompt("Metal guitar with a lot of distortion", &metas);
101    let ids = tokenizer.encode(&text_prompt, false).ids;
102    let mut ids: Vec<i64> = ids.into_iter().map(i64::from).collect();
103    ids.push(151643); // explicit EOS (official add_eos=true)
104    let n_text = ids.len();
105    let ids_t = Tensor::<B, 2, Int>::from_data(TensorData::new(ids, [1, n_text]), &device);
106    let text_hidden = text_encoder.forward(ids_t, true);
107    let text_hidden_vec = vec_of(text_hidden.clone());
108    dump_tensor(
109        &out_dir,
110        "text_hidden",
111        &text_hidden_vec,
112        &[n_text as i32, 1024],
113    );
114
115    let lyric_ids = tokenizer
116        .encode(
117            maolan_generate::acestep::pipeline::INSTRUMENTAL_LYRIC_PROMPT,
118            false,
119        )
120        .ids;
121    let mut lyric_ids: Vec<i64> = lyric_ids.into_iter().map(i64::from).collect();
122    lyric_ids.push(151643); // explicit EOS (official add_eos=true)
123    let n_lyric = lyric_ids.len();
124    let lyric_ids_t =
125        Tensor::<B, 2, Int>::from_data(TensorData::new(lyric_ids, [1, n_lyric]), &device);
126    // Official: lyric branch is a raw embed_tokens lookup (no transformer).
127    let lyric_hidden = text_encoder.embed_tokens.forward(lyric_ids_t);
128    let lyric_embed_vec = vec_of(lyric_hidden.clone());
129    dump_tensor(
130        &out_dir,
131        "lyric_embed",
132        &lyric_embed_vec,
133        &[n_lyric as i32, 1024],
134    );
135    let lyric_mask = Tensor::<B, 2, Int>::ones([1, n_lyric], &device);
136
137    // ---- conditioning: use the oracle's enc_hidden/context when provided ----
138    let mut oracle_context: Option<Tensor<B, 3>> = None;
139    let enc = if let Some(dir) = std::env::var_os("MAOLAN_ORACLE_DUMP_DIR") {
140        let dir = PathBuf::from(dir);
141        let (enc_shape, enc_data) = load_bin(&dir.join("enc_hidden.bin"));
142        eprintln!("oracle enc_hidden: {enc_shape:?}");
143        let (ctx_shape, ctx_data) = load_bin(&dir.join("context.bin"));
144        eprintln!("oracle context: {ctx_shape:?}");
145        let enc = Tensor::<B, 3>::from_data(
146            TensorData::new(enc_data, [1, enc_shape[0] as usize, enc_shape[1] as usize]),
147            &device,
148        );
149        let ctx = Tensor::<B, 3>::from_data(TensorData::new(ctx_data, [1, 96, 128]), &device);
150        oracle_context = Some(ctx);
151        enc
152    } else {
153        condition.encode(
154            text_hidden,
155            lyric_hidden,
156            lyric_mask,
157            silence.timbre_reference(),
158        )
159    };
160    let [_, enc_len, enc_dim] = enc.dims();
161    let enc_vec = vec_of(enc.clone());
162    dump_tensor(
163        &out_dir,
164        "enc_hidden",
165        &enc_vec,
166        &[enc_len as i32, enc_dim as i32],
167    );
168
169    // ---- FSQ hints from the given codes (env MAOLAN_ACESTEP_CODES or the
170    // built-in oracle sequence) ----
171    let codes: Vec<u32> = std::env::var_os("MAOLAN_ACESTEP_CODES")
172        .map(|raw| {
173            raw.to_string_lossy()
174                .split(',')
175                .filter_map(|part| part.trim().parse::<u32>().ok())
176                .collect()
177        })
178        .unwrap_or_else(|| ORACLE_CODES.to_vec());
179    let n_codes = codes.len();
180    let codes_tensor =
181        Tensor::<B, 2, Int>::from_data(TensorData::new(codes, [1, n_codes]), &device);
182    let hints = condition.codes_to_hints(codes_tensor);
183    let [_, hint_len, _] = hints.dims();
184    let hints_vec = vec_of(hints.clone());
185    dump_tensor(&out_dir, "detok_output", &hints_vec, &[hint_len as i32, 64]);
186
187    // ---- noise ----
188    let noise = if let Some(path) = noise_path {
189        let (shape, data) = load_bin(&path);
190        let frames = shape[0] as usize;
191        Tensor::<B, 3>::from_data(TensorData::new(data, [1, frames, 64]), &device)
192    } else {
193        maolan_generate::acestep::pipeline::seeded_latent_noise(0, hint_len, 64, &device)
194    };
195    let [_, frames, _] = noise.dims();
196
197    // src latents: hints cropped or silence-padded to the noise's frame count
198    let src = if hint_len >= frames {
199        hints.narrow(1, 0, frames)
200    } else {
201        let padding = silence.slice(frames - hint_len);
202        Tensor::cat(vec![hints, padding], 1)
203    };
204    let chunk_mask = Tensor::ones([1, frames, 64], &device);
205    let context = Tensor::cat(vec![src.clone(), chunk_mask], 2);
206    let context = oracle_context.unwrap_or(context);
207    let context_vec = vec_of(context.clone());
208    dump_tensor(&out_dir, "context", &context_vec, &[frames as i32, 128]);
209
210    let noise_vec = vec_of(noise.clone());
211    dump_tensor(&out_dir, "noise", &noise_vec, &[frames as i32, 64]);
212
213    // ---- DiT loop with per-step dumps ----
214    let kv = dit.prepare_cross_kv(enc);
215    let mut xt = noise;
216    let total = TURBO_TIMESTEPS.len();
217    for (index, &t_cur) in TURBO_TIMESTEPS.iter().enumerate() {
218        let xt_vec = vec_of(xt.clone());
219        dump_tensor(
220            &out_dir,
221            &format!("dit_step{index}_xt"),
222            &xt_vec,
223            &[frames as i32, 64],
224        );
225        let v = dit.forward_with_kv(xt.clone(), t_cur, context.clone(), &kv);
226        let v_vec = vec_of(v.clone());
227        dump_tensor(
228            &out_dir,
229            &format!("dit_step{index}_vt"),
230            &v_vec,
231            &[frames as i32, 64],
232        );
233        let dt = if index + 1 == total {
234            t_cur
235        } else {
236            t_cur - TURBO_TIMESTEPS[index + 1]
237        };
238        xt = xt - v * dt;
239    }
240    let x0_vec = vec_of(xt.clone());
241    dump_tensor(&out_dir, "dit_x0", &x0_vec, &[frames as i32, 64]);
242
243    // ---- VAE decode ----
244    let audio = vae.decode(xt);
245    let [_, channels, samples] = audio.dims();
246    let audio_vec = vec_of(audio);
247    dump_tensor(
248        &out_dir,
249        "vae_audio",
250        &audio_vec,
251        &[channels as i32, samples as i32],
252    );
253    println!("dumps written to {}", out_dir.display());
254    Ok(())
255}
examples/dit_dump.rs (line 264)
63fn main() -> anyhow::Result<()> {
64    let mut args = std::env::args().skip(1);
65    let model_dir = PathBuf::from(args.next().unwrap_or_else(|| "/home/meka/repos/ace".into()));
66    let out_dir = PathBuf::from(args.next().unwrap_or_else(|| "/var/tmp/dump_ours".into()));
67    let noise_path = args.next().map(PathBuf::from);
68    std::fs::create_dir_all(&out_dir)?;
69
70    let device = burn::backend::wgpu::WgpuDevice::default();
71    burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
72        &device,
73        burn::backend::wgpu::RuntimeOptions {
74            memory_config: burn::backend::wgpu::MemoryConfiguration::ExclusivePages,
75            ..Default::default()
76        },
77    );
78
79    let variant_sft = std::env::var("MAOLAN_ACESTEP_VARIANT")
80        .map(|v| v == "sft")
81        .unwrap_or(false);
82    let prefix = if variant_sft { "sft-" } else { "" };
83
84    let dit_config = AceStepConfig::load(&model_dir.join(format!("{prefix}dit_config.json")))?;
85    let text_config = Qwen3Config::load(&model_dir.join("qwen3_config.json"))?;
86    let vae_config = OobleckVaeConfig::load(&model_dir.join("vae_config.json"))?;
87    eprintln!("loading components (prefix '{prefix}')...");
88    let text_encoder = Qwen3Model::<B>::from_burnpack(
89        &text_config,
90        &model_dir.join("qwen3-encoder.bpk"),
91        &device,
92    )?;
93    let condition = AceStepCondition::<B>::from_burnpack(
94        &dit_config,
95        &model_dir.join(format!("{prefix}acestep-condition.bpk")),
96        &device,
97    )?;
98    let dit = AceStepDiT::<B>::from_burnpack(
99        &dit_config,
100        &model_dir.join(format!("{prefix}acestep-dit.bpk")),
101        &device,
102    )?;
103    let vae = OobleckDecoder::<B>::from_burnpack(
104        &vae_config,
105        &model_dir.join("acestep-vae.bpk"),
106        &device,
107    )?;
108    let silence =
109        SilenceLatent::<B>::from_burnpack(&model_dir.join("silence_latent.bpk"), &device)?;
110    let tokenizer = tokie::Tokenizer::from_json(model_dir.join("tokenizer.json"))
111        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
112
113    // ---- text + lyric encoding (causal, official prompts) ----
114    let metas = build_metas_block(Some(120.0), Some("A minor"), Some("4/4"), 4);
115    let text_prompt = build_dit_text_prompt("Metal guitar with a lot of distortion", &metas);
116    let ids = tokenizer.encode(&text_prompt, false).ids;
117    let mut ids: Vec<i64> = ids.into_iter().map(i64::from).collect();
118    ids.push(151643); // explicit EOS (official add_eos=true)
119    let n_text = ids.len();
120    let ids_t = Tensor::<B, 2, Int>::from_data(TensorData::new(ids, [1, n_text]), &device);
121    let text_hidden = text_encoder.forward(ids_t, true);
122    let text_hidden_vec = vec_of(text_hidden.clone());
123    dump_tensor(
124        &out_dir,
125        "text_hidden",
126        &text_hidden_vec,
127        &[n_text as i32, 1024],
128    );
129
130    let lyric_ids = tokenizer
131        .encode(
132            maolan_generate::acestep::pipeline::INSTRUMENTAL_LYRIC_PROMPT,
133            false,
134        )
135        .ids;
136    let mut lyric_ids: Vec<i64> = lyric_ids.into_iter().map(i64::from).collect();
137    lyric_ids.push(151643); // explicit EOS (official add_eos=true)
138    let n_lyric = lyric_ids.len();
139    let lyric_ids_t =
140        Tensor::<B, 2, Int>::from_data(TensorData::new(lyric_ids, [1, n_lyric]), &device);
141    // Official: lyric branch is a raw embed_tokens lookup (no transformer).
142    let lyric_hidden = text_encoder.embed_tokens.forward(lyric_ids_t);
143    let lyric_embed_vec = vec_of(lyric_hidden.clone());
144    dump_tensor(
145        &out_dir,
146        "lyric_embed",
147        &lyric_embed_vec,
148        &[n_lyric as i32, 1024],
149    );
150    let lyric_mask = Tensor::<B, 2, Int>::ones([1, n_lyric], &device);
151
152    // ---- conditioning: use the oracle's enc_hidden/context when provided ----
153    let mut oracle_context: Option<Tensor<B, 3>> = None;
154    let enc = if let Some(dir) = std::env::var_os("MAOLAN_ORACLE_DUMP_DIR") {
155        let dir = PathBuf::from(dir);
156        let (enc_shape, enc_data) = load_bin(&dir.join("enc_hidden.bin"));
157        eprintln!("oracle enc_hidden: {enc_shape:?}");
158        let (ctx_shape, ctx_data) = load_bin(&dir.join("context.bin"));
159        eprintln!("oracle context: {ctx_shape:?}");
160        let enc = Tensor::<B, 3>::from_data(
161            TensorData::new(enc_data, [1, enc_shape[0] as usize, enc_shape[1] as usize]),
162            &device,
163        );
164        let ctx = Tensor::<B, 3>::from_data(TensorData::new(ctx_data, [1, 96, 128]), &device);
165        oracle_context = Some(ctx);
166        enc
167    } else {
168        condition.encode(
169            text_hidden,
170            lyric_hidden,
171            lyric_mask,
172            silence.timbre_reference(),
173        )
174    };
175    let [_, enc_len, enc_dim] = enc.dims();
176    let enc_vec = vec_of(enc.clone());
177    dump_tensor(
178        &out_dir,
179        "enc_hidden",
180        &enc_vec,
181        &[enc_len as i32, enc_dim as i32],
182    );
183
184    // ---- FSQ hints from the given codes (env MAOLAN_ACESTEP_CODES or the
185    // built-in oracle sequence) ----
186    let codes: Vec<u32> = std::env::var_os("MAOLAN_ACESTEP_CODES")
187        .map(|raw| {
188            raw.to_string_lossy()
189                .split(',')
190                .filter_map(|part| part.trim().parse::<u32>().ok())
191                .collect()
192        })
193        .unwrap_or_else(|| ORACLE_CODES.to_vec());
194    let n_codes = codes.len();
195    let codes_tensor =
196        Tensor::<B, 2, Int>::from_data(TensorData::new(codes, [1, n_codes]), &device);
197    let hints = condition.codes_to_hints(codes_tensor);
198    let [_, hint_len, _] = hints.dims();
199    let hints_vec = vec_of(hints.clone());
200    dump_tensor(&out_dir, "detok_output", &hints_vec, &[hint_len as i32, 64]);
201
202    // ---- noise ----
203    let noise = if let Some(path) = noise_path {
204        let (shape, data) = load_bin(&path);
205        let frames = shape[0] as usize;
206        Tensor::<B, 3>::from_data(TensorData::new(data, [1, frames, 64]), &device)
207    } else {
208        maolan_generate::acestep::pipeline::seeded_latent_noise(0, hint_len, 64, &device)
209    };
210    let [_, frames, _] = noise.dims();
211
212    // src latents: hints cropped or silence-padded to the noise's frame count
213    let src = if hint_len >= frames {
214        hints.narrow(1, 0, frames)
215    } else {
216        let padding = silence.slice(frames - hint_len);
217        Tensor::cat(vec![hints, padding], 1)
218    };
219    let chunk_mask = Tensor::ones([1, frames, 64], &device);
220    let context = Tensor::cat(vec![src.clone(), chunk_mask], 2);
221    let context = oracle_context.unwrap_or(context);
222    let context_vec = vec_of(context.clone());
223    dump_tensor(&out_dir, "context", &context_vec, &[frames as i32, 128]);
224
225    let noise_vec = vec_of(noise.clone());
226    dump_tensor(&out_dir, "noise", &noise_vec, &[frames as i32, 64]);
227
228    // ---- DiT loop with per-step dumps (schedule from is_turbo) ----
229    let timesteps: &[f32] = if dit_config.is_turbo {
230        &TURBO_TIMESTEPS
231    } else {
232        &SFT_TIMESTEPS
233    };
234    let kv = dit.prepare_cross_kv(enc);
235    let mut xt = noise;
236    let total = timesteps.len();
237    for (index, &t_cur) in timesteps.iter().enumerate() {
238        let xt_vec = vec_of(xt.clone());
239        dump_tensor(
240            &out_dir,
241            &format!("dit_step{index}_xt"),
242            &xt_vec,
243            &[frames as i32, 64],
244        );
245        let v = dit.forward_with_kv(xt.clone(), t_cur, context.clone(), &kv);
246        let v_vec = vec_of(v.clone());
247        dump_tensor(
248            &out_dir,
249            &format!("dit_step{index}_vt"),
250            &v_vec,
251            &[frames as i32, 64],
252        );
253        let dt = if index + 1 == total {
254            t_cur
255        } else {
256            t_cur - timesteps[index + 1]
257        };
258        xt = xt - v * dt;
259    }
260    let x0_vec = vec_of(xt.clone());
261    dump_tensor(&out_dir, "dit_x0", &x0_vec, &[frames as i32, 64]);
262
263    // ---- VAE decode ----
264    let audio = vae.decode(xt);
265    let [_, channels, samples] = audio.dims();
266    let audio_vec = vec_of(audio);
267    dump_tensor(
268        &out_dir,
269        "vae_audio",
270        &audio_vec,
271        &[channels as i32, samples as i32],
272    );
273    println!("dumps written to {}", out_dir.display());
274    Ok(())
275}

Trait Implementations§

Source§

impl<B> AutodiffModule<B> for OobleckDecoder<B>
where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

Source§

type InnerModule = OobleckDecoder<<B as AutodiffBackend>::InnerBackend>

Inner module without auto-differentiation.
Source§

fn valid(&self) -> Self::InnerModule

Returns the same module, but on the inner backend without auto-differentiation.
Source§

fn from_inner(module: Self::InnerModule) -> Self

Wraps an inner module back into an auto-diff module.
Source§

impl<B: Backend> Clone for OobleckDecoder<B>

Source§

fn clone(&self) -> Self

Returns a duplicate of the value. Read more
1.0.0 (const: unstable) · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more
Source§

impl<B: Debug + Backend> Debug for OobleckDecoder<B>

Source§

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more
Source§

impl<B: Backend> Display for OobleckDecoder<B>

Source§

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more
Source§

impl<B> HasAutodiffModule<B> for OobleckDecoder<B::InnerBackend>
where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

Source§

type TrainModule = OobleckDecoder<B>

The module with auto-differentiation.
Source§

impl<B: Backend> Module<B> for OobleckDecoder<B>

Source§

type Record = OobleckDecoderRecord<B>

Type to save and load the module.
Source§

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.
Source§

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.
Source§

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.
Source§

fn visit<Visitor: ModuleVisitor<B>>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.
Source§

fn map<Mapper: ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.
Source§

fn collect_devices(&self, devices: Devices<B>) -> Devices<B>

Return all the devices found in the underneath module tree added to the given vector without duplicates.
Source§

fn to_device(self, device: &B::Device) -> Self

Move the module and all of its sub-modules to the given device. Read more
Source§

fn fork(self, device: &B::Device) -> Self

Fork the module and all of its sub-modules to the given device. Read more
Source§

fn devices(&self) -> Vec<<B as BackendTypes>::Device>

Return all the devices found in the underneath module tree without duplicates.
Source§

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad. Read more
Source§

fn train<AB>(self) -> Self::TrainModule
where AB: AutodiffBackend<InnerBackend = B>, Self: HasAutodiffModule<AB>,

Move the module and all of its sub-modules to the autodiff backend. Read more
Source§

fn save_file<FR, PB>( self, file_path: PB, recorder: &FR, ) -> Result<(), RecorderError>
where FR: FileRecorder<B>, PB: Into<PathBuf>,

Save the module to a file using the provided file recorder. Read more
Source§

fn load_file<FR, PB>( self, file_path: PB, recorder: &FR, device: &<B as BackendTypes>::Device, ) -> Result<Self, RecorderError>
where FR: FileRecorder<B>, PB: Into<PathBuf>,

Load the module from a file using the provided file recorder. Read more
Source§

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.
Source§

impl<B: Backend> ModuleDisplay for OobleckDecoder<B>

Source§

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings. Read more
Source§

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module. Read more
Source§

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module. Read more
Source§

impl<B: Backend> ModuleDisplayDefault for OobleckDecoder<B>

Source§

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes. Read more
Source§

fn num_params(&self) -> usize

Gets the number of the parameters of the module.

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
where ST: ?Sized, DT: ?Sized,

Source§

impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
where ST: ?Sized, DT: ?Sized,

Source§

impl<C> CloneExpand for C
where C: Clone,

Source§

fn __expand_clone_method(&self, _scope: &mut Scope) -> C

Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dest. Read more
Source§

impl<T> Downcast<T> for T

Source§

fn downcast(&self) -> &T

Source§

impl<T> ErasedDestructor for T
where T: 'static,

Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> IntoComptime for T

Source§

fn comptime(self) -> Self

Source§

impl<T> IntoEither for T

Source§

fn into_either(self, into_left: bool) -> Either<Self, Self>

Converts self into a Left variant of Either<Self, Self> if into_left is true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
where F: FnOnce(&Self) -> bool,

Converts self into a Left variant of Either<Self, Self> if into_left(&self) returns true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

impl<T> MaybeSendSync for T

Source§

impl<B, M> ModuleSnapshot<B> for M
where B: Backend, M: Module<B>,

Source§

fn collect( &self, filter: Option<PathFilter>, adapter: Option<Box<dyn ModuleAdapter>>, skip_enum_variants: bool, ) -> Vec<TensorSnapshot>

Collects tensor snapshots for inspection without copying data. Read more
Source§

fn apply( &mut self, snapshots: Vec<TensorSnapshot>, filter: Option<PathFilter>, adapter: Option<Box<dyn ModuleAdapter>>, skip_enum_variants: bool, ) -> ApplyResult
where Self: Sized,

Applies tensor snapshots to the module. Read more
Source§

fn save_into<P>(&self, store: &mut P) -> Result<(), <P as ModuleStore>::Error>
where P: ModuleStore,

Saves tensor snapshots into a ModuleStore. Read more
Source§

fn load_from<P>( &mut self, store: &mut P, ) -> Result<ApplyResult, <P as ModuleStore>::Error>
where P: ModuleStore,

Loads tensor data from a ModuleStore. Read more
Source§

impl<T> Pointable for T

Source§

const ALIGN: usize

The alignment of pointer.
Source§

type Init = T

The type for initializers.
Source§

unsafe fn init(init: <T as Pointable>::Init) -> usize

Initializes a with the given initializer. Read more
Source§

unsafe fn deref<'a>(ptr: usize) -> &'a T

Dereferences the given pointer. Read more
Source§

unsafe fn deref_mut<'a>(ptr: usize) -> &'a mut T

Mutably dereferences the given pointer. Read more
Source§

unsafe fn drop(ptr: usize)

Drops the object pointed to by the given pointer. Read more
Source§

impl<T> Read<Exclusive, BecauseExclusive> for T
where T: ?Sized,

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T> ToCompactString for T
where T: Display,

Source§

fn try_to_compact_string(&self) -> Result<CompactString, ToCompactStringError>

Fallible version of ToCompactString::to_compact_string() Read more
Source§

fn to_compact_string(&self) -> CompactString

Converts the given value to a CompactString. Read more
Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T> ToString for T
where T: Display + ?Sized,

Source§

fn to_string(&self) -> String

Converts the given value to a String. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> TuneInputs for T
where T: Clone + Send + Sync + 'static,

Source§

type At<'a> = T

The concrete input type at lifetime 'a.
Source§

impl<T> Upcast<T> for T

Source§

fn upcast(&self) -> Option<&T>

Source§

impl<V, T> VZip<V> for T
where V: MultiLane<T>,

Source§

fn vzip(self) -> V

Source§

impl<T> WasmNotSend for T
where T: Send,

Source§

impl<T> WasmNotSendSync for T
where T: WasmNotSend + WasmNotSync,

Source§

impl<T> WasmNotSync for T
where T: Sync,