diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index e3b977db..bfe353bd 100644 --- a/fastembed/image/transform/functional.py +++ b/fastembed/image/transform/functional.py @@ -186,7 +186,8 @@ def resize_longest_edge( if new_width % 2 != 0: new_width += 1 - return image.resize((new_width, new_height), resample) + # A very narrow image can round down to zero along its shorter edge. + return image.resize((max(1, new_width), max(1, new_height)), resample) def crop_ndarray( diff --git a/fastembed/image/transform/operators.py b/fastembed/image/transform/operators.py index 5c747b85..669cfcdc 100644 --- a/fastembed/image/transform/operators.py +++ b/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 diff --git a/fastembed/late_interaction_multimodal/colmodernvbert.py b/fastembed/late_interaction_multimodal/colmodernvbert.py index f8fe0154..f64b14fa 100644 --- a/fastembed/late_interaction_multimodal/colmodernvbert.py +++ b/fastembed/late_interaction_multimodal/colmodernvbert.py @@ -3,6 +3,7 @@ import json import numpy as np +import onnxruntime as ort from tokenizers import Encoding from PIL import Image @@ -27,9 +28,10 @@ description="The late-interaction version of ModernVBERT, CPU friendly, English, 2025.", license="mit", size_in_GB=1.0, + # TEMPORARY: test copy of the re-exported model, switch back to Qdrant/colmodernvbert before merging sources=ModelSource(hf="Qdrant/colmodernvbert"), additional_files=["processor_config.json"], - model_file="model.onnx", + model_file="model_v2.onnx", ), ] @@ -125,6 +127,13 @@ def _list_supported_models(cls) -> list[DenseModelDescription]: return supported_colmodernvbert_models def load_onnx_model(self) -> None: + # onnxruntime 1.18 crashes the whole process while optimizing this graph (an access violation + # on Windows, a failed bounds check on Linux), and 1.17 can't read its IR version + if tuple(int(part) for part in ort.__version__.split(".")[:2]) < (1, 19): + raise RuntimeError( + f"Could not load {self.model_name}: it requires onnxruntime>=1.19, but " + f"onnxruntime {ort.__version__} is installed. Please upgrade onnxruntime." + ) self._load_onnx_model( model_dir=self._model_dir, model_file=self.model_description.model_file, @@ -168,12 +177,13 @@ def _preprocess_onnx_text_input( Returns: Iterable[NumpyArray]: Post-processed output as NumPy arrays. """ - batch_size, seq_length = onnx_input["input_ids"].shape - empty_image_placeholder: NumpyArray = np.zeros( - (batch_size, seq_length, 3, self.image_size, self.image_size), + # one blank tile for the whole batch: the model drops blank tiles and skips the vision + # encoder; zero tiles would also skip it, but CUDA can't reduce the resulting empty tensor + blank_image_placeholder: NumpyArray = np.zeros( + (1, 1, 3, self.image_size, self.image_size), dtype=np.float32, # type: ignore[type-var,arg-type,assignment] ) - onnx_input["pixel_values"] = empty_image_placeholder + onnx_input["pixel_values"] = blank_image_placeholder return onnx_input def _post_process_onnx_text_output( @@ -189,7 +199,10 @@ def _post_process_onnx_text_output( Returns: Iterable[NumpyArray]: Post-processed output as NumPy arrays. """ - return output.model_output + assert output.attention_mask is not None + # drop the padding rows, so a query gets the same vectors whatever else is in its batch + for embedding, attention_mask in zip(output.model_output, output.attention_mask): + yield embedding[attention_mask == 1] def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]: # Add query augmentation tokens (matching process_queries logic from colpali-engine) @@ -222,16 +235,21 @@ 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( model_output=model_output[0], - attention_mask=attention_mask, # type: ignore[arg-type] + # the token mask, not the tile mask: post-processing drops the padding rows with it + attention_mask=onnx_input["attention_mask"], # type: ignore[arg-type] metadata=metadata, ) @@ -276,17 +294,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 +321,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 +366,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 ( "" @@ -427,9 +428,13 @@ def _post_process_onnx_image_output( Iterable[NumpyArray]: Post-processed output as NumPy arrays. """ assert self.model_description.dim is not None, "Model dim is not defined" - return output.model_output.reshape( + assert output.attention_mask is not None + embeddings: NumpyArray = output.model_output.reshape( output.model_output.shape[0], -1, self.model_description.dim ) + # drop the padding rows, so an image gets the same vectors whatever else is in its batch + for embedding, attention_mask in zip(embeddings, output.attention_mask): + yield embedding[attention_mask == 1] def embed_text( self, @@ -470,7 +475,7 @@ def embed_text( def embed_image( self, images: ImageInput | Iterable[ImageInput], - batch_size: int = 16, + batch_size: int = 2, parallel: Optional[int] = None, **kwargs: Any, ) -> Iterable[NumpyArray]: @@ -479,7 +484,8 @@ def embed_image( Args: images: Iterator of image paths or single image path to embed - batch_size: Batch size for encoding -- higher values will use more memory, but be faster + batch_size: Batch size for encoding -- every image is split into up to 17 tiles, so + memory grows quickly with it, while on CPU larger batches are not faster parallel: If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. If 0, use all available cores. diff --git a/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py b/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py index 4318409b..689b646f 100644 --- a/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py +++ b/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py @@ -1,3 +1,4 @@ +import warnings from typing import Any, Iterable, Sequence, Type from dataclasses import asdict @@ -67,6 +68,14 @@ def __init__( **kwargs: Any, ): super().__init__(model_name, cache_dir, threads, **kwargs) + # keep for a couple of releases after the 2026-10 re-export of the model + if model_name.lower() == "qdrant/colmodernvbert": + warnings.warn( + "The model 'Qdrant/colmodernvbert' has been updated: re-embed images that were " + "embedded in batches with images of other aspect ratios.", + UserWarning, + stacklevel=2, + ) for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY: supported_models = EMBEDDING_MODEL_TYPE._list_supported_models() if any(model_name.lower() == model.model.lower() for model in supported_models): @@ -147,7 +156,7 @@ def embed_text( def embed_image( self, images: ImageInput | Iterable[ImageInput], - batch_size: int = 16, + batch_size: int = 2, parallel: int | None = None, **kwargs: Any, ) -> Iterable[NumpyArray]: diff --git a/tests/test_late_interaction_multimodal.py b/tests/test_late_interaction_multimodal.py index 5237475a..78e7f874 100644 --- a/tests/test_late_interaction_multimodal.py +++ b/tests/test_late_interaction_multimodal.py @@ -130,11 +130,27 @@ def test_single_embedding(model_cache): assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3) -def test_single_embedding_query(model_cache): - # text inputs get a zero image placeholder per token, so this is as slow as the image tests +def test_batch_embedding_mixed_tile_counts(model_cache): + # image splitting makes colmodernvbert slow on CI runners, so only linux runs it if platform.system() != "Linux" and os.getenv("CI"): pytest.skip("too slow on macOS and Windows CI runners") + model_name = "Qdrant/colmodernvbert" + image = Image.open(TEST_MISC_DIR / "image.jpeg") + # the 4:1 strip is split into 5 tiles and the full image into 13, so the strip gets padded + # with empty tiles, which must not leak into the image after it + strip = image.crop((0, 0, image.width, image.width // 4)) + mixed_images = [strip, image] + with model_cache(model_name) as model: + batch_result = list(model.embed_image(mixed_images, batch_size=len(mixed_images))) + single_result = list(model.embed_image(mixed_images, batch_size=1)) + + for batch_value, single_value in zip(batch_result, single_result): + assert batch_value.shape == single_value.shape + assert np.allclose(batch_value, single_value, atol=1e-3) + + +def test_single_embedding_query(model_cache): for model_name, expected_result in CANONICAL_QUERY_VALUES.items(): if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"): continue # colpali is too large for ci