1use crate::audiovae::AudioVae;
20use crate::mmh3::{Layout, MiniMaxH3, time_shift_sigma};
21use crate::qwen3te::{ImageSpan, Qwen3Encoder};
22use crate::qwen3vis::{self, VisionTower};
23use crate::sampler::SplitMix64;
24use crate::tokenizer::Tokenizer;
25use crate::vae3d::VideoVae;
26use crate::vae3d::VideoVaeEncoder;
27use std::path::Path;
28use std::sync::Arc;
29
30pub const FPS: usize = 24;
31pub const AUDIO_LATENT_FPS: usize = 40;
32
33pub struct AnimParams {
34 pub width: usize,
35 pub height: usize,
36 pub frames: usize,
38 pub steps: usize,
39 pub seed: u64,
40 pub stock_sampler: bool,
43 pub max_tokens: usize,
44 pub first_frame: Option<(Vec<f32>, usize, usize)>,
47 pub last_frame: Option<(Vec<f32>, usize, usize)>,
48 pub mid_frames: Vec<((Vec<f32>, usize, usize), usize)>,
53 pub lora: Option<String>,
55 pub lora_strength: f32,
56 pub upscale: Option<String>,
61 pub upscale_by: f32,
62}
63
64const VISION_START: u32 = 151_652;
66const VISION_END: u32 = 151_653;
67
68pub fn fit_to_canvas(
72 rgb: &[f32],
73 h: usize,
74 w: usize,
75 out_h: usize,
76 out_w: usize,
77 crop: bool,
78) -> Vec<f32> {
79 let (sx0, sy0, sw, sh) = if crop {
82 let (tw, th) = (out_w as f64, out_h as f64);
83 let scale = (w as f64 / tw).min(h as f64 / th);
84 let (cw, ch) = ((tw * scale).round() as usize, (th * scale).round() as usize);
85 ((w - cw) / 2, (h - ch) / 2, cw.max(1), ch.max(1))
86 } else {
87 (0, 0, w, h)
88 };
89 let mut out = vec![0f32; 3 * out_h * out_w];
90 for c in 0..3 {
91 for y in 0..out_h {
92 let sy = ((y as f64 + 0.5) * sh as f64 / out_h as f64 - 0.5).max(0.0);
93 let y0 = sy.floor() as usize;
94 let y1 = (y0 + 1).min(sh - 1);
95 let fy = (sy - y0 as f64) as f32;
96 for x in 0..out_w {
97 let sx = ((x as f64 + 0.5) * sw as f64 / out_w as f64 - 0.5).max(0.0);
98 let x0 = sx.floor() as usize;
99 let x1 = (x0 + 1).min(sw - 1);
100 let fx = (sx - x0 as f64) as f32;
101 let p = |yy: usize, xx: usize| rgb[(c * h + sy0 + yy) * w + sx0 + xx];
102 let top = p(y0, x0) * (1.0 - fx) + p(y0, x1) * fx;
103 let bot = p(y1, x0) * (1.0 - fx) + p(y1, x1) * fx;
104 out[(c * out_h + y) * out_w + x] = top * (1.0 - fy) + bot * fy;
105 }
106 }
107 }
108 out
109}
110
111impl Default for AnimParams {
112 fn default() -> Self {
113 Self {
114 width: 512,
115 height: 288,
116 frames: 39,
117 steps: 4,
118 seed: 42,
119 stock_sampler: false,
120 max_tokens: 512,
121 first_frame: None,
122 last_frame: None,
123 mid_frames: Vec::new(),
124 lora: None,
125 lora_strength: 1.0,
126 upscale: None,
127 upscale_by: 2.0,
128 }
129 }
130}
131
132pub struct Anim {
135 pub rgb: Vec<f32>,
136 pub frames: usize,
137 pub height: usize,
138 pub width: usize,
139 pub audio: Vec<f32>,
140 pub samples: usize,
141 pub sample_rate: usize,
142}
143
144pub fn align_frames(n: usize) -> usize {
147 let mut n = n.max(5);
148 while n % 17 != 5 {
149 n += 1;
150 }
151 n
152}
153
154pub fn video_latent_t(frames: usize) -> usize {
155 if frames <= 5 {
156 2
157 } else {
158 (frames - 5) / 17 * 5 + 2
159 }
160}
161
162pub fn temporal_shape(len: usize) -> (usize, usize, usize) {
164 let frames = align_frames(len);
165 let audio_t = ((frames as f64 / FPS as f64) * AUDIO_LATENT_FPS as f64).round() as usize;
166 (frames, video_latent_t(frames), audio_t)
167}
168
169pub fn sigmas(steps: usize, shift: f64) -> Vec<f64> {
172 let table = 1000usize;
173 let mut out: Vec<f64> = (0..steps)
174 .map(|x| {
175 let idx = table - 1 - x * table / steps;
176 let t = (idx + 1) as f64 / table as f64;
177 shift * t / (1.0 + (shift - 1.0) * t)
178 })
179 .collect();
180 out.push(0.0);
181 out
182}
183
184pub fn gauss_pub(n: usize, seed: u64) -> Vec<f32> {
186 gauss(n, seed)
187}
188
189fn gauss(n: usize, seed: u64) -> Vec<f32> {
190 let mut rng = SplitMix64::new(seed);
191 let mut u = || (rng.next_u64() >> 11) as f64 / (1u64 << 53) as f64;
192 let mut out = Vec::with_capacity(n);
193 while out.len() < n {
194 let (a, b) = (u().max(1e-300), u());
195 let r = (-2.0 * a.ln()).sqrt();
196 let ang = 2.0 * std::f64::consts::PI * b;
197 out.push((r * ang.cos()) as f32);
198 if out.len() < n {
199 out.push((r * ang.sin()) as f32);
200 }
201 }
202 out
203}
204
205pub fn generate(
207 path: &Path,
208 prompt: &str,
209 p: &AnimParams,
210 mut progress: impl FnMut(&str, usize, usize),
211) -> Result<Anim, String> {
212 if p.width % 32 != 0 || p.height % 32 != 0 {
213 return Err("width/height must be multiples of 32".into());
214 }
215 let use_gpu = match std::env::var("CMF_MMH3_GPU").ok().as_deref() {
223 Some("1") => true,
224 Some("0") => false,
225 _ => mmh3_gpu_parity_probe(path).unwrap_or(false),
226 };
227 if use_gpu {
228 generate_inner(path, prompt, p, &mut progress)
229 } else {
230 crate::gpu::cpu_scope(|| generate_inner(path, prompt, p, &mut progress))
231 }
232}
233
234fn mmh3_gpu_parity_probe(path: &Path) -> Result<bool, String> {
241 let model = Arc::new(
242 cortiq_core::CmfModel::open(path).map_err(|e| format!("{}: {e}", path.display()))?,
243 );
244 let Some(idx) = model.tensors.iter().position(|t| {
249 t.name.starts_with("dit.") && t.name.ends_with("attn.qkv_proj.weight")
250 }) else {
251 tracing::info!("mmh3 GPU parity probe: no qkv weight — host path");
252 return Ok(false);
253 };
254 let entry = &model.tensors[idx];
255 let (rows, cols) = (entry.shape[0], entry.shape[1]);
256 let b = 64usize;
257 let mut xs = vec![0f32; b * cols];
258 for (i, v) in xs.iter_mut().enumerate() {
259 let base = ((i * 37 + 11) % 1009) as f32 / 1009.0 - 0.5;
260 *v = base * if i % 7 == 0 { 2000.0 } else { 2.0 };
261 }
262 let mut gpu = vec![0f32; b * rows];
263 if std::env::var("CMF_GPU_DEBUG").is_ok() {
264 eprintln!(
265 "mmh3 probe env: CMF_GPU={:?} enabled={} avail={}",
266 std::env::var("CMF_GPU").ok(),
267 crate::gpu::enabled(),
268 crate::gpu::backend_available(),
269 );
270 }
271 let qt = crate::qtensor::QTensor::from_model(&model, &entry.name.clone())?;
272 if !qt.device_matmat(&xs, b, &mut gpu) {
273 tracing::info!(
274 "mmh3 GPU parity probe: {:?} device GEMM refused ({rows}x{cols}) — host path",
275 entry.dtype
276 );
277 if std::env::var("CMF_GPU_DEBUG").is_ok() {
278 eprintln!("mmh3 probe: device GEMM refused {rows}x{cols} for {:?}", entry.dtype);
279 }
280 return Ok(false);
281 }
282 let host = {
283 let name = entry.name.clone();
284 let proj = crate::dit::Proj::from_model(&model, &name)?;
285 let mut out = vec![0f32; b * rows];
286 crate::gpu::cpu_scope(|| proj.matmat(&xs, b, &mut out, None));
287 out
288 };
289 let mut num = 0f64;
290 let mut den = 0f64;
291 for (g, h) in gpu.iter().zip(&host) {
292 num += ((g - h) as f64).powi(2);
293 den += (*h as f64).powi(2);
294 }
295 let rel = (num / den.max(1e-30)).sqrt();
296 let ok = rel < 1e-2;
297 tracing::info!(
298 "mmh3 GPU parity probe: rel rms {rel:.2e} → {}",
299 if ok { "device" } else { "host" }
300 );
301 if std::env::var("CMF_GPU_DEBUG").is_ok() {
302 eprintln!("mmh3 GPU parity probe: rel rms {rel:.2e}");
303 }
304 Ok(ok)
305}
306
307fn generate_inner(
308 path: &Path,
309 prompt: &str,
310 p: &AnimParams,
311 progress: &mut dyn FnMut(&str, usize, usize),
312) -> Result<Anim, String> {
313 let model = Arc::new(
314 cortiq_core::CmfModel::open(path).map_err(|e| format!("{}: {e}", path.display()))?,
315 );
316 let (frames_total, latent_t, audio_t) = temporal_shape(p.frames);
317 let (lat_h, lat_w) = (p.height / 16, p.width / 16);
318
319 let t_stage = std::time::Instant::now();
326 let mut marks: Vec<(&str, f32)> = Vec::new();
327 let lap = |marks: &mut Vec<(&'static str, f32)>, name: &'static str| {
328 let prev: f32 = marks.iter().map(|(_, v)| v).sum();
329 marks.push((name, t_stage.elapsed().as_secs_f32() - prev));
330 };
331 let vocab = model
332 .vocab
333 .as_deref()
334 .ok_or("packaged .cmf has no embedded tokenizer")?;
335 let tok = Tokenizer::from_bytes(vocab).map_err(|e| format!("tokenizer: {e}"))?;
336 let mut keyframes: Vec<(&(Vec<f32>, usize, usize), usize)> = p
340 .first_frame
341 .iter()
342 .map(|f| (f, 0usize))
343 .chain(p.mid_frames.iter().map(|(f, i)| (f, (*i).min(frames_total - 1))))
344 .chain(p.last_frame.iter().map(|f| (f, frames_total - 1)))
345 .collect();
346 keyframes.sort_by_key(|(_, i)| *i);
350 keyframes.dedup_by_key(|(_, i)| *i);
351 let mut ids: Vec<u32> = Vec::new();
352 let mut spans: Vec<ImageSpan> = Vec::new();
353 let mut embeds: Vec<Vec<f32>> = Vec::new();
354 let mut deepstack: Vec<Vec<f32>> = Vec::new();
355 let mut cond: Vec<Vec<f32>> = Vec::new();
356 let mut tags: Vec<u8> = Vec::new();
357
358 if !keyframes.is_empty() {
359 let tower = VisionTower::from_cmf(&model)?;
360 let te_only = std::env::var("CMF_TE_ONLY").as_deref() == Ok("1");
364 let venc = if te_only {
365 None
366 } else {
367 Some(VideoVaeEncoder::from_cmf(&model)?)
368 };
369 for (i, (frame, _)) in keyframes.iter().enumerate() {
370 let (src, sh, sw) = *frame;
371 let fitted = fit_to_canvas(src, *sh, *sw, p.height, p.width, i > 0);
374 if let Some(venc) = &venc {
375 let (z, _, _) = venc.encode_frame(
376 &fitted.iter().map(|&v| v * 2.0 - 1.0).collect::<Vec<_>>(),
377 p.height,
378 p.width,
379 );
380 cond.push(z);
381 }
382
383 for t in tok.encode(&format!("<Picture {}>: ", i + 1)) {
384 ids.push(t);
385 tags.push(1);
386 }
387 let (patches, gh, gw) = qwen3vis::preprocess(
388 &fitted,
389 p.height,
390 p.width,
391 tower.patch_size,
392 tower.temporal_patch,
393 tower.merge,
394 );
395 let (merged, deep) = tower.forward(&patches, gh, gw);
396 let n_img = merged.len() / tower.out_hidden;
397 ids.push(VISION_START);
400 tags.push(0);
401 let start = ids.len();
402 for _ in 0..n_img {
403 ids.push(VISION_START); tags.push(0);
405 }
406 ids.push(VISION_END);
407 tags.push(0);
408 spans.push(ImageSpan {
409 start,
410 len: n_img,
411 merged_h: gh / tower.merge,
412 merged_w: gw / tower.merge,
413 });
414 embeds.push(merged);
415 if deepstack.is_empty() {
416 deepstack = deep;
417 } else {
418 for (a, b) in deepstack.iter_mut().zip(deep) {
419 a.extend_from_slice(&b);
420 }
421 }
422 }
423 }
424 for t in tok.encode(prompt) {
425 ids.push(t);
426 tags.push(1);
427 }
428 if ids.is_empty() {
429 ids.push(151643); tags.push(1);
431 }
432 ids.truncate(p.max_tokens);
433 tags.truncate(ids.len());
434 lap(&mut marks, "prepare");
435 crate::gpu::mm_kill_arm(false);
441 progress("encode", 0, 1);
442 let states = {
443 let enc = Qwen3Encoder::from_cmf(&model)?;
444 enc.encode_with_images(&ids, &spans, &embeds, &deepstack)
445 };
446 if let Ok(p) = std::env::var("CMF_TE_DUMP") {
451 let w = states.len() / ids.len().max(1);
452 let mut b = Vec::with_capacity(16 + states.len() * 4);
453 b.extend_from_slice(&(ids.len() as u64).to_le_bytes());
454 b.extend_from_slice(&(w as u64).to_le_bytes());
455 for v in &states {
456 b.extend_from_slice(&v.to_le_bytes());
457 }
458 std::fs::write(&p, &b).map_err(|e| format!("CMF_TE_DUMP {p}: {e}"))?;
459 eprintln!("te dump: {} tokens x {w} -> {p}", ids.len());
460 }
461 progress("encode", 1, 1);
462 lap(&mut marks, "text encode");
463 crate::gpu::mm_kill_arm(true);
464 if std::env::var("CMF_TE_ONLY").as_deref() == Ok("1") {
467 return Err("CMF_TE_ONLY: encode dumped, render skipped".into());
468 }
469 {
475 let dropped = model.advise_done(|n| n.starts_with("model.") || n.starts_with("vis."));
476 if dropped > 0 {
477 tracing::info!(
478 "encoder pages released after prompt encode: {} MB",
479 dropped / (1024 * 1024)
480 );
481 }
482 }
483
484 let (mut video, audio) = {
486 let bank = match p.lora.as_deref() {
491 None => None,
492 Some(path) => {
493 let k = crate::ltxlora::LoraBank::load(std::path::Path::new(path), p.lora_strength)?;
494 Some(k)
495 }
496 };
497 let dit = MiniMaxH3::from_cmf_lora(&model, bank.as_ref())?;
498 if let Some(k) = &bank {
499 let bound = dit.lora_bound();
500 let mut skipped: std::collections::BTreeMap<String, usize> = Default::default();
504 for name in k.keys() {
505 if !dit.lora_binds(name) {
506 let fam = name
507 .rsplit_once('.')
508 .map(|(_, t)| {
509 let head = name.split('.').next().unwrap_or("");
510 format!("{head}…{t}")
511 })
512 .unwrap_or_else(|| name.to_string());
513 *skipped.entry(fam).or_default() += 1;
514 }
515 }
516 let tail = if skipped.is_empty() {
517 String::new()
518 } else {
519 let parts: Vec<String> =
520 skipped.iter().map(|(k, v)| format!("{k} ×{v}")).collect();
521 format!("; not applied: {}", parts.join(", "))
522 };
523 tracing::info!(
524 "lora: rank {}, {} branches, {} bound at strength {}{}",
525 k.rank(),
526 k.len(),
527 bound,
528 p.lora_strength,
529 tail
530 );
531 println!(
532 "lora: rank {}, {}/{} branches bound{}",
533 k.rank(),
534 bound,
535 k.len(),
536 tail
537 );
538 }
539 let kf: Vec<(usize, usize)> = keyframes
540 .iter()
541 .map(|&(_, idx)| (idx, frames_total))
542 .collect();
543 let layout = if kf.is_empty() {
544 Layout::t2va(ids.len(), latent_t, lat_h, lat_w, audio_t)
545 } else {
546 Layout::fl2va(ids.len(), latent_t, lat_h, lat_w, audio_t, &kf, &tags)
547 };
548 let text = dit.refine_text(&states, ids.len());
549 let mut v = gauss(dit.latents_dim * latent_t * lat_h * lat_w, p.seed);
550 let mut a = gauss(dit.audio_dim * 2 * audio_t, p.seed ^ 0x9E37_79B9_7F4A_7C15);
551 let sg = sigmas(p.steps, dit.shift_video);
552 let prof = std::env::var_os("CMF_ANIM_PROF").is_some();
556 let rms = |x: &[f32]| {
557 (x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64).sqrt()
558 };
559 if prof {
560 eprintln!(
561 " text {} tok, refined rms {:.4}, sigmas {:?}",
562 ids.len(),
563 rms(&text),
564 sg.iter()
565 .map(|v| (v * 1e4).round() / 1e4)
566 .collect::<Vec<_>>()
567 );
568 }
569 for i in 0..p.steps {
570 let (sv, sv_n) = (sg[i], sg[i + 1]);
571 let (dv, da) = dit.forward(&layout, &text, &v, &a, sv, &cond);
572 let step_v = (sv_n - sv) as f32;
573 for (x, &d) in v.iter_mut().zip(&dv) {
574 *x += step_v * d;
575 }
576 let step_a = if p.stock_sampler {
577 step_v
578 } else {
579 (time_shift_sigma(sv_n, dit.shift_video, dit.shift_audio)
580 - time_shift_sigma(sv.max(1e-6), dit.shift_video, dit.shift_audio))
581 as f32
582 };
583 for (x, &d) in a.iter_mut().zip(&da) {
584 *x += step_a * d;
585 }
586 if prof {
587 eprintln!(
588 " step {i}: sv {sv:.4}->{sv_n:.4} v_vel {:.4} a_vel {:.4} | video {:.4} audio {:.4}",
589 rms(&dv),
590 rms(&da),
591 rms(&v),
592 rms(&a)
593 );
594 }
595 progress("denoise", i + 1, p.steps);
596 }
597 if let Some(rep) = dit.lora_report() {
598 eprint!("{rep}");
599 }
600 (v, a)
601 };
602
603 lap(&mut marks, "denoise");
604
605 let (mut out_h, mut out_w) = (p.height, p.width);
607 let (mut lat_h, mut lat_w) = (lat_h, lat_w);
608 if let Some(path) = p.upscale.as_deref() {
609 progress("upscale", 0, 1);
610 let t = std::time::Instant::now();
611 let ups = crate::mmh3ups::LatentUpscaler::load(std::path::Path::new(path))?;
612 let z = crate::mmh3ups::Vol {
613 c: video.len() / (latent_t * lat_h * lat_w),
614 t: latent_t,
615 h: lat_h,
616 w: lat_w,
617 data: video,
618 };
619 let f = p.upscale_by.max(1.0);
622 let nh = ((lat_h as f32 * f).round() as usize).max(lat_h);
623 let nw = ((lat_w as f32 * f).round() as usize).max(lat_w);
624 let big = ups.upscale(&z, nh, nw, None);
625 tracing::info!(
626 "latent upscale {}x{} -> {}x{} in {:.1}s",
627 lat_h,
628 lat_w,
629 nh,
630 nw,
631 t.elapsed().as_secs_f64()
632 );
633 out_h = out_h * nh / lat_h;
634 out_w = out_w * nw / lat_w;
635 lat_h = nh;
636 lat_w = nw;
637 video = big.data;
638 progress("upscale", 1, 1);
639 lap(&mut marks, "upscale");
640 }
641
642 progress("video vae", 0, 1);
644 let (rgb, out_frames) = {
645 let vae = VideoVae::from_cmf(&model)?;
646 vae.decode(&video, latent_t, lat_h, lat_w)
647 };
648 progress("video vae", 1, 1);
649 lap(&mut marks, "video vae");
650 progress("audio vae", 0, 1);
651 let (wave, samples, sr) = {
652 let vae = AudioVae::from_cmf(&model)?;
653 let c = audio.len() / (2 * audio_t);
654 let (w, n) = vae.decode(&audio, c, audio_t);
655 (w, n, vae.sample_rate)
656 };
657 progress("audio vae", 1, 1);
658 lap(&mut marks, "audio vae");
659 tracing::info!(
660 "stages: {}",
661 marks
662 .iter()
663 .map(|(n, v)| format!("{n} {v:.1}s"))
664 .collect::<Vec<_>>()
665 .join(" · ")
666 );
667 #[cfg(target_os = "macos")]
671 if std::env::var("CMF_METAL_MMPROF").is_ok() {
672 use std::sync::atomic::Ordering::Relaxed;
673 let n = crate::gpu_metal::MM_N.load(Relaxed);
674 eprintln!(
675 " q4tp mm x{n}: upload {:.1}s · submit+wait {:.1}s · readback {:.1}s",
676 crate::gpu_metal::MM_UP.load(Relaxed) as f64 / 1e6,
677 crate::gpu_metal::MM_GPU.load(Relaxed) as f64 / 1e6,
678 crate::gpu_metal::MM_DN.load(Relaxed) as f64 / 1e6,
679 );
680 }
681
682 let keep = out_frames.min(frames_total);
686 Ok(Anim {
687 rgb: trim_frames(&rgb, out_frames, keep, out_h, out_w),
688 frames: keep,
689 height: out_h,
690 width: out_w,
691 audio: wave,
692 samples,
693 sample_rate: sr,
694 })
695}
696
697fn trim_frames(rgb: &[f32], have: usize, keep: usize, h: usize, w: usize) -> Vec<f32> {
698 if keep == have {
699 return rgb.to_vec();
700 }
701 let mut out = vec![0f32; 3 * keep * h * w];
702 for c in 0..3 {
703 let s = c * have * h * w;
704 let d = c * keep * h * w;
705 out[d..d + keep * h * w].copy_from_slice(&rgb[s..s + keep * h * w]);
706 }
707 out
708}
709
710#[cfg(test)]
711mod tests {
712 use super::*;
713
714 #[test]
715 fn the_four_step_schedule_is_the_references() {
716 let s = sigmas(4, 12.0);
717 let want = [1.0, 0.972_973, 0.923_077, 0.8, 0.0];
718 assert_eq!(s.len(), want.len());
719 for (g, w) in s.iter().zip(&want) {
720 assert!((g - w).abs() < 1e-6, "{s:?}");
721 }
722 }
723
724 #[test]
725 fn a_stretch_keeps_the_corners_and_a_crop_takes_the_middle() {
726 let (h, w) = (2usize, 4usize);
729 let mut rgb = vec![0f32; 3 * h * w];
730 for c in 0..3 {
731 for y in 0..h {
732 for x in 0..w {
733 rgb[(c * h + y) * w + x] = x as f32 / (w - 1) as f32;
734 }
735 }
736 }
737 let s = fit_to_canvas(&rgb, h, w, 4, 4, false);
739 assert!((s[0] - 0.0).abs() < 1e-6, "left edge");
740 assert!((s[3] - 1.0).abs() < 1e-6, "right edge");
741 let c = fit_to_canvas(&rgb, h, w, 4, 4, true);
744 let (lo, hi) = c[..16]
745 .iter()
746 .fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
747 assert!(lo > 0.05, "crop kept the left edge: {lo}");
748 assert!(hi < 0.95, "crop kept the right edge: {hi}");
749 }
750
751 #[test]
752 fn frame_counts_snap_to_the_models_grid() {
753 assert_eq!(align_frames(1), 5);
755 assert_eq!(align_frames(39), 39);
756 assert_eq!(align_frames(41), 56);
757 assert_eq!(align_frames(124), 124);
758 assert_eq!(video_latent_t(124), 37);
759 assert_eq!(video_latent_t(39), 12);
760 let (f, lt, at) = temporal_shape(124);
761 assert_eq!((f, lt, at), (124, 37, 207));
762 }
763}