inference-models 0.18.3__py3-none-any.whl

inference_models/utils/imports.py

@@ -0,0 +1,48 @@
+import importlib
+import os
+import sys
+from typing import Any, Optional
+
+
+class LazyClass:
+
+    def __init__(self, module_name: str, class_name: str):
+        self._module_name = module_name
+        self._class_name = class_name
+        self._symbol: Optional[type] = None
+
+    def resolve(self) -> type:
+        if self._symbol is None:
+            module = importlib.import_module(self._module_name)
+            self._symbol = getattr(module, self._class_name)
+        return self._symbol
+
+
+def import_class_from_file(
+    file_path: str, class_name: str, alias_name: Optional[str] = None
+) -> type:
+    """
+    Emulates what huggingface transformers does to load remote code with trust_remote_code=True,
+    but allows us to use the class directly so that we don't have to load untrusted code.
+    """
+    file_path = os.path.abspath(file_path)
+    module_name = os.path.splitext(os.path.basename(file_path))[0]
+    module_dir = os.path.dirname(file_path)
+    parent_dir = os.path.dirname(module_dir)
+
+    sys.path.insert(0, parent_dir)
+
+    try:
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        module = importlib.util.module_from_spec(spec)
+
+        # Manually set the __package__ attribute to the parent package
+        module.__package__ = os.path.basename(module_dir)
+
+        spec.loader.exec_module(module)
+        cls = getattr(module, class_name)
+        if alias_name:
+            globals()[alias_name] = cls
+        return cls
+    finally:
+        sys.path.pop(0)

inference_models/utils/onnx_introspection.py

@@ -0,0 +1,17 @@
+from functools import cache
+from typing import List
+
+from inference_models.configuration import ONNXRUNTIME_EXECUTION_PROVIDERS
+
+
+@cache
+def get_selected_onnx_execution_providers() -> List[str]:
+    try:
+        import onnxruntime
+
+        available_providers = set(onnxruntime.get_available_providers())
+        return [
+            ep for ep in ONNXRUNTIME_EXECUTION_PROVIDERS if ep in available_providers
+        ]
+    except ImportError:
+        return []

inference_models/weights_providers/core.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+from inference_models.errors import ModelRetrievalError
+from inference_models.weights_providers.entities import ModelMetadata
+from inference_models.weights_providers.roboflow import get_roboflow_model
+
+WEIGHTS_PROVIDERS = {
+    "roboflow": get_roboflow_model,
+}
+
+
+def get_model_from_provider(
+    model_id: str, provider: str, api_key: Optional[str] = None
+) -> ModelMetadata:
+    if provider not in WEIGHTS_PROVIDERS:
+        raise ModelRetrievalError(
+            message=f"Requested model to be retrieved using '{provider}' provider which is not implemented.",
+            help_url="https://todo",
+        )
+    return WEIGHTS_PROVIDERS[provider](model_id=model_id, api_key=api_key)

inference_models/weights_providers/entities.py

@@ -0,0 +1,159 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import List, Optional, Set, Tuple, Union
+
+from packaging.version import Version
+from pydantic import BaseModel, Field
+
+from inference_models.models.auto_loaders.entities import BackendType
+
+
+class Quantization(str, Enum):
+    FP32 = "fp32"
+    FP16 = "fp16"
+    BF16 = "bf16"
+    INT8 = "int8"
+    UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True)
+class FileDownloadSpecs:
+    download_url: str
+    file_handle: str
+    md5_hash: Optional[str] = field(default=None)
+
+
+@dataclass(frozen=True)
+class ServerEnvironmentRequirements:
+    cuda_device_cc: Version
+    cuda_device_name: str
+    driver_version: Optional[Version]
+    cuda_version: Optional[Version]
+    trt_version: Optional[Version]
+    os_version: Optional[str]
+
+    def __str__(self) -> str:
+        return (
+            f"Server(device={self.cuda_device_name}, cc={self.cuda_device_cc}, "
+            f"os={self.os_version}, driver={self.driver_version}, cu={self.cuda_version}, "
+            f"trt={self.trt_version})"
+        )
+
+
+@dataclass(frozen=True)
+class JetsonEnvironmentRequirements:
+    cuda_device_cc: Version
+    cuda_device_name: str
+    l4t_version: Version
+    jetson_product_name: Optional[str]
+    cuda_version: Optional[Version]
+    trt_version: Optional[Version]
+    driver_version: Optional[Version]
+
+    def __str__(self) -> str:
+        return (
+            f"Jetson(device={self.cuda_device_name}, cc={self.cuda_device_cc}, "
+            f"l4t={self.l4t_version}, product={self.jetson_product_name}, cu={self.cuda_version}, "
+            f"trt={self.trt_version}, driver={self.driver_version})"
+        )
+
+
+@dataclass(frozen=True)
+class TRTPackageDetails:
+    min_dynamic_batch_size: Optional[int] = field(default=None)
+    opt_dynamic_batch_size: Optional[int] = field(default=None)
+    max_dynamic_batch_size: Optional[int] = field(default=None)
+    same_cc_compatible: bool = field(default=False)
+    trt_forward_compatible: bool = field(default=False)
+    trt_lean_runtime_excluded: bool = field(default=False)
+
+    def __str__(self):
+        return (
+            f"TRTPackageDetails("
+            f"dynamic_batch=({self.min_dynamic_batch_size}/{self.opt_dynamic_batch_size}/{self.max_dynamic_batch_size}), "
+            f"same_cc_compatible={self.same_cc_compatible}, "
+            f"trt_forward_compatible={self.trt_forward_compatible}, "
+            f"trt_lean_runtime_excluded={self.trt_lean_runtime_excluded})"
+        )
+
+
+@dataclass(frozen=True)
+class ONNXPackageDetails:
+    opset: int
+    incompatible_providers: Optional[List[str]] = field(default=None)
+
+
+@dataclass(frozen=True)
+class TorchScriptPackageDetails:
+    supported_device_types: Set[str]
+    torch_version: Version
+    torch_vision_version: Optional[Version] = field(default=None)
+
+
+@dataclass(frozen=True)
+class ModelPackageMetadata:
+    package_id: str
+    backend: BackendType
+    package_artefacts: List[FileDownloadSpecs]
+    quantization: Optional[Quantization] = field(default=None)
+    dynamic_batch_size_supported: Optional[bool] = field(default=None)
+    static_batch_size: Optional[int] = field(default=None)
+    trt_package_details: Optional[TRTPackageDetails] = field(default=None)
+    onnx_package_details: Optional[ONNXPackageDetails] = field(default=None)
+    torch_script_package_details: Optional[TorchScriptPackageDetails] = field(
+        default=None
+    )
+    trusted_source: bool = field(default=False)
+    environment_requirements: Optional[
+        Union[ServerEnvironmentRequirements, JetsonEnvironmentRequirements]
+    ] = field(default=None)
+    model_features: Optional[dict] = field(default=None)
+
+    def get_summary(self) -> str:
+        return (
+            f"ModelPackageMetadata(package_id={self.package_id}, backend={self.backend.value}, quantization={self.quantization} "
+            f"dynamic_batch_size_supported={self.dynamic_batch_size_supported}, "
+            f"static_batch_size={self.static_batch_size}, trt_package_details={self.trt_package_details}, "
+            f"environment_requirements={self.environment_requirements}, model_features={self.model_features})"
+        )
+
+    def get_dynamic_batch_boundaries(self) -> Tuple[int, int]:
+        if not self.specifies_dynamic_batch_boundaries():
+            raise RuntimeError(
+                "Requested dynamic batch boundaries from model package that does not support dynamic batches."
+            )
+        values = []
+        if self.trt_package_details.min_dynamic_batch_size is not None:
+            values.append(self.trt_package_details.min_dynamic_batch_size)
+        if self.trt_package_details.opt_dynamic_batch_size is not None:
+            values.append(self.trt_package_details.opt_dynamic_batch_size)
+        if self.trt_package_details.max_dynamic_batch_size is not None:
+            values.append(self.trt_package_details.max_dynamic_batch_size)
+        return min(values), max(values)
+
+    def specifies_dynamic_batch_boundaries(self) -> bool:
+        if not self.dynamic_batch_size_supported:
+            return False
+        if self.trt_package_details is None:
+            return False
+        return (
+            self.trt_package_details.min_dynamic_batch_size is not None
+            or self.trt_package_details.opt_dynamic_batch_size is not None
+            or self.trt_package_details.max_dynamic_batch_size is not None
+        )
+
+
+class ModelDependency(BaseModel):
+    name: str
+    model_id: str
+    model_package_id: Optional[str] = Field(default=None)
+
+
+@dataclass(frozen=True)
+class ModelMetadata:
+    model_id: str
+    model_architecture: str
+    model_packages: List[ModelPackageMetadata]
+    task_type: Optional[str] = field(default=None)
+    model_variant: Optional[str] = field(default=None)
+    model_dependencies: Optional[List[ModelDependency]] = field(default=None)