mimo_vis_calib/
mimo_vis_calib.rs

1//! Dev-only: capture the MiMo ViT's input Hessians over calibration images
2//! and write them as a `cortiq quantize-gptq --hessians` cache, so the GPTQ
3//! q4tp driver can fold the tower without its text calibration pass.
4//!
5//!     mimo_vis_calib TOWER.cmf OUT.hess MAX_PIXELS IMG_OR_DIR...
6//!
7//! TOWER must hold dense (F32/F16/BF16) `visual.*` matrices: the tower's own
8//! hook folds the inputs of dense projections only. The cache layout is the
9//! one `cortiq-cli/src/gptq.rs::save_hessians` writes (magic `CMFHESS1`,
10//! then per group: names, cols, count, H length, Σx², upper triangle of H),
11//! here with one name per group.
12
13use cortiq_core::CmfModel;
14use cortiq_engine::gptq_capture;
15use cortiq_engine::media;
16use cortiq_engine::mimo_vision::{MimoProcessorConfig, MimoVit, prepare_image};
17use std::io::Write;
18use std::path::PathBuf;
19use std::sync::Arc;
20
21fn main() {
22    let args: Vec<String> = std::env::args().collect();
23    if args.len() < 5 {
24        eprintln!("usage: mimo_vis_calib TOWER.cmf OUT.hess MAX_PIXELS IMG_OR_DIR...");
25        std::process::exit(2);
26    }
27    let model = Arc::new(CmfModel::open(&args[1]).expect("open tower"));
28    let vit = MimoVit::from_model(&model).expect("load tower");
29    let max_px: usize = args[3].parse().expect("MAX_PIXELS");
30    let mut paths: Vec<PathBuf> = Vec::new();
31    for a in &args[4..] {
32        let p = PathBuf::from(a);
33        if p.is_dir() {
34            let mut v: Vec<PathBuf> = std::fs::read_dir(&p)
35                .unwrap()
36                .filter_map(|e| e.ok().map(|e| e.path()))
37                .filter(|p| {
38                    p.extension()
39                        .and_then(|e| e.to_str())
40                        .is_some_and(|e| matches!(e, "png" | "jpg" | "jpeg" | "webp"))
41                })
42                .collect();
43            v.sort();
44            paths.extend(v);
45        } else {
46            paths.push(p);
47        }
48    }
49    let cfg = MimoProcessorConfig::default();
50    let t0 = std::time::Instant::now();
51    gptq_capture::begin(true);
52    let mut rows = 0usize;
53    for p in &paths {
54        let frame = match media::read_rgb(p) {
55            Ok(f) => f,
56            Err(e) => {
57                eprintln!("skip {}: {e}", p.display());
58                continue;
59            }
60        };
61        let input = match prepare_image(&frame, &cfg, Some(max_px)) {
62            Ok(i) => i,
63            Err(e) => {
64                eprintln!("skip {}: {e}", p.display());
65                continue;
66            }
67        };
68        vit.forward(&input).expect("forward");
69        rows += input.patches();
70        eprintln!(
71            "  {} {}x{} → {} patches ({:.0} s)",
72            p.display(),
73            frame.width,
74            frame.height,
75            input.patches(),
76            t0.elapsed().as_secs_f64()
77        );
78    }
79    let hess = gptq_capture::end();
80    let mut names: Vec<&String> = hess.keys().collect();
81    names.sort();
82    let tmp = format!("{}.tmp", args[2]);
83    let mut f = std::io::BufWriter::with_capacity(1 << 22, std::fs::File::create(&tmp).unwrap());
84    f.write_all(b"CMFHESS1").unwrap();
85    f.write_all(&(names.len() as u64).to_le_bytes()).unwrap();
86    for n in &names {
87        let a = &hess[*n];
88        f.write_all(&1u32.to_le_bytes()).unwrap();
89        f.write_all(&(n.len() as u32).to_le_bytes()).unwrap();
90        f.write_all(n.as_bytes()).unwrap();
91        f.write_all(&(a.cols as u64).to_le_bytes()).unwrap();
92        f.write_all(&(a.count as u64).to_le_bytes()).unwrap();
93        f.write_all(&(a.h.len() as u64).to_le_bytes()).unwrap();
94        for v in &a.sumsq {
95            f.write_all(&v.to_le_bytes()).unwrap();
96        }
97        let c = a.cols;
98        if a.h.len() == c * c {
99            for i in 0..c {
100                for v in &a.h[i * c + i..i * c + c] {
101                    f.write_all(&v.to_le_bytes()).unwrap();
102                }
103            }
104        }
105    }
106    f.flush().unwrap();
107    drop(f);
108    std::fs::rename(&tmp, &args[2]).unwrap();
109    println!(
110        "wrote {}: {} linears, {rows} patches from {} images, {:.0} s",
111        args[2],
112        names.len(),
113        paths.len(),
114        t0.elapsed().as_secs_f64()
115    );
116}
mimo_vis_calib/mimo_vis_calib.rs

mimo_vis_calib/
mimo_vis_calib.rs
