Skip to content

predictors

sleap_nn.export.predictors

Predictors for exported models.

Modules:

Name Description
base

Predictor base class for exported models.
onnx

ONNX Runtime predictor.
tensorrt

TensorRT predictor for exported models.

Classes:

Name Description
ExportPredictor

Base interface for exported model inference.
ONNXPredictor

ONNX Runtime inference with provider selection.
TensorRTPredictor

TensorRT inference for exported models.

Functions:

Name Description
detect_runtime

Auto-detect runtime from file extension or folder contents.
load_exported_model

Load an exported model and return a predictor instance.

ExportPredictor

Bases: ABC

Base interface for exported model inference.

Methods:

Name Description
benchmark

Benchmark inference latency and throughput.
predict

Run inference on a batch of images.
Source code in sleap_nn/export/predictors/base.py 
11
12
13
14
15
16
17
18
19
20
21
22
class ExportPredictor(ABC):
    """Base interface for exported model inference."""

    @abstractmethod
    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run inference on a batch of images."""

    @abstractmethod
    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference latency and throughput."""

benchmark(image, n_warmup=50, n_runs=200) abstractmethod

Benchmark inference latency and throughput.

Source code in sleap_nn/export/predictors/base.py 
18
19
20
21
22
@abstractmethod
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference latency and throughput."""

predict(image) abstractmethod

Run inference on a batch of images.

Source code in sleap_nn/export/predictors/base.py 
14
15
16
@abstractmethod
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run inference on a batch of images."""

ONNXPredictor

Bases: ExportPredictor

ONNX Runtime inference with provider selection.

Methods:

Name Description
__init__

Initialize ONNX predictor with execution providers.
benchmark

Benchmark inference latency and throughput.
predict

Run inference on a batch of images.
Source code in sleap_nn/export/predictors/onnx.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class ONNXPredictor(ExportPredictor):
    """ONNX Runtime inference with provider selection."""

    def __init__(
        self,
        model_path: str,
        device: str = "auto",
        providers: Optional[Iterable[str]] = None,
    ) -> None:
        """Initialize ONNX predictor with execution providers.

        Args:
            model_path: Path to the ONNX model file.
            device: Device for inference ("auto", "cpu", or "cuda").
            providers: ONNX Runtime execution providers. Auto-selected if None.
        """
        try:
            import onnxruntime as ort
        except ImportError as exc:
            raise ImportError(
                "onnxruntime is required for ONNXPredictor. Install with "
                "`pip install onnxruntime` or `onnxruntime-gpu`."
            ) from exc

        # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
        # This is required for onnxruntime-gpu to find the CUDA libraries
        if hasattr(ort, "preload_dlls"):
            ort.preload_dlls()

        self.ort = ort
        if providers is None:
            providers = _select_providers(device, ort.get_available_providers())

        self.session = ort.InferenceSession(model_path, providers=list(providers))
        input_info = self.session.get_inputs()[0]
        self.input_name = input_info.name
        self.input_type = input_info.type
        self.input_dtype = _onnx_type_to_numpy(self.input_type)
        self.output_names = [out.name for out in self.session.get_outputs()]

    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run inference on a batch of images."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        outputs = self.session.run(None, {self.input_name: image})
        return dict(zip(self.output_names, outputs))

    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference latency and throughput."""
        image = _as_numpy(image, expected_dtype=self.input_dtype)
        for _ in range(n_warmup):
            self.session.run(None, {self.input_name: image})

        times = []
        for _ in range(n_runs):
            start = time.perf_counter()
            self.session.run(None, {self.input_name: image})
            times.append(time.perf_counter() - start)

        times_ms = np.array(times) * 1000.0
        mean_ms = float(times_ms.mean())
        p50_ms = float(np.percentile(times_ms, 50))
        p95_ms = float(np.percentile(times_ms, 95))
        fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

        return {
            "latency_ms_mean": mean_ms,
            "latency_ms_p50": p50_ms,
            "latency_ms_p95": p95_ms,
            "fps": fps,
        }

__init__(model_path, device='auto', providers=None)

Initialize ONNX predictor with execution providers.

Parameters:

Name Type Description Default
model_path str

Path to the ONNX model file.

 required
device str

Device for inference ("auto", "cpu", or "cuda").

 'auto'
providers Optional[Iterable[str]]

ONNX Runtime execution providers. Auto-selected if None.

 None
Source code in sleap_nn/export/predictors/onnx.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def __init__(
    self,
    model_path: str,
    device: str = "auto",
    providers: Optional[Iterable[str]] = None,
) -> None:
    """Initialize ONNX predictor with execution providers.

    Args:
        model_path: Path to the ONNX model file.
        device: Device for inference ("auto", "cpu", or "cuda").
        providers: ONNX Runtime execution providers. Auto-selected if None.
    """
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "onnxruntime is required for ONNXPredictor. Install with "
            "`pip install onnxruntime` or `onnxruntime-gpu`."
        ) from exc

    # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
    # This is required for onnxruntime-gpu to find the CUDA libraries
    if hasattr(ort, "preload_dlls"):
        ort.preload_dlls()

    self.ort = ort
    if providers is None:
        providers = _select_providers(device, ort.get_available_providers())

    self.session = ort.InferenceSession(model_path, providers=list(providers))
    input_info = self.session.get_inputs()[0]
    self.input_name = input_info.name
    self.input_type = input_info.type
    self.input_dtype = _onnx_type_to_numpy(self.input_type)
    self.output_names = [out.name for out in self.session.get_outputs()]

benchmark(image, n_warmup=50, n_runs=200)

Benchmark inference latency and throughput.

Source code in sleap_nn/export/predictors/onnx.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference latency and throughput."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    for _ in range(n_warmup):
        self.session.run(None, {self.input_name: image})

    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        self.session.run(None, {self.input_name: image})
        times.append(time.perf_counter() - start)

    times_ms = np.array(times) * 1000.0
    mean_ms = float(times_ms.mean())
    p50_ms = float(np.percentile(times_ms, 50))
    p95_ms = float(np.percentile(times_ms, 95))
    fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0

    return {
        "latency_ms_mean": mean_ms,
        "latency_ms_p50": p50_ms,
        "latency_ms_p95": p95_ms,
        "fps": fps,
    }

predict(image)

Run inference on a batch of images.

Source code in sleap_nn/export/predictors/onnx.py 
53
54
55
56
57
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run inference on a batch of images."""
    image = _as_numpy(image, expected_dtype=self.input_dtype)
    outputs = self.session.run(None, {self.input_name: image})
    return dict(zip(self.output_names, outputs))

TensorRTPredictor

Bases: ExportPredictor

TensorRT inference for exported models.

This predictor loads a native TensorRT engine file (.trt) and provides inference capabilities using CUDA for high-throughput predictions.

Parameters:

Name Type Description Default
engine_path str | Path

Path to the TensorRT engine file (.trt).

 required
device str

Device to run inference on (only "cuda" supported for TRT).

 'cuda'
Example

predictor = TensorRTPredictor("model.trt") outputs = predictor.predict(images) # uint8 [B, C, H, W]

Methods:

Name Description
__init__

Initialize TensorRT predictor with a serialized engine.
benchmark

Benchmark inference performance.
predict

Run TensorRT inference.
Source code in sleap_nn/export/predictors/tensorrt.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
class TensorRTPredictor(ExportPredictor):
    """TensorRT inference for exported models.

    This predictor loads a native TensorRT engine file (.trt) and provides
    inference capabilities using CUDA for high-throughput predictions.

    Args:
        engine_path: Path to the TensorRT engine file (.trt).
        device: Device to run inference on (only "cuda" supported for TRT).

    Example:
        >>> predictor = TensorRTPredictor("model.trt")
        >>> outputs = predictor.predict(images)  # uint8 [B, C, H, W]
    """

    def __init__(
        self,
        engine_path: str | Path,
        device: str = "cuda",
    ) -> None:
        """Initialize TensorRT predictor with a serialized engine.

        Args:
            engine_path: Path to the TensorRT engine file.
            device: Device for inference ("cuda" or "auto"). TensorRT requires CUDA.
        """
        import tensorrt as trt

        if device not in ("cuda", "auto"):
            raise ValueError(f"TensorRT only supports CUDA devices, got: {device}")

        self.engine_path = Path(engine_path)
        if not self.engine_path.exists():
            raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")

        self.logger = trt.Logger(trt.Logger.WARNING)

        # Load engine
        with open(self.engine_path, "rb") as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())

        if self.engine is None:
            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

        # Create execution context
        self.context = self.engine.create_execution_context()

        # Get input/output info
        self.input_names: List[str] = []
        self.output_names: List[str] = []
        self.input_shapes: Dict[str, tuple] = {}
        self.output_shapes: Dict[str, tuple] = {}

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            shape = tuple(self.engine.get_tensor_shape(name))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.input_names.append(name)
                self.input_shapes[name] = shape
            else:
                self.output_names.append(name)
                self.output_shapes[name] = shape

        self.device = torch.device("cuda")

    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Run TensorRT inference.

        Args:
            image: Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

        Returns:
            Dict mapping output names to numpy arrays.
        """
        import tensorrt as trt

        # Convert input to torch tensor on GPU
        input_tensor = torch.from_numpy(image).to(self.device)

        # Check if engine expects uint8 or float32 and convert if needed
        input_name = self.input_names[0]
        expected_dtype = self.engine.get_tensor_dtype(input_name)
        if expected_dtype == trt.DataType.UINT8:
            # Engine expects uint8 - keep as uint8
            if input_tensor.dtype != torch.uint8:
                input_tensor = input_tensor.to(torch.uint8)
        else:
            # Engine expects float - convert uint8 to float32
            if input_tensor.dtype == torch.uint8:
                input_tensor = input_tensor.to(torch.float32)

        # Ensure contiguous memory
        input_tensor = input_tensor.contiguous()

        # Set input shape for dynamic dimensions
        input_name = self.input_names[0]
        self.context.set_input_shape(input_name, tuple(input_tensor.shape))

        # Allocate output tensors
        outputs: Dict[str, torch.Tensor] = {}
        bindings: Dict[str, int] = {}

        # Set input binding
        bindings[input_name] = input_tensor.data_ptr()

        # Allocate outputs
        for name in self.output_names:
            shape = self.context.get_tensor_shape(name)
            dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
            outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=self.device)
            bindings[name] = outputs[name].data_ptr()

        # Set tensor addresses
        for name, ptr in bindings.items():
            self.context.set_tensor_address(name, ptr)

        # Run inference
        stream = torch.cuda.current_stream().cuda_stream
        success = self.context.execute_async_v3(stream)
        torch.cuda.current_stream().synchronize()

        if not success:
            raise RuntimeError("TensorRT inference failed")

        # Convert outputs to numpy
        return {name: tensor.cpu().numpy() for name, tensor in outputs.items()}

    def benchmark(
        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
    ) -> Dict[str, float]:
        """Benchmark inference performance.

        Args:
            image: Input image(s) as numpy array [B, C, H, W].
            n_warmup: Number of warmup runs (not timed).
            n_runs: Number of timed runs.

        Returns:
            Dict with timing statistics:
                - mean_ms: Mean inference time in milliseconds
                - std_ms: Standard deviation of inference time
                - min_ms: Minimum inference time
                - max_ms: Maximum inference time
                - fps: Frames per second (based on mean time and batch size)
        """
        batch_size = image.shape[0]

        # Warmup
        for _ in range(n_warmup):
            _ = self.predict(image)

        # Timed runs
        times = []
        for _ in range(n_runs):
            start = time.perf_counter()
            _ = self.predict(image)
            times.append((time.perf_counter() - start) * 1000)

        times_arr = np.array(times)
        mean_ms = float(np.mean(times_arr))

        return {
            "mean_ms": mean_ms,
            "std_ms": float(np.std(times_arr)),
            "min_ms": float(np.min(times_arr)),
            "max_ms": float(np.max(times_arr)),
            "fps": (batch_size * 1000) / mean_ms if mean_ms > 0 else 0.0,
        }

    @staticmethod
    def _trt_dtype_to_torch(trt_dtype):
        """Convert TensorRT dtype to PyTorch dtype."""
        import tensorrt as trt

        mapping = {
            trt.DataType.FLOAT: torch.float32,
            trt.DataType.HALF: torch.float16,
            trt.DataType.INT32: torch.int32,
            trt.DataType.INT8: torch.int8,
            trt.DataType.BOOL: torch.bool,
        }
        return mapping.get(trt_dtype, torch.float32)

__init__(engine_path, device='cuda')

Initialize TensorRT predictor with a serialized engine.

Parameters:

Name Type Description Default
engine_path str | Path

Path to the TensorRT engine file.

 required
device str

Device for inference ("cuda" or "auto"). TensorRT requires CUDA.

 'cuda'
Source code in sleap_nn/export/predictors/tensorrt.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(
    self,
    engine_path: str | Path,
    device: str = "cuda",
) -> None:
    """Initialize TensorRT predictor with a serialized engine.

    Args:
        engine_path: Path to the TensorRT engine file.
        device: Device for inference ("cuda" or "auto"). TensorRT requires CUDA.
    """
    import tensorrt as trt

    if device not in ("cuda", "auto"):
        raise ValueError(f"TensorRT only supports CUDA devices, got: {device}")

    self.engine_path = Path(engine_path)
    if not self.engine_path.exists():
        raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")

    self.logger = trt.Logger(trt.Logger.WARNING)

    # Load engine
    with open(self.engine_path, "rb") as f:
        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())

    if self.engine is None:
        raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

    # Create execution context
    self.context = self.engine.create_execution_context()

    # Get input/output info
    self.input_names: List[str] = []
    self.output_names: List[str] = []
    self.input_shapes: Dict[str, tuple] = {}
    self.output_shapes: Dict[str, tuple] = {}

    for i in range(self.engine.num_io_tensors):
        name = self.engine.get_tensor_name(i)
        shape = tuple(self.engine.get_tensor_shape(name))
        if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
            self.input_names.append(name)
            self.input_shapes[name] = shape
        else:
            self.output_names.append(name)
            self.output_shapes[name] = shape

    self.device = torch.device("cuda")

benchmark(image, n_warmup=50, n_runs=200)

Benchmark inference performance.

Parameters:

Name Type Description Default
image ndarray

Input image(s) as numpy array [B, C, H, W].

 required
n_warmup int

Number of warmup runs (not timed).

 50
n_runs int

Number of timed runs.

 200

Returns:

Type Description
Dict[str, float]

Dict with timing statistics: - mean_ms: Mean inference time in milliseconds - std_ms: Standard deviation of inference time - min_ms: Minimum inference time - max_ms: Maximum inference time - fps: Frames per second (based on mean time and batch size)
Source code in sleap_nn/export/predictors/tensorrt.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def benchmark(
    self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
) -> Dict[str, float]:
    """Benchmark inference performance.

    Args:
        image: Input image(s) as numpy array [B, C, H, W].
        n_warmup: Number of warmup runs (not timed).
        n_runs: Number of timed runs.

    Returns:
        Dict with timing statistics:
            - mean_ms: Mean inference time in milliseconds
            - std_ms: Standard deviation of inference time
            - min_ms: Minimum inference time
            - max_ms: Maximum inference time
            - fps: Frames per second (based on mean time and batch size)
    """
    batch_size = image.shape[0]

    # Warmup
    for _ in range(n_warmup):
        _ = self.predict(image)

    # Timed runs
    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        _ = self.predict(image)
        times.append((time.perf_counter() - start) * 1000)

    times_arr = np.array(times)
    mean_ms = float(np.mean(times_arr))

    return {
        "mean_ms": mean_ms,
        "std_ms": float(np.std(times_arr)),
        "min_ms": float(np.min(times_arr)),
        "max_ms": float(np.max(times_arr)),
        "fps": (batch_size * 1000) / mean_ms if mean_ms > 0 else 0.0,
    }

predict(image)

Run TensorRT inference.

Parameters:

Name Type Description Default
image ndarray

Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

 required

Returns:

Type Description
Dict[str, ndarray]

Dict mapping output names to numpy arrays.
Source code in sleap_nn/export/predictors/tensorrt.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
    """Run TensorRT inference.

    Args:
        image: Input image(s) as numpy array [B, C, H, W] with uint8 dtype.

    Returns:
        Dict mapping output names to numpy arrays.
    """
    import tensorrt as trt

    # Convert input to torch tensor on GPU
    input_tensor = torch.from_numpy(image).to(self.device)

    # Check if engine expects uint8 or float32 and convert if needed
    input_name = self.input_names[0]
    expected_dtype = self.engine.get_tensor_dtype(input_name)
    if expected_dtype == trt.DataType.UINT8:
        # Engine expects uint8 - keep as uint8
        if input_tensor.dtype != torch.uint8:
            input_tensor = input_tensor.to(torch.uint8)
    else:
        # Engine expects float - convert uint8 to float32
        if input_tensor.dtype == torch.uint8:
            input_tensor = input_tensor.to(torch.float32)

    # Ensure contiguous memory
    input_tensor = input_tensor.contiguous()

    # Set input shape for dynamic dimensions
    input_name = self.input_names[0]
    self.context.set_input_shape(input_name, tuple(input_tensor.shape))

    # Allocate output tensors
    outputs: Dict[str, torch.Tensor] = {}
    bindings: Dict[str, int] = {}

    # Set input binding
    bindings[input_name] = input_tensor.data_ptr()

    # Allocate outputs
    for name in self.output_names:
        shape = self.context.get_tensor_shape(name)
        dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
        outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=self.device)
        bindings[name] = outputs[name].data_ptr()

    # Set tensor addresses
    for name, ptr in bindings.items():
        self.context.set_tensor_address(name, ptr)

    # Run inference
    stream = torch.cuda.current_stream().cuda_stream
    success = self.context.execute_async_v3(stream)
    torch.cuda.current_stream().synchronize()

    if not success:
        raise RuntimeError("TensorRT inference failed")

    # Convert outputs to numpy
    return {name: tensor.cpu().numpy() for name, tensor in outputs.items()}

detect_runtime(model_path)

Auto-detect runtime from file extension or folder contents.

Source code in sleap_nn/export/predictors/__init__.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def detect_runtime(model_path: str | Path) -> str:
    """Auto-detect runtime from file extension or folder contents."""
    model_path = Path(model_path)
    if model_path.is_dir():
        onnx_path = model_path / "exported" / "model.onnx"
        trt_path = model_path / "exported" / "model.trt"
        if trt_path.exists():
            return "tensorrt"
        if onnx_path.exists():
            return "onnx"
        raise ValueError(f"No exported model found in {model_path}")

    ext = model_path.suffix.lower()
    if ext == ".onnx":
        return "onnx"
    if ext in (".trt", ".engine"):
        return "tensorrt"

    raise ValueError(f"Unknown model format: {ext}")

load_exported_model(model_path, runtime='auto', device='auto', **kwargs)

Load an exported model and return a predictor instance.

Source code in sleap_nn/export/predictors/__init__.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def load_exported_model(
    model_path: str | Path,
    runtime: str = "auto",
    device: str = "auto",
    **kwargs,
) -> ExportPredictor:
    """Load an exported model and return a predictor instance."""
    if runtime == "auto":
        runtime = detect_runtime(model_path)

    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if runtime == "onnx":
        return ONNXPredictor(str(model_path), device=device, **kwargs)
    if runtime == "tensorrt":
        return TensorRTPredictor(str(model_path), device=device, **kwargs)

    raise ValueError(f"Unknown runtime: {runtime}")