diff --git a/crates/diarize-rs/Cargo.toml b/crates/diarize-rs/Cargo.toml index 3985be5..4f40efc 100644 --- a/crates/diarize-rs/Cargo.toml +++ b/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"] } + diff --git a/crates/diarize-rs/src/lib.rs b/crates/diarize-rs/src/lib.rs index 957f353..60f361a 100644 --- a/crates/diarize-rs/src/lib.rs +++ b/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 = std::result::Result; @@ -23,12 +26,15 @@ pub struct Diarizer { impl Diarizer { pub fn new(model_path: impl AsRef) -> Result { 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,