inference-models 0.18.3__py3-none-any.whl

inference_models/models/clip/clip_onnx.py

@@ -0,0 +1,148 @@
+from threading import Lock
+from typing import List, Optional, Union
+
+import clip
+import numpy as np
+import torch
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import (
+    EnvironmentConfigurationError,
+    MissingDependencyError,
+)
+from inference_models.models.base.embeddings import TextImageEmbeddingModel
+from inference_models.models.clip.preprocessing import create_clip_preprocessor
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.onnx import (
+    run_session_with_batch_size_limit,
+    set_execution_provider_defaults,
+)
+from inference_models.utils.onnx_introspection import (
+    get_selected_onnx_execution_providers,
+)
+
+try:
+    import onnxruntime
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import CLIP model with ONNX backend - this error means that some additional dependencies "
+        f"are not installed in the environment. If you run the `inference-models` library directly in your Python "
+        f"program, make sure the following extras of the package are installed: \n"
+        f"\t* `onnx-cpu` - when you wish to use library with CPU support only\n"
+        f"\t* `onnx-cu12` - for running on GPU with Cuda 12 installed\n"
+        f"\t* `onnx-cu118` - for running on GPU with Cuda 11.8 installed\n"
+        f"\t* `onnx-jp6-cu126` - for running on Jetson with Jetpack 6\n"
+        f"If you see this error using Roboflow infrastructure, make sure the service you use does support the model. "
+        f"You can also contact Roboflow to get support.",
+        help_url="https://todo",
+    ) from import_error
+
+
+MEAN = (0.48145466, 0.4578275, 0.40821073)
+STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+class ClipOnnx(TextImageEmbeddingModel):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        onnx_execution_providers: Optional[List[Union[str, tuple]]] = None,
+        default_onnx_trt_options: bool = True,
+        device: torch.device = DEFAULT_DEVICE,
+        max_batch_size: int = 32,
+        **kwargs,
+    ) -> "ClipOnnx":
+        if onnx_execution_providers is None:
+            onnx_execution_providers = get_selected_onnx_execution_providers()
+        if not onnx_execution_providers:
+            raise EnvironmentConfigurationError(
+                message=f"Could not initialize model - selected backend is ONNX which requires execution provider to "
+                f"be specified - explicitly in `from_pretrained(...)` method or via env variable "
+                f"`ONNXRUNTIME_EXECUTION_PROVIDERS`. If you run model locally - adjust your setup, otherwise "
+                f"contact the platform support.",
+                help_url="https://todo",
+            )
+        onnx_execution_providers = set_execution_provider_defaults(
+            providers=onnx_execution_providers,
+            model_package_path=model_name_or_path,
+            device=device,
+            default_onnx_trt_options=default_onnx_trt_options,
+        )
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=[
+                "textual.onnx",
+                "visual.onnx",
+            ],
+        )
+        visual_onnx_session = onnxruntime.InferenceSession(
+            path_or_bytes=model_package_content["visual.onnx"],
+            providers=onnx_execution_providers,
+        )
+        textual_onnx_session = onnxruntime.InferenceSession(
+            path_or_bytes=model_package_content["textual.onnx"],
+            providers=onnx_execution_providers,
+        )
+        image_size = visual_onnx_session.get_inputs()[0].shape[2]
+        visual_input_name = visual_onnx_session.get_inputs()[0].name
+        textual_input_name = textual_onnx_session.get_inputs()[0].name
+        return cls(
+            visual_onnx_session=visual_onnx_session,
+            textual_onnx_session=textual_onnx_session,
+            image_size=image_size,
+            visual_input_name=visual_input_name,
+            textual_input_name=textual_input_name,
+            device=device,
+            max_batch_size=max_batch_size,
+        )
+
+    def __init__(
+        self,
+        visual_onnx_session: onnxruntime.InferenceSession,
+        textual_onnx_session: onnxruntime.InferenceSession,
+        image_size: int,
+        visual_input_name: str,
+        textual_input_name: str,
+        device: torch.device,
+        max_batch_size: int,
+    ):
+        self._visual_onnx_session = visual_onnx_session
+        self._textual_onnx_session = textual_onnx_session
+        self._image_size = image_size
+        self._visual_input_name = visual_input_name
+        self._textual_input_name = textual_input_name
+        self._device = device
+        self._max_batch_size = max_batch_size
+        self._visual_session_thread_lock = Lock()
+        self._textual_session_thread_lock = Lock()
+        self._preprocessor = create_clip_preprocessor(image_size=image_size)
+
+    def embed_images(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        pre_processed_images = self._preprocessor(
+            images, input_color_format, self._device
+        )
+        with self._visual_session_thread_lock:
+            return run_session_with_batch_size_limit(
+                session=self._visual_onnx_session,
+                inputs={self._visual_input_name: pre_processed_images},
+                max_batch_size=self._max_batch_size,
+            )[0]
+
+    def embed_text(self, texts: Union[str, List[str]], **kwargs) -> torch.Tensor:
+        if not isinstance(texts, list):
+            texts = [texts]
+        tokenized_batch = clip.tokenize(texts)
+        with self._textual_session_thread_lock:
+            return run_session_with_batch_size_limit(
+                session=self._textual_onnx_session,
+                inputs={self._textual_input_name: tokenized_batch},
+                max_batch_size=self._max_batch_size,
+            )[0]

inference_models/models/clip/clip_pytorch.py

@@ -0,0 +1,104 @@
+from typing import Callable, List, Optional, Union
+
+import clip
+import numpy as np
+import torch
+from clip.model import CLIP, build_model
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import CorruptedModelPackageError
+from inference_models.models.base.embeddings import TextImageEmbeddingModel
+from inference_models.models.clip.preprocessing import create_clip_preprocessor
+from inference_models.models.common.model_packages import get_model_package_contents
+
+
+class ClipTorch(TextImageEmbeddingModel):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        max_batch_size: int = 32,
+        **kwargs,
+    ) -> "ClipTorch":
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["model.pt"],
+        )
+        model_weights_file = model_package_content["model.pt"]
+        model = build_clip_model(model_weights_file=model_weights_file, device=device)
+        model.eval()
+        return cls(
+            model=model,
+            tokenizer=clip.tokenize,
+            device=device,
+            max_batch_size=max_batch_size,
+        )
+
+    def __init__(
+        self,
+        model: CLIP,
+        tokenizer: Callable,
+        device: torch.device,
+        max_batch_size: int,
+    ):
+        self._model = model
+        self._tokenizer = tokenizer
+        self._device = device
+        self._preprocessor = create_clip_preprocessor(
+            image_size=model.visual.input_resolution
+        )
+        self._max_batch_size = max_batch_size
+
+    @torch.no_grad()
+    def embed_images(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        tensor_batch = self._preprocessor(images, input_color_format, self._device)
+        if tensor_batch.shape[0] <= self._max_batch_size:
+            return self._model.encode_image(tensor_batch.to(self._device))
+        results = []
+        for i in range(0, tensor_batch.shape[0], self._max_batch_size):
+            batch_input = tensor_batch[i : i + self._max_batch_size].contiguous()
+            batch_results = self._model.encode_image(batch_input.to(self._device))
+            results.append(batch_results)
+        return torch.cat(results, dim=0)
+
+    @torch.no_grad()
+    def embed_text(
+        self,
+        texts: Union[str, List[str]],
+        **kwargs,
+    ) -> torch.Tensor:
+        if isinstance(texts, str):
+            texts = [texts]
+        text_tokens = self._tokenizer(texts).to(self._device)
+        if text_tokens.shape[0] <= self._max_batch_size:
+            return self._model.encode_text(text_tokens)
+        results = []
+        for i in range(0, text_tokens.shape[0], self._max_batch_size):
+            batch_input = text_tokens[i : i + self._max_batch_size].contiguous()
+            batch_results = self._model.encode_text(batch_input)
+            results.append(batch_results)
+        return torch.cat(results, dim=0)
+
+
+def build_clip_model(model_weights_file: str, device: torch.device) -> CLIP:
+    try:
+        # The model file is a JIT archive, so we load it as such
+        # and then build a new model from its state dict.
+        jit_model = torch.jit.load(model_weights_file, map_location="cpu").eval()
+        state_dict = jit_model.state_dict()
+        model = build_model(state_dict).to(device)
+        if device.type == "cpu":
+            model.float()
+        return model
+    except Exception as e:
+        raise CorruptedModelPackageError(
+            f"Could not load TorchScript model from {model_weights_file}. Details: {e}"
+        ) from e

inference_models/models/clip/preprocessing.py

@@ -0,0 +1,162 @@
+from functools import partial
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import torch
+from torchvision.transforms import (
+    CenterCrop,
+    Compose,
+    InterpolationMode,
+    Normalize,
+    Resize,
+)
+
+from inference_models.entities import ColorFormat
+from inference_models.errors import ModelRuntimeError
+
+MEAN = (0.48145466, 0.4578275, 0.40821073)
+STD = (0.26862954, 0.26130258, 0.27577711)
+
+PreprocessorFun = Callable[
+    [
+        Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        Optional[ColorFormat],
+        torch.device,
+    ],
+    torch.Tensor,
+]
+
+
+def create_clip_preprocessor(image_size: int) -> PreprocessorFun:
+    """
+    Creates a preprocessor for CLIP models that operates on tensors.
+
+    This implementation replicates the logic of the original CLIP preprocessing pipeline
+    but is designed to work directly with torch.Tensors and np.ndarrays, avoiding
+    the need to convert to and from PIL.Image objects.
+
+    Note: Due to differences in the underlying resizing algorithms (torchvision vs. PIL),
+    the output of this preprocessor may have minor numerical differences compared to
+    the original. These differences have been tested and are known to produce
+    embeddings with very high cosine similarity, making them functionally equivalent.
+
+    Args:
+        image_size (int): The target size for the input images.`
+        device (torch.device): The device to move the tensors to.
+
+    Returns:
+        A callable function that preprocesses images.
+    """
+    # This pre-processing pipeline matches the original CLIP implementation.
+    # 1. Resize to `image_size`
+    # 2. Center crop to `image_size`
+    # 3. Scale pixel values to [0, 1]
+    # 4. Normalize with CLIP's specific mean and standard deviation.
+    transforms = Compose(
+        [
+            Resize(image_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
+            CenterCrop(image_size),
+            lambda x: x.to(torch.float32) / 255.0,
+            Normalize(MEAN, STD),
+        ]
+    )
+    return partial(pre_process_image, transforms=transforms)
+
+
+def pre_process_image(
+    images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+    input_color_format: Optional[ColorFormat],
+    device: torch.device,
+    transforms: Compose,
+) -> torch.Tensor:
+    images = inputs_to_tensor(
+        images=images, device=device, input_color_format=input_color_format
+    )
+    if isinstance(images, torch.Tensor):
+        return transforms(images)
+    return torch.cat([transforms(i) for i in images], dim=0).contiguous()
+
+
+def inputs_to_tensor(
+    images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+    device: torch.device,
+    input_color_format: Optional[ColorFormat] = None,
+) -> Union[torch.Tensor, List[torch.Tensor]]:
+    if not isinstance(images, (list, np.ndarray, torch.Tensor)):
+        raise ModelRuntimeError(
+            f"Unsupported input type: {type(images)}. Must be one of list, np.ndarray, or torch.Tensor."
+        )
+    if isinstance(images, list):
+        if not images:
+            raise ModelRuntimeError("Input image list cannot be empty.")
+        return [
+            input_to_tensor(
+                image=image,
+                device=device,
+                input_color_format=input_color_format,
+                batched_tensors_allowed=False,
+            )
+            for image in images
+        ]
+    return input_to_tensor(
+        image=images,
+        device=device,
+        input_color_format=input_color_format,
+    ).contiguous()
+
+
+def input_to_tensor(
+    image: Union[torch.Tensor, np.ndarray],
+    device: torch.device,
+    input_color_format: Optional[ColorFormat] = None,
+    batched_tensors_allowed: bool = True,
+) -> torch.Tensor:
+    if not isinstance(image, (np.ndarray, torch.Tensor)):
+        raise ModelRuntimeError(
+            f"Unsupported input type: {type(image)}. Each element must be one of np.ndarray, or torch.Tensor."
+        )
+    is_numpy = isinstance(image, np.ndarray)
+    if is_numpy:
+        if len(image.shape) != 3:
+            raise ModelRuntimeError(
+                f"Unsupported input type: detected np.ndarray image of shape {image.shape} which has "
+                f"number of dimensions different than 3. This input is invalid."
+            )
+        if image.shape[-1] != 3:
+            raise ModelRuntimeError(
+                f"Unsupported input type: detected np.ndarray image of shape {image.shape} which has "
+                f"incorrect number of color channels (expected: 3)."
+            )
+        # HWC -> CHW
+        tensor_image = torch.from_numpy(image).to(device).permute(2, 0, 1).unsqueeze(0)
+    else:
+        expected_dimensions_str = (
+            "expected: 3 or 4" if batched_tensors_allowed else "expected: 3"
+        )
+        if len(image.shape) == 4 and not batched_tensors_allowed:
+            raise ModelRuntimeError(
+                f"Unsupported input type: detected torch.Tensor image of shape {image.shape} which has "
+                f"incorrect number of dimensions ({expected_dimensions_str})."
+            )
+        if len(image.shape) != 3 and len(image.shape) != 4:
+            raise ModelRuntimeError(
+                f"Unsupported input type: detected torch.Tensor image of shape {image.shape} which has "
+                f"incorrect number of dimensions ({expected_dimensions_str})."
+            )
+        if (len(image.shape) == 3 and image.shape[0] != 3) or (
+            len(image.shape) == 4 and image.shape[1] != 3
+        ):
+            raise ModelRuntimeError(
+                f"Unsupported input type: detected torch.Tensor image of shape {image.shape} which has "
+                f"incorrect number of color channels (expected: 3)."
+            )
+        if len(image.shape) == 3:
+            image = image.unsqueeze(0)
+        tensor_image = image.to(device)
+    effective_color_format = input_color_format
+    if effective_color_format is None:
+        effective_color_format = "bgr" if is_numpy else "rgb"
+    if effective_color_format == "bgr":
+        # BGR -> RGB
+        tensor_image = tensor_image[:, [2, 1, 0], :, :]
+    return tensor_image

inference_models/models/common/cuda.py

@@ -0,0 +1,30 @@
+import contextlib
+from typing import Generator
+
+from inference_models.errors import MissingDependencyError
+
+try:
+    import pycuda.driver as cuda
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message="TODO",
+        help_url="https://todo",
+    ) from import_error
+
+
+@contextlib.contextmanager
+def use_primary_cuda_context(
+    cuda_device: cuda.Device,
+) -> Generator[cuda.Context, None, None]:
+    context = cuda_device.retain_primary_context()
+    with use_cuda_context(context) as ctx:
+        yield ctx
+
+
+@contextlib.contextmanager
+def use_cuda_context(context: cuda.Context) -> Generator[cuda.Context, None, None]:
+    context.push()
+    try:
+        yield context
+    finally:
+        context.pop()

inference_models/models/common/model_packages.py

@@ -0,0 +1,25 @@
+import os.path
+from typing import Dict, List
+
+from inference_models.errors import CorruptedModelPackageError
+
+
+def get_model_package_contents(
+    model_package_dir: str,
+    elements: List[str],
+) -> Dict[str, str]:
+    result = {}
+    for element in elements:
+        element_path = os.path.join(model_package_dir, element)
+        if not os.path.exists(element_path):
+            raise CorruptedModelPackageError(
+                message=f"Model package is incomplete. Could not find element {element}. "
+                f"If you attempt to run `inference-models` locally - inspect the contents of local directory to check for "
+                f"completeness of model package download - lack of files may indicate network issues. Verification "
+                f"of connectivity may be a good first step. If you prepared the model package manually - examine the "
+                f"correctness of the setup. If you run on managed serving - contact support if the issue is "
+                f"not ephemeral.",
+                help_url="https://todo",
+            )
+        result[element] = element_path
+    return result