mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-09-28 19:06:52 +03:00
Nick RogersandGitHub
c959df32c9
* Run ONNX models on a Mac's Neural Engine through lighter's plugin provider lighter's lighter.sh/ane device places an ONNX Runtime plugin execution provider in the container. When it is present, the ONNX session setup registers it once and opens sessions on its Neural Engine device, the same place CUDA, ROCm and OpenVINO are chosen, so the onnx detector (and any model that is not pinned to the CPU) runs there with no configuration. The hardware probe reports it as an onnx unit. * docs: hardware decode on an Apple Silicon Mac under lighter A community section on the video decoding page: lighter's lighter.sh/video device, hwaccel_args -c:v h264_v4l2m2m, and why the Raspberry Pi presets decode a single-stream camera in software. The detector docs link to it. * docs: set the lighter decoder per camera when codecs are mixed * docs: the ONNX detector on a Mac's Neural Engine under lighter * Format the Neural Engine provider setup * Fall back to the default providers when the Neural Engine cannot load a model * Apple Silicon ffmpeg presets for lighter's media engine, recommended when it is present
799 lines
29 KiB
Python
799 lines
29 KiB
Python
"""Base runner implementation for ONNX models."""
|
|
|
|
import logging
|
|
import os
|
|
import platform
|
|
import threading
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
from frigate.util.model import get_ort_providers
|
|
from frigate.util.rknn_converter import auto_convert_model, is_rknn_compatible
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Process-wide lock serializing all OpenVINO compile/inference calls
|
|
_OPENVINO_LOCK = threading.Lock()
|
|
|
|
# model file path -> (model type, device the runner actually loaded on); the
|
|
# embeddings maintainer folds this per enrichment for the stats endpoint.
|
|
# Models load on several threads (reindex, lazy first use), so writes and
|
|
# snapshots go through the lock.
|
|
loaded_devices: dict[str, tuple[str, str]] = {}
|
|
_loaded_devices_lock = threading.Lock()
|
|
|
|
|
|
def record_loaded_device(model_path: str, model_type: str, device: str) -> None:
|
|
"""Record the device a model loaded on, for the enrichment stats."""
|
|
with _loaded_devices_lock:
|
|
loaded_devices[model_path] = (model_type, device)
|
|
|
|
logger.info("Loaded %s model on %s", model_type, device)
|
|
|
|
|
|
def snapshot_loaded_devices() -> dict[str, tuple[str, str]]:
|
|
"""A copy that is safe to iterate while other threads load models."""
|
|
with _loaded_devices_lock:
|
|
return dict(loaded_devices)
|
|
|
|
|
|
_PROVIDER_LABELS = {
|
|
"CUDAExecutionProvider": "CUDA",
|
|
"TensorrtExecutionProvider": "TensorRT",
|
|
"MIGraphXExecutionProvider": "MIGraphX",
|
|
"OpenVINOExecutionProvider": "OpenVINO",
|
|
"CPUExecutionProvider": "CPU",
|
|
"LighterANE": "Neural Engine",
|
|
}
|
|
|
|
# lighter (https://github.com/fieldwork-ai/lighter) places an ONNX Runtime plugin
|
|
# execution provider in a container started with --device lighter.sh/ane=all,
|
|
# which runs models on a Mac's Neural Engine; LIGHTER_ANE_EP names where it is
|
|
LIGHTER_ANE_EP_NAME = "LighterANE"
|
|
LIGHTER_ANE_LIBRARY = "/usr/lib/lighter/liblighter_ane_ep.so"
|
|
|
|
|
|
def get_lighter_ane_devices() -> list[Any]:
|
|
"""Get the Neural Engine devices lighter's provider offers, registering it once.
|
|
|
|
Returns:
|
|
The provider's ONNX Runtime devices, or an empty list without lighter's device
|
|
"""
|
|
library = os.environ.get("LIGHTER_ANE_EP", LIGHTER_ANE_LIBRARY)
|
|
|
|
if not os.path.exists(library):
|
|
return []
|
|
|
|
devices = [d for d in ort.get_ep_devices() if d.ep_name == LIGHTER_ANE_EP_NAME]
|
|
|
|
if not devices:
|
|
try:
|
|
ort.register_execution_provider_library(LIGHTER_ANE_EP_NAME, library)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Failed to load the Neural Engine provider from {library}: {e}"
|
|
)
|
|
return []
|
|
|
|
devices = [d for d in ort.get_ep_devices() if d.ep_name == LIGHTER_ANE_EP_NAME]
|
|
|
|
return devices
|
|
|
|
|
|
def is_arm64_platform() -> bool:
|
|
"""Check if we're running on an ARM platform."""
|
|
machine = platform.machine().lower()
|
|
return machine in ("aarch64", "arm64", "armv8", "armv7l")
|
|
|
|
|
|
def get_ort_session_options(model_type: str | None = None) -> ort.SessionOptions | None:
|
|
"""Get ONNX Runtime session options with appropriate settings.
|
|
|
|
Args:
|
|
model_type: Model being loaded, used to pin its graph optimization level.
|
|
|
|
Returns:
|
|
SessionOptions with a pinned optimization level, or None for default settings.
|
|
"""
|
|
# Import here to avoid circular imports
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
if model_type == EnrichmentModelTypeEnum.jina_v2.value:
|
|
# below EXTENDED the CUDA EP returns an identical vector for every image,
|
|
# and ORT_ENABLE_ALL fails to build on CPU with a SimplifiedLayerNormFusion error
|
|
level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
|
|
elif model_type == EnrichmentModelTypeEnum.jina_v1.value:
|
|
# aggressive optimizations create or expect nodes that don't exist
|
|
level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
|
|
else:
|
|
return None
|
|
|
|
sess_options = ort.SessionOptions()
|
|
sess_options.graph_optimization_level = level
|
|
return sess_options
|
|
|
|
|
|
# Import OpenVINO only when needed to avoid circular dependencies
|
|
try:
|
|
import openvino as ov
|
|
except ImportError:
|
|
ov = None
|
|
|
|
|
|
def get_openvino_available_devices() -> list[str]:
|
|
"""Get available OpenVINO devices without using ONNX Runtime.
|
|
|
|
Returns:
|
|
List of available OpenVINO device names (e.g., ['CPU', 'GPU', 'MYRIAD'])
|
|
"""
|
|
if ov is None:
|
|
logger.debug("OpenVINO is not available")
|
|
return []
|
|
|
|
try:
|
|
core = ov.Core()
|
|
available_devices = core.available_devices
|
|
logger.debug(f"OpenVINO available devices: {available_devices}")
|
|
return available_devices
|
|
except Exception as e:
|
|
logger.warning(f"Failed to get OpenVINO available devices: {e}")
|
|
return []
|
|
|
|
|
|
def is_openvino_gpu_npu_available() -> bool:
|
|
"""Check if OpenVINO GPU or NPU devices are available.
|
|
|
|
Returns:
|
|
True if GPU or NPU devices are available, False otherwise
|
|
"""
|
|
available_devices = get_openvino_available_devices()
|
|
# Check for GPU, NPU, or other acceleration devices (excluding CPU)
|
|
acceleration_devices = ["GPU", "MYRIAD", "NPU", "GNA", "HDDL"]
|
|
return any(
|
|
avail_dev == accel_dev or avail_dev.startswith(accel_dev + ".")
|
|
for avail_dev in available_devices
|
|
for accel_dev in acceleration_devices
|
|
)
|
|
|
|
|
|
class BaseModelRunner(ABC):
|
|
"""Abstract base class for model runners."""
|
|
|
|
def __init__(self, model_path: str, device: str, **kwargs):
|
|
self.model_path = model_path
|
|
self.device = device
|
|
|
|
@abstractmethod
|
|
def get_input_names(self) -> list[str]:
|
|
"""Get input names for the model."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_input_width(self) -> int:
|
|
"""Get the input width of the model."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def run(self, input: dict[str, Any]) -> Any | None:
|
|
"""Run inference with the model."""
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def device_name(self) -> str:
|
|
"""Short label of the device the model actually loaded on."""
|
|
pass
|
|
|
|
|
|
class ONNXModelRunner(BaseModelRunner):
|
|
"""Run ONNX models using ONNX Runtime."""
|
|
|
|
@staticmethod
|
|
def is_migraphx_complex_model(model_type: str) -> bool:
|
|
# Import here to avoid circular imports
|
|
from frigate.detectors.detector_config import ModelTypeEnum
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
return model_type in [
|
|
EnrichmentModelTypeEnum.paddleocr.value,
|
|
EnrichmentModelTypeEnum.jina_v2.value,
|
|
ModelTypeEnum.rfdetr.value,
|
|
ModelTypeEnum.dfine.value,
|
|
]
|
|
|
|
@staticmethod
|
|
def is_concurrent_model(model_type: str | None) -> bool:
|
|
"""Check if model requires thread locking for concurrent inference.
|
|
|
|
Some models (like JinaV2) share one runner between text and vision embeddings
|
|
called from different threads, requiring thread synchronization.
|
|
"""
|
|
if not model_type:
|
|
return False
|
|
|
|
# Import here to avoid circular imports
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
return model_type == EnrichmentModelTypeEnum.jina_v2.value
|
|
|
|
def __init__(self, ort: ort.InferenceSession, model_type: str | None = None):
|
|
self.ort = ort
|
|
self.model_type = model_type
|
|
|
|
# Thread lock to prevent concurrent inference (needed for JinaV2 which shares
|
|
# one runner between text and vision embeddings called from different threads)
|
|
if self.is_concurrent_model(model_type):
|
|
self._inference_lock = threading.Lock()
|
|
else:
|
|
self._inference_lock = None
|
|
|
|
def get_input_names(self) -> list[str]:
|
|
return [input.name for input in self.ort.get_inputs()]
|
|
|
|
def get_input_width(self) -> int:
|
|
"""Get the input width of the model."""
|
|
return self.ort.get_inputs()[0].shape[3]
|
|
|
|
def run(self, input: dict[str, Any]) -> Any | None:
|
|
if self._inference_lock:
|
|
with self._inference_lock:
|
|
return self.ort.run(None, input)
|
|
|
|
return self.ort.run(None, input)
|
|
|
|
@property
|
|
def device_name(self) -> str:
|
|
providers = self.ort.get_providers()
|
|
|
|
if not providers:
|
|
return "CPU"
|
|
|
|
provider = providers[0]
|
|
return _PROVIDER_LABELS.get(
|
|
provider, provider.removesuffix("ExecutionProvider")
|
|
)
|
|
|
|
|
|
class CudaGraphRunner(BaseModelRunner):
|
|
"""Encapsulates CUDA Graph capture and replay using ONNX Runtime IOBinding.
|
|
|
|
This runner assumes a single tensor input and binds all model outputs.
|
|
|
|
NOTE: CUDA Graphs limit supported model operations, so they are not usable
|
|
for more complex models like CLIP or PaddleOCR.
|
|
"""
|
|
|
|
@staticmethod
|
|
def is_model_supported(model_type: str) -> bool:
|
|
# Import here to avoid circular imports
|
|
from frigate.detectors.detector_config import ModelTypeEnum
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
return model_type not in [
|
|
ModelTypeEnum.yolonas.value,
|
|
ModelTypeEnum.dfine.value,
|
|
EnrichmentModelTypeEnum.paddleocr.value,
|
|
EnrichmentModelTypeEnum.jina_v1.value,
|
|
EnrichmentModelTypeEnum.jina_v2.value,
|
|
EnrichmentModelTypeEnum.yolov9_license_plate.value,
|
|
]
|
|
|
|
# ORT performs two regular runs before it starts capturing, but on some
|
|
# driver / cuDNN combinations the arena still has to extend on the run that
|
|
# captures, and cudaMalloc is not allowed during capture. Running with
|
|
# capture disabled first keeps those allocations outside of the capture.
|
|
GRAPH_FREE_WARMUP_RUNS = 2
|
|
|
|
def __init__(self, session: ort.InferenceSession, cuda_device_id: int):
|
|
self._session = session
|
|
self._cuda_device_id = cuda_device_id
|
|
self._prepared = False
|
|
self._io_binding: ort.IOBinding | None = None
|
|
self._input_name: str | None = None
|
|
self._output_names: list[str] | None = None
|
|
self._input_ortvalue: ort.OrtValue | None = None
|
|
|
|
def get_input_names(self) -> list[str]:
|
|
"""Get input names for the model."""
|
|
return [input.name for input in self._session.get_inputs()]
|
|
|
|
def get_input_width(self) -> int:
|
|
"""Get the input width of the model."""
|
|
return self._session.get_inputs()[0].shape[3]
|
|
|
|
def _prepare(self, input_name: str, tensor_input: np.ndarray) -> None:
|
|
"""Bind CUDA buffers and warm the session up with capture disabled."""
|
|
self._io_binding = self._session.io_binding()
|
|
self._input_name = input_name
|
|
self._output_names = [o.name for o in self._session.get_outputs()]
|
|
|
|
self._input_ortvalue = ort.OrtValue.ortvalue_from_numpy(
|
|
tensor_input, "cuda", self._cuda_device_id
|
|
)
|
|
self._io_binding.bind_ortvalue_input(self._input_name, self._input_ortvalue)
|
|
|
|
for name in self._output_names:
|
|
# Bind outputs to CUDA and allow ORT to allocate appropriately
|
|
self._io_binding.bind_output(name, "cuda", self._cuda_device_id)
|
|
|
|
# gpu_graph_id -1 disables capture and replay for the run
|
|
warmup_options = ort.RunOptions()
|
|
warmup_options.add_run_config_entry("gpu_graph_id", "-1")
|
|
|
|
for _ in range(self.GRAPH_FREE_WARMUP_RUNS):
|
|
self._session.run_with_iobinding(self._io_binding, warmup_options)
|
|
|
|
self._prepared = True
|
|
|
|
def run(self, input: dict[str, Any]):
|
|
# Extract the single tensor input (assuming one input)
|
|
input_name = list(input.keys())[0]
|
|
tensor_input = np.ascontiguousarray(input[input_name])
|
|
|
|
if not self._prepared:
|
|
self._prepare(input_name, tensor_input)
|
|
else:
|
|
# Replay using updated input
|
|
self._input_ortvalue.update_inplace(tensor_input)
|
|
|
|
ro = ort.RunOptions()
|
|
self._session.run_with_iobinding(self._io_binding, ro)
|
|
return self._io_binding.copy_outputs_to_cpu()
|
|
|
|
@property
|
|
def device_name(self) -> str:
|
|
return "CUDA"
|
|
|
|
|
|
class OpenVINOModelRunner(BaseModelRunner):
|
|
"""OpenVINO model runner that handles inference efficiently."""
|
|
|
|
@staticmethod
|
|
def is_complex_model(model_type: str) -> bool:
|
|
# Import here to avoid circular imports
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
return model_type in [
|
|
EnrichmentModelTypeEnum.paddleocr.value,
|
|
EnrichmentModelTypeEnum.jina_v2.value,
|
|
]
|
|
|
|
@staticmethod
|
|
def is_model_npu_supported(model_type: str) -> bool:
|
|
# Import here to avoid circular imports
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
return model_type not in [
|
|
EnrichmentModelTypeEnum.paddleocr.value,
|
|
EnrichmentModelTypeEnum.jina_v1.value,
|
|
EnrichmentModelTypeEnum.jina_v2.value,
|
|
EnrichmentModelTypeEnum.arcface.value,
|
|
]
|
|
|
|
@staticmethod
|
|
def is_detection_model(model_type: str) -> bool:
|
|
# Import here to avoid circular imports
|
|
from frigate.detectors.detector_config import ModelTypeEnum
|
|
|
|
return model_type in [m.value for m in ModelTypeEnum]
|
|
|
|
def __init__(self, model_path: str, device: str, model_type: str, **kwargs):
|
|
self.model_path = model_path
|
|
self.device = device
|
|
self.model_type = model_type
|
|
|
|
if device == "NPU" and not OpenVINOModelRunner.is_model_npu_supported(
|
|
model_type
|
|
):
|
|
logger.warning(
|
|
f"OpenVINO model {model_type} is not supported on NPU, using GPU instead"
|
|
)
|
|
device = "GPU"
|
|
|
|
self.complex_model = OpenVINOModelRunner.is_complex_model(model_type)
|
|
|
|
if not os.path.isfile(model_path):
|
|
raise FileNotFoundError(f"OpenVINO model file {model_path} not found.")
|
|
|
|
if ov is None:
|
|
raise ImportError(
|
|
"OpenVINO is not available. Please install openvino package."
|
|
)
|
|
|
|
self.ov_core = ov.Core()
|
|
|
|
# Apply performance optimization
|
|
self.ov_core.set_property(device, {"PERF_COUNT": "NO"})
|
|
|
|
if device in ["GPU", "AUTO", "NPU"]:
|
|
self.ov_core.set_property(device, {"PERFORMANCE_HINT": "LATENCY"})
|
|
|
|
if device in ["GPU", "AUTO"]:
|
|
try:
|
|
self.ov_core.set_property("GPU", {"GPU_QUEUE_THROTTLE": "LOW"})
|
|
except Exception as e:
|
|
logger.debug(f"GPU_QUEUE_THROTTLE not supported: {e}")
|
|
|
|
# Some keys must be passed as compile-time config so that it can be caught
|
|
compile_config = {}
|
|
|
|
if device == "NPU" and OpenVINOModelRunner.is_detection_model(model_type):
|
|
compile_config["NPU_TURBO"] = "YES"
|
|
|
|
self.compiled_device = device
|
|
|
|
# Compile model under the shared lock
|
|
with _OPENVINO_LOCK:
|
|
try:
|
|
self.compiled_model = self.ov_core.compile_model(
|
|
model=model_path, device_name=device, config=compile_config
|
|
)
|
|
except RuntimeError as e:
|
|
if not compile_config:
|
|
raise
|
|
|
|
logger.debug(
|
|
f"Failed to compile with {compile_config}, retrying without: {e}"
|
|
)
|
|
self.compiled_model = self.ov_core.compile_model(
|
|
model=model_path, device_name=device
|
|
)
|
|
|
|
# Create reusable inference request
|
|
self.infer_request = self.compiled_model.create_infer_request()
|
|
|
|
self.input_tensor: ov.Tensor | None = None
|
|
|
|
if not self.complex_model:
|
|
try:
|
|
input_shape = self.compiled_model.inputs[0].get_shape()
|
|
input_element_type = self.compiled_model.inputs[0].get_element_type()
|
|
self.input_tensor = ov.Tensor(input_element_type, input_shape)
|
|
except RuntimeError:
|
|
# model is complex and has dynamic shape
|
|
pass
|
|
|
|
@property
|
|
def device_name(self) -> str:
|
|
device = self.compiled_device
|
|
|
|
if device == "AUTO":
|
|
try:
|
|
resolved = self.compiled_model.get_property("EXECUTION_DEVICES")
|
|
|
|
if resolved:
|
|
device = ",".join(str(d) for d in resolved)
|
|
except Exception:
|
|
# older OpenVINO builds do not expose the property
|
|
pass
|
|
|
|
return f"OpenVINO {device}"
|
|
|
|
def get_input_names(self) -> list[str]:
|
|
"""Get input names for the model."""
|
|
return [input.get_any_name() for input in self.compiled_model.inputs]
|
|
|
|
def get_input_width(self) -> int:
|
|
"""Get the input width of the model."""
|
|
input_info = self.compiled_model.inputs
|
|
first_input = input_info[0]
|
|
|
|
try:
|
|
partial_shape = first_input.get_partial_shape()
|
|
# width dimension
|
|
if len(partial_shape) >= 4 and partial_shape[3].is_static:
|
|
return partial_shape[3].get_length()
|
|
|
|
# If width is dynamic or we can't determine it
|
|
return -1
|
|
except Exception:
|
|
try:
|
|
# gemini says some ov versions might still allow this
|
|
input_shape = first_input.shape
|
|
return input_shape[3] if len(input_shape) >= 4 else -1
|
|
except Exception:
|
|
return -1
|
|
|
|
def run(self, inputs: dict[str, Any]) -> list[np.ndarray]:
|
|
"""Run inference with the model.
|
|
|
|
Args:
|
|
inputs: Dictionary mapping input names to input data
|
|
|
|
Returns:
|
|
List of output tensors
|
|
"""
|
|
# Shared lock serializes inference across every OpenVINO runner in this
|
|
# process — both the shared-runner JinaV2 case (genai text thread +
|
|
# embeddings vision thread) and distinct runners running on separate
|
|
# threads (e.g. the ArcFace face-model build vs the LPR detector).
|
|
with _OPENVINO_LOCK:
|
|
from frigate.embeddings.types import EnrichmentModelTypeEnum
|
|
|
|
if self.model_type in [EnrichmentModelTypeEnum.arcface.value]:
|
|
# For face recognition models, create a fresh infer_request
|
|
# for each inference to avoid state pollution that causes incorrect results.
|
|
self.infer_request = self.compiled_model.create_infer_request()
|
|
|
|
# Handle single input case for backward compatibility
|
|
if (
|
|
len(inputs) == 1
|
|
and len(self.compiled_model.inputs) == 1
|
|
and self.input_tensor is not None
|
|
):
|
|
# Single input case - use the pre-allocated tensor for efficiency
|
|
input_data = list(inputs.values())[0]
|
|
np.copyto(self.input_tensor.data, input_data)
|
|
self.infer_request.infer(self.input_tensor)
|
|
else:
|
|
if self.complex_model:
|
|
try:
|
|
# This ensures the model starts with a clean state for each sequence
|
|
# Important for RNN models like PaddleOCR recognition
|
|
self.infer_request.reset_state()
|
|
except Exception:
|
|
# this will raise an exception for models with AUTO set as the device
|
|
pass
|
|
|
|
# Multiple inputs case - set each input by name
|
|
for input_name, input_data in inputs.items():
|
|
# Find the input by name and its index
|
|
input_port = None
|
|
input_index = None
|
|
for idx, port in enumerate(self.compiled_model.inputs):
|
|
if port.get_any_name() == input_name:
|
|
input_port = port
|
|
input_index = idx
|
|
break
|
|
|
|
if input_port is None:
|
|
raise ValueError(f"Input '{input_name}' not found in model")
|
|
|
|
# Create tensor with the correct element type
|
|
input_element_type = input_port.get_element_type()
|
|
|
|
# Ensure input data matches the expected dtype to prevent type mismatches
|
|
# that can occur with models like Jina-CLIP v2 running on OpenVINO
|
|
expected_dtype = input_element_type.to_dtype()
|
|
if input_data.dtype != expected_dtype:
|
|
logger.debug(
|
|
f"Converting input '{input_name}' from {input_data.dtype} to {expected_dtype}"
|
|
)
|
|
input_data = input_data.astype(expected_dtype)
|
|
|
|
input_tensor = ov.Tensor(input_element_type, input_data.shape)
|
|
np.copyto(input_tensor.data, input_data)
|
|
|
|
# Set the input tensor for the specific port index
|
|
self.infer_request.set_input_tensor(input_index, input_tensor)
|
|
|
|
# Run inference
|
|
try:
|
|
self.infer_request.infer()
|
|
except Exception as e:
|
|
logger.error(f"Error during OpenVINO inference: {e}")
|
|
return []
|
|
|
|
# Get all output tensors
|
|
outputs = []
|
|
for i in range(len(self.compiled_model.outputs)):
|
|
outputs.append(self.infer_request.get_output_tensor(i).data)
|
|
|
|
return outputs
|
|
|
|
|
|
class RKNNModelRunner(BaseModelRunner):
|
|
"""Run RKNN models for embeddings."""
|
|
|
|
def __init__(self, model_path: str, model_type: str = None, core_mask: int = 0):
|
|
self.model_path = model_path
|
|
self.model_type = model_type
|
|
self.core_mask = core_mask
|
|
self.rknn = None
|
|
self._load_model()
|
|
|
|
def _load_model(self):
|
|
"""Load the RKNN model."""
|
|
try:
|
|
from rknnlite.api import RKNNLite
|
|
|
|
self.rknn = RKNNLite(verbose=False)
|
|
|
|
if self.rknn.load_rknn(self.model_path) != 0:
|
|
logger.error(f"Failed to load RKNN model: {self.model_path}")
|
|
raise RuntimeError("Failed to load RKNN model")
|
|
|
|
if self.rknn.init_runtime(core_mask=self.core_mask) != 0:
|
|
logger.error("Failed to initialize RKNN runtime")
|
|
raise RuntimeError("Failed to initialize RKNN runtime")
|
|
|
|
logger.info(f"Successfully loaded RKNN model: {self.model_path}")
|
|
|
|
except ImportError:
|
|
logger.error("RKNN Lite not available")
|
|
raise ImportError("RKNN Lite not available") from None
|
|
except Exception as e:
|
|
logger.error(f"Error loading RKNN model: {e}")
|
|
raise
|
|
|
|
@property
|
|
def device_name(self) -> str:
|
|
return "RKNN"
|
|
|
|
def get_input_names(self) -> list[str]:
|
|
"""Get input names for the model."""
|
|
# For detection models, we typically use "input" as the default input name
|
|
# For CLIP models, we need to determine the model type from the path
|
|
model_name = os.path.basename(self.model_path).lower()
|
|
|
|
if "vision" in model_name:
|
|
return ["pixel_values"]
|
|
elif "arcface" in model_name:
|
|
return ["data"]
|
|
else:
|
|
# Default fallback - try to infer from model type
|
|
if self.model_type and "jina-clip" in self.model_type:
|
|
if "vision" in self.model_type:
|
|
return ["pixel_values"]
|
|
|
|
# Generic fallback
|
|
return ["input"]
|
|
|
|
def get_input_width(self) -> int:
|
|
"""Get the input width of the model."""
|
|
# For CLIP vision models, this is typically 224
|
|
model_name = os.path.basename(self.model_path).lower()
|
|
if "vision" in model_name:
|
|
return 224 # CLIP V1 uses 224x224
|
|
elif "arcface" in model_name:
|
|
return 112
|
|
# For detection models, we can't easily determine this from the RKNN model
|
|
# The calling code should provide this information
|
|
return -1
|
|
|
|
def run(self, inputs: dict[str, Any]) -> Any:
|
|
"""Run inference with the RKNN model."""
|
|
if not self.rknn:
|
|
raise RuntimeError("RKNN model not loaded")
|
|
|
|
try:
|
|
input_names = self.get_input_names()
|
|
rknn_inputs = []
|
|
|
|
for name in input_names:
|
|
if name in inputs:
|
|
if name == "pixel_values":
|
|
# RKNN expects NHWC format, but ONNX typically provides NCHW
|
|
# Transpose from [batch, channels, height, width] to [batch, height, width, channels]
|
|
pixel_data = inputs[name]
|
|
if len(pixel_data.shape) == 4 and pixel_data.shape[1] == 3:
|
|
# Transpose from NCHW to NHWC
|
|
pixel_data = np.transpose(pixel_data, (0, 2, 3, 1))
|
|
rknn_inputs.append(pixel_data)
|
|
elif name == "data":
|
|
# ArcFace: undo Python normalisation to uint8 [0,255]
|
|
# RKNN runtime applies mean=127.5/std=127.5 internally before first layer
|
|
face_data = inputs[name]
|
|
if len(face_data.shape) == 4 and face_data.shape[1] == 3:
|
|
# Transpose from NCHW to NHWC
|
|
face_data = np.transpose(face_data, (0, 2, 3, 1))
|
|
face_data = (
|
|
((face_data + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
|
|
)
|
|
rknn_inputs.append(face_data)
|
|
else:
|
|
rknn_inputs.append(inputs[name])
|
|
|
|
outputs = self.rknn.inference(inputs=rknn_inputs)
|
|
return outputs
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during RKNN inference: {e}")
|
|
raise
|
|
|
|
def __del__(self):
|
|
"""Cleanup when the runner is destroyed."""
|
|
if self.rknn:
|
|
try:
|
|
self.rknn.release()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _record_runner(
|
|
model_path: str, model_type: str, runner: BaseModelRunner
|
|
) -> BaseModelRunner:
|
|
"""Record the device a freshly loaded runner ended up on."""
|
|
record_loaded_device(model_path, model_type, runner.device_name)
|
|
return runner
|
|
|
|
|
|
def get_optimized_runner(
|
|
model_path: str, device: str | None, model_type: str, **kwargs
|
|
) -> BaseModelRunner:
|
|
"""Get an optimized runner for the hardware."""
|
|
device = device or "AUTO"
|
|
|
|
if device != "CPU" and is_rknn_compatible(model_path):
|
|
rknn_path = auto_convert_model(model_path)
|
|
|
|
if rknn_path:
|
|
return _record_runner(model_path, model_type, RKNNModelRunner(rknn_path))
|
|
|
|
if device != "CPU" and (ane_devices := get_lighter_ane_devices()):
|
|
sess_options = get_ort_session_options(model_type) or ort.SessionOptions()
|
|
sess_options.add_provider_for_devices(ane_devices, {})
|
|
|
|
try:
|
|
session = ort.InferenceSession(model_path, sess_options=sess_options)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Failed to load {model_path} on the Neural Engine, using the default providers: {e}"
|
|
)
|
|
else:
|
|
return _record_runner(
|
|
model_path,
|
|
model_type,
|
|
ONNXModelRunner(session, model_type=model_type),
|
|
)
|
|
|
|
providers, options = get_ort_providers(device == "CPU", device, **kwargs)
|
|
|
|
if providers[0] == "CPUExecutionProvider":
|
|
# In the default image, ONNXRuntime is used so we will only get CPUExecutionProvider
|
|
# In other images we will get CUDA / ROCm which are preferred over OpenVINO
|
|
# There is currently no way to prioritize OpenVINO over CUDA / ROCm in these images
|
|
if device != "CPU" and is_openvino_gpu_npu_available():
|
|
return _record_runner(
|
|
model_path,
|
|
model_type,
|
|
OpenVINOModelRunner(model_path, device, model_type, **kwargs),
|
|
)
|
|
|
|
if (
|
|
CudaGraphRunner.is_model_supported(model_type)
|
|
and providers[0] == "CUDAExecutionProvider"
|
|
):
|
|
options[0] = {
|
|
**options[0],
|
|
"enable_cuda_graph": True,
|
|
}
|
|
return _record_runner(
|
|
model_path,
|
|
model_type,
|
|
CudaGraphRunner(
|
|
ort.InferenceSession(
|
|
model_path,
|
|
providers=providers,
|
|
provider_options=options,
|
|
),
|
|
options[0]["device_id"],
|
|
),
|
|
)
|
|
|
|
if (
|
|
providers
|
|
and providers[0] == "MIGraphXExecutionProvider"
|
|
and ONNXModelRunner.is_migraphx_complex_model(model_type)
|
|
):
|
|
# Don't use MIGraphX for models that are not supported
|
|
providers.pop(0)
|
|
options.pop(0)
|
|
|
|
return _record_runner(
|
|
model_path,
|
|
model_type,
|
|
ONNXModelRunner(
|
|
ort.InferenceSession(
|
|
model_path,
|
|
sess_options=get_ort_session_options(model_type),
|
|
providers=providers,
|
|
provider_options=options,
|
|
),
|
|
model_type=model_type,
|
|
),
|
|
)
|