diff --git a/crates/segment/src/mobilesam.rs b/crates/segment/src/mobilesam.rs index 6599b018a..abb12fb0f 100644 --- a/crates/segment/src/mobilesam.rs +++ b/crates/segment/src/mobilesam.rs @@ -19,6 +19,9 @@ const IMG: usize = 1024; const EMB: usize = 64; const DIM: usize = 256; const LOW: usize = 256; +/// Most prompt points the decoder takes (each is an attention token, and self-attention is +/// quadratic in them); SAM is used with a handful. More are thinned evenly, keeping both ends. +const MAX_POINTS: usize = 256; const MEAN: [f32; 3] = [123.675, 116.28, 103.53]; const STD: [f32; 3] = [58.395, 57.12, 57.375]; @@ -609,7 +612,7 @@ impl MobileSam { /// Foreground probability (`w×h`) for a prompt on an embedded image, and the predicted IoU. pub fn predict(&self, e: &Embedding, prompt: &Prompt) -> (Vec, f32) { let s = e.scale; - let points: Vec<([f32; 2], bool)> = prompt.points.iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect(); + let points: Vec<([f32; 2], bool)> = thin(&prompt.points).iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect(); let bbox = prompt.bbox.map(|b| b.map(|v| v * s)); // A prior mask (probabilities at image size) → logits on the low-res grid. let mask: Option> = prompt.mask.as_ref().filter(|m| m.len() == e.w * e.h).map(|m| { @@ -675,3 +678,28 @@ fn sample(rgb: &[[f32; 3]], w: usize, h: usize, x: f32, y: f32, foot: f32) -> [f } s.map(|v| v / f32::max(n, 1.0)) } + +/// `points`, or `MAX_POINTS` evenly spaced ones (first and last kept) when there are more. +fn thin(points: &[([f32; 2], bool)]) -> Vec<([f32; 2], bool)> { + let n = points.len(); + if n <= MAX_POINTS { + return points.to_vec(); + } + (0..MAX_POINTS).filter_map(|i| points.get(i * (n - 1) / (MAX_POINTS - 1)).copied()).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn many_prompt_points_are_thinned_to_the_budget() { + let few: Vec<_> = (0..MAX_POINTS).map(|i| ([i as f32, 0.0], i % 2 == 0)).collect(); + assert_eq!(thin(&few), few, "a prompt within the budget is untouched"); + let many: Vec<_> = (0..1_000_000).map(|i| ([i as f32, 0.0], i % 2 == 0)).collect(); + let t = thin(&many); + assert_eq!(t.len(), MAX_POINTS); + assert_eq!((t[0].0, t[MAX_POINTS - 1].0), (many[0].0, many[many.len() - 1].0)); + assert!(t.windows(2).all(|w| w[0].0[0] < w[1].0[0]), "spread in order"); + } +} diff --git a/crates/segment/src/tflite.rs b/crates/segment/src/tflite.rs index 34e85dfab..6aafd056f 100644 --- a/crates/segment/src/tflite.rs +++ b/crates/segment/src/tflite.rs @@ -131,6 +131,8 @@ const MAX_VALUES: u64 = 1 << 30; const MAX_LIVE: u64 = 1 << 29; /// The most constant values decoded and laid out at load (a weight shared by several tensors or /// operators counts once per use). +/// Bytes of tensor names and shapes a model may carry in all (the official ones take well under 1 MB). +const MAX_META: usize = 1 << 24; const MAX_CONSTS: usize = 1 << 27; /// Tensor types (`TensorType`). @@ -261,6 +263,7 @@ impl Model { let mut ints: Vec>> = vec![None; n]; let mut data: Vec<(usize, u8, usize, &[u8])> = vec![]; let mut consts = 0usize; + let mut meta = 0usize; for (i, t) in tensors.iter().enumerate() { let shape: Vec = fb .ints(*t, 0)? @@ -269,7 +272,13 @@ impl Model { .collect::>()?; let ty = fb.byte(*t, 1, FLOAT32)?; let buf = fb.int(*t, 2, 0)? as u32 as usize; - names.push(fb.string(*t, 3)?); + let name = fb.string(*t, 3)?; + // Names and shapes are copied per tensor: bound them in all. + meta = meta.saturating_add(name.len()).saturating_add(shape.len().saturating_mul(8)); + if meta > MAX_META { + return Err("the model's tensor metadata is too large".into()); + } + names.push(name); // Every size computed from shapes later stays in range, even a product of some of the // dimensions (a zero dimension does not let the others grow without bound). let count = shape.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).ok_or("a tensor is too large")?; @@ -329,6 +338,11 @@ impl Model { && let Some(i) = ins.first().copied().flatten() && let Some(len) = floats[i].as_ref().map(Vec::len) { + // The folded constant is what the graph outputs and the run's memory estimate count: + // it must be the declared size. + if shapes[output].iter().try_fold(1usize, |a, &d| a.checked_mul(d)) != Some(len) { + return Err(format!("{}: DEQUANTIZE output size disagrees with its input", names[output])); + } spend(len)?; floats[output] = floats[i].clone(); continue; diff --git a/crates/segment/src/tflite/tests.rs b/crates/segment/src/tflite/tests.rs index 29a781d34..9ba96971f 100644 --- a/crates/segment/src/tflite/tests.rs +++ b/crates/segment/src/tflite/tests.rs @@ -334,6 +334,26 @@ fn shapes_computed_at_run_time_must_match_the_declared_ones() { refused(&zero, "tensor is too large"); } +#[test] +fn a_folded_dequantize_has_its_declared_size() { + let half = |n: usize| vec![0u8; 2 * n]; + let fold = |declared: Vec| { + model(vec![half(4)], vec![(vec![1, 4], FLOAT16, Some(0)), (declared, FLOAT32, None)], vec![(op::DEQUANTIZE, vec![0], vec![1], vec![])], vec![], vec![1]) + }; + // A declared output of one value would let the folded four-value constant be copied out per reference. + refused(&fold(vec![1, 1]), "DEQUANTIZE output size"); + assert!(Model::read(&fold(vec![1, 4])).is_ok()); +} + +#[test] +fn tensor_metadata_is_bounded_in_all() { + // A rank of 4 million dimensions of 1 is a valid shape by size, but copied per tensor. + let m = model(vec![], vec![(vec![1; 1 << 22], FLOAT32, None)], vec![], vec![], vec![]); + refused(&m, "metadata is too large"); + let ok = model(vec![], (0..70).map(|_| (vec![1, 2, 3], FLOAT32, None)).collect(), vec![], vec![], vec![]); + assert!(Model::read(&ok).is_ok()); +} + #[test] fn shared_weights_count_once_per_use() { // One 128 KB buffer decoded for 2100 tensors would be 2100 copies of 2¹⁶ floats.