feat: support CUDA

.github/workflows/build.yml

@@ -5,7 +5,7 @@ on:
       - '*'
 env:
   ORT_STRATEGY: system
-  ORT_LIB_LOCATION: ./onnxruntime-linux-x64-1.14.1/lib/
+  ORT_LIB_LOCATION: ./onnxruntime-linux-x64-gpu-1.14.1/lib/
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -21,7 +21,7 @@ jobs:
     - name: Install requirement
       run: sudo apt-get update && sudo apt-get install -y protobuf-compiler
     - name: Download Onnx runtime
-      run: wget https://github.com/microsoft/onnxruntime/releases/download/v1.14.1/onnxruntime-linux-x64-1.14.1.tgz && tar -xzvf onnxruntime-linux-x64-1.14.1.tgz
+      run: wget https://github.com/microsoft/onnxruntime/releases/download/v1.14.1/onnxruntime-linux-x64-gpu-1.14.1.tgz && tar -xzvf onnxruntime-linux-x64-gpu-1.14.1.tgz
     - name: Build
       run: cargo build --bin clip-as-service-server --release
     - name: Upload artifact

Cargo.toml

@@ -10,13 +10,14 @@ license = "MIT"
 ndarray = "0.15.0"
 itertools = "0.10.5"
 num_cpus = "1.15.0"
-ort = { git = "https://github.com/pykeio/ort.git", default-features = false }
+ort = { git = "https://github.com/pykeio/ort.git", features = ["cuda"] }
 prost = "0.11.8"
 tokenizers = "0.13.2"
 tokio = { version = "1.26.0", features = ["macros", "rt-multi-thread"] }
 tonic = "0.8.3"
 image = "0.24.6"
 clap = { version = "4.2.4", features = ["derive"] }
+tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
 
 [build-dependencies]
 tonic-build = "0.8.4"
@@ -31,3 +32,9 @@ rustflags = ["-C", "linker-flavor=ld.lld"]
 
 [target.x86_64-pc-windows-msvc]
 rustflags = ["-C", "target-feature=+crt-static"]
+
+[profile.dev]
+rpath = true
+
+[profile.release]
+rpath = true

src/server.rs

@@ -186,13 +186,17 @@ impl Encoder for EncoderService {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tracing_subscriber::fmt::init();
+
     let args = Args::parse();
 
     let addr: &String = &args.listen;
+
     let environment =
         Arc::new(Environment::builder()
-            .with_name("clip")
-            .build().unwrap());
+                .with_name("clip")
+                .with_execution_providers([ExecutionProvider::cuda()])
+                .build().unwrap());
 
     let server = EncoderService::new(&environment, &args);