onnx: fix provider detection

koush · Nov 17, 2024 · ea4922d · ea4922d
1 parent 2ec6760
commit ea4922d
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/plugins/onnx/package-lock.json b/plugins/onnx/package-lock.json
diff --git a/plugins/onnx/package.json b/plugins/onnx/package.json
@@ -42,5 +42,5 @@
    "devDependencies": {
       "@scrypted/sdk": "file:../../sdk"
    },
-   "version": "0.1.116"
+   "version": "0.1.117"
 }
diff --git a/plugins/onnx/src/ort/__init__.py b/plugins/onnx/src/ort/__init__.py
@@ -122,7 +122,14 @@ def executor_initializer():
             thread_name = threading.current_thread().name
             interpreter = compiled_models.pop()
             self.compiled_models[thread_name] = interpreter
-            self.provider = "CUDAExecutionProvider" if "CUDAExecutionProvider" in interpreter.get_providers() else "CPUExecutionProvider"
+            # remove CPUExecutionProider from providers
+            providers = interpreter.get_providers()
+            if not len(providers):
+                providers = ["CPUExecutionProvider"]
+            if "CPUExecutionProvider" in providers:
+                providers.remove("CPUExecutionProvider")
+            # join the remaining providers string
+            self.provider = ", ".join(providers)
             print('Runtime initialized on thread {}'.format(thread_name))
 
         self.executor = concurrent.futures.ThreadPoolExecutor(