fastembed-gpu 0.5.1__py3-none-any.whl

fastembed/image/image_embedding.py

@@ -0,0 +1,97 @@
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
+
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+
+
+class ImageEmbedding(ImageEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "Qdrant/clip-ViT-B-32-vision",
+                        "dim": 512,
+                        "description": "CLIP vision encoder based on ViT-B/32",
+                        "license": "mit",
+                        "size_in_GB": 0.33,
+                        "sources": {
+                            "hf": "Qdrant/clip-ViT-B-32-vision",
+                        },
+                        "model_file": "model.onnx",
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        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):
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in ImageEmbedding."
+            "Please check the supported models using `ImageEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        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
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self.model.embed(images, batch_size, parallel, **kwargs)

fastembed/image/image_embedding_base.py

@@ -0,0 +1,44 @@
+from typing import Iterable, Optional
+
+import numpy as np
+
+from fastembed.common.model_management import ModelManagement
+from fastembed.common.types import ImageInput
+
+
+class ImageEmbeddingBase(ModelManagement):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Embeds a list of images into a list of embeddings.
+
+        Args:
+            images: The list of image paths to preprocess and embed.
+            batch_size: Batch size for encoding
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[np.ndarray]: The embeddings.
+        """
+        raise NotImplementedError()

fastembed/image/onnx_embedding.py

@@ -0,0 +1,211 @@
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
+
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
+
+supported_onnx_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-vision",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-vision",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/resnet50-onnx",
+        "dim": 2048,
+        "description": "Image embeddings, Unimodal (image), 2016 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.1,
+        "sources": {
+            "hf": "Qdrant/resnet50-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Image embeddings, Multimodal (text&image), 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/vision_model.onnx",
+    },
+]
+
+
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        """
+        Load the onnx model.
+        """
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+        )
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            list[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_onnx_models
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of images into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        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
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+
+        yield from self._embed_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
+        return OnnxImageEmbeddingWorker
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+
+        return onnx_input
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        return normalize(output.model_output).astype(np.float32)
+
+
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_image_model.py

@@ -0,0 +1,131 @@
+import contextlib
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
+from PIL import Image
+
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
+from fastembed.common.preprocessor_utils import load_preprocessor
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool
+
+# Holds type of the embedding result
+
+
+class OnnxImageModel(OnnxModel[T]):
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.processor = None
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+        )
+        self.processor = load_preprocessor(model_dir=model_dir)
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
+        return {node.name: encoded for node in self.model.get_inputs()}
+
+    def onnx_embed(self, images: list[ImageInput], **kwargs) -> OnnxOutputContext:
+        with contextlib.ExitStack():
+            image_files = [
+                Image.open(image) if not isinstance(image, Image.Image) else image
+                for image in images
+            ]
+            encoded = self.processor(image_files)
+        onnx_input = self._build_onnx_input(encoded)
+        onnx_input = self._preprocess_onnx_input(onnx_input)
+        model_output = self.model.run(None, onnx_input)
+        embeddings = model_output[0].reshape(len(images), -1)
+        return OnnxOutputContext(model_output=embeddings)
+
+    def _embed_images(
+        self,
+        model_name: str,
+        cache_dir: str,
+        images: ImageInput,
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        **kwargs,
+    ) -> Iterable[T]:
+        is_small = False
+
+        if isinstance(images, (str, Path, Image.Image)):
+            images = [images]
+            is_small = True
+
+        if isinstance(images, list) and len(images) < batch_size:
+            is_small = True
+
+        if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+
+            for batch in iter_batch(images, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                **kwargs,
+            }
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)
+
+
+class ImageEmbeddingWorker(EmbeddingWorker):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            embeddings = self.model.onnx_embed(batch)
+            yield idx, embeddings

fastembed/image/transform/functional.py

@@ -0,0 +1,150 @@
+from typing import Sized, Union
+
+import numpy as np
+from PIL import Image
+
+
+def convert_to_rgb(image: Image.Image) -> Image.Image:
+    if image.mode == "RGB":
+        return image
+
+    image = image.convert("RGB")
+    return image
+
+
+def center_crop(
+    image: Union[Image.Image, np.ndarray],
+    size: tuple[int, int],
+) -> np.ndarray:
+    if isinstance(image, np.ndarray):
+        _, orig_height, orig_width = image.shape
+    else:
+        orig_height, orig_width = image.height, image.width
+        # (H, W, C) -> (C, H, W)
+        image = np.array(image).transpose((2, 0, 1))
+
+    crop_height, crop_width = size
+
+    # left upper corner (0, 0)
+    top = (orig_height - crop_height) // 2
+    bottom = top + crop_height
+    left = (orig_width - crop_width) // 2
+    right = left + crop_width
+
+    # Check if cropped area is within image boundaries
+    if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width:
+        image = image[..., top:bottom, left:right]
+        return image
+
+    # Padding with zeros
+    new_height = max(crop_height, orig_height)
+    new_width = max(crop_width, orig_width)
+    new_shape = image.shape[:-2] + (new_height, new_width)
+    new_image = np.zeros_like(image, shape=new_shape)
+
+    top_pad = (new_height - orig_height) // 2
+    bottom_pad = top_pad + orig_height
+    left_pad = (new_width - orig_width) // 2
+    right_pad = left_pad + orig_width
+    new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image
+
+    top += top_pad
+    bottom += top_pad
+    left += left_pad
+    right += left_pad
+
+    new_image = new_image[
+        ..., max(0, top) : min(new_height, bottom), max(0, left) : min(new_width, right)
+    ]
+
+    return new_image
+
+
+def normalize(
+    image: np.ndarray,
+    mean: Union[float, np.ndarray],
+    std: Union[float, np.ndarray],
+) -> np.ndarray:
+    if not isinstance(image, np.ndarray):
+        raise ValueError("image must be a numpy array")
+
+    num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
+
+    if not np.issubdtype(image.dtype, np.floating):
+        image = image.astype(np.float32)
+
+    if isinstance(mean, Sized):
+        if len(mean) != num_channels:
+            raise ValueError(
+                f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}"
+            )
+    else:
+        mean = [mean] * num_channels
+    mean = np.array(mean, dtype=image.dtype)
+
+    if isinstance(std, Sized):
+        if len(std) != num_channels:
+            raise ValueError(
+                f"std must have {num_channels} elements if it is an iterable, got {len(std)}"
+            )
+    else:
+        std = [std] * num_channels
+    std = np.array(std, dtype=image.dtype)
+
+    image = ((image.T - mean) / std).T
+    return image
+
+
+def resize(
+    image: Image.Image,
+    size: Union[int, tuple[int, int]],
+    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
+) -> Image.Image:
+    if isinstance(size, tuple):
+        return image.resize(size, resample)
+
+    height, width = image.height, image.width
+    short, long = (width, height) if width <= height else (height, width)
+
+    new_short, new_long = size, int(size * long / short)
+    if width <= height:
+        new_size = (new_short, new_long)
+    else:
+        new_size = (new_long, new_short)
+    return image.resize(new_size, resample)
+
+
+def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
+    return (image * scale).astype(dtype)
+
+
+def pil2ndarray(image: Union[Image.Image, np.ndarray]):
+    if isinstance(image, Image.Image):
+        return np.asarray(image).transpose((2, 0, 1))
+    return image
+
+
+def pad2square(
+    image: Image.Image,
+    size: int,
+    fill_color: Union[str, int, tuple[int, ...]] = 0,
+) -> Image.Image:
+    height, width = image.height, image.width
+
+    left, right = 0, width
+    top, bottom = 0, height
+
+    crop_required = False
+    if width > size:
+        left = (width - size) // 2
+        right = left + size
+        crop_required = True
+
+    if height > size:
+        top = (height - size) // 2
+        bottom = top + size
+        crop_required = True
+
+    new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
+    new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
+    return new_image