Run Sortformer diarization on GPU via DirectML on Windows

crates/diarize-rs/Cargo.toml

@@ -11,7 +11,13 @@ name = "diarize_rs"
 path = "src/lib.rs"
 
 [dependencies]
-parakeet-rs = { version = "0.3.2", features = ["sortformer"] }
 serde = { version = "1", features = ["derive"] }
 thiserror = "2"
 
+# The DirectML execution provider is Windows-only; other platforms stay on CPU.
+[target.'cfg(windows)'.dependencies]
+parakeet-rs = { version = "0.3.6", features = ["sortformer", "directml"] }
+
+[target.'cfg(not(windows))'.dependencies]
+parakeet-rs = { version = "0.3.6", features = ["sortformer"] }
+

crates/diarize-rs/src/lib.rs

@@ -1,6 +1,9 @@
 use std::path::{Path, PathBuf};
 
 use parakeet_rs::sortformer::{DiarizationConfig, Sortformer};
+use parakeet_rs::ExecutionConfig;
+#[cfg(windows)]
+use parakeet_rs::ExecutionProvider;
 use serde::{Deserialize, Serialize};
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -23,12 +26,15 @@ pub struct Diarizer {
 impl Diarizer {
     pub fn new(model_path: impl AsRef<Path>) -> Result<Self> {
         let path = model_path.as_ref();
-        let sortformer =
-            Sortformer::with_config(path.to_string_lossy().as_ref(), None, DiarizationConfig::callhome())
-                .map_err(|source| Error::Init {
-                    path: path.to_path_buf(),
-                    source: Box::new(source),
-                })?;
+        let sortformer = Sortformer::with_config(
+            path.to_string_lossy().as_ref(),
+            Some(execution_config()),
+            DiarizationConfig::callhome(),
+        )
+        .map_err(|source| Error::Init {
+            path: path.to_path_buf(),
+            source: Box::new(source),
+        })?;
         Ok(Self { sortformer })
     }
 
@@ -54,6 +60,24 @@ impl Diarizer {
     }
 }
 
+/// On Windows, run the Sortformer graph on GPU via DirectML (ort registers a
+/// CPU fallback automatically if the provider is unavailable); set
+/// SONA_DIARIZE_EP=cpu to opt out. Elsewhere, stay on CPU but use all
+/// available cores instead of the crate default of 4.
+fn execution_config() -> ExecutionConfig {
+    let threads = std::thread::available_parallelism()
+        .map(|count| count.get())
+        .unwrap_or(4);
+    let config = ExecutionConfig::new().with_intra_threads(threads);
+    #[cfg(windows)]
+    let config = if std::env::var("SONA_DIARIZE_EP").as_deref() == Ok("cpu") {
+        config
+    } else {
+        config.with_execution_provider(ExecutionProvider::DirectML)
+    };
+    config
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Segment {
     pub start: f64,