diff --git a/frigate/test/test_runner_devices.py b/frigate/test/test_runner_devices.py index aaee9f7ad5..cd6c76c69d 100644 --- a/frigate/test/test_runner_devices.py +++ b/frigate/test/test_runner_devices.py @@ -321,6 +321,36 @@ class TestLighterANE(unittest.TestCase): ane.assert_not_called() +class TestTensorRTProvider(unittest.TestCase): + def _trt_options(self, env): + with ( + patch.object( + model_util.ort, + "get_available_providers", + return_value=["TensorrtExecutionProvider", "CPUExecutionProvider"], + ), + patch.object(model_util.os, "makedirs"), + patch.dict(os.environ, env, clear=False), + ): + if "USE_FP16" not in env: + os.environ.pop("USE_FP16", None) + providers, options = model_util.get_ort_providers(device="Tensorrt") + + return options[providers.index("TensorrtExecutionProvider")] + + def test_fp16_is_disabled_by_default(self): + self.assertFalse(self._trt_options({})["trt_fp16_enable"]) + + def test_use_fp16_enables_fp16(self): + for value in ("true", "True", "TRUE"): + with self.subTest(value=value): + options = self._trt_options({"USE_FP16": value}) + self.assertTrue(options["trt_fp16_enable"]) + + def test_use_fp16_false_keeps_fp16_disabled(self): + self.assertFalse(self._trt_options({"USE_FP16": "False"})["trt_fp16_enable"]) + + class TestMIGraphX(unittest.TestCase): """AMD's MIGraphX plugin provider, selected through its devices in the ROCm image.""" diff --git a/frigate/util/model.py b/frigate/util/model.py index 8d34d422a5..fbf21ae432 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -318,7 +318,6 @@ def post_process_yolox( def get_ort_providers( force_cpu: bool = False, device: str | None = "AUTO", - requires_fp16: bool = False, ) -> tuple[list[str], list[dict[str, Any]]]: if force_cpu: return ( @@ -357,8 +356,8 @@ def get_ort_providers( options.append( { "device_id": device_id, - "trt_fp16_enable": requires_fp16 - and os.environ.get("USE_FP16", "True") != "False", + "trt_fp16_enable": os.environ.get("USE_FP16", "False").lower() + == "true", "trt_timing_cache_enable": True, "trt_engine_cache_enable": True, "trt_timing_cache_path": os.path.join(