spark_model/vision_item.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The unit of vision input handed to the model.
4
5/// One vision item: a still image, or a video, ready for the encoder.
6///
7/// `groups` holds the TEMPORAL GROUPS. A still image has exactly one; a video
8/// has `frames / temporal_patch_size`, each group a full `grid_h x grid_w`
9/// patch plane built from `temporal_patch_size` consecutive frames. Every
10/// group is shaped identically — which is why the ViT consumes them on the
11/// same path and needs no notion of time.
12///
13/// Grouping lives in the TYPE rather than in a parallel `groups_per_item`
14/// vector carried alongside the pixels. Two vectors that must agree is
15/// precisely the failure this feature is exposed to: the pad run, the encoder
16/// rows and the MRoPE position stream all have to describe the same item, and
17/// a desync between them is silent — fluent output, wrong answer.
18#[derive(Debug, Clone, PartialEq)]
19pub struct VisionItem {
20 /// Per temporal group: `[grid_h * grid_w, C * temporal_patch_size * patch^2]`.
21 pub groups: Vec<Vec<f32>>,
22 /// Pre-merge patch grid, identical across this item's groups.
23 pub grid_h: usize,
24 pub grid_w: usize,
25}
26
27impl VisionItem {
28 /// A still image: one temporal group.
29 pub fn image(pixels: Vec<f32>, grid_h: usize, grid_w: usize) -> Self {
30 Self {
31 groups: vec![pixels],
32 grid_h,
33 grid_w,
34 }
35 }
36
37 /// Temporal extent, in groups. 1 for a still.
38 pub fn t_len(&self) -> usize {
39 self.groups.len().max(1)
40 }
41
42 /// Merged tokens this item occupies in the prompt — the length of its pad
43 /// run, and the number of embedding rows the encoder will return for it.
44 /// The two are the same number by construction, which is the invariant
45 /// the whole splice depends on.
46 pub fn pad_count(&self, spatial_merge_size: usize) -> usize {
47 let sms = spatial_merge_size.max(1);
48 self.t_len() * (self.grid_h / sms) * (self.grid_w / sms)
49 }
50}