Enrichments: Allow targeting a specific GPU ID (#19342)

frigate/embeddings/onnx/face_embedding.py

@@ -9,6 +9,7 @@ from frigate.const import MODEL_CACHE_DIR
 from frigate.log import redirect_output_to_logger
 from frigate.util.downloader import ModelDownloader
 
+from ...config import FaceRecognitionConfig
 from .base_embedding import BaseEmbedding
 from .runner import ONNXModelRunner
 
@@ -111,7 +112,7 @@ class FaceNetEmbedding(BaseEmbedding):
 
 
 class ArcfaceEmbedding(BaseEmbedding):
-    def __init__(self):
+    def __init__(self, config: FaceRecognitionConfig):
         super().__init__(
             model_name="facedet",
             model_file="arcface.onnx",
@@ -119,6 +120,7 @@ class ArcfaceEmbedding(BaseEmbedding):
                 "arcface.onnx": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/arcface.onnx",
             },
         )
+        self.config = config
         self.download_path = os.path.join(MODEL_CACHE_DIR, self.model_name)
         self.tokenizer = None
         self.feature_extractor = None
@@ -148,7 +150,7 @@ class ArcfaceEmbedding(BaseEmbedding):
 
             self.runner = ONNXModelRunner(
                 os.path.join(self.download_path, self.model_file),
-                "GPU",
+                device=self.config.device or "GPU",
             )
 
     def _preprocess_inputs(self, raw_inputs):