metrale_bench/benchmarks/vision/
geometry.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The vision-token counts the geometry leg predicts, from an
4//! image's size, the patch and merge sizes, and the serve's declared area
5//! bound.
6//!
7//! Owner: bench, vision.
8//! Invariants: none beyond the types.
9
10/// 2026-09-26: Snap a side to the grid as the engine's `target_size_for` does
11/// at scale 1: `max(round(side / grid), 1) * grid`. `f32::round` rounds half
12/// away from zero, so 336 snaps to 352 at grid 32
13/// (`the_rounding_mode_is_pinned`).
14pub fn snap(side: u32, grid: u32) -> u32 {
15    (((side as f32) / (grid as f32)).round() as u32).max(1) * grid
16}
17
18/// 2026-09-26: Merged vision tokens for a `w × h` image at native size.
19///
20/// Both sides snap to `patch × merge`, then each `patch × patch` square is one
21/// patch and each `merge × merge` block of patches one token. At patch 16,
22/// merge 2 a 448×448 image is 784 patches and 196 tokens; 196 was measured
23/// against a live server on 2026-08-14.
24pub fn expected_vision_tokens(w: u32, h: u32, patch: u32, merge: u32) -> u32 {
25    let grid = patch * merge;
26    let sw = snap(w, grid);
27    let sh = snap(h, grid);
28    (sw / patch) * (sh / patch) / (merge * merge)
29}
30
31/// 2026-09-26: The engine's long-side clamp when no area bound is given; a copy
32/// of `FALLBACK_MAX_DIM` in `model-layers/src/vision_preprocess.rs`.
33pub const FALLBACK_MAX_DIM: u32 = 1280;
34
35/// 2026-09-26: The long-side ceiling applied with or without an area bound; a
36/// copy of `ABS_MAX_DIM` in `model-layers/src/vision_preprocess.rs`. It binds
37/// when an area bound admits a long side over 4096, such as a 64×8192 strip.
38pub const ABS_MAX_DIM: u32 = 4096;
39
40/// 2026-09-26: The `(w, h)` the server encodes for a `w × h` source.
41///
42/// Copies the scale and grid rounding of
43/// `metrale_model_layers::vision_preprocess::target_size_for`, which is
44/// crate-private and in a crate metrale-bench does not link. It omits the
45/// engine's last step, which shrinks a side one grid unit at a time while the
46/// rounded area is over `max_pixels`. `max_pixels` is an area, as with
47/// `--vision-max-pixels`; `None` applies the 1280 px fallback clamp.
48pub fn served_size(w: u32, h: u32, grid: u32, max_pixels: Option<u64>) -> (u32, u32) {
49    let long_side = w.max(h) as f32;
50    let area = (w as f32) * (h as f32);
51    let bound_scale = match max_pixels.filter(|&p| p > 0) {
52        // 2026-09-26: A declared area bound replaces the fallback clamp.
53        Some(p) => ((p as f32) / area).sqrt(),
54        None => (FALLBACK_MAX_DIM as f32) / long_side,
55    };
56    let abs_scale = (ABS_MAX_DIM as f32) / long_side;
57    let scale = bound_scale.min(abs_scale).min(1.0);
58    let tw = ((w as f32 * scale / grid as f32).round() as u32).max(1) * grid;
59    let th = ((h as f32 * scale / grid as f32).round() as u32).max(1) * grid;
60    (tw, th)
61}
62
63/// 2026-09-26: Merged vision tokens for `w × h` under a declared area bound.
64///
65/// A serve with a bound downscales a larger image and answers
66/// (`preprocess_image_with_max_pixels`), so the prediction takes the bound as
67/// input. `max_pixels == 0` means no bound was declared and returns
68/// [`expected_vision_tokens`], on the premise that the checkpoint's own bound
69/// (16777216 px in the Qwen3.6 preprocessor configs) is above every fixture.
70/// It is not `served_size(.., None)`, which applies the 1280 px fallback clamp.
71pub fn expected_vision_tokens_bounded(
72    w: u32,
73    h: u32,
74    patch: u32,
75    merge: u32,
76    max_pixels: u64,
77) -> u32 {
78    if max_pixels == 0 {
79        return expected_vision_tokens(w, h, patch, merge);
80    }
81    let (tw, th) = served_size(w, h, patch * merge, Some(max_pixels));
82    (tw / patch) * (th / patch) / (merge * merge)
83}
84
85#[cfg(test)]
86#[path = "geometry_tests.rs"]
87mod geometry_tests;
metrale_bench/benchmarks/vision/geometry.rs

metrale_bench/benchmarks/vision/
geometry.rs
