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}