inference-models 0.18.3__py3-none-any.whl

inference_models/models/auto_loaders/ranking.py

@@ -0,0 +1,413 @@
+from typing import List, Optional, Union
+
+import torch
+
+from inference_models.models.auto_loaders.constants import (
+    NMS_CLASS_AGNOSTIC_KEY,
+    NMS_CONFIDENCE_THRESHOLD_KEY,
+    NMS_FUSED_FEATURE,
+    NMS_IOU_THRESHOLD_KEY,
+    NMS_MAX_DETECTIONS_KEY,
+)
+from inference_models.models.auto_loaders.entities import BackendType
+from inference_models.models.auto_loaders.utils import (
+    filter_available_devices_with_selected_device,
+)
+from inference_models.runtime_introspection.core import x_ray_runtime_environment
+from inference_models.weights_providers.entities import (
+    JetsonEnvironmentRequirements,
+    ModelPackageMetadata,
+    Quantization,
+    ServerEnvironmentRequirements,
+)
+
+BACKEND_PRIORITY = {
+    BackendType.TRT: 7,
+    BackendType.TORCH: 6,
+    BackendType.HF: 5,
+    BackendType.ONNX: 3,
+    BackendType.TORCH_SCRIPT: 3,
+    BackendType.MEDIAPIPE: 2,
+    BackendType.ULTRALYTICS: 1,
+    BackendType.CUSTOM: 0,
+}
+QUANTIZATION_PRIORITY = {
+    Quantization.INT8: 4,
+    Quantization.FP16: 3,
+    Quantization.FP32: 2,
+    Quantization.UNKNOWN: 1,
+}
+DYNAMIC_BATCH_SIZE_KEY = "dynamic"
+STATIC_BATCH_SIZE_KEY = "static"
+BATCH_SIZE_PRIORITY = {
+    DYNAMIC_BATCH_SIZE_KEY: 2,
+    STATIC_BATCH_SIZE_KEY: 1,
+}
+
+
+def rank_model_packages(
+    model_packages: List[ModelPackageMetadata],
+    selected_device: Optional[torch.device] = None,
+    nms_fusion_preferences: Optional[Union[bool, dict]] = None,
+) -> List[ModelPackageMetadata]:
+    # I feel like this will be the biggest liability of new inference :))
+    # Some dimensions are just hard to rank arbitrarily and reasonably
+    sorting_features = []
+    # ordering TRT and Cu versions from older to newest -
+    # with the assumption that incompatible versions are eliminated earlier, and
+    # ranking implicitly attempts to match version closes to the current -
+    # in come cases we would rank high versions below current one, but for that
+    # it is assumed such versions are compatible, otherwise should be
+    # discarded in the previous stage.
+    cuda_ranking = rank_cuda_versions(model_packages=model_packages)
+    trt_ranking = rank_trt_versions(model_packages=model_packages)
+    # this is to ensure determinism when other methods fail
+    identifiers_ranking = rank_packages_ids(model_packages=model_packages)
+    for model_package, package_cu_rank, package_trt_rank, package_id_rank in zip(
+        model_packages, cuda_ranking, trt_ranking, identifiers_ranking
+    ):
+        batch_mode = (
+            DYNAMIC_BATCH_SIZE_KEY
+            if model_package.dynamic_batch_size_supported
+            else STATIC_BATCH_SIZE_KEY
+        )
+        static_batch_size_score = (
+            0
+            if model_package.static_batch_size is None
+            else -1 * model_package.static_batch_size
+        )
+        sorting_features.append(
+            (
+                BACKEND_PRIORITY.get(model_package.backend, 0),
+                QUANTIZATION_PRIORITY.get(model_package.quantization, 0),
+                model_package.trusted_source,
+                BATCH_SIZE_PRIORITY[batch_mode],
+                static_batch_size_score,  # the bigger statis batch size, the worse - requires padding
+                retrieve_onnx_opset_score(
+                    model_package
+                ),  # the higher opset, the better
+                retrieve_trt_forward_compatible_match_score(
+                    model_package
+                ),  # exact matches first
+                retrieve_same_trt_cc_compatibility_score(model_package),
+                retrieve_cuda_device_match_score(
+                    model_package, selected_device
+                ),  # we like more direct matches
+                package_cu_rank,
+                package_trt_rank,
+                retrieve_onnx_incompatible_providers_score(model_package),
+                retrieve_trt_dynamic_batch_size_score(model_package),
+                retrieve_fused_nms_rank(
+                    model_package, nms_fusion_preferences=nms_fusion_preferences
+                ),
+                retrieve_trt_lean_runtime_excluded_score(model_package),
+                retrieve_jetson_device_name_match_score(model_package),
+                retrieve_os_version_match_score(model_package),
+                retrieve_l4t_version_match_score(model_package),
+                retrieve_driver_version_match_score(model_package),
+                package_id_rank,
+                model_package,
+            )
+        )
+    sorted_features = sorted(sorting_features, key=lambda x: x[:-1], reverse=True)
+    return [f[-1] for f in sorted_features]
+
+
+def retrieve_onnx_opset_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.onnx_package_details is None:
+        return -1
+    return model_package.onnx_package_details.opset
+
+
+def retrieve_cuda_device_match_score(
+    model_package: ModelPackageMetadata,
+    selected_device: Optional[torch.device] = None,
+) -> int:
+    if model_package.backend is not BackendType.TRT:
+        return 0
+    if model_package.environment_requirements is None:
+        return 0
+    if not isinstance(
+        model_package.environment_requirements,
+        (JetsonEnvironmentRequirements, ServerEnvironmentRequirements),
+    ):
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    all_available_cuda_devices, _ = filter_available_devices_with_selected_device(
+        selected_device=selected_device,
+        all_available_cuda_devices=runtime_x_ray.gpu_devices,
+        all_available_devices_cc=runtime_x_ray.gpu_devices_cc,
+    )
+    compilation_device = model_package.environment_requirements.cuda_device_name
+    return sum(dev == compilation_device for dev in all_available_cuda_devices)
+
+
+def retrieve_same_trt_cc_compatibility_score(
+    model_package: ModelPackageMetadata,
+) -> int:
+    if model_package.trt_package_details is None:
+        return 1
+    return int(not model_package.trt_package_details.same_cc_compatible)
+
+
+def retrieve_trt_forward_compatible_match_score(
+    model_package: ModelPackageMetadata,
+) -> int:
+    if model_package.trt_package_details is None:
+        return 1
+    return int(not model_package.trt_package_details.trt_forward_compatible)
+
+
+def retrieve_onnx_incompatible_providers_score(
+    model_package: ModelPackageMetadata,
+) -> int:
+    if model_package.onnx_package_details is None:
+        return 0
+    if not model_package.onnx_package_details.incompatible_providers:
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    available_onnx_execution_providers = set(
+        runtime_x_ray.available_onnx_execution_providers or []
+    )
+    return -len(
+        available_onnx_execution_providers.intersection(
+            model_package.onnx_package_details.incompatible_providers
+        )
+    )
+
+
+def retrieve_trt_dynamic_batch_size_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.trt_package_details is None:
+        return 0
+    if any(
+        bs is None
+        for bs in [
+            model_package.trt_package_details.min_dynamic_batch_size,
+            model_package.trt_package_details.max_dynamic_batch_size,
+        ]
+    ):
+        return 0
+    return (
+        model_package.trt_package_details.max_dynamic_batch_size
+        - model_package.trt_package_details.min_dynamic_batch_size
+    )
+
+
+def retrieve_trt_lean_runtime_excluded_score(
+    model_package: ModelPackageMetadata,
+) -> int:
+    if model_package.trt_package_details is None:
+        return 0
+    return int(not model_package.trt_package_details.trt_lean_runtime_excluded)
+
+
+def retrieve_os_version_match_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.backend is not BackendType.TRT:
+        # irrelevant for not trt
+        return 0
+    if model_package.environment_requirements is None:
+        return 0
+    if not isinstance(
+        model_package.environment_requirements, ServerEnvironmentRequirements
+    ):
+        return 0
+    if not model_package.environment_requirements.os_version:
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    return int(
+        runtime_x_ray.os_version == model_package.environment_requirements.os_version
+    )
+
+
+def retrieve_l4t_version_match_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.backend is not BackendType.TRT:
+        # irrelevant for not trt
+        return 0
+    if model_package.environment_requirements is None:
+        return 0
+    if not isinstance(
+        model_package.environment_requirements, JetsonEnvironmentRequirements
+    ):
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    return int(
+        runtime_x_ray.l4t_version == model_package.environment_requirements.l4t_version
+    )
+
+
+def retrieve_driver_version_match_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.trt_package_details is None:
+        # irrelevant for not trt
+        return 0
+    if model_package.environment_requirements is None:
+        return 0
+    if not isinstance(
+        model_package.environment_requirements, JetsonEnvironmentRequirements
+    ) and not isinstance(
+        model_package.environment_requirements, ServerEnvironmentRequirements
+    ):
+        return 0
+    if not model_package.environment_requirements.driver_version:
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    return int(
+        runtime_x_ray.driver_version
+        == model_package.environment_requirements.driver_version
+    )
+
+
+def retrieve_jetson_device_name_match_score(model_package: ModelPackageMetadata) -> int:
+    if model_package.trt_package_details is None:
+        # irrelevant for not trt
+        return 0
+    if model_package.environment_requirements is None:
+        return 0
+    if not isinstance(
+        model_package.environment_requirements, JetsonEnvironmentRequirements
+    ):
+        return 0
+    runtime_x_ray = x_ray_runtime_environment()
+    return int(
+        runtime_x_ray.jetson_type
+        == model_package.environment_requirements.jetson_product_name
+    )
+
+
+def retrieve_fused_nms_rank(
+    model_package: ModelPackageMetadata,
+    nms_fusion_preferences: Optional[Union[bool, dict]],
+) -> Union[float, int]:
+    if nms_fusion_preferences is None or nms_fusion_preferences is False:
+        return 0
+    if not model_package.model_features:
+        return 0
+    nms_fused = model_package.model_features.get(NMS_FUSED_FEATURE)
+    if not isinstance(nms_fused, dict):
+        return 0
+    if nms_fusion_preferences is True:
+        # default values should be passed by filter, so here we treat every package equally good
+        return 1
+    actual_max_detections = nms_fused[NMS_MAX_DETECTIONS_KEY]
+    actual_confidence_threshold = nms_fused[NMS_CONFIDENCE_THRESHOLD_KEY]
+    actual_iou_threshold = nms_fused[NMS_IOU_THRESHOLD_KEY]
+    actual_class_agnostic = nms_fused[NMS_CLASS_AGNOSTIC_KEY]
+    final_score = 0.0
+    if NMS_MAX_DETECTIONS_KEY in nms_fusion_preferences:
+        requested_max_detections = nms_fusion_preferences[NMS_MAX_DETECTIONS_KEY]
+        if isinstance(requested_max_detections, (list, tuple)):
+            min_detections, max_detections = requested_max_detections
+        else:
+            min_detections, max_detections = (
+                requested_max_detections,
+                requested_max_detections,
+            )
+        final_score += score_distance_from_mean(
+            min_value=min_detections,
+            max_value=max_detections,
+            examined_value=actual_max_detections,
+        )
+    if NMS_CONFIDENCE_THRESHOLD_KEY in nms_fusion_preferences:
+        requested_confidence = nms_fusion_preferences[NMS_CONFIDENCE_THRESHOLD_KEY]
+        if isinstance(requested_confidence, (list, tuple)):
+            min_confidence, max_confidence = requested_confidence
+        else:
+            min_confidence, max_confidence = (
+                requested_confidence,
+                requested_confidence,
+            )
+        final_score += score_distance_from_mean(
+            min_value=min_confidence,
+            max_value=max_confidence,
+            examined_value=actual_confidence_threshold,
+        )
+    if NMS_IOU_THRESHOLD_KEY in nms_fusion_preferences:
+        requested_iou_threshold = nms_fusion_preferences[NMS_IOU_THRESHOLD_KEY]
+        if isinstance(requested_iou_threshold, (list, tuple)):
+            min_iou_threshold, max_iou_threshold = requested_iou_threshold
+        else:
+            min_iou_threshold, max_iou_threshold = (
+                requested_iou_threshold,
+                requested_iou_threshold,
+            )
+        final_score += score_distance_from_mean(
+            min_value=min_iou_threshold,
+            max_value=max_iou_threshold,
+            examined_value=actual_iou_threshold,
+        )
+    if NMS_CLASS_AGNOSTIC_KEY in nms_fusion_preferences:
+        final_score += float(
+            actual_class_agnostic == nms_fusion_preferences[NMS_CLASS_AGNOSTIC_KEY]
+        )
+    return final_score
+
+
+def score_distance_from_mean(
+    min_value: float, max_value: float, examined_value: float
+) -> float:
+    min_value, max_value = min(min_value, max_value), max(min_value, max_value)
+    if min_value == max_value:
+        return float(abs(examined_value - max_value) < 1e-5)
+    if examined_value < min_value or examined_value > max_value:
+        return 0.0
+    span = max_value - min_value
+    examined_value_scaled = min(max((examined_value - min_value) / (span + 1e-6), 0), 1)
+    return 1.0 - abs(0.5 - examined_value_scaled)
+
+
+def rank_cuda_versions(model_packages: List[ModelPackageMetadata]) -> List[int]:
+    cuda_versions = []
+    package_id_to_cuda_version = {}
+    last_ranking = -len(model_packages) + 1
+    for package in model_packages:
+        if isinstance(package.environment_requirements, ServerEnvironmentRequirements):
+            cuda_versions.append(package.environment_requirements.cuda_version)
+            package_id_to_cuda_version[package.package_id] = (
+                package.environment_requirements.cuda_version
+            )
+        elif isinstance(
+            package.environment_requirements, JetsonEnvironmentRequirements
+        ):
+            cuda_versions.append(package.environment_requirements.cuda_version)
+            package_id_to_cuda_version[package.package_id] = (
+                package.environment_requirements.cuda_version
+            )
+    cuda_versions = sorted(set(cuda_versions))
+    cuda_versions_ranking = {version: -idx for idx, version in enumerate(cuda_versions)}
+    results = []
+    for package in model_packages:
+        package_cu_version = package_id_to_cuda_version.get(package.package_id)
+        result = cuda_versions_ranking.get(package_cu_version, last_ranking)
+        results.append(result)
+    return results
+
+
+def rank_trt_versions(model_packages: List[ModelPackageMetadata]) -> List[int]:
+    trt_versions = []
+    package_id_to_trt_version = {}
+    last_ranking = -len(model_packages) + 1
+    for package in model_packages:
+        if isinstance(package.environment_requirements, ServerEnvironmentRequirements):
+            trt_versions.append(package.environment_requirements.trt_version)
+            package_id_to_trt_version[package.package_id] = (
+                package.environment_requirements.trt_version
+            )
+        elif isinstance(
+            package.environment_requirements, JetsonEnvironmentRequirements
+        ):
+            trt_versions.append(package.environment_requirements.trt_version)
+            package_id_to_trt_version[package.package_id] = (
+                package.environment_requirements.trt_version
+            )
+    trt_versions = sorted(set(trt_versions))
+    trt_versions_ranking = {version: -idx for idx, version in enumerate(trt_versions)}
+    results = []
+    for package in model_packages:
+        package_trt_version = package_id_to_trt_version.get(package.package_id)
+        result = trt_versions_ranking.get(package_trt_version, last_ranking)
+        results.append(result)
+    return results
+
+
+def rank_packages_ids(model_packages: List[ModelPackageMetadata]) -> List[int]:
+    package_ids = [p.package_id for p in model_packages]
+    return sorted(range(len(package_ids)), key=lambda i: package_ids[i])

inference_models/models/auto_loaders/utils.py

@@ -0,0 +1,31 @@
+from typing import List, Optional, Tuple
+
+import torch
+from packaging.version import Version
+
+from inference_models.errors import ModelPackageNegotiationError
+
+
+def filter_available_devices_with_selected_device(
+    selected_device: Optional[torch.device],
+    all_available_cuda_devices: List[str],
+    all_available_devices_cc: List[Version],
+) -> Tuple[List[str], List[Version]]:
+    if selected_device is not None and selected_device.type != "cuda":
+        return [], []
+    if selected_device is not None and selected_device.type == "cuda":
+        index = selected_device.index or 0
+        if index >= len(all_available_cuda_devices) or index >= len(
+            all_available_devices_cc
+        ):
+            raise ModelPackageNegotiationError(
+                message=f"Model Package Negotiation algorithm received selected device: {selected_device} which "
+                f"does not match runtime introspection results. If you selected device to run the model "
+                f"manually - verify your choice. Otherwise, this error most likely is a bug. Create new "
+                f"issue: https://github.com/roboflow/inference/issues",
+                help_url="https://todo",
+            )
+        all_available_cuda_devices = [all_available_cuda_devices[index]]
+        all_available_devices_cc = [all_available_devices_cc[index]]
+        return all_available_cuda_devices, all_available_devices_cc
+    return all_available_cuda_devices, all_available_devices_cc

inference_models/models/base/classification.py

@@ -0,0 +1,123 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, List, Optional, Union
+
+import numpy as np
+import torch
+
+from inference_models.models.base.types import PreprocessedInputs, RawPrediction
+
+
+@dataclass
+class ClassificationPrediction:
+    class_id: torch.Tensor  # (bs, )
+    confidence: torch.Tensor  # (bs, )
+    images_metadata: Optional[List[dict]] = None  # if given, list of size equal to bs
+
+
+class ClassificationModel(ABC, Generic[PreprocessedInputs, RawPrediction]):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "ClassificationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> ClassificationPrediction:
+        pre_processed_images = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> PreprocessedInputs:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: RawPrediction, **kwargs
+    ) -> ClassificationPrediction:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> ClassificationPrediction:
+        return self.infer(images, **kwargs)
+
+
+@dataclass
+class MultiLabelClassificationPrediction:
+    class_ids: torch.Tensor  # (predicted_labels_ids, )
+    confidence: torch.Tensor  # (predicted_labels_confidence, )
+    image_metadata: Optional[dict] = None
+
+
+class MultiLabelClassificationModel(ABC, Generic[PreprocessedInputs, RawPrediction]):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "MultiLabelClassificationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[MultiLabelClassificationPrediction]:
+        pre_processed_images = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> PreprocessedInputs:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: RawPrediction, **kwargs
+    ) -> List[MultiLabelClassificationPrediction]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[MultiLabelClassificationPrediction]:
+        return self.infer(images, **kwargs)

inference_models/models/base/depth_estimation.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from typing import Generic, List, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+class DepthEstimationModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "DepthEstimationModel":
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[torch.Tensor]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessingMetadata,
+        **kwargs,
+    ) -> List[torch.Tensor]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[torch.Tensor]:
+        return self.infer(images, **kwargs)