Skip to main content

cortiq_engine/
videogen.rs

1//! End-to-end MiniMax-H3 text→(video + synchronized stereo audio):
2//! Qwen3-VL prompt encode → 4-step dual-schedule flow sampling over the
3//! packed DiT → ViT3D video decode and BigVGAN audio decode.
4//!
5//! Stages load and drop one at a time, as the image pipeline next door
6//! does: peak resident is one component, not their sum.
7//!
8//! ## Two clocks, four steps
9//!
10//! The sampler walks the VIDEO sigma grid — `simple` at shift 12, which
11//! at four steps is 1, 0.973, 0.923, 0.8, 0 — and the audio stream is
12//! integrated on its own remap of that grid (shift 3). Stepping both on
13//! the video grid is what a stock sampler does, and it is fine at
14//! twenty steps and audibly wrong at four: `Δσ_a` and `Δσ_v` differ by
15//! a factor of three over the last interval, and no per-step slope
16//! correction fixes a step that large. Hence `--stock-sampler`, which
17//! reproduces the broken behaviour on purpose, for comparison.
18
19use 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    /// Frames at 24 fps; snapped up to the model's 17k+5 grid.
37    pub frames: usize,
38    pub steps: usize,
39    pub seed: u64,
40    /// Integrate the audio on the video's grid, as a stock sampler
41    /// would. Wrong at four steps; kept for A/B.
42    pub stock_sampler: bool,
43    pub max_tokens: usize,
44    /// RGB in [0, 1] as `[3, h, w]` with its size — the clip's first
45    /// frame, and/or its last.
46    pub first_frame: Option<(Vec<f32>, usize, usize)>,
47    pub last_frame: Option<(Vec<f32>, usize, usize)>,
48    /// Extra reference frames with the pixel index each one stands for:
49    /// video-to-video as this architecture actually takes it — every
50    /// frame is a condition row pinned to its own time coordinate, the
51    /// same machinery `--first-frame`/`--last-frame` use for two.
52    pub mid_frames: Vec<((Vec<f32>, usize, usize), usize)>,
53    /// A LoRA adapter (.safetensors) applied at runtime, and how hard.
54    pub lora: Option<String>,
55    pub lora_strength: f32,
56    /// The published latent upscaler (.safetensors) and the factor to
57    /// apply: the denoised latent is resized by the learned net and the
58    /// VAE decodes at the larger size, so the 5 B-parameter decode →
59    /// pixel resize → encode round trip never happens.
60    pub upscale: Option<String>,
61    pub upscale_by: f32,
62}
63
64/// The vision-block token ids the H3 presentation flanks a picture with.
65const VISION_START: u32 = 151_652;
66const VISION_END: u32 = 151_653;
67
68/// Resize `[3, h, w]` RGB to the canvas. The first frame is a geometry
69/// anchor and is stretched; the last one follows and is cover-cropped,
70/// which is what the reference node does with each.
71pub 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    // Cover-crop picks the largest centred rectangle of the source with
80    // the target's aspect; a stretch takes the whole thing.
81    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
132/// The rendered result: RGB in [0, 1] as `[3, frames, h, w]`, and
133/// stereo f32 in [-1, 1] as `[2, samples]`.
134pub 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
144/// Frame counts snap UP to 17k+5 — the grid the temporal VAE and the
145/// DiT's frame-span pattern agree on.
146pub 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
162/// `(frames, latent_t, audio_t)` for a requested length.
163pub 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
169/// The `simple` scheduler over `ModelSamplingDiscreteFlow(shift)`: the
170/// 1000-entry sigma table sampled at even strides, terminal 0 appended.
171pub 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
184/// Shared with the DiT's condition-row noise blend.
185pub 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
205/// Text→(video, audio) from a packaged `.cmf`.
206pub 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    // The wgpu wide-GEMM arm was measured WRONG on one driver stack
216    // (RTX PRO 6000: step-1 velocity rms off, step-2 NaN) and byte-
217    // healthy on another (2×RTX 5090, coop and plain arms within 0.5%
218    // of each other and 3.5% of the host render). Trust is therefore
219    // PER-STACK, decided by a parity probe on this file's own first
220    // qkv weight at DiT-scale activations — not by a hardcoded verdict
221    // either way. CMF_MMH3_GPU=1/0 still forces.
222    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
234/// GPU-vs-host parity on the packed DiT's first attention projection:
235/// the real q4tp bytes, activations spanning the modulation range
236/// (±2000 mixed with ±2), rms gate at 1e-2 — the measured failure was
237/// ~24%, honest drift is ~1e-5, so the gate has a decade of margin on
238/// each side. Any refusal (no adapter, dtype outside the kernel) is a
239/// clean "no": the host path is never wrong, only slower.
240fn 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    // Any codec, not just q4tp: the probe's job is to check THIS file's
245    // device arm against the host, and a container packed as q8_2f has one
246    // too. Matching on dtype sent every non-q4tp build to the CPU for the
247    // whole render — measured at 357 s against 60 on the same card.
248    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    // ── prompt ──
320    // The H3 presentation is raw text: no chat template, no BOS, no
321    // special tokens at all.
322    // Stage clock. Half of a render is not the DiT — at 8 steps the
323    // denoiser is 63 s of 116 — and until this line existed there was
324    // no way to see which half anything went to.
325    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    // fl2va: every keyframe is presented as "<Picture i>: " and a
337    // vision block BEFORE the prompt, and separately conditions the DiT
338    // as a latent. Both halves come from the same picture.
339    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    // Condition rows are placed by time coordinate, so they have to arrive
347    // in time order and only once per frame — two references pinned to the
348    // same pixel index would be two rows claiming one moment.
349    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        // An activation harvest (CMF_TE_ONLY) never denoises, and the
361        // VAE latent is the DiT's food alone — the frame's 3-D conv
362        // encode is 99.5 s of a 102.5 s M4 run. Skip it.
363        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            // The picture the DiT sees is on the generation canvas; the
372            // one Qwen sees keeps its own resolution policy.
373            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            // The whole block carries the VIDEO tag, the flanking
398            // markers included.
399            ids.push(VISION_START);
400            tags.push(0);
401            let start = ids.len();
402            for _ in 0..n_img {
403                ids.push(VISION_START); // a placeholder the embed replaces
404                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); // the pad id, as the reference does for ""
430        tags.push(1);
431    }
432    ids.truncate(p.max_tokens);
433    tags.truncate(ids.len());
434    lap(&mut marks, "prepare");
435    // The prompt encoder is a one-shot pass over 12 GB of weights; on a
436    // machine the file does not fit it streams from disk and its GEMMs
437    // run over any contention budget for reasons that are not contention.
438    // The kill stays disarmed until the encode is done (users on 24 GB
439    // Macs had to patch it out to keep the denoise loop on the GPU).
440    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    // `CMF_TE_DUMP=<path>`: the conditioning as `[u64 n][u64 width]`
447    // then f32 rows. A stand-in encoder is only as good as the stream
448    // it hands the DiT, and that is measurable against the teacher's
449    // dump on the same prompt WITHOUT rendering a frame.
450    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    // `CMF_TE_ONLY=1`: stop after the dump — an activation-harvest run
465    // (the ClipProj refit) wants hundreds of encodes and zero renders.
466    if std::env::var("CMF_TE_ONLY").as_deref() == Ok("1") {
467        return Err("CMF_TE_ONLY: encode dumped, render skipped".into());
468    }
469    // The prompt encoder and vision tower ran their once-per-generation
470    // pass; release their page cache so the denoise loop's DiT does not
471    // fight 12+ GB of dead weights for RAM. On a 24 GB Mac with the
472    // 25.7 GB full-encoder fl2va file this is the difference between
473    // denoise steps at DiT speed and 320 s/step of SSD thrash.
474    {
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    // ── denoise ──
485    let (mut video, audio) = {
486        // The adapter is read here, after the encoder's pages are gone:
487        // a rank-32 file for this DiT is 130 MB of f32 once expanded and
488        // there is no reason for it to share a peak with 12 GB of text
489        // tower on a 24 GB machine.
490        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            // Say what did NOT land. An adaLN branch on a curve-form
501            // pack is the one real gap, and a user who sees "applied"
502            // while half the adapter sat out has been lied to.
503            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        // `CMF_ANIM_PROF=1`: the per-step rms of both streams and of
553        // their velocities. A run that is not denoising shows it here
554        // long before anything is written out.
555        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    // ── the learned latent resize, when one was handed in ──
606    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        // Snap to the VAE's own 16-pixel grid: the net takes any target,
620        // the decoder does not.
621        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    // ── decode ──
643    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    // Where a GEMM's wall time goes on unified memory: copies or kernel.
668    // The answer decides whether fusing blocks or tuning the kernel is
669    // the optimization worth doing.
670    #[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    // The VAE emits latent_t·4 frames; the request snapped to 17k+5,
683    // which is one fewer than a multiple of four plus the leading key
684    // frame, so trim rather than pad.
685    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        // A 4x2 ramp: value rises left to right, so the corners name
727        // themselves.
728        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        // Stretch to a square: the far edges survive.
738        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        // Cover-crop to a square takes the centre 2x2, so the extremes
742        // are gone and the span is narrower.
743        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        // The grid is 5 + 17k: 5, 22, 39, 56, … 124.
754        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}