From ba647a101b63614ff56acf8664ad81cebbdfbdcc Mon Sep 17 00:00:00 2001 From: Neko Ayaka Date: Wed, 2 Jul 2025 21:45:34 +0800 Subject: [PATCH] feat(stage-tamagotchi): use coreml (metal) for macOS, directml, or cuda for windows, cuda for linux --- Cargo.lock | 216 ++++++++++++++++-- apps/stage-tamagotchi/src-tauri/Cargo.toml | 5 + .../src/app/models/silero_vad/mod.rs | 173 ++++++++++---- .../src/pages/settings/modules/hearing.vue | 3 +- cspell.config.yaml | 5 + 5 files changed, 337 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 606ed6a49..cca115c6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -191,8 +191,10 @@ dependencies = [ "crossbeam-channel", "hf-hub", "log", + "ndarray", "objc2 0.6.1", "objc2-foundation 0.3.1", + "ort", "rubato", "serde", "serde_json", @@ -295,6 +297,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64ct" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba" + [[package]] name = "bindgen_cuda" version = "0.1.5" @@ -711,7 +719,7 @@ version = "0.15.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02" dependencies = [ - "smallvec", + "smallvec 1.15.1", "target-lexicon", ] @@ -1024,7 +1032,7 @@ dependencies = [ "phf 0.10.1", "proc-macro2", "quote", - "smallvec", + "smallvec 1.15.1", "syn 1.0.109", ] @@ -1108,6 +1116,16 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c87e182de0887fd5361989c677c4e8f5000cd9491d6d563161a8f3a5519fc7f" +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "pem-rfc7468", + "zeroize", +] + [[package]] name = "deranged" version = "0.4.0" @@ -1458,6 +1476,18 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "filetime" +version = "0.2.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35c0522e981e68cbfa8c3f978441a5f34b30b96e146b33cd3359176b50fe8586" +dependencies = [ + "cfg-if", + "libc", + "libredox", + "windows-sys 0.59.0", +] + [[package]] name = "fixedbitset" version = "0.4.2" @@ -2066,7 +2096,7 @@ dependencies = [ "libc", "once_cell", "pin-project-lite", - "smallvec", + "smallvec 1.15.1", "thiserror 1.0.69", ] @@ -2102,7 +2132,7 @@ dependencies = [ "libc", "memchr", "once_cell", - "smallvec", + "smallvec 1.15.1", "thiserror 1.0.69", ] @@ -2309,7 +2339,7 @@ dependencies = [ "serde_json", "thiserror 2.0.12", "tokio", - "ureq", + "ureq 2.12.1", "windows-sys 0.60.2", ] @@ -2380,7 +2410,7 @@ dependencies = [ "httparse", "itoa", "pin-project-lite", - "smallvec", + "smallvec 1.15.1", "tokio", "want", ] @@ -2514,7 +2544,7 @@ dependencies = [ "icu_normalizer_data", "icu_properties", "icu_provider", - "smallvec", + "smallvec 1.15.1", "zerovec", ] @@ -2576,7 +2606,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" dependencies = [ "idna_adapter", - "smallvec", + "smallvec 1.15.1", "utf8_iter", ] @@ -2862,6 +2892,7 @@ checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d" dependencies = [ "bitflags 2.9.1", "libc", + "redox_syscall", ] [[package]] @@ -2972,6 +3003,16 @@ version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.7.5" @@ -3125,6 +3166,21 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "ndk" version = "0.9.0" @@ -3678,6 +3734,31 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "ort" +version = "2.0.0-rc.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa7e49bd669d32d7bc2a15ec540a527e7764aec722a45467814005725bcd721" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec 2.0.0-alpha.10", + "tracing", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2aba9f5c7c479925205799216e7e5d07cc1d4fa76ea8058c60a9a30f6a4e890" +dependencies = [ + "flate2", + "pkg-config", + "sha2", + "tar", + "ureq 3.0.12", +] + [[package]] name = "os_info" version = "3.12.0" @@ -3734,7 +3815,7 @@ dependencies = [ "cfg-if", "libc", "redox_syscall", - "smallvec", + "smallvec 1.15.1", "windows-targets 0.52.6", ] @@ -3744,6 +3825,15 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.1" @@ -3944,6 +4034,15 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483" +[[package]] +name = "portable-atomic-util" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.2" @@ -4333,6 +4432,12 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.10.0" @@ -4673,6 +4778,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "rustls-pki-types" version = "1.12.0" @@ -4822,7 +4936,7 @@ dependencies = [ "phf_codegen 0.8.0", "precomputed-hash", "servo_arc", - "smallvec", + "smallvec 1.15.1", ] [[package]] @@ -5060,6 +5174,12 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "smallvec" +version = "2.0.0-alpha.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d44cfb396c3caf6fbfd0ab422af02631b69ddd96d2eff0b0f0724f9024051b" + [[package]] name = "socket2" version = "0.5.10" @@ -5531,6 +5651,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" +[[package]] +name = "tar" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d863878d212c87a19c1a610eb53bb01fe12951c0501cf5a0d65f724914a667a" +dependencies = [ + "filetime", + "libc", + "xattr", +] + [[package]] name = "target-lexicon" version = "0.12.16" @@ -6388,7 +6519,7 @@ version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" dependencies = [ - "smallvec", + "smallvec 1.15.1", ] [[package]] @@ -6399,9 +6530,9 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" -version = "0.2.1" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" +checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" [[package]] name = "unicode_categories" @@ -6435,6 +6566,37 @@ dependencies = [ "webpki-roots 0.26.11", ] +[[package]] +name = "ureq" +version = "3.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f0fde9bc91026e381155f8c67cb354bcd35260b2f4a29bcc84639f762760c39" +dependencies = [ + "base64 0.22.1", + "der", + "log", + "native-tls", + "percent-encoding", + "rustls-pemfile", + "rustls-pki-types", + "socks", + "ureq-proto", + "utf-8", + "webpki-root-certs 0.26.11", +] + +[[package]] +name = "ureq-proto" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59db78ad1923f2b1be62b6da81fe80b173605ca0d57f85da2e005382adf693f7" +dependencies = [ + "base64 0.22.1", + "http", + "httparse", + "log", +] + [[package]] name = "url" version = "2.5.4" @@ -6727,6 +6889,24 @@ dependencies = [ "system-deps", ] +[[package]] +name = "webpki-root-certs" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75c7f0ef91146ebfb530314f5f1d24528d7f0767efbfd31dce919275413e393e" +dependencies = [ + "webpki-root-certs 1.0.1", +] + +[[package]] +name = "webpki-root-certs" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86138b15b2b7d561bc4469e77027b8dd005a43dc502e9031d1f5afc8ce1f280e" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "0.26.11" @@ -7399,6 +7579,16 @@ version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec107c4503ea0b4a98ef47356329af139c0a4f7750e621cf2973cd3385ebcb3d" +[[package]] +name = "xattr" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af3a19837351dc82ba89f8a125e22a3c475f05aba604acc023d62b2739ae2909" +dependencies = [ + "libc", + "rustix 1.0.7", +] + [[package]] name = "xkeysym" version = "0.2.1" diff --git a/apps/stage-tamagotchi/src-tauri/Cargo.toml b/apps/stage-tamagotchi/src-tauri/Cargo.toml index 8597ad3a9..f81c98b1d 100644 --- a/apps/stage-tamagotchi/src-tauri/Cargo.toml +++ b/apps/stage-tamagotchi/src-tauri/Cargo.toml @@ -34,6 +34,8 @@ tauri-plugin-prevent-default = "1.3" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.140" tokio = "1.45.1" +ort = { version = "2.0.0-rc.10", features = ["ndarray", "copy-dylibs"] } +ndarray = "0.16.1" candle-core = { version = "0.9.1" } candle-nn = { version = "0.9.1" } candle-transformers = { version = "0.9.1" } @@ -54,17 +56,20 @@ tauri-plugin-positioner = "2.3.0" [target.'cfg(target_os = "macos")'.dependencies] objc2 = "0.6.1" objc2-foundation = "0.3.1" +ort = { version = "2.0.0-rc.10", features = ["ndarray", "coreml", "copy-dylibs"] } candle-core = { version = "0.9.1", features = ["metal"] } candle-nn = { version = "0.9.1", features = ["metal"] } candle-transformers = { version = "0.9.1", features = ["metal"] } [target.'cfg(target_os = "windows")'.dependencies] windows = { version = "0.61.1", features = ["Win32_UI_WindowsAndMessaging", "Win32_Foundation"] } +ort = { version = "2.0.0-rc.10", features = ["ndarray", "directml", "cuda", "copy-dylibs"] } candle-core = { version = "0.9.1", features = ["cuda"] } candle-nn = { version = "0.9.1", features = ["cuda"] } candle-transformers = { version = "0.9.1", features = ["cuda"] } [target.'cfg(target_os = "linux")'.dependencies] +ort = { version = "2.0.0-rc.10", features = ["ndarray", "cuda", "copy-dylibs"] } candle-core = { version = "0.9.1", features = ["cuda"] } candle-nn = { version = "0.9.1", features = ["cuda"] } candle-transformers = { version = "0.9.1", features = ["cuda"] } diff --git a/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs b/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs index d5d9382ae..a411e86b0 100644 --- a/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs +++ b/apps/stage-tamagotchi/src-tauri/src/app/models/silero_vad/mod.rs @@ -1,32 +1,42 @@ -use std::collections::HashMap; +use std::path::PathBuf; use anyhow::Result; -use candle_core::{DType, Device, Tensor}; -use candle_onnx::simple_eval; -use hf_hub::{Repo, RepoType}; +use log::{debug, info}; +use ndarray::{Array2, Array3}; +use ort::{ + execution_providers::{ + CPUExecutionProvider, + CUDAExecutionProvider, + CoreMLExecutionProvider, + DirectMLExecutionProvider, + }, + session::{Session, builder::GraphOptimizationLevel}, + value::Tensor, +}; use tauri::Runtime; use crate::helpers::huggingface::create_progress_emitter; +/// Main Silero VAD model wrapper with hardware acceleration support pub struct Processor { - model: candle_onnx::onnx::ModelProto, - frame_size: usize, - context_size: usize, - sample_rate: Tensor, - state: Tensor, - context: Tensor, - device: Device, + session: Session, + context: Array2, + state: ndarray::Array3, + last_batch_size: usize, + frame_size: usize, + context_size: usize, + sample_rate: i64, } impl Processor { pub fn new( - device: Device, + _device: candle_core::Device, window: tauri::WebviewWindow, ) -> Result { - let api = hf_hub::api::sync::Api::new()?; - let repo = api.repo(Repo::with_revision( + let api = hf_hub::api::sync::ApiBuilder::new().build()?; + let repo = api.repo(hf_hub::Repo::with_revision( "onnx-community/silero-vad".into(), - RepoType::Model, + hf_hub::RepoType::Model, "main".into(), )); @@ -35,61 +45,124 @@ impl Processor { create_progress_emitter(window.clone(), "onnx/model.onnx".to_string()), )?; - let model = candle_onnx::read_file(model_path)?; - - let sample_rate_value = 16000i64; + let session = Self::create_optimized_session(model_path.clone())?; let (frame_size, context_size) = (512, 64); Ok(Self { - model, + session, + context: Array2::zeros((1, context_size)), + state: Array3::zeros((2, 1, 128)), + last_batch_size: 0, frame_size, context_size, - sample_rate: Tensor::new(sample_rate_value, &device)?, - state: Tensor::zeros((2, 1, 128), DType::F32, &device)?, - context: Tensor::zeros((1, context_size), DType::F32, &device)?, - device, + sample_rate: 16000, }) } + /// Create an optimized ONNX session with hardware acceleration + fn create_optimized_session(model_path: PathBuf) -> Result { + let builder = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_intra_threads(1)?; + + let session = builder + .with_execution_providers(vec![ + CUDAExecutionProvider::default() + .with_device_id(0) + .build(), + CoreMLExecutionProvider::default().build(), + DirectMLExecutionProvider::default() + .with_device_id(0) + .build(), + CPUExecutionProvider::default().build(), + ])? + .commit_from_file(model_path)?; + info!("VAD model loaded successfully"); + + Ok(session) + } + + /// Reset the model's internal state + fn reset_states( + &mut self, + batch_size: usize, + ) { + self.context = Array2::zeros((batch_size, self.context_size)); + self.state = Array3::zeros((2, batch_size, 128)); + } + + /// Validate input audio chunk + fn validate_input( + &self, + x: &[f32], + ) -> Result<()> { + if x.len() != self.frame_size { + return Err(anyhow::anyhow!( + "Input chunk must be {} samples, got {}", + self.frame_size, + x.len() + )); + } + Ok(()) + } + + /// Process a single audio chunk and return speech probability + /// This is the main API used by hearing.vue via the audio_vad plugin function pub fn process_chunk( &mut self, chunk: &[f32], ) -> Result { - if chunk.len() != self.frame_size { - return Ok(0.0); + self.validate_input(chunk)?; + + let batch_size = 1; + if self.last_batch_size != batch_size { + self.reset_states(batch_size); } - let next_context = Tensor::from_slice( - &chunk[self.frame_size - self.context_size..], - (1, self.context_size), - &self.device, - )?; - let chunk_tensor = Tensor::from_vec(chunk.to_vec(), (1, self.frame_size), &self.device)?; + // Prepare input tensor by concatenating context and new audio data + let mut input_data = Vec::with_capacity(self.context_size + chunk.len()); + input_data.extend_from_slice(self.context.row(0).as_slice().unwrap()); + input_data.extend_from_slice(chunk); - let input = Tensor::cat(&[&self.context, &chunk_tensor], 1)?; - let inputs: HashMap = HashMap::from_iter([ - ("input".to_string(), input), - ("sr".to_string(), self.sample_rate.clone()), - ("state".to_string(), self.state.clone()), - ]); + let input_shape = vec![batch_size, input_data.len()]; - let outputs = simple_eval(&self.model, inputs)?; - let graph = self.model.graph.as_ref().unwrap(); - let out_names = &graph.output; + // Create input tensors for the ONNX model + // Silero VAD requires audio input, sample rate, and state + let inputs = vec![ + ( + "input", + Tensor::from_array((input_shape.clone(), input_data.clone()))?.into_dyn(), + ), + ( + "sr", + Tensor::from_array(([1], vec![self.sample_rate]))?.into_dyn(), + ), + ( + "state", + Tensor::from_array((vec![2, 1, 128], self.state.as_slice().unwrap().to_vec()))?.into_dyn(), + ), + ]; - let output = outputs - .get(&out_names[0].name) - .ok_or_else(|| anyhow::anyhow!("Missing VAD output tensor: {}", &out_names[0].name))? - .clone(); + // Run inference + let outputs = self.session.run(inputs)?; - self.state = outputs - .get(&out_names[1].name) - .ok_or_else(|| anyhow::anyhow!("Missing VAD state tensor: {}", &out_names[1].name))? - .clone(); + // Update context from the last portion of the input + let context_start = input_data.len() - self.context_size; + let new_context_data = input_data[context_start..].to_vec(); + self.context = Array2::from_shape_vec((batch_size, self.context_size), new_context_data) + .map_err(|e| anyhow::anyhow!("Failed to update context: {}", e))?; - self.context = next_context; + // Update state from model output + let (_state_shape, state_data) = outputs[1].try_extract_tensor::()?; + self.state = Array3::from_shape_vec((2, 1, 128), state_data.to_vec()) + .map_err(|e| anyhow::anyhow!("Failed to update state: {}", e))?; + + self.last_batch_size = batch_size; + + // Extract speech probability + let (_shape, data) = outputs[0].try_extract_tensor::()?; + let speech_prob = data[0]; - let speech_prob = output.flatten_all()?.to_vec1::()?[0]; Ok(speech_prob) } } diff --git a/apps/stage-tamagotchi/src/pages/settings/modules/hearing.vue b/apps/stage-tamagotchi/src/pages/settings/modules/hearing.vue index 7303ffc56..ecf3f5254 100644 --- a/apps/stage-tamagotchi/src/pages/settings/modules/hearing.vue +++ b/apps/stage-tamagotchi/src/pages/settings/modules/hearing.vue @@ -466,8 +466,7 @@ onUnmounted(() => { Loading... -
-
+
Inference error: {{ vadModelError }}
diff --git a/cspell.config.yaml b/cspell.config.yaml index 0f5ef2fb8..e5fc71a5a 100644 --- a/cspell.config.yaml +++ b/cspell.config.yaml @@ -25,6 +25,7 @@ words: - Bitstream - browserbasehq - bumpp + - byteorder - catppuccin - cdylib - cientos @@ -33,6 +34,7 @@ words: - collectblock - composables - cooldown + - coreml - cosyvoice - cozyvoice - crossws @@ -57,6 +59,7 @@ words: - demi - demodel - devlogs + - directml - dotenvx - DownloadLive2DSDK - dtolnay @@ -118,6 +121,7 @@ words: - moeru - Morioki - Myriam + - ndarray - Neko - nekomeowww - neuri @@ -205,6 +209,7 @@ words: - wavefile - webgpu - weeb + - wgpu - wlipsync - worklet - xast