onnx: publish

koush · May 3, 2024 · 6fd6c7a · 6fd6c7a
1 parent d69ec69
commit 6fd6c7a
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 7 deletions.
diff --git a/plugins/onnx/.vscode/settings.json b/plugins/onnx/.vscode/settings.json
@@ -1,9 +1,9 @@
 
 {
     // docker installation
-    // "scrypted.debugHost": "koushik-ubuntu",
+    // "scrypted.debugHost": "koushik-ubuntuvm",
     // "scrypted.serverRoot": "/server",
-
+ 
     // pi local installation
     // "scrypted.debugHost": "192.168.2.119",
     // "scrypted.serverRoot": "/home/pi/.scrypted",

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
@@ -41,5 +41,5 @@
    "devDependencies": {
       "@scrypted/sdk": "file:../../sdk"
    },
-   "version": "0.1.79"
+   "version": "0.1.80"
 }
diff --git a/plugins/onnx/src/ort/__init__.py b/plugins/onnx/src/ort/__init__.py
@@ -3,6 +3,8 @@
 import asyncio
 from typing import Any, Tuple
 
+import sys
+import platform
 import numpy as np
 import onnxruntime
 import scrypted_sdk
@@ -65,7 +67,18 @@ def __init__(self, nativeId: str | None = None):
         print(onnxfile)
 
         try:
-            self.compiled_model = onnxruntime.InferenceSession(onnxfile)
+            sess_options = onnxruntime.SessionOptions()
+
+            providers: list[str] = []
+            if sys.platform == 'darwin':
+                providers.append("CoreMLExecutionProvider")
+
+            if sys.platform == 'linux' and platform.machine() == 'x86_64':
+                providers.append("CUDAExecutionProvider")
+
+            providers.append('CPUExecutionProvider')
+
+            self.compiled_model = onnxruntime.InferenceSession(onnxfile, sess_options=sess_options, providers=providers)
         except:
             import traceback
 

diff --git a/plugins/onnx/src/requirements.txt b/plugins/onnx/src/requirements.txt
@@ -1,4 +1,9 @@
-onnxruntime
+--extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
+# cuda 12
+onnxruntime-gpu; sys_platform == 'linux' and platform_machine == 'x86_64'
+# cpu and coreml execution provider
+onnxruntime; sys_platform != 'linux' or platform_machine != 'x86_64'
+# ort-nightly-gpu==1.17.3.dev20240409002
 
 # pillow-simd is available on x64 linux
 # pillow-simd confirmed not building with arm64 linux or apple silicon