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}
49
50/// The vision-block token ids the H3 presentation flanks a picture with.
51const VISION_START: u32 = 151_652;
52const VISION_END: u32 = 151_653;
53
54/// Resize `[3, h, w]` RGB to the canvas. The first frame is a geometry
55/// anchor and is stretched; the last one follows and is cover-cropped,
56/// which is what the reference node does with each.
57pub fn fit_to_canvas(
58    rgb: &[f32],
59    h: usize,
60    w: usize,
61    out_h: usize,
62    out_w: usize,
63    crop: bool,
64) -> Vec<f32> {
65    // Cover-crop picks the largest centred rectangle of the source with
66    // the target's aspect; a stretch takes the whole thing.
67    let (sx0, sy0, sw, sh) = if crop {
68        let (tw, th) = (out_w as f64, out_h as f64);
69        let scale = (w as f64 / tw).min(h as f64 / th);
70        let (cw, ch) = ((tw * scale).round() as usize, (th * scale).round() as usize);
71        ((w - cw) / 2, (h - ch) / 2, cw.max(1), ch.max(1))
72    } else {
73        (0, 0, w, h)
74    };
75    let mut out = vec![0f32; 3 * out_h * out_w];
76    for c in 0..3 {
77        for y in 0..out_h {
78            let sy = ((y as f64 + 0.5) * sh as f64 / out_h as f64 - 0.5).max(0.0);
79            let y0 = sy.floor() as usize;
80            let y1 = (y0 + 1).min(sh - 1);
81            let fy = (sy - y0 as f64) as f32;
82            for x in 0..out_w {
83                let sx = ((x as f64 + 0.5) * sw as f64 / out_w as f64 - 0.5).max(0.0);
84                let x0 = sx.floor() as usize;
85                let x1 = (x0 + 1).min(sw - 1);
86                let fx = (sx - x0 as f64) as f32;
87                let p = |yy: usize, xx: usize| rgb[(c * h + sy0 + yy) * w + sx0 + xx];
88                let top = p(y0, x0) * (1.0 - fx) + p(y0, x1) * fx;
89                let bot = p(y1, x0) * (1.0 - fx) + p(y1, x1) * fx;
90                out[(c * out_h + y) * out_w + x] = top * (1.0 - fy) + bot * fy;
91            }
92        }
93    }
94    out
95}
96
97impl Default for AnimParams {
98    fn default() -> Self {
99        Self {
100            width: 512,
101            height: 288,
102            frames: 39,
103            steps: 4,
104            seed: 42,
105            stock_sampler: false,
106            max_tokens: 512,
107            first_frame: None,
108            last_frame: None,
109        }
110    }
111}
112
113/// The rendered result: RGB in [0, 1] as `[3, frames, h, w]`, and
114/// stereo f32 in [-1, 1] as `[2, samples]`.
115pub struct Anim {
116    pub rgb: Vec<f32>,
117    pub frames: usize,
118    pub height: usize,
119    pub width: usize,
120    pub audio: Vec<f32>,
121    pub samples: usize,
122    pub sample_rate: usize,
123}
124
125/// Frame counts snap UP to 17k+5 — the grid the temporal VAE and the
126/// DiT's frame-span pattern agree on.
127pub fn align_frames(n: usize) -> usize {
128    let mut n = n.max(5);
129    while n % 17 != 5 {
130        n += 1;
131    }
132    n
133}
134
135pub fn video_latent_t(frames: usize) -> usize {
136    if frames <= 5 {
137        2
138    } else {
139        (frames - 5) / 17 * 5 + 2
140    }
141}
142
143/// `(frames, latent_t, audio_t)` for a requested length.
144pub fn temporal_shape(len: usize) -> (usize, usize, usize) {
145    let frames = align_frames(len);
146    let audio_t = ((frames as f64 / FPS as f64) * AUDIO_LATENT_FPS as f64).round() as usize;
147    (frames, video_latent_t(frames), audio_t)
148}
149
150/// The `simple` scheduler over `ModelSamplingDiscreteFlow(shift)`: the
151/// 1000-entry sigma table sampled at even strides, terminal 0 appended.
152pub fn sigmas(steps: usize, shift: f64) -> Vec<f64> {
153    let table = 1000usize;
154    let mut out: Vec<f64> = (0..steps)
155        .map(|x| {
156            let idx = table - 1 - x * table / steps;
157            let t = (idx + 1) as f64 / table as f64;
158            shift * t / (1.0 + (shift - 1.0) * t)
159        })
160        .collect();
161    out.push(0.0);
162    out
163}
164
165/// Shared with the DiT's condition-row noise blend.
166pub fn gauss_pub(n: usize, seed: u64) -> Vec<f32> {
167    gauss(n, seed)
168}
169
170fn gauss(n: usize, seed: u64) -> Vec<f32> {
171    let mut rng = SplitMix64::new(seed);
172    let mut u = || (rng.next_u64() >> 11) as f64 / (1u64 << 53) as f64;
173    let mut out = Vec::with_capacity(n);
174    while out.len() < n {
175        let (a, b) = (u().max(1e-300), u());
176        let r = (-2.0 * a.ln()).sqrt();
177        let ang = 2.0 * std::f64::consts::PI * b;
178        out.push((r * ang.cos()) as f32);
179        if out.len() < n {
180            out.push((r * ang.sin()) as f32);
181        }
182    }
183    out
184}
185
186/// Text→(video, audio) from a packaged `.cmf`.
187pub fn generate(
188    path: &Path,
189    prompt: &str,
190    p: &AnimParams,
191    mut progress: impl FnMut(&str, usize, usize),
192) -> Result<Anim, String> {
193    if p.width % 32 != 0 || p.height % 32 != 0 {
194        return Err("width/height must be multiples of 32".into());
195    }
196    // The wgpu wide-GEMM arm was measured WRONG on one driver stack
197    // (RTX PRO 6000: step-1 velocity rms off, step-2 NaN) and byte-
198    // healthy on another (2×RTX 5090, coop and plain arms within 0.5%
199    // of each other and 3.5% of the host render). Trust is therefore
200    // PER-STACK, decided by a parity probe on this file's own first
201    // qkv weight at DiT-scale activations — not by a hardcoded verdict
202    // either way. CMF_MMH3_GPU=1/0 still forces.
203    let use_gpu = match std::env::var("CMF_MMH3_GPU").ok().as_deref() {
204        Some("1") => true,
205        Some("0") => false,
206        _ => mmh3_gpu_parity_probe(path).unwrap_or(false),
207    };
208    if use_gpu {
209        generate_inner(path, prompt, p, &mut progress)
210    } else {
211        crate::gpu::cpu_scope(|| generate_inner(path, prompt, p, &mut progress))
212    }
213}
214
215/// GPU-vs-host parity on the packed DiT's first attention projection:
216/// the real q4tp bytes, activations spanning the modulation range
217/// (±2000 mixed with ±2), rms gate at 1e-2 — the measured failure was
218/// ~24%, honest drift is ~1e-5, so the gate has a decade of margin on
219/// each side. Any refusal (no adapter, dtype outside the kernel) is a
220/// clean "no": the host path is never wrong, only slower.
221fn mmh3_gpu_parity_probe(path: &Path) -> Result<bool, String> {
222    let model = Arc::new(
223        cortiq_core::CmfModel::open(path).map_err(|e| format!("{}: {e}", path.display()))?,
224    );
225    let Some(idx) = model.tensors.iter().position(|t| {
226        t.name.starts_with("dit.")
227            && t.name.ends_with("attn.qkv_proj.weight")
228            && t.dtype == cortiq_core::TensorDtype::Q4TiledP
229    }) else {
230        tracing::info!("mmh3 GPU parity probe: no q4tp qkv tensor — host path");
231        return Ok(false);
232    };
233    let entry = &model.tensors[idx];
234    let (rows, cols) = (entry.shape[0], entry.shape[1]);
235    let b = 64usize;
236    let mut xs = vec![0f32; b * cols];
237    for (i, v) in xs.iter_mut().enumerate() {
238        let base = ((i * 37 + 11) % 1009) as f32 / 1009.0 - 0.5;
239        *v = base * if i % 7 == 0 { 2000.0 } else { 2.0 };
240    }
241    let mut gpu = vec![0f32; b * rows];
242    if std::env::var("CMF_GPU_DEBUG").is_ok() {
243        eprintln!(
244            "mmh3 probe env: CMF_GPU={:?} enabled={} avail={}",
245            std::env::var("CMF_GPU").ok(),
246            crate::gpu::enabled(),
247            crate::gpu::backend_available(),
248        );
249    }
250    if !crate::gpu::q4tp_matmat(&model, idx, &xs, b, rows, cols, &mut gpu) {
251        tracing::info!("mmh3 GPU parity probe: q4tp_matmat refused ({rows}x{cols}) — host path");
252        if std::env::var("CMF_GPU_DEBUG").is_ok() {
253            eprintln!("mmh3 probe: q4tp_matmat refused {rows}x{cols}");
254        }
255        return Ok(false);
256    }
257    let host = {
258        let name = entry.name.clone();
259        let proj = crate::dit::Proj::from_model(&model, &name)?;
260        let mut out = vec![0f32; b * rows];
261        crate::gpu::cpu_scope(|| proj.matmat(&xs, b, &mut out, None));
262        out
263    };
264    let mut num = 0f64;
265    let mut den = 0f64;
266    for (g, h) in gpu.iter().zip(&host) {
267        num += ((g - h) as f64).powi(2);
268        den += (*h as f64).powi(2);
269    }
270    let rel = (num / den.max(1e-30)).sqrt();
271    let ok = rel < 1e-2;
272    tracing::info!(
273        "mmh3 GPU parity probe: rel rms {rel:.2e} → {}",
274        if ok { "device" } else { "host" }
275    );
276    if std::env::var("CMF_GPU_DEBUG").is_ok() {
277        eprintln!("mmh3 GPU parity probe: rel rms {rel:.2e}");
278    }
279    Ok(ok)
280}
281
282fn generate_inner(
283    path: &Path,
284    prompt: &str,
285    p: &AnimParams,
286    progress: &mut dyn FnMut(&str, usize, usize),
287) -> Result<Anim, String> {
288    let model = Arc::new(
289        cortiq_core::CmfModel::open(path).map_err(|e| format!("{}: {e}", path.display()))?,
290    );
291    let (frames_total, latent_t, audio_t) = temporal_shape(p.frames);
292    let (lat_h, lat_w) = (p.height / 16, p.width / 16);
293
294    // ── prompt ──
295    // The H3 presentation is raw text: no chat template, no BOS, no
296    // special tokens at all.
297    // Stage clock. Half of a render is not the DiT — at 8 steps the
298    // denoiser is 63 s of 116 — and until this line existed there was
299    // no way to see which half anything went to.
300    let t_stage = std::time::Instant::now();
301    let mut marks: Vec<(&str, f32)> = Vec::new();
302    let lap = |marks: &mut Vec<(&'static str, f32)>, name: &'static str| {
303        let prev: f32 = marks.iter().map(|(_, v)| v).sum();
304        marks.push((name, t_stage.elapsed().as_secs_f32() - prev));
305    };
306    let vocab = model
307        .vocab
308        .as_deref()
309        .ok_or("packaged .cmf has no embedded tokenizer")?;
310    let tok = Tokenizer::from_bytes(vocab).map_err(|e| format!("tokenizer: {e}"))?;
311    // fl2va: every keyframe is presented as "<Picture i>: " and a
312    // vision block BEFORE the prompt, and separately conditions the DiT
313    // as a latent. Both halves come from the same picture.
314    let keyframes: Vec<(&(Vec<f32>, usize, usize), usize)> = p
315        .first_frame
316        .iter()
317        .map(|f| (f, 0usize))
318        .chain(p.last_frame.iter().map(|f| (f, frames_total - 1)))
319        .collect();
320    let mut ids: Vec<u32> = Vec::new();
321    let mut spans: Vec<ImageSpan> = Vec::new();
322    let mut embeds: Vec<Vec<f32>> = Vec::new();
323    let mut deepstack: Vec<Vec<f32>> = Vec::new();
324    let mut cond: Vec<Vec<f32>> = Vec::new();
325    let mut tags: Vec<u8> = Vec::new();
326
327    if !keyframes.is_empty() {
328        let tower = VisionTower::from_cmf(&model)?;
329        // An activation harvest (CMF_TE_ONLY) never denoises, and the
330        // VAE latent is the DiT's food alone — the frame's 3-D conv
331        // encode is 99.5 s of a 102.5 s M4 run. Skip it.
332        let te_only = std::env::var("CMF_TE_ONLY").as_deref() == Ok("1");
333        let venc = if te_only {
334            None
335        } else {
336            Some(VideoVaeEncoder::from_cmf(&model)?)
337        };
338        for (i, (frame, _)) in keyframes.iter().enumerate() {
339            let (src, sh, sw) = *frame;
340            // The picture the DiT sees is on the generation canvas; the
341            // one Qwen sees keeps its own resolution policy.
342            let fitted = fit_to_canvas(src, *sh, *sw, p.height, p.width, i > 0);
343            if let Some(venc) = &venc {
344                let (z, _, _) = venc.encode_frame(
345                    &fitted.iter().map(|&v| v * 2.0 - 1.0).collect::<Vec<_>>(),
346                    p.height,
347                    p.width,
348                );
349                cond.push(z);
350            }
351
352            for t in tok.encode(&format!("<Picture {}>: ", i + 1)) {
353                ids.push(t);
354                tags.push(1);
355            }
356            let (patches, gh, gw) = qwen3vis::preprocess(
357                &fitted,
358                p.height,
359                p.width,
360                tower.patch_size,
361                tower.temporal_patch,
362                tower.merge,
363            );
364            let (merged, deep) = tower.forward(&patches, gh, gw);
365            let n_img = merged.len() / tower.out_hidden;
366            // The whole block carries the VIDEO tag, the flanking
367            // markers included.
368            ids.push(VISION_START);
369            tags.push(0);
370            let start = ids.len();
371            for _ in 0..n_img {
372                ids.push(VISION_START); // a placeholder the embed replaces
373                tags.push(0);
374            }
375            ids.push(VISION_END);
376            tags.push(0);
377            spans.push(ImageSpan {
378                start,
379                len: n_img,
380                merged_h: gh / tower.merge,
381                merged_w: gw / tower.merge,
382            });
383            embeds.push(merged);
384            if deepstack.is_empty() {
385                deepstack = deep;
386            } else {
387                for (a, b) in deepstack.iter_mut().zip(deep) {
388                    a.extend_from_slice(&b);
389                }
390            }
391        }
392    }
393    for t in tok.encode(prompt) {
394        ids.push(t);
395        tags.push(1);
396    }
397    if ids.is_empty() {
398        ids.push(151643); // the pad id, as the reference does for ""
399        tags.push(1);
400    }
401    ids.truncate(p.max_tokens);
402    tags.truncate(ids.len());
403    lap(&mut marks, "prepare");
404    // The prompt encoder is a one-shot pass over 12 GB of weights; on a
405    // machine the file does not fit it streams from disk and its GEMMs
406    // run over any contention budget for reasons that are not contention.
407    // The kill stays disarmed until the encode is done (users on 24 GB
408    // Macs had to patch it out to keep the denoise loop on the GPU).
409    crate::gpu::mm_kill_arm(false);
410    progress("encode", 0, 1);
411    let states = {
412        let enc = Qwen3Encoder::from_cmf(&model)?;
413        enc.encode_with_images(&ids, &spans, &embeds, &deepstack)
414    };
415    // `CMF_TE_DUMP=<path>`: the conditioning as `[u64 n][u64 width]`
416    // then f32 rows. A stand-in encoder is only as good as the stream
417    // it hands the DiT, and that is measurable against the teacher's
418    // dump on the same prompt WITHOUT rendering a frame.
419    if let Ok(p) = std::env::var("CMF_TE_DUMP") {
420        let w = states.len() / ids.len().max(1);
421        let mut b = Vec::with_capacity(16 + states.len() * 4);
422        b.extend_from_slice(&(ids.len() as u64).to_le_bytes());
423        b.extend_from_slice(&(w as u64).to_le_bytes());
424        for v in &states {
425            b.extend_from_slice(&v.to_le_bytes());
426        }
427        std::fs::write(&p, &b).map_err(|e| format!("CMF_TE_DUMP {p}: {e}"))?;
428        eprintln!("te dump: {} tokens x {w} -> {p}", ids.len());
429    }
430    progress("encode", 1, 1);
431    lap(&mut marks, "text encode");
432    crate::gpu::mm_kill_arm(true);
433    // `CMF_TE_ONLY=1`: stop after the dump — an activation-harvest run
434    // (the ClipProj refit) wants hundreds of encodes and zero renders.
435    if std::env::var("CMF_TE_ONLY").as_deref() == Ok("1") {
436        return Err("CMF_TE_ONLY: encode dumped, render skipped".into());
437    }
438    // The prompt encoder and vision tower ran their once-per-generation
439    // pass; release their page cache so the denoise loop's DiT does not
440    // fight 12+ GB of dead weights for RAM. On a 24 GB Mac with the
441    // 25.7 GB full-encoder fl2va file this is the difference between
442    // denoise steps at DiT speed and 320 s/step of SSD thrash.
443    {
444        let dropped = model.advise_done(|n| n.starts_with("model.") || n.starts_with("vis."));
445        if dropped > 0 {
446            tracing::info!(
447                "encoder pages released after prompt encode: {} MB",
448                dropped / (1024 * 1024)
449            );
450        }
451    }
452
453    // ── denoise ──
454    let (video, audio) = {
455        let dit = MiniMaxH3::from_cmf(&model)?;
456        let kf: Vec<(usize, usize)> = keyframes
457            .iter()
458            .map(|&(_, idx)| (idx, frames_total))
459            .collect();
460        let layout = if kf.is_empty() {
461            Layout::t2va(ids.len(), latent_t, lat_h, lat_w, audio_t)
462        } else {
463            Layout::fl2va(ids.len(), latent_t, lat_h, lat_w, audio_t, &kf, &tags)
464        };
465        let text = dit.refine_text(&states, ids.len());
466        let mut v = gauss(dit.latents_dim * latent_t * lat_h * lat_w, p.seed);
467        let mut a = gauss(dit.audio_dim * 2 * audio_t, p.seed ^ 0x9E37_79B9_7F4A_7C15);
468        let sg = sigmas(p.steps, dit.shift_video);
469        // `CMF_ANIM_PROF=1`: the per-step rms of both streams and of
470        // their velocities. A run that is not denoising shows it here
471        // long before anything is written out.
472        let prof = std::env::var_os("CMF_ANIM_PROF").is_some();
473        let rms = |x: &[f32]| {
474            (x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64).sqrt()
475        };
476        if prof {
477            eprintln!(
478                "  text {} tok, refined rms {:.4}, sigmas {:?}",
479                ids.len(),
480                rms(&text),
481                sg.iter()
482                    .map(|v| (v * 1e4).round() / 1e4)
483                    .collect::<Vec<_>>()
484            );
485        }
486        for i in 0..p.steps {
487            let (sv, sv_n) = (sg[i], sg[i + 1]);
488            let (dv, da) = dit.forward(&layout, &text, &v, &a, sv, &cond);
489            let step_v = (sv_n - sv) as f32;
490            for (x, &d) in v.iter_mut().zip(&dv) {
491                *x += step_v * d;
492            }
493            let step_a = if p.stock_sampler {
494                step_v
495            } else {
496                (time_shift_sigma(sv_n, dit.shift_video, dit.shift_audio)
497                    - time_shift_sigma(sv.max(1e-6), dit.shift_video, dit.shift_audio))
498                    as f32
499            };
500            for (x, &d) in a.iter_mut().zip(&da) {
501                *x += step_a * d;
502            }
503            if prof {
504                eprintln!(
505                    "  step {i}: sv {sv:.4}->{sv_n:.4} v_vel {:.4} a_vel {:.4} | video {:.4} audio {:.4}",
506                    rms(&dv),
507                    rms(&da),
508                    rms(&v),
509                    rms(&a)
510                );
511            }
512            progress("denoise", i + 1, p.steps);
513        }
514        (v, a)
515    };
516
517    lap(&mut marks, "denoise");
518    // ── decode ──
519    progress("video vae", 0, 1);
520    let (rgb, out_frames) = {
521        let vae = VideoVae::from_cmf(&model)?;
522        vae.decode(&video, latent_t, lat_h, lat_w)
523    };
524    progress("video vae", 1, 1);
525    lap(&mut marks, "video vae");
526    progress("audio vae", 0, 1);
527    let (wave, samples, sr) = {
528        let vae = AudioVae::from_cmf(&model)?;
529        let c = audio.len() / (2 * audio_t);
530        let (w, n) = vae.decode(&audio, c, audio_t);
531        (w, n, vae.sample_rate)
532    };
533    progress("audio vae", 1, 1);
534    lap(&mut marks, "audio vae");
535    tracing::info!(
536        "stages: {}",
537        marks
538            .iter()
539            .map(|(n, v)| format!("{n} {v:.1}s"))
540            .collect::<Vec<_>>()
541            .join(" · ")
542    );
543    // Where a GEMM's wall time goes on unified memory: copies or kernel.
544    // The answer decides whether fusing blocks or tuning the kernel is
545    // the optimization worth doing.
546    #[cfg(target_os = "macos")]
547    if std::env::var("CMF_METAL_MMPROF").is_ok() {
548        use std::sync::atomic::Ordering::Relaxed;
549        let n = crate::gpu_metal::MM_N.load(Relaxed);
550        eprintln!(
551            "  q4tp mm x{n}: upload {:.1}s · submit+wait {:.1}s · readback {:.1}s",
552            crate::gpu_metal::MM_UP.load(Relaxed) as f64 / 1e6,
553            crate::gpu_metal::MM_GPU.load(Relaxed) as f64 / 1e6,
554            crate::gpu_metal::MM_DN.load(Relaxed) as f64 / 1e6,
555        );
556    }
557
558    // The VAE emits latent_t·4 frames; the request snapped to 17k+5,
559    // which is one fewer than a multiple of four plus the leading key
560    // frame, so trim rather than pad.
561    let keep = out_frames.min(frames_total);
562    Ok(Anim {
563        rgb: trim_frames(&rgb, out_frames, keep, p.height, p.width),
564        frames: keep,
565        height: p.height,
566        width: p.width,
567        audio: wave,
568        samples,
569        sample_rate: sr,
570    })
571}
572
573fn trim_frames(rgb: &[f32], have: usize, keep: usize, h: usize, w: usize) -> Vec<f32> {
574    if keep == have {
575        return rgb.to_vec();
576    }
577    let mut out = vec![0f32; 3 * keep * h * w];
578    for c in 0..3 {
579        let s = c * have * h * w;
580        let d = c * keep * h * w;
581        out[d..d + keep * h * w].copy_from_slice(&rgb[s..s + keep * h * w]);
582    }
583    out
584}
585
586#[cfg(test)]
587mod tests {
588    use super::*;
589
590    #[test]
591    fn the_four_step_schedule_is_the_references() {
592        let s = sigmas(4, 12.0);
593        let want = [1.0, 0.972_973, 0.923_077, 0.8, 0.0];
594        assert_eq!(s.len(), want.len());
595        for (g, w) in s.iter().zip(&want) {
596            assert!((g - w).abs() < 1e-6, "{s:?}");
597        }
598    }
599
600    #[test]
601    fn a_stretch_keeps_the_corners_and_a_crop_takes_the_middle() {
602        // A 4x2 ramp: value rises left to right, so the corners name
603        // themselves.
604        let (h, w) = (2usize, 4usize);
605        let mut rgb = vec![0f32; 3 * h * w];
606        for c in 0..3 {
607            for y in 0..h {
608                for x in 0..w {
609                    rgb[(c * h + y) * w + x] = x as f32 / (w - 1) as f32;
610                }
611            }
612        }
613        // Stretch to a square: the far edges survive.
614        let s = fit_to_canvas(&rgb, h, w, 4, 4, false);
615        assert!((s[0] - 0.0).abs() < 1e-6, "left edge");
616        assert!((s[3] - 1.0).abs() < 1e-6, "right edge");
617        // Cover-crop to a square takes the centre 2x2, so the extremes
618        // are gone and the span is narrower.
619        let c = fit_to_canvas(&rgb, h, w, 4, 4, true);
620        let (lo, hi) = c[..16]
621            .iter()
622            .fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
623        assert!(lo > 0.05, "crop kept the left edge: {lo}");
624        assert!(hi < 0.95, "crop kept the right edge: {hi}");
625    }
626
627    #[test]
628    fn frame_counts_snap_to_the_models_grid() {
629        // The grid is 5 + 17k: 5, 22, 39, 56, … 124.
630        assert_eq!(align_frames(1), 5);
631        assert_eq!(align_frames(39), 39);
632        assert_eq!(align_frames(41), 56);
633        assert_eq!(align_frames(124), 124);
634        assert_eq!(video_latent_t(124), 37);
635        assert_eq!(video_latent_t(39), 12);
636        let (f, lt, at) = temporal_shape(124);
637        assert_eq!((f, lt, at), (124, 37, 207));
638    }
639}