fix: preserve ColModernVBERT image patch geometry

fastembed/image/transform/operators.py

@@ -191,7 +191,15 @@ def __init__(
         self.resample = resample
 
     def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        patches, _ = self.split_with_grid(images)
+        return patches
+
+    def split_with_grid(
+        self, images: list[NumpyArray]
+    ) -> tuple[list[list[NumpyArray]], list[tuple[int, int]]]:
+        """Return patches and their (rows, cols), or (0, 0) for unsplit images."""
         result = []
+        image_grid = []
 
         for image in images:
             # Assume (C, H, W) format
@@ -204,6 +212,7 @@ def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type:
                 # Calculate the number of splits needed
                 num_splits_h = math.ceil(height / max_height)
                 num_splits_w = math.ceil(width / max_width)
+                image_grid.append((num_splits_h, num_splits_w))
 
                 # Calculate optimal patch dimensions
                 optimal_height = math.ceil(height / num_splits_h)
@@ -234,12 +243,13 @@ def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type:
                 frames.append(global_view)
             else:
                 # Image is small enough, no splitting needed
+                image_grid.append((0, 0))
                 frames.append(image)
 
             # Append (not extend) to preserve per-image grouping
             result.append(frames)
 
-        return result
+        return result, image_grid
 
 
 class SquareResize(Transform):
@@ -272,10 +282,23 @@ def __init__(self, transforms: list[Transform]):
         self.transforms = transforms
 
     def __call__(
-        self, images: list[Image.Image] | list[NumpyArray]
+        self,
+        images: list[Image.Image] | list[NumpyArray],
+        *,
+        metadata: dict[str, Any] | None = None,
     ) -> list[NumpyArray] | list[Image.Image]:
+        """Apply transforms without changing their usual return type.
+
+        If supplied, metadata["image_grid"] records (rows, cols) for each image,
+        or (0, 0) for images that are not split.
+        """
+        if metadata is not None:
+            metadata["image_grid"] = [(0, 0) for _ in images]
         for transform in self.transforms:
-            images = transform(images)
+            if metadata is not None and isinstance(transform, ImageSplitter):
+                images, metadata["image_grid"] = transform.split_with_grid(images)  # type: ignore[assignment,arg-type]
+            else:
+                images = transform(images)
         return images
 
     @classmethod

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -222,11 +222,15 @@ def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutpu
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
-            processed = self.processor(image_files)
+            processor_metadata: dict[str, Any] = {}
+            processed = self.processor(image_files, metadata=processor_metadata)
             encoded, attention_mask, metadata = self._process_nested_patches(processed)  # type: ignore[arg-type]
+            metadata.update(processor_metadata)
 
         onnx_input = {"pixel_values": encoded, "attention_mask": attention_mask}
-        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        onnx_input = self._preprocess_onnx_image_input(
+            onnx_input, image_grid=metadata["image_grid"], **kwargs
+        )
         model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
 
         return OnnxOutputContext(
@@ -276,17 +280,22 @@ def _process_nested_patches(
         return encoded, attention_mask, metadata  # type: ignore[return-value]
 
     def _preprocess_onnx_image_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+        self,
+        onnx_input: dict[str, np.ndarray],
+        *,
+        image_grid: list[tuple[int, int]] | None = None,
+        **kwargs: Any,
     ) -> dict[str, NumpyArray]:
         """
         Add text input placeholders for image data, following Idefics3 processing logic.
 
-        Constructs input_ids dynamically based on the actual number of image patches,
-        using the same token expansion logic as Idefics3Processor.
+        Constructs input_ids from the actual image patch grid, using the same
+        token expansion logic as Idefics3Processor.
 
         Args:
             onnx_input: Dict with 'pixel_values' (batch, num_patches, C, H, W)
                         and 'attention_mask' (batch, num_patches) indicating real patches
+            image_grid: Actual (rows, cols) from image splitting, (0, 0) if unsplit
             **kwargs: Additional arguments
 
         Returns:
@@ -298,15 +307,18 @@ def _preprocess_onnx_image_input(
         pixel_values = onnx_input["pixel_values"]
 
         batch_size = pixel_values.shape[0]
+        if image_grid is None or len(image_grid) != batch_size:
+            raise ValueError("The image patch grid is required for each image in the batch")
         batch_input_ids = []
 
-        # Build input_ids for each image based on its actual patch count
-        for i in range(batch_size):
+        # A patch count cannot distinguish portrait, landscape, and square grids.
+        for i, (rows, cols) in enumerate(image_grid):
             # Count real patches (non-padded) from attention mask
             patch_count = int(np.sum(patch_attention_mask[i]))
 
-            # Compute rows/cols from patch count
-            rows, cols = self._compute_rows_cols_from_patches(patch_count)
+            valid_grid = (rows == 0 and cols == 0) or (rows > 0 and cols > 0)
+            if not valid_grid or rows * cols + 1 != patch_count:
+                raise ValueError("The image patch grid does not match the number of patches")
 
             # Build input_ids for this image
             input_ids = self._build_input_ids_for_image(rows, cols)
@@ -340,31 +352,6 @@ def _preprocess_onnx_image_input(
         onnx_input["attention_mask"] = attention_mask
         return onnx_input
 
-    @staticmethod
-    def _compute_rows_cols_from_patches(patch_count: int) -> tuple[int, int]:
-        if patch_count <= 1:
-            return 0, 0
-
-        # Subtract 1 for the global image
-        grid_patches = patch_count - 1
-
-        # Find rows and cols (assume square or near-square grid)
-        rows = int(grid_patches**0.5)
-        cols = grid_patches // rows
-
-        # Verify the calculation
-        if rows * cols + 1 != patch_count:
-            # Handle non-square grids
-            for r in range(1, grid_patches + 1):
-                if grid_patches % r == 0:
-                    c = grid_patches // r
-                    if r * c + 1 == patch_count:
-                        return r, c
-            # Fallback: treat as unsplit
-            return 0, 0
-
-        return rows, cols
-
     def _create_single_image_prompt_string(self) -> str:
         return (
             "<fake_token_around_image>"

tests/test_colmodernvbert_preprocessing.py

@@ -0,0 +1,159 @@
+"""Offline preprocessing regressions: real transforms/tokenizer, no model weights."""
+
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+from PIL import Image
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.trainers import BpeTrainer
+
+from fastembed.image.transform.operators import Compose
+from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT
+
+
+def expected_prompt(rows: int, cols: int, image_seq_len: int = 2) -> str:
+    """Idefics3's row-major local crops followed by the global image."""
+    image_tokens = "<image>" * image_seq_len
+    fake = "<fake_token_around_image>"
+    local_rows = [
+        "".join(f"{fake}<row_{row}_col_{col}>{image_tokens}" for col in range(1, cols + 1))
+        for row in range(1, rows + 1)
+    ]
+    local = "\n".join(local_rows) + "\n\n" if local_rows else ""
+    image = local + f"{fake}<global-img>{image_tokens}{fake}"
+    return f"<|begin_of_text|>User:{image}Describe the image.<end_of_utterance>\nAssistant:"
+
+
+@pytest.fixture
+def model() -> ColModernVBERT:
+    # Bypass downloading weights and constructing an ONNX session. Every preprocessing
+    # component below is real; only the final inference call is a recording test double.
+    model = object.__new__(ColModernVBERT)
+    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    special_tokens = [
+        "[UNK]",
+        "[PAD]",
+        "<image>",
+        "<fake_token_around_image>",
+        "<global-img>",
+        "<|begin_of_text|>",
+        "<end_of_utterance>",
+    ] + [f"<row_{row}_col_{col}>" for row in range(1, 5) for col in range(1, 5)]
+    tokenizer.train_from_iterator(
+        [expected_prompt(4, 4)], BpeTrainer(vocab_size=256, special_tokens=special_tokens)
+    )
+    tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
+    model.tokenizer = tokenizer
+    model.image_seq_len = 2
+    model.model = Mock()
+    model.model.run.return_value = [np.empty((0,), dtype=np.float32)]
+    return model
+
+
+def configure(model: ColModernVBERT, *, splitting: bool = True, resize: bool = False) -> None:
+    model.processor = Compose.from_config(
+        {
+            "image_processor_type": "Idefics3ImageProcessor",
+            "do_resize": resize,
+            "size": {"longest_edge": 16},
+            "do_image_splitting": splitting,
+            "max_image_size": {"longest_edge": 4},
+            "do_rescale": True,
+            "do_normalize": True,
+            "image_mean": [0.5, 0.5, 0.5],
+            "image_std": [0.5, 0.5, 0.5],
+        }
+    )
+
+
+def assert_prompt(
+    model: ColModernVBERT, onnx_input: dict, index: int, grid: tuple[int, int]
+) -> None:
+    expected = model.tokenizer.encode(expected_prompt(*grid)).ids
+    active = onnx_input["attention_mask"][index].astype(bool)
+    actual = onnx_input["input_ids"][index][active]
+    # Assert actual special-token content as well as exact IDs and row separators.
+    actual_tokens = [model.tokenizer.id_to_token(int(token)) for token in actual]
+    expected_tokens = [model.tokenizer.id_to_token(token) for token in expected]
+    assert actual_tokens == expected_tokens
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("size", "grid"),
+    [
+        ((4, 8), (2, 1)),
+        ((8, 4), (1, 2)),
+        ((16, 4), (1, 4)),
+        ((4, 16), (4, 1)),
+        ((8, 8), (2, 2)),
+        ((4, 4), (0, 0)),
+        ((5, 9), (3, 2)),
+    ],
+    ids=["portrait", "landscape", "wide", "tall", "square", "unsplit", "rounded"],
+)
+def test_image_prompt_retains_patch_grid(model, size, grid) -> None:
+    configure(model)
+    output = model.onnx_embed_image([Image.new("RGB", size, (32, 64, 128))])
+    onnx_input = model.model.run.call_args.args[1]
+    assert_prompt(model, onnx_input, 0, grid)
+    assert set(onnx_input) == {"pixel_values", "input_ids", "attention_mask"}
+    assert output.metadata["image_grid"] == [grid]
+    count = grid[0] * grid[1] + 1
+    assert output.metadata["patch_counts"] == [count]
+    assert onnx_input["pixel_values"].shape == (1, count, 3, 4, 4)
+
+
+@pytest.mark.parametrize("direction", ["left", "right"])
+def test_mixed_image_batch_preserves_grids_and_padding(model, direction) -> None:
+    configure(model)
+    model.tokenizer.enable_padding(direction=direction, pad_id=1, pad_token="[PAD]")
+    sizes = [(4, 8), (8, 4), (16, 4), (8, 8), (4, 4)]
+    grids = [(2, 1), (1, 2), (1, 4), (2, 2), (0, 0)]
+    output = model.onnx_embed_image([Image.new("RGB", size, (32, 64, 128)) for size in sizes])
+    onnx_input = model.model.run.call_args.args[1]
+    assert output.metadata["patch_counts"] == [3, 3, 5, 5, 1]
+    assert onnx_input["pixel_values"].shape == (5, 5, 3, 4, 4)
+    for i, grid in enumerate(grids):
+        assert_prompt(model, onnx_input, i, grid)
+        count = grid[0] * grid[1] + 1
+        assert np.count_nonzero(onnx_input["pixel_values"][i, count:]) == 0
+        mask = onnx_input["attention_mask"][i]
+        assert np.all(onnx_input["input_ids"][i][mask == 0] == 1)
+        padding_count = int(np.sum(mask == 0))
+        if padding_count:
+            padded = mask[:padding_count] if direction == "left" else mask[-padding_count:]
+            assert not np.any(padded)
+
+
+@pytest.mark.parametrize("size", [(4, 8), (8, 4)])
+def test_grid_is_retained_after_longest_edge_resize(model, size) -> None:
+    configure(model, resize=True)
+    model.onnx_embed_image([Image.new("RGB", size)])
+    grid = (4, 2) if size[1] > size[0] else (2, 4)
+    assert_prompt(model, model.model.run.call_args.args[1], 0, grid)
+
+
+def test_disabled_splitting_uses_only_global_prompt(model) -> None:
+    configure(model, splitting=False)
+    model.onnx_embed_image([Image.new("RGB", (4, 16))])
+    onnx_input = model.model.run.call_args.args[1]
+    assert_prompt(model, onnx_input, 0, (0, 0))
+    assert onnx_input["pixel_values"].shape == (1, 1, 3, 4, 4)
+
+
+@pytest.mark.parametrize(
+    ("grid", "count"),
+    [(None, 3), ([], 3), ([(2, 2)], 3), ([(-1, -2)], 3), ([(0, 2)], 1), ([(2, 0)], 1)],
+)
+def test_missing_or_inconsistent_grid_is_rejected(model, grid, count) -> None:
+    inputs = {
+        "pixel_values": np.zeros((1, count, 3, 4, 4), dtype=np.float32),
+        "attention_mask": np.ones((1, count), dtype=np.int64),
+    }
+    with pytest.raises(ValueError, match="image patch grid"):
+        model._preprocess_onnx_image_input(inputs, image_grid=grid)

tests/test_image_transform.py

@@ -109,3 +109,71 @@ def test_normalize_channel_count_mismatch_raises(
     image = np.zeros((3, 4, 4), dtype=np.float32)
     with pytest.raises(ValueError, match=expected):
         normalize(image, mean=mean, std=std)
+
+
+@pytest.mark.parametrize("splitting", [False, True])
+@pytest.mark.parametrize("rescale_and_normalize", [False, True])
+def test_compose_grid_metadata_keeps_default_outputs(splitting, rescale_and_normalize) -> None:
+    processor = Compose.from_config(
+        {
+            "image_processor_type": "Idefics3ImageProcessor",
+            "do_resize": False,
+            "do_image_splitting": splitting,
+            "max_image_size": {"longest_edge": 4},
+            "do_rescale": rescale_and_normalize,
+            "do_normalize": rescale_and_normalize,
+            "image_mean": [0.5, 0.5, 0.5],
+            "image_std": [0.5, 0.5, 0.5],
+        }
+    )
+    images = [Image.new("RGB", size, (32, 64, 128)) for size in [(4, 8), (8, 4), (4, 4)]]
+    ordinary = processor(images)
+    metadata = {"existing": "preserved"}
+    with_metadata = processor(images, metadata=metadata)
+
+    assert isinstance(ordinary, list) and isinstance(with_metadata, list)
+    assert metadata["image_grid"] == ([(2, 1), (1, 2), (0, 0)] if splitting else [(0, 0)] * 3)
+    assert metadata["existing"] == "preserved"
+    for expected, actual in zip(ordinary, with_metadata):
+        assert isinstance(expected, list) and isinstance(actual, list)
+        assert len(expected) == len(actual)
+        for expected_patch, actual_patch in zip(expected, actual):
+            assert isinstance(expected_patch, np.ndarray) and isinstance(actual_patch, np.ndarray)
+            np.testing.assert_array_equal(actual_patch, expected_patch)
+
+    # Geometry belongs to each call, including repeated use of a metadata dictionary.
+    processor([images[-1]], metadata=metadata)
+    assert metadata["image_grid"] == [(0, 0)]
+
+
+def test_splitter_grid_keeps_row_major_pixels_and_global_last() -> None:
+    from fastembed.image.transform.operators import ImageSplitter
+
+    image = np.arange(3 * 8 * 12, dtype=np.float32).reshape(3, 8, 12)
+    splitter = ImageSplitter(max_size=4)
+    ordinary = splitter([image])
+    patches, grids = splitter.split_with_grid([image])
+
+    assert grids == [(2, 3)]
+    assert isinstance(ordinary, list) and isinstance(ordinary[0], list)
+    assert len(patches[0]) == len(ordinary[0]) == 7
+    for row in range(2):
+        for col in range(3):
+            index = row * 3 + col
+            np.testing.assert_array_equal(
+                patches[0][index], image[:, row * 4 : (row + 1) * 4, col * 4 : (col + 1) * 4]
+            )
+    for actual, expected in zip(patches[0], ordinary[0]):
+        np.testing.assert_array_equal(actual, expected)
+    assert patches[0][-1].shape == (3, 4, 4)
+
+
+def test_compose_metadata_does_not_change_flat_image_batches() -> None:
+    processor = Compose.from_config({"do_resize": False, "do_rescale": False})
+    images = [Image.new("RGB", (4, 4))]
+    metadata = {}
+    actual = processor(images, metadata=metadata)
+
+    assert isinstance(actual, list) and isinstance(actual[0], np.ndarray)
+    assert metadata == {"image_grid": [(0, 0)]}
+    np.testing.assert_array_equal(actual, processor(images))