Struct MimoVit

Source

pub struct MimoVit {
    pub cfg: MimoVisionConfig,
    /* private fields */
}

Expand description

The MiMo vision transformer and its patch merger.

Fields§

§cfg: MimoVisionConfig

Implementations§

Source §

impl MimoVit

Source

pub fn from_model(model: &Arc<CmfModel>) -> Result<Self, String>

Load from a CMF that carries visual.* and the mm.config_json blob.

Examples found in repository ?

examples/mimo_vis_calib.rs (line 28)

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}

Source

pub fn from_model_with_config( model: &Arc<CmfModel>, config: &Value, ) -> Result<Self, String>

Load visual.* from model with an explicit config (config.json or its vision_config).

Source

pub fn set_sink_mode(&mut self, mode: SinkMode)

Source

pub fn sink_mode(&self) -> SinkMode

Source

pub fn set_gpu_attention(&mut self, on: bool)

Allow (default) or forbid the device attention for the full blocks.

Source

pub fn forward(&self, input: &VisualInput) -> Result<Vec<f32>, String>

Encode one image or video: [tokens, out_hidden], in the order the item’s placeholders take them (raster (t, block row, block col)).

Examples found in repository ?

examples/mimo_vis_calib.rs (line 68)

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}