blakeblackshear.frigate/frigate/util/model.py

"""Model Utils"""

import os

import onnxruntime as ort


def get_ort_providers(
    force_cpu: bool = False, openvino_device: str = "AUTO", requires_fp16: bool = False
) -> tuple[list[str], list[dict[str, any]]]:
    if force_cpu:
        return (["CPUExecutionProvider"], [{}])

    providers = ort.get_available_providers()
    options = []

    for provider in providers:
        if provider == "TensorrtExecutionProvider":
            os.makedirs("/config/model_cache/tensorrt/ort/trt-engines", exist_ok=True)

            if not requires_fp16 or os.environ.get("USE_FP_16", "True") != "False":
                options.append(
                    {
                        "trt_fp16_enable": requires_fp16,
                        "trt_timing_cache_enable": True,
                        "trt_engine_cache_enable": True,
                        "trt_timing_cache_path": "/config/model_cache/tensorrt/ort",
                        "trt_engine_cache_path": "/config/model_cache/tensorrt/ort/trt-engines",
                    }
                )
            else:
                options.append({})
        elif provider == "OpenVINOExecutionProvider":
            os.makedirs("/config/model_cache/openvino/ort", exist_ok=True)
            options.append(
                {
                    "cache_dir": "/config/model_cache/openvino/ort",
                    "device_type": openvino_device,
                }
            )
        else:
            options.append({})

    return (providers, options)