From 30b6d5507459f2235d3a9a5cc45a94ad5b91c69f Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Sun, 22 Mar 2026 14:25:15 -0700 Subject: [PATCH 1/4] Switch from mlx-rs to quill-mlx fork Replace mlx-rs/mlx-macros with quill-mlx/quill-mlx-macros across all workspace crates. The fork adds Metal memory management APIs that were missing from upstream, fixing unbounded GPU buffer cache growth (102GB+) in long-running MCP server sessions. - Cap Metal buffer cache at 2GB on MCP server startup - Update scaled_dot_product_attention call for new sinks parameter --- Cargo.lock | 163 +++++++++------------- crates/voice-cli/Cargo.toml | 1 + crates/voice-cli/src/mcp.rs | 6 + crates/voice-dsp/Cargo.toml | 2 +- crates/voice-dsp/README.md | 2 +- crates/voice-dsp/src/lib.rs | 50 +++---- crates/voice-nn/Cargo.toml | 4 +- crates/voice-nn/README.md | 2 +- crates/voice-nn/src/ada_norm.rs | 16 +-- crates/voice-nn/src/albert.rs | 18 +-- crates/voice-nn/src/conv_weighted.rs | 16 +-- crates/voice-nn/src/lstm.rs | 26 ++-- crates/voice-nn/src/prosody.rs | 20 +-- crates/voice-nn/src/text_encoder.rs | 22 +-- crates/voice-nn/src/vocoder/decoder.rs | 14 +- crates/voice-nn/src/vocoder/generator.rs | 16 +-- crates/voice-nn/src/vocoder/source.rs | 26 ++-- crates/voice-stt/Cargo.toml | 4 +- crates/voice-stt/src/error.rs | 2 +- crates/voice-stt/src/lib.rs | 4 +- crates/voice-stt/src/moonshine/model.rs | 29 ++-- crates/voice-tts/Cargo.toml | 4 +- crates/voice-tts/src/builtin.rs | 2 +- crates/voice-tts/src/error.rs | 2 +- crates/voice-tts/src/lib.rs | 2 +- crates/voice-tts/src/model.rs | 20 +-- crates/voice-tts/src/voice.rs | 2 +- crates/voice-tts/src/weights.rs | 6 +- crates/voice-tts/tests/istft_roundtrip.rs | 4 +- 29 files changed, 231 insertions(+), 254 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e810ac9..096f824 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -194,26 +194,6 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" -[[package]] -name = "bindgen" -version = "0.70.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f" -dependencies = [ - "bitflags 2.11.0", - "cexpr", - "clang-sys", - "itertools 0.13.0", - "log", - "prettyplease", - "proc-macro2", - "quote", - "regex", - "rustc-hash 1.1.0", - "shlex", - "syn", -] - [[package]] name = "bindgen" version = "0.72.1" @@ -224,10 +204,12 @@ dependencies = [ "cexpr", "clang-sys", "itertools 0.13.0", + "log", + "prettyplease", "proc-macro2", "quote", "regex", - "rustc-hash 2.1.1", + "rustc-hash", "shlex", "syn", ] @@ -553,7 +535,7 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ceec7a6067e62d6f931a2baf6f3a751f4a892595bcec1461a3c94ef9949864b6" dependencies = [ - "bindgen 0.72.1", + "bindgen", ] [[package]] @@ -1727,68 +1709,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "mlx-internal-macros" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a7c4444d624bf6b93db5cc22ebff4fdfa13593fd56154fe33b1f302a557c2c6" -dependencies = [ - "darling 0.21.3", - "itertools 0.14.0", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "mlx-macros" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a819ee8b4434690572b6feb9c3ef0b6e90137e4190b340cf00150703b410aaf9" -dependencies = [ - "darling 0.21.3", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "mlx-rs" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3a0f592c5839b0237b3072530b1d3c503923579a2009ee0761df3edd1b1f27b" -dependencies = [ - "bytemuck", - "dyn-clone", - "half", - "itertools 0.14.0", - "libc", - "mach-sys", - "mlx-internal-macros", - "mlx-macros", - "mlx-sys", - "num-complex", - "num-traits", - "num_enum", - "parking_lot", - "paste", - "safetensors", - "smallvec", - "strum", - "thiserror 2.0.18", -] - -[[package]] -name = "mlx-sys" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e3bc3880111918b2d5018f845d48fd995f9901f16efc81d1fcfd2f4210b8219" -dependencies = [ - "bindgen 0.70.1", - "cc", - "cmake", -] - [[package]] name = "monostate" version = "0.1.18" @@ -2350,6 +2270,60 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae" +[[package]] +name = "quill-mlx" +version = "0.25.3" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.14.0", + "libc", + "mach-sys", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "quill-mlx-internal-macros", + "quill-mlx-macros", + "quill-mlx-sys", + "safetensors", + "smallvec", + "strum", + "thiserror 2.0.18", +] + +[[package]] +name = "quill-mlx-internal-macros" +version = "0.25.3" +dependencies = [ + "darling 0.21.3", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quill-mlx-macros" +version = "0.25.3" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quill-mlx-sys" +version = "0.2.0" +dependencies = [ + "bindgen", + "cc", + "cmake", +] + [[package]] name = "quote" version = "1.0.45" @@ -2611,12 +2585,6 @@ dependencies = [ "windowfunctions", ] -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.1" @@ -3614,6 +3582,7 @@ dependencies = [ "ctrlc", "hound", "pulldown-cmark", + "quill-mlx", "rodio", "serde", "serde_json", @@ -3627,7 +3596,7 @@ dependencies = [ name = "voice-dsp" version = "0.2.0" dependencies = [ - "mlx-rs", + "quill-mlx", ] [[package]] @@ -3645,8 +3614,8 @@ dependencies = [ name = "voice-nn" version = "0.2.0" dependencies = [ - "mlx-macros", - "mlx-rs", + "quill-mlx", + "quill-mlx-macros", "serde", "voice-dsp", ] @@ -3658,8 +3627,8 @@ dependencies = [ "audioadapter-buffers", "hf-hub", "hound", - "mlx-macros", - "mlx-rs", + "quill-mlx", + "quill-mlx-macros", "rubato", "serde", "serde_json", @@ -3673,8 +3642,8 @@ version = "0.2.1" dependencies = [ "hf-hub", "hound", - "mlx-macros", - "mlx-rs", + "quill-mlx", + "quill-mlx-macros", "serde", "serde_json", "thiserror 2.0.18", diff --git a/crates/voice-cli/Cargo.toml b/crates/voice-cli/Cargo.toml index 002c8db..fe88c93 100644 --- a/crates/voice-cli/Cargo.toml +++ b/crates/voice-cli/Cargo.toml @@ -24,6 +24,7 @@ pubkey = "RWRmHACtt9SpjQB/l+dAoCe+3iqxq/Xe/90P6dyCW70axzy3UkWl+Ave" [dependencies] voice-tts = { version = "0.2.0", path = "../voice-tts" } +quill-mlx = { path = "../../../mlx-rs/mlx-rs" } voice-g2p = { version = "0.2.0", path = "../voice-g2p" } voice-stt = { version = "0.1.0", path = "../voice-stt" } clap = { version = "4", features = ["derive"] } diff --git a/crates/voice-cli/src/mcp.rs b/crates/voice-cli/src/mcp.rs index d8eefdb..9cabdd5 100644 --- a/crates/voice-cli/src/mcp.rs +++ b/crates/voice-cli/src/mcp.rs @@ -204,6 +204,12 @@ pub fn run(config: ServerConfig) { stt_tokenizer: None, }; + // Cap the Metal buffer cache at 2 GB to prevent unbounded memory growth + // across repeated inference calls. + if let Err(e) = quill_mlx::metal::set_cache_limit(2 * 1024 * 1024 * 1024) { + eprintln!("warning: failed to set Metal cache limit: {e}"); + } + let mut stdout = io::stdout(); if !QUIET.load(Ordering::Relaxed) { diff --git a/crates/voice-dsp/Cargo.toml b/crates/voice-dsp/Cargo.toml index 31a90b0..b3044b1 100644 --- a/crates/voice-dsp/Cargo.toml +++ b/crates/voice-dsp/Cargo.toml @@ -8,4 +8,4 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -mlx-rs = { version = "0.25" } +quill-mlx = { path = "../../../mlx-rs/mlx-rs" } diff --git a/crates/voice-dsp/README.md b/crates/voice-dsp/README.md index 060bccf..4d4a15d 100644 --- a/crates/voice-dsp/README.md +++ b/crates/voice-dsp/README.md @@ -28,7 +28,7 @@ let (magnitude, phase) = stft.transform(&audio_batch)?; let reconstructed = stft.inverse(&magnitude, &phase)?; ``` -All functions operate on `mlx_rs::Array` and return `Result<_, mlx_rs::error::Exception>`. +All functions operate on `quill_mlx::Array` and return `Result<_, quill_mlx::error::Exception>`. ## Requirements diff --git a/crates/voice-dsp/src/lib.rs b/crates/voice-dsp/src/lib.rs index 4e1dc21..5e1a94e 100644 --- a/crates/voice-dsp/src/lib.rs +++ b/crates/voice-dsp/src/lib.rs @@ -1,6 +1,6 @@ -use mlx_rs::error::Exception; -use mlx_rs::ops::indexing::{ArrayIndex, Ellipsis, IndexOp, TryIndexOp}; -use mlx_rs::Array; +use quill_mlx::error::Exception; +use quill_mlx::ops::indexing::{ArrayIndex, Ellipsis, IndexOp, TryIndexOp}; +use quill_mlx::Array; /// Perform overlap-add on CPU for istft. /// @@ -82,7 +82,7 @@ pub fn stft( let w = if win_length < n_fft { let pad_len = n_fft - win_length; let pad = Array::zeros::(&[pad_len])?; - mlx_rs::ops::concatenate_axis(&[&w, &pad], 0)? + quill_mlx::ops::concatenate_axis(&[&w, &pad], 0)? } else { w }; @@ -102,7 +102,7 @@ pub fn stft( let suffix_idx = Array::from_slice(&suffix_indices, &[half]); let suffix = x.index(&suffix_idx); - mlx_rs::ops::concatenate_axis(&[&prefix, x, &suffix], 0)? + quill_mlx::ops::concatenate_axis(&[&prefix, x, &suffix], 0)? } else { x.clone() }; @@ -117,7 +117,7 @@ pub fn stft( // Window and FFT (rfft along last axis by default) let windowed = frames * &w; - mlx_rs::fft::rfft(&windowed, None, None) + quill_mlx::fft::rfft(&windowed, None, None) } /// Compute the inverse Short-Time Fourier Transform. @@ -151,7 +151,7 @@ pub fn istft( let t = (num_frames - 1) * hop_length + win_length; // irfft along axis 0, then transpose to (num_frames, win_length) - let frames_time = mlx_rs::fft::irfft(x, None, 0)?; + let frames_time = quill_mlx::fft::irfft(x, None, 0)?; let frames_time = frames_time.transpose_axes(&[1, 0])?; // Overlap-add on CPU to avoid massive graph @@ -169,7 +169,7 @@ pub fn istft( let eps = Array::from_slice(&[1e-10f32], &[1]); let mask = window_sum.gt(&eps)?; let ones = Array::from_slice(&[1.0f32], &[1]); - let safe_window = mlx_rs::ops::r#where(&mask, &window_sum, &ones)?; + let safe_window = quill_mlx::ops::r#where(&mask, &window_sum, &ones)?; let reconstructed = reconstructed / safe_window; let reconstructed = if center && length.is_none() { @@ -245,8 +245,8 @@ impl MlxStft { let mag_refs: Vec<&Array> = mags.iter().collect(); let phase_refs: Vec<&Array> = phases.iter().collect(); - let mag_batch = mlx_rs::ops::stack_axis(&mag_refs, 0)?; - let phase_batch = mlx_rs::ops::stack_axis(&phase_refs, 0)?; + let mag_batch = quill_mlx::ops::stack_axis(&mag_refs, 0)?; + let phase_batch = quill_mlx::ops::stack_axis(&phase_refs, 0)?; Ok((mag_batch, phase_batch)) } @@ -266,8 +266,8 @@ impl MlxStft { // Unwrap phase along the time axis (axis 1) before reconstruction let phase_cont = mlx_unwrap(&phase_b, None, Some(1), None)?; - let cos_phase = mlx_rs::ops::cos(&phase_cont)?; - let sin_phase = mlx_rs::ops::sin(&phase_cont)?; + let cos_phase = quill_mlx::ops::cos(&phase_cont)?; + let sin_phase = quill_mlx::ops::sin(&phase_cont)?; let r = &spec_b * &cos_phase; let im = &spec_b * &sin_phase; @@ -279,13 +279,13 @@ impl MlxStft { // complex64 = [f32 real, f32 imag] per element. let r_flat = r_t.reshape(&[-1, 1])?; let im_flat = im_t.reshape(&[-1, 1])?; - let interleaved = mlx_rs::ops::concatenate_axis(&[&r_flat, &im_flat], -1)?; + let interleaved = quill_mlx::ops::concatenate_axis(&[&r_flat, &im_flat], -1)?; let interleaved = interleaved.reshape(&[r_t.shape()[0], r_t.shape()[1], 2])?; // View as complex64: (frames, freq, 2) float32 -> (frames, freq) complex64 let complex_spec = interleaved .reshape(&[r_t.shape()[0], r_t.shape()[1] * 2])? - .view::()?; + .view::()?; // istft operates on (freq, frames) so transpose let complex_spec_t = complex_spec.transpose_axes(&[1, 0])?; @@ -303,7 +303,7 @@ impl MlxStft { } let out_refs: Vec<&Array> = outputs.iter().collect(); - mlx_rs::ops::stack_axis(&out_refs, 0) + quill_mlx::ops::stack_axis(&out_refs, 0) } } @@ -361,7 +361,7 @@ pub fn interpolate1d( let indices = indices.as_type::()?; let zero = Array::from_slice(&[0i32], &[1]); let max_idx = Array::from_slice(&[in_width - 1], &[1]); - let indices = mlx_rs::ops::clip(&indices, (&zero, &max_idx))?; + let indices = quill_mlx::ops::clip(&indices, (&zero, &max_idx))?; Ok(input.index((Ellipsis, &indices))) } else { // Linear interpolation @@ -381,7 +381,7 @@ pub fn interpolate1d( let x_low = x.floor()?.as_type::()?; let x_low_plus_one = &x_low + Array::from_slice(&[1i32], &[1]); let max_idx = Array::from_slice(&[in_width - 1], &[1]); - let x_high = mlx_rs::ops::minimum(&x_low_plus_one, &max_idx)?; + let x_high = quill_mlx::ops::minimum(&x_low_plus_one, &max_idx)?; let x_low_f = x_low.as_type::()?; let x_frac = &x - &x_low_f; @@ -417,14 +417,14 @@ pub fn check_array_shape(arr: &Array) -> bool { /// Compute the angle (argument) of a complex array, or zero for real arrays. pub fn mlx_angle(z: &Array) -> Result { - use mlx_rs::Dtype; + use quill_mlx::Dtype; if z.dtype() == Dtype::Complex64 { let zimag = z.imag()?; let zreal = z.real()?; - mlx_rs::ops::atan2(&zimag, &zreal) + quill_mlx::ops::atan2(&zimag, &zreal) } else { - let zimag = mlx_rs::ops::zeros_like(z)?; - mlx_rs::ops::atan2(&zimag, z) + let zimag = quill_mlx::ops::zeros_like(z)?; + quill_mlx::ops::atan2(&zimag, z) } } @@ -469,21 +469,21 @@ pub fn mlx_unwrap( let close_to_high = diff_from_high.lt(&eps)?; let dd_positive = dd.gt(&zero)?; let condition = close_to_high.logical_and(&dd_positive)?; - let ddmod = mlx_rs::ops::r#where(&condition, &interval_high_arr, &ddmod)?; + let ddmod = quill_mlx::ops::r#where(&condition, &interval_high_arr, &ddmod)?; let ph_correct = &ddmod - ⅆ // Zero out corrections where jump is small let abs_dd = dd.abs()?; let small_jump = abs_dd.lt(&discont_arr)?; - let ph_correct = mlx_rs::ops::r#where(&small_jump, &zero, &ph_correct)?; + let ph_correct = quill_mlx::ops::r#where(&small_jump, &zero, &ph_correct)?; // Pad with zeros along axis, then cumulative sum let mut pad_shape: Vec = p.shape().to_vec(); pad_shape[resolved_axis as usize] = 1; let zero_padding = Array::zeros::(&pad_shape)?; - let padded = mlx_rs::ops::concatenate_axis(&[&zero_padding, &ph_correct], resolved_axis)?; + let padded = quill_mlx::ops::concatenate_axis(&[&zero_padding, &ph_correct], resolved_axis)?; let cumulative = padded.cumsum(resolved_axis, None, None)?; Ok(p + cumulative) @@ -498,7 +498,7 @@ fn slice_along_axis(arr: &Array, axis: i32, start: i32, stop: i32) -> Result = Vec::with_capacity(ndim); for i in 0..ndim { if i == axis_usize { diff --git a/crates/voice-nn/Cargo.toml b/crates/voice-nn/Cargo.toml index c181108..912df63 100644 --- a/crates/voice-nn/Cargo.toml +++ b/crates/voice-nn/Cargo.toml @@ -8,7 +8,7 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -mlx-rs = { version = "0.25" } -mlx-macros = "0.25.3" +quill-mlx = { path = "../../../mlx-rs/mlx-rs" } +quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } serde = { version = "1", features = ["derive"] } voice-dsp = { version = "0.2.0", path = "../voice-dsp" } diff --git a/crates/voice-nn/README.md b/crates/voice-nn/README.md index c6826e4..31ae48b 100644 --- a/crates/voice-nn/README.md +++ b/crates/voice-nn/README.md @@ -29,7 +29,7 @@ use voice_nn::albert; use voice_nn::vocoder; ``` -All modules operate on `mlx_rs::Array` tensors and use `mlx_macros::ModuleParameters` for weight loading from safetensors checkpoints. +All modules operate on `quill_mlx::Array` tensors and use `quill_mlx_macros::ModuleParameters` for weight loading from safetensors checkpoints. ## License diff --git a/crates/voice-nn/src/ada_norm.rs b/crates/voice-nn/src/ada_norm.rs index 6d0879f..e0e4361 100644 --- a/crates/voice-nn/src/ada_norm.rs +++ b/crates/voice-nn/src/ada_norm.rs @@ -1,9 +1,9 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{Linear, LinearBuilder}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{Linear, LinearBuilder}; +use quill_mlx::Array; // --------------------------------------------------------------------------- // InstanceNorm1d @@ -80,7 +80,7 @@ impl AdaIN1d { let h = self.fc.forward(s)?; let num_features = self.norm.num_features; - let parts = mlx_rs::ops::split(&h, 2, -1)?; + let parts = quill_mlx::ops::split(&h, 2, -1)?; let gamma = &parts[0]; // (batch, num_features) let beta = &parts[1]; @@ -152,7 +152,7 @@ impl AdaLayerNorm { let h = self.fc.forward(s)?; let d_model = self.d_model; - let parts = mlx_rs::ops::split(&h, 2, -1)?; + let parts = quill_mlx::ops::split(&h, 2, -1)?; let gamma = &parts[0]; let beta = &parts[1]; diff --git a/crates/voice-nn/src/albert.rs b/crates/voice-nn/src/albert.rs index aec563f..17ea648 100644 --- a/crates/voice-nn/src/albert.rs +++ b/crates/voice-nn/src/albert.rs @@ -1,12 +1,12 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{ +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{ Dropout, DropoutBuilder, Embedding, Gelu, LayerNorm, LayerNormBuilder, Linear, LinearBuilder, }; -use mlx_rs::ops::{expand_dims, softmax_axis, zeros_like}; -use mlx_rs::Array; +use quill_mlx::ops::{expand_dims, softmax_axis, zeros_like}; +use quill_mlx::Array; use serde::Deserialize; #[derive(Debug, Clone, Deserialize)] @@ -517,10 +517,10 @@ impl<'a> Module> for CustomAlbert { })?; // Pool: take first token, apply pooler + tanh - use mlx_rs::ops::indexing::IndexOp; + use quill_mlx::ops::indexing::IndexOp; let first_token = encoder_output.index((.., 0)); let pooled = self.pooler.forward(&first_token)?; - let pooled = mlx_rs::ops::tanh(&pooled)?; + let pooled = quill_mlx::ops::tanh(&pooled)?; Ok(CustomAlbertOutput { encoder_output, diff --git a/crates/voice-nn/src/conv_weighted.rs b/crates/voice-nn/src/conv_weighted.rs index f7b06d2..2f0bad4 100644 --- a/crates/voice-nn/src/conv_weighted.rs +++ b/crates/voice-nn/src/conv_weighted.rs @@ -1,7 +1,7 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::error::Exception; -use mlx_rs::module::Param; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::error::Exception; +use quill_mlx::module::Param; +use quill_mlx::Array; /// Which convolution operation to use. #[derive(Debug, Clone, Copy)] @@ -99,11 +99,11 @@ impl ConvWeighted { let g = &*self.weight_g; // L2 norm of v along dims [1, 2], keeping dims for broadcast - let v_norm = mlx_rs::linalg::norm_l2(v, &[1, 2], true)?; + let v_norm = quill_mlx::linalg::norm_l2(v, &[1, 2], true)?; // Avoid division by zero let eps = Array::from_slice(&[1e-12f32], &[1]); - let v_norm = mlx_rs::ops::maximum(&v_norm, &eps)?; + let v_norm = quill_mlx::ops::maximum(&v_norm, &eps)?; let weight = g * v / v_norm; Ok(weight) @@ -136,7 +136,7 @@ impl ConvWeighted { /// Input: `(batch, length, channels)` -- MLX channel-last layout. pub fn forward_conv1d(&self, x: &Array) -> Result { let weight = self.resolve_weight(x)?; - let result = mlx_rs::ops::conv1d( + let result = quill_mlx::ops::conv1d( x, &weight, self.stride, @@ -168,7 +168,7 @@ impl ConvWeighted { weight.transpose_axes(&[2, 1, 0])? }; - let result = mlx_rs::ops::conv_transpose1d( + let result = quill_mlx::ops::conv_transpose1d( x, &weight, self.stride, diff --git a/crates/voice-nn/src/lstm.rs b/crates/voice-nn/src/lstm.rs index e6c0637..33498e1 100644 --- a/crates/voice-nn/src/lstm.rs +++ b/crates/voice-nn/src/lstm.rs @@ -1,11 +1,11 @@ -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::nn::{Lstm, LstmBuilder}; -use mlx_rs::ops::indexing::IndexOp; -use mlx_rs::ops::{addmm, split, stack_axis}; -use mlx_rs::Array; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::nn::{Lstm, LstmBuilder}; +use quill_mlx::ops::indexing::IndexOp; +use quill_mlx::ops::{addmm, split, stack_axis}; +use quill_mlx::Array; -use mlx_macros::ModuleParameters; +use quill_mlx_macros::ModuleParameters; /// Run a single-direction LSTM with proper hidden state propagation. /// @@ -40,17 +40,17 @@ fn lstm_forward_recurrent( } let pieces = split(&ifgo, 4, -1)?; - let i = mlx_rs::ops::sigmoid(&pieces[0])?; - let f = mlx_rs::ops::sigmoid(&pieces[1])?; - let g = mlx_rs::ops::tanh(&pieces[2])?; - let o = mlx_rs::ops::sigmoid(&pieces[3])?; + let i = quill_mlx::ops::sigmoid(&pieces[0])?; + let f = quill_mlx::ops::sigmoid(&pieces[1])?; + let g = quill_mlx::ops::tanh(&pieces[2])?; + let o = quill_mlx::ops::sigmoid(&pieces[3])?; let new_cell = match &cell { Some(c) => f.multiply(c)?.add(i.multiply(&g)?)?, None => i.multiply(&g)?, }; - let new_hidden = o.multiply(mlx_rs::ops::tanh(&new_cell)?)?; + let new_hidden = o.multiply(quill_mlx::ops::tanh(&new_cell)?)?; all_hidden.push(new_hidden.clone()); all_cell.push(new_cell.clone()); @@ -120,7 +120,7 @@ impl BiLstm { let bwd_out = bwd_out_rev.index((0.., &rev_idx, 0..)); // Concatenate forward and backward outputs along the feature dimension (axis 2) - let output = mlx_rs::ops::concatenate_axis(&[&fwd_out, &bwd_out], 2)?; + let output = quill_mlx::ops::concatenate_axis(&[&fwd_out, &bwd_out], 2)?; Ok((output, ())) } diff --git a/crates/voice-nn/src/prosody.rs b/crates/voice-nn/src/prosody.rs index 66ed2cb..2afe7c4 100644 --- a/crates/voice-nn/src/prosody.rs +++ b/crates/voice-nn/src/prosody.rs @@ -1,10 +1,10 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::{Module, ModuleParameters}; -use mlx_rs::nn::{Conv1d, Conv1dBuilder, Linear, LinearBuilder}; -use mlx_rs::ops::{broadcast_to, concatenate_axis, squeeze_axes, which, zeros}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::{Module, ModuleParameters}; +use quill_mlx::nn::{Conv1d, Conv1dBuilder, Linear, LinearBuilder}; +use quill_mlx::ops::{broadcast_to, concatenate_axis, squeeze_axes, which, zeros}; +use quill_mlx::Array; use super::ada_norm::AdaLayerNorm; use super::lstm::BiLstm; @@ -55,21 +55,21 @@ impl ModuleParameters for DurationEncoderBlock { } } - fn parameters(&self) -> mlx_rs::module::ModuleParamRef<'_> { + fn parameters(&self) -> quill_mlx::module::ModuleParamRef<'_> { match self { DurationEncoderBlock::Lstm(l) => l.parameters(), DurationEncoderBlock::Norm(n) => n.parameters(), } } - fn parameters_mut(&mut self) -> mlx_rs::module::ModuleParamMut<'_> { + fn parameters_mut(&mut self) -> quill_mlx::module::ModuleParamMut<'_> { match self { DurationEncoderBlock::Lstm(l) => l.parameters_mut(), DurationEncoderBlock::Norm(n) => n.parameters_mut(), } } - fn trainable_parameters(&self) -> mlx_rs::module::ModuleParamRef<'_> { + fn trainable_parameters(&self) -> quill_mlx::module::ModuleParamRef<'_> { match self { DurationEncoderBlock::Lstm(l) => l.trainable_parameters(), DurationEncoderBlock::Norm(n) => n.trainable_parameters(), diff --git a/crates/voice-nn/src/text_encoder.rs b/crates/voice-nn/src/text_encoder.rs index 517b228..adab0d5 100644 --- a/crates/voice-nn/src/text_encoder.rs +++ b/crates/voice-nn/src/text_encoder.rs @@ -1,10 +1,10 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{Dropout, DropoutBuilder, Embedding, LayerNorm, LayerNormBuilder}; -use mlx_rs::ops::{concatenate_axis, expand_dims, zeros}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{Dropout, DropoutBuilder, Embedding, LayerNorm, LayerNormBuilder}; +use quill_mlx::ops::{concatenate_axis, expand_dims, zeros}; +use quill_mlx::Array; use super::conv_weighted::ConvWeighted; use super::lstm::BiLstm; @@ -51,7 +51,7 @@ impl ConvBlock { /// Apply leaky ReLU with negative slope 0.2. fn leaky_relu_02(x: &Array) -> Result { - mlx_rs::nn::leaky_relu(x, 0.2) + quill_mlx::nn::leaky_relu(x, 0.2) } // --------------------------------------------------------------------------- @@ -112,7 +112,7 @@ impl<'a> Module> for TextEncoder { // Apply mask: zero out masked positions let zero = Array::from_f32(0.0); - x = mlx_rs::ops::r#where(&m, &zero, &x)?; + x = quill_mlx::ops::r#where(&m, &zero, &x)?; // CNN blocks for block in self.cnn.iter_mut() { @@ -133,7 +133,7 @@ impl<'a> Module> for TextEncoder { x = block.dropout.forward(&x)?; // Re-apply mask - x = mlx_rs::ops::r#where(&m, &zero, &x)?; + x = quill_mlx::ops::r#where(&m, &zero, &x)?; } // LSTM: needs (B, seq, channels) @@ -156,7 +156,7 @@ impl<'a> Module> for TextEncoder { } // Final mask application - x = mlx_rs::ops::r#where(&m, &zero, &x)?; + x = quill_mlx::ops::r#where(&m, &zero, &x)?; Ok(x) } diff --git a/crates/voice-nn/src/vocoder/decoder.rs b/crates/voice-nn/src/vocoder/decoder.rs index 3e93002..9cc71e4 100644 --- a/crates/voice-nn/src/vocoder/decoder.rs +++ b/crates/voice-nn/src/vocoder/decoder.rs @@ -1,10 +1,10 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{leaky_relu, Dropout, DropoutBuilder, Upsample, UpsampleMode}; -use mlx_rs::ops::{concatenate_axis, pad, sqrt, PadWidth}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{leaky_relu, Dropout, DropoutBuilder, Upsample, UpsampleMode}; +use quill_mlx::ops::{concatenate_axis, pad, sqrt, PadWidth}; +use quill_mlx::Array; use super::generator::Generator; use crate::ada_norm::AdaIN1d; diff --git a/crates/voice-nn/src/vocoder/generator.rs b/crates/voice-nn/src/vocoder/generator.rs index 9c6ae27..5d5a704 100644 --- a/crates/voice-nn/src/vocoder/generator.rs +++ b/crates/voice-nn/src/vocoder/generator.rs @@ -1,11 +1,11 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{leaky_relu, Conv1d, Conv1dBuilder, Upsample, UpsampleMode}; -use mlx_rs::ops::indexing::IndexOp; -use mlx_rs::ops::{concatenate_axis, exp, pad, sin, squeeze_axes, PadWidth}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{leaky_relu, Conv1d, Conv1dBuilder, Upsample, UpsampleMode}; +use quill_mlx::ops::indexing::IndexOp; +use quill_mlx::ops::{concatenate_axis, exp, pad, sin, squeeze_axes, PadWidth}; +use quill_mlx::Array; use super::source::{AdaINResBlock1, SourceModuleHnNSF}; use crate::conv_weighted::ConvWeighted; diff --git a/crates/voice-nn/src/vocoder/source.rs b/crates/voice-nn/src/vocoder/source.rs index 5b31513..58aa2b9 100644 --- a/crates/voice-nn/src/vocoder/source.rs +++ b/crates/voice-nn/src/vocoder/source.rs @@ -1,11 +1,11 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{Linear, LinearBuilder}; -use mlx_rs::ops::indexing::IndexOp; -use mlx_rs::ops::{cumsum, sin, tanh}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{Linear, LinearBuilder}; +use quill_mlx::ops::indexing::IndexOp; +use quill_mlx::ops::{cumsum, sin, tanh}; +use quill_mlx::Array; use crate::ada_norm::AdaIN1d; use crate::conv_weighted::ConvWeighted; @@ -79,8 +79,8 @@ impl SineGen { let dim = f0_values.dim(2); let zero_col = Array::zeros::(&[b, 1])?; let rest_cols = if dim > 1 { - let rest = mlx_rs::random::normal::(&[b, dim - 1], None, None, None)?; - mlx_rs::ops::concatenate_axis(&[&zero_col, &rest], -1)? + let rest = quill_mlx::random::normal::(&[b, dim - 1], None, None, None)?; + quill_mlx::ops::concatenate_axis(&[&zero_col, &rest], -1)? } else { zero_col }; @@ -96,7 +96,7 @@ impl SineGen { let rad_values = if t > 1 { let rest_frames = rad_values.index((.., 1_i32.., ..)); - mlx_rs::ops::concatenate_axis(&[&first_frame, &rest_frames], 1)? + quill_mlx::ops::concatenate_axis(&[&first_frame, &rest_frames], 1)? } else { first_frame }; @@ -154,7 +154,7 @@ impl SineGen { let noise_amp = (&uv * &noise_std_arr).add(one_minus_uv.multiply(&sine_amp_third)?)?; // Generate noise - let noise_raw = mlx_rs::random::normal::(sine_waves.shape(), None, None, None)?; + let noise_raw = quill_mlx::random::normal::(sine_waves.shape(), None, None, None)?; let noise = noise_amp.multiply(&noise_raw)?; // Mix: voiced regions get sine, unvoiced get noise @@ -239,7 +239,7 @@ impl Module<&Array> for SourceModuleHnNSF { // Generate noise for unvoiced regions let noise_amp = Array::from_f32(self.sine_amp / 3.0); let noise = - mlx_rs::random::normal::(uv.shape(), None, None, None)?.multiply(&noise_amp)?; + quill_mlx::random::normal::(uv.shape(), None, None, None)?.multiply(&noise_amp)?; Ok(SourceModuleOutput { sine_merge, diff --git a/crates/voice-stt/Cargo.toml b/crates/voice-stt/Cargo.toml index bb10ca2..1263635 100644 --- a/crates/voice-stt/Cargo.toml +++ b/crates/voice-stt/Cargo.toml @@ -8,8 +8,8 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -mlx-rs = { version = "0.25", features = ["safetensors"] } -mlx-macros = "0.25.3" +quill-mlx = { path = "../../../mlx-rs/mlx-rs", features = ["safetensors"] } +quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" diff --git a/crates/voice-stt/src/error.rs b/crates/voice-stt/src/error.rs index 7385ce6..8d86c11 100644 --- a/crates/voice-stt/src/error.rs +++ b/crates/voice-stt/src/error.rs @@ -5,7 +5,7 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum SttError { #[error("MLX error: {0}")] - Mlx(#[from] mlx_rs::error::Exception), + Mlx(#[from] quill_mlx::error::Exception), #[error("IO error: {0}")] Io(#[from] std::io::Error), diff --git a/crates/voice-stt/src/lib.rs b/crates/voice-stt/src/lib.rs index 43f967d..3494a40 100644 --- a/crates/voice-stt/src/lib.rs +++ b/crates/voice-stt/src/lib.rs @@ -29,7 +29,7 @@ pub mod moonshine; use std::path::{Path, PathBuf}; pub use error::{Result, SttError}; -pub use mlx_rs::Array; +pub use quill_mlx::Array; pub use moonshine::{MoonshineConfig, MoonshineModel}; pub use tokenizers; @@ -87,7 +87,7 @@ pub fn load_model(path_or_repo: &str) -> Result { // Apply weights to model parameters { - use mlx_rs::module::ModuleParameters; + use quill_mlx::module::ModuleParameters; let mut params = model.parameters_mut().flatten(); let mut loaded = 0; let mut missing = Vec::new(); diff --git a/crates/voice-stt/src/moonshine/model.rs b/crates/voice-stt/src/moonshine/model.rs index cb84ce4..e820f0e 100644 --- a/crates/voice-stt/src/moonshine/model.rs +++ b/crates/voice-stt/src/moonshine/model.rs @@ -13,14 +13,14 @@ //! gradient flow. //! - **Tied embeddings**: Output projection reuses the token embedding matrix. -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{self, Embedding, GroupNorm, LayerNorm, LayerNormBuilder, Linear, LinearBuilder}; -use mlx_rs::ops::indexing::{IndexOp, IntoStrideBy}; -use mlx_rs::ops::{concatenate_axis, expand_dims, stack_axis}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{self, Embedding, GroupNorm, LayerNorm, LayerNormBuilder, Linear, LinearBuilder}; +use quill_mlx::ops::indexing::{IndexOp, IntoStrideBy}; +use quill_mlx::ops::{concatenate_axis, expand_dims, stack_axis}; +use quill_mlx::Array; use super::config::MoonshineConfig; @@ -54,7 +54,7 @@ impl MoonshineRotaryEmbedding { pub fn forward(&self, position_ids: &Array) -> Result<(Array, Array), Exception> { // position_ids: [B, T] -> [B, T, 1] let pos = position_ids.reshape(&[position_ids.shape()[0], position_ids.shape()[1], 1])?; - let pos = pos.as_dtype(mlx_rs::Dtype::Float32)?; + let pos = pos.as_dtype(quill_mlx::Dtype::Float32)?; // inv_freq: [dim/2] -> [1, 1, dim/2] let inv = self.inv_freq.reshape(&[1, 1, (self.dim / 2) as i32])?; @@ -312,14 +312,15 @@ impl MoonshineAttention { // Scaled dot-product attention let o = if let Some(ref m) = mask { - mlx_rs::fast::scaled_dot_product_attention(&q, &k_exp, &v_exp, self.scale, m)? + quill_mlx::fast::scaled_dot_product_attention(&q, &k_exp, &v_exp, self.scale, m, None)? } else { - mlx_rs::fast::scaled_dot_product_attention( + quill_mlx::fast::scaled_dot_product_attention( &q, &k_exp, &v_exp, self.scale, - None::, + None::, + None, )? }; @@ -654,7 +655,7 @@ impl MoonshineEncoder { // Conv frontend with progressive downsampling let x = self.conv1.forward(&x)?; - let x = mlx_rs::ops::tanh(&x)?; + let x = quill_mlx::ops::tanh(&x)?; let x = self.groupnorm.forward(&x)?; let x = self.conv2.forward(&x)?; let x = nn::gelu(&x)?; @@ -860,7 +861,7 @@ impl MoonshineModel { let logits = self.get_logits(&last_hidden)?; // Greedy: argmax - let next_token = mlx_rs::ops::indexing::argmax_axis(&logits, -1, false)?; + let next_token = quill_mlx::ops::indexing::argmax_axis(&logits, -1, false)?; let next_token_val: u32 = next_token.item(); if next_token_val == eos { diff --git a/crates/voice-tts/Cargo.toml b/crates/voice-tts/Cargo.toml index b9e4f2c..45eac2c 100644 --- a/crates/voice-tts/Cargo.toml +++ b/crates/voice-tts/Cargo.toml @@ -8,12 +8,12 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -mlx-rs = { version = "0.25", features = ["safetensors"] } +quill-mlx = { path = "../../../mlx-rs/mlx-rs", features = ["safetensors"] } serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" hound = "3" thiserror = "2" -mlx-macros = "0.25.3" +quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } voice-dsp = { version = "0.2.0", path = "../voice-dsp" } voice-nn = { version = "0.2.0", path = "../voice-nn" } diff --git a/crates/voice-tts/src/builtin.rs b/crates/voice-tts/src/builtin.rs index 417f824..a295482 100644 --- a/crates/voice-tts/src/builtin.rs +++ b/crates/voice-tts/src/builtin.rs @@ -4,7 +4,7 @@ //! `include_str!`/`include_bytes!`, enabling instant loading with no //! network or filesystem access (model weights still require HF Hub). -use mlx_rs::Array; +use quill_mlx::Array; use crate::config::ModelConfig; use crate::error::{Result, VoicersError}; diff --git a/crates/voice-tts/src/error.rs b/crates/voice-tts/src/error.rs index cdf8671..8da3795 100644 --- a/crates/voice-tts/src/error.rs +++ b/crates/voice-tts/src/error.rs @@ -1,4 +1,4 @@ -use mlx_rs::error::Exception; +use quill_mlx::error::Exception; #[derive(Debug, thiserror::Error)] pub enum VoicersError { diff --git a/crates/voice-tts/src/lib.rs b/crates/voice-tts/src/lib.rs index 0b28ce8..10d730e 100644 --- a/crates/voice-tts/src/lib.rs +++ b/crates/voice-tts/src/lib.rs @@ -8,7 +8,7 @@ pub mod weights; use std::path::Path; -pub use mlx_rs::Array; +pub use quill_mlx::Array; pub use builtin::BUILTIN_VOICES; pub use config::ModelConfig; diff --git a/crates/voice-tts/src/model.rs b/crates/voice-tts/src/model.rs index 3ef3d40..c310642 100644 --- a/crates/voice-tts/src/model.rs +++ b/crates/voice-tts/src/model.rs @@ -1,11 +1,11 @@ -use mlx_macros::ModuleParameters; -use mlx_rs::builder::Builder; -use mlx_rs::error::Exception; -use mlx_rs::module::Module; -use mlx_rs::nn::{Linear, LinearBuilder}; -use mlx_rs::ops::indexing::IndexOp; -use mlx_rs::ops::{clip, sigmoid}; -use mlx_rs::Array; +use quill_mlx_macros::ModuleParameters; +use quill_mlx::builder::Builder; +use quill_mlx::error::Exception; +use quill_mlx::module::Module; +use quill_mlx::nn::{Linear, LinearBuilder}; +use quill_mlx::ops::indexing::IndexOp; +use quill_mlx::ops::{clip, sigmoid}; +use quill_mlx::Array; use std::collections::HashMap; use crate::config::{AlbertConfig, ModelConfig}; @@ -136,7 +136,7 @@ impl KokoroModel { let text_mask = arange_plus_one.gt(&lengths_expanded)?; // ALBERT encoder - let mask_int = text_mask.logical_not()?.as_dtype(mlx_rs::Dtype::Int32)?; + let mask_int = text_mask.logical_not()?.as_dtype(quill_mlx::Dtype::Int32)?; let bert_output = self.bert.forward(CustomAlbertInput { input_ids: &input_ids_arr, token_type_ids: None, @@ -182,7 +182,7 @@ impl KokoroModel { // Round and clip durations let rounded = duration_scaled.round(None)?; let pred_dur = clip(&rounded, (1.0f32, ()))?; - let pred_dur = pred_dur.as_dtype(mlx_rs::Dtype::Int32)?; + let pred_dur = pred_dur.as_dtype(quill_mlx::Dtype::Int32)?; // Take first batch element let pred_dur = pred_dur.index(0); pred_dur.eval()?; diff --git a/crates/voice-tts/src/voice.rs b/crates/voice-tts/src/voice.rs index 5a1d0cd..5f8b231 100644 --- a/crates/voice-tts/src/voice.rs +++ b/crates/voice-tts/src/voice.rs @@ -1,7 +1,7 @@ use std::path::Path; use hf_hub::api::sync::Api; -use mlx_rs::Array; +use quill_mlx::Array; use crate::error::{Result, VoicersError}; diff --git a/crates/voice-tts/src/weights.rs b/crates/voice-tts/src/weights.rs index f47940c..ca07aba 100644 --- a/crates/voice-tts/src/weights.rs +++ b/crates/voice-tts/src/weights.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use hf_hub::api::sync::Api; -use mlx_rs::Array; +use quill_mlx::Array; use crate::builtin; use crate::config::ModelConfig; @@ -326,7 +326,7 @@ pub fn load_model(path_or_repo: &str) -> Result { // Load weights into model by matching parameter keys { - use mlx_rs::module::ModuleParameters; + use quill_mlx::module::ModuleParameters; let mut params = model.parameters_mut().flatten(); let mut loaded = 0; let mut missing = Vec::new(); @@ -356,7 +356,7 @@ pub fn load_model(path_or_repo: &str) -> Result { } // Evaluate loaded params - use mlx_rs::module::ModuleParametersExt; + use quill_mlx::module::ModuleParametersExt; model .eval() .map_err(|e| VoicersError::Weight(e.to_string()))?; diff --git a/crates/voice-tts/tests/istft_roundtrip.rs b/crates/voice-tts/tests/istft_roundtrip.rs index 8c97e61..3247f7d 100644 --- a/crates/voice-tts/tests/istft_roundtrip.rs +++ b/crates/voice-tts/tests/istft_roundtrip.rs @@ -1,5 +1,5 @@ -use mlx_rs::ops::indexing::IndexOp; -use mlx_rs::Array; +use quill_mlx::ops::indexing::IndexOp; +use quill_mlx::Array; use voice_dsp::{mlx_angle, stft, MlxStft}; #[test] From 182c8c5f68061ff459545b8c6bc9cc16cec0c5e5 Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Sun, 22 Mar 2026 14:45:21 -0700 Subject: [PATCH 2/4] Switch to quill-mlx from crates.io and fix metallib discovery Update all workspace crates to use quill-mlx 0.25.4 from crates.io instead of path dependencies. Update build.rs metallib discovery to also match quill-mlx-sys-* build directories. --- Cargo.lock | 16 ++++++++++++---- crates/voice-cli/Cargo.toml | 2 +- crates/voice-cli/build.rs | 2 +- crates/voice-dsp/Cargo.toml | 2 +- crates/voice-nn/Cargo.toml | 4 ++-- crates/voice-stt/Cargo.toml | 4 ++-- crates/voice-tts/Cargo.toml | 4 ++-- 7 files changed, 21 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 096f824..54207b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2272,7 +2272,9 @@ checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae" [[package]] name = "quill-mlx" -version = "0.25.3" +version = "0.25.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21ced7db6db4a4ac552781bc286c16f1ccb0f3877645471d8327f8ee36ac93e7" dependencies = [ "bytemuck", "dyn-clone", @@ -2296,7 +2298,9 @@ dependencies = [ [[package]] name = "quill-mlx-internal-macros" -version = "0.25.3" +version = "0.25.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f2b1f27315134b6f7b8c572d2cbd5653992c07f5d1aaa0c03ab6baccb0b9631" dependencies = [ "darling 0.21.3", "itertools 0.14.0", @@ -2307,7 +2311,9 @@ dependencies = [ [[package]] name = "quill-mlx-macros" -version = "0.25.3" +version = "0.25.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0922d0dedc44d56a61b5f79a6db20ecafe1eef3f8ffd510c163adff258ec7684" dependencies = [ "darling 0.21.3", "proc-macro2", @@ -2317,7 +2323,9 @@ dependencies = [ [[package]] name = "quill-mlx-sys" -version = "0.2.0" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dfbb6718cd0197626ebf2da1dac0bc24171a602db06cd3f2ebb6c2c7bec2213" dependencies = [ "bindgen", "cc", diff --git a/crates/voice-cli/Cargo.toml b/crates/voice-cli/Cargo.toml index fe88c93..f9ee0f9 100644 --- a/crates/voice-cli/Cargo.toml +++ b/crates/voice-cli/Cargo.toml @@ -24,7 +24,7 @@ pubkey = "RWRmHACtt9SpjQB/l+dAoCe+3iqxq/Xe/90P6dyCW70axzy3UkWl+Ave" [dependencies] voice-tts = { version = "0.2.0", path = "../voice-tts" } -quill-mlx = { path = "../../../mlx-rs/mlx-rs" } +quill-mlx = "0.25.4" voice-g2p = { version = "0.2.0", path = "../voice-g2p" } voice-stt = { version = "0.1.0", path = "../voice-stt" } clap = { version = "4", features = ["derive"] } diff --git a/crates/voice-cli/build.rs b/crates/voice-cli/build.rs index 59226fb..941f025 100644 --- a/crates/voice-cli/build.rs +++ b/crates/voice-cli/build.rs @@ -76,7 +76,7 @@ fn find_metallib(build_dir: &Path) -> Option { for entry in entries.flatten() { let name = entry.file_name(); let name_str = name.to_string_lossy(); - if name_str.starts_with("mlx-sys-") { + if name_str.starts_with("quill-mlx-sys-") || name_str.starts_with("mlx-sys-") { println!( "cargo:warning=[metallib] found mlx-sys dir: {}", entry.path().display() diff --git a/crates/voice-dsp/Cargo.toml b/crates/voice-dsp/Cargo.toml index b3044b1..d23bfc2 100644 --- a/crates/voice-dsp/Cargo.toml +++ b/crates/voice-dsp/Cargo.toml @@ -8,4 +8,4 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { path = "../../../mlx-rs/mlx-rs" } +quill-mlx = "0.25.4" diff --git a/crates/voice-nn/Cargo.toml b/crates/voice-nn/Cargo.toml index 912df63..fcf27f2 100644 --- a/crates/voice-nn/Cargo.toml +++ b/crates/voice-nn/Cargo.toml @@ -8,7 +8,7 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { path = "../../../mlx-rs/mlx-rs" } -quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } +quill-mlx = "0.25.4" +quill-mlx-macros = "0.25.4" serde = { version = "1", features = ["derive"] } voice-dsp = { version = "0.2.0", path = "../voice-dsp" } diff --git a/crates/voice-stt/Cargo.toml b/crates/voice-stt/Cargo.toml index 1263635..c5f613e 100644 --- a/crates/voice-stt/Cargo.toml +++ b/crates/voice-stt/Cargo.toml @@ -8,8 +8,8 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { path = "../../../mlx-rs/mlx-rs", features = ["safetensors"] } -quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } +quill-mlx = { version = "0.25.4", features = ["safetensors"] } +quill-mlx-macros = "0.25.4" serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" diff --git a/crates/voice-tts/Cargo.toml b/crates/voice-tts/Cargo.toml index 45eac2c..cb88990 100644 --- a/crates/voice-tts/Cargo.toml +++ b/crates/voice-tts/Cargo.toml @@ -8,12 +8,12 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { path = "../../../mlx-rs/mlx-rs", features = ["safetensors"] } +quill-mlx = { version = "0.25.4", features = ["safetensors"] } serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" hound = "3" thiserror = "2" -quill-mlx-macros = { path = "../../../mlx-rs/mlx-macros" } +quill-mlx-macros = "0.25.4" voice-dsp = { version = "0.2.0", path = "../voice-dsp" } voice-nn = { version = "0.2.0", path = "../voice-nn" } From 890aef47c19177de25b7a0f0f0369dcd586a9743 Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Sun, 22 Mar 2026 14:51:51 -0700 Subject: [PATCH 3/4] Fix import ordering for cargo fmt --- crates/voice-nn/src/ada_norm.rs | 2 +- crates/voice-nn/src/albert.rs | 2 +- crates/voice-nn/src/conv_weighted.rs | 2 +- crates/voice-nn/src/prosody.rs | 2 +- crates/voice-nn/src/text_encoder.rs | 2 +- crates/voice-nn/src/vocoder/decoder.rs | 2 +- crates/voice-nn/src/vocoder/generator.rs | 2 +- crates/voice-nn/src/vocoder/source.rs | 2 +- crates/voice-stt/src/lib.rs | 2 +- crates/voice-stt/src/moonshine/model.rs | 6 ++++-- crates/voice-tts/src/model.rs | 2 +- 11 files changed, 14 insertions(+), 12 deletions(-) diff --git a/crates/voice-nn/src/ada_norm.rs b/crates/voice-nn/src/ada_norm.rs index e0e4361..c1bf925 100644 --- a/crates/voice-nn/src/ada_norm.rs +++ b/crates/voice-nn/src/ada_norm.rs @@ -1,9 +1,9 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; use quill_mlx::nn::{Linear, LinearBuilder}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; // --------------------------------------------------------------------------- // InstanceNorm1d diff --git a/crates/voice-nn/src/albert.rs b/crates/voice-nn/src/albert.rs index 17ea648..3df8b50 100644 --- a/crates/voice-nn/src/albert.rs +++ b/crates/voice-nn/src/albert.rs @@ -1,4 +1,3 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; @@ -7,6 +6,7 @@ use quill_mlx::nn::{ }; use quill_mlx::ops::{expand_dims, softmax_axis, zeros_like}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use serde::Deserialize; #[derive(Debug, Clone, Deserialize)] diff --git a/crates/voice-nn/src/conv_weighted.rs b/crates/voice-nn/src/conv_weighted.rs index 2f0bad4..aa755f9 100644 --- a/crates/voice-nn/src/conv_weighted.rs +++ b/crates/voice-nn/src/conv_weighted.rs @@ -1,7 +1,7 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::error::Exception; use quill_mlx::module::Param; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; /// Which convolution operation to use. #[derive(Debug, Clone, Copy)] diff --git a/crates/voice-nn/src/prosody.rs b/crates/voice-nn/src/prosody.rs index 2afe7c4..245403a 100644 --- a/crates/voice-nn/src/prosody.rs +++ b/crates/voice-nn/src/prosody.rs @@ -1,10 +1,10 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::{Module, ModuleParameters}; use quill_mlx::nn::{Conv1d, Conv1dBuilder, Linear, LinearBuilder}; use quill_mlx::ops::{broadcast_to, concatenate_axis, squeeze_axes, which, zeros}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use super::ada_norm::AdaLayerNorm; use super::lstm::BiLstm; diff --git a/crates/voice-nn/src/text_encoder.rs b/crates/voice-nn/src/text_encoder.rs index adab0d5..77716c2 100644 --- a/crates/voice-nn/src/text_encoder.rs +++ b/crates/voice-nn/src/text_encoder.rs @@ -1,10 +1,10 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; use quill_mlx::nn::{Dropout, DropoutBuilder, Embedding, LayerNorm, LayerNormBuilder}; use quill_mlx::ops::{concatenate_axis, expand_dims, zeros}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use super::conv_weighted::ConvWeighted; use super::lstm::BiLstm; diff --git a/crates/voice-nn/src/vocoder/decoder.rs b/crates/voice-nn/src/vocoder/decoder.rs index 9cc71e4..9899f82 100644 --- a/crates/voice-nn/src/vocoder/decoder.rs +++ b/crates/voice-nn/src/vocoder/decoder.rs @@ -1,10 +1,10 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; use quill_mlx::nn::{leaky_relu, Dropout, DropoutBuilder, Upsample, UpsampleMode}; use quill_mlx::ops::{concatenate_axis, pad, sqrt, PadWidth}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use super::generator::Generator; use crate::ada_norm::AdaIN1d; diff --git a/crates/voice-nn/src/vocoder/generator.rs b/crates/voice-nn/src/vocoder/generator.rs index 5d5a704..a280829 100644 --- a/crates/voice-nn/src/vocoder/generator.rs +++ b/crates/voice-nn/src/vocoder/generator.rs @@ -1,4 +1,3 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; @@ -6,6 +5,7 @@ use quill_mlx::nn::{leaky_relu, Conv1d, Conv1dBuilder, Upsample, UpsampleMode}; use quill_mlx::ops::indexing::IndexOp; use quill_mlx::ops::{concatenate_axis, exp, pad, sin, squeeze_axes, PadWidth}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use super::source::{AdaINResBlock1, SourceModuleHnNSF}; use crate::conv_weighted::ConvWeighted; diff --git a/crates/voice-nn/src/vocoder/source.rs b/crates/voice-nn/src/vocoder/source.rs index 58aa2b9..da066fb 100644 --- a/crates/voice-nn/src/vocoder/source.rs +++ b/crates/voice-nn/src/vocoder/source.rs @@ -1,4 +1,3 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; @@ -6,6 +5,7 @@ use quill_mlx::nn::{Linear, LinearBuilder}; use quill_mlx::ops::indexing::IndexOp; use quill_mlx::ops::{cumsum, sin, tanh}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use crate::ada_norm::AdaIN1d; use crate::conv_weighted::ConvWeighted; diff --git a/crates/voice-stt/src/lib.rs b/crates/voice-stt/src/lib.rs index 3494a40..b3705e5 100644 --- a/crates/voice-stt/src/lib.rs +++ b/crates/voice-stt/src/lib.rs @@ -29,8 +29,8 @@ pub mod moonshine; use std::path::{Path, PathBuf}; pub use error::{Result, SttError}; -pub use quill_mlx::Array; pub use moonshine::{MoonshineConfig, MoonshineModel}; +pub use quill_mlx::Array; pub use tokenizers; /// Result of a transcription. diff --git a/crates/voice-stt/src/moonshine/model.rs b/crates/voice-stt/src/moonshine/model.rs index e820f0e..a3d24ea 100644 --- a/crates/voice-stt/src/moonshine/model.rs +++ b/crates/voice-stt/src/moonshine/model.rs @@ -13,14 +13,16 @@ //! gradient flow. //! - **Tied embeddings**: Output projection reuses the token embedding matrix. -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; -use quill_mlx::nn::{self, Embedding, GroupNorm, LayerNorm, LayerNormBuilder, Linear, LinearBuilder}; +use quill_mlx::nn::{ + self, Embedding, GroupNorm, LayerNorm, LayerNormBuilder, Linear, LinearBuilder, +}; use quill_mlx::ops::indexing::{IndexOp, IntoStrideBy}; use quill_mlx::ops::{concatenate_axis, expand_dims, stack_axis}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use super::config::MoonshineConfig; diff --git a/crates/voice-tts/src/model.rs b/crates/voice-tts/src/model.rs index c310642..2d2ae8e 100644 --- a/crates/voice-tts/src/model.rs +++ b/crates/voice-tts/src/model.rs @@ -1,4 +1,3 @@ -use quill_mlx_macros::ModuleParameters; use quill_mlx::builder::Builder; use quill_mlx::error::Exception; use quill_mlx::module::Module; @@ -6,6 +5,7 @@ use quill_mlx::nn::{Linear, LinearBuilder}; use quill_mlx::ops::indexing::IndexOp; use quill_mlx::ops::{clip, sigmoid}; use quill_mlx::Array; +use quill_mlx_macros::ModuleParameters; use std::collections::HashMap; use crate::config::{AlbertConfig, ModelConfig}; From 4bb263022a37864b57be8e7d22284152f89eecee Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Sun, 22 Mar 2026 15:28:27 -0700 Subject: [PATCH 4/4] Switch to quill-mlx 0.25.5 (based on MLX v0.25.1) Previous 0.25.4 was based on upstream main which had mlx-c v0.5.0 (MLX v0.30.6), causing kernel load failures at runtime. 0.25.5 is based on the v0.25.3 tag with mlx-c v0.2.0 (MLX v0.25.1). Revert scaled_dot_product_attention sinks parameter since the v0.25.x API doesn't have it. --- crates/voice-cli/Cargo.toml | 2 +- crates/voice-dsp/Cargo.toml | 2 +- crates/voice-nn/Cargo.toml | 4 ++-- crates/voice-stt/Cargo.toml | 4 ++-- crates/voice-stt/src/moonshine/model.rs | 3 +-- crates/voice-tts/Cargo.toml | 4 ++-- 6 files changed, 9 insertions(+), 10 deletions(-) diff --git a/crates/voice-cli/Cargo.toml b/crates/voice-cli/Cargo.toml index f9ee0f9..eea5d45 100644 --- a/crates/voice-cli/Cargo.toml +++ b/crates/voice-cli/Cargo.toml @@ -24,7 +24,7 @@ pubkey = "RWRmHACtt9SpjQB/l+dAoCe+3iqxq/Xe/90P6dyCW70axzy3UkWl+Ave" [dependencies] voice-tts = { version = "0.2.0", path = "../voice-tts" } -quill-mlx = "0.25.4" +quill-mlx = "0.25.5" voice-g2p = { version = "0.2.0", path = "../voice-g2p" } voice-stt = { version = "0.1.0", path = "../voice-stt" } clap = { version = "4", features = ["derive"] } diff --git a/crates/voice-dsp/Cargo.toml b/crates/voice-dsp/Cargo.toml index d23bfc2..af2e0b2 100644 --- a/crates/voice-dsp/Cargo.toml +++ b/crates/voice-dsp/Cargo.toml @@ -8,4 +8,4 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = "0.25.4" +quill-mlx = "0.25.5" diff --git a/crates/voice-nn/Cargo.toml b/crates/voice-nn/Cargo.toml index fcf27f2..f34fd61 100644 --- a/crates/voice-nn/Cargo.toml +++ b/crates/voice-nn/Cargo.toml @@ -8,7 +8,7 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = "0.25.4" -quill-mlx-macros = "0.25.4" +quill-mlx = "0.25.5" +quill-mlx-macros = "0.25.5" serde = { version = "1", features = ["derive"] } voice-dsp = { version = "0.2.0", path = "../voice-dsp" } diff --git a/crates/voice-stt/Cargo.toml b/crates/voice-stt/Cargo.toml index c5f613e..2f5313e 100644 --- a/crates/voice-stt/Cargo.toml +++ b/crates/voice-stt/Cargo.toml @@ -8,8 +8,8 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { version = "0.25.4", features = ["safetensors"] } -quill-mlx-macros = "0.25.4" +quill-mlx = { version = "0.25.5", features = ["safetensors"] } +quill-mlx-macros = "0.25.5" serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" diff --git a/crates/voice-stt/src/moonshine/model.rs b/crates/voice-stt/src/moonshine/model.rs index a3d24ea..df4f480 100644 --- a/crates/voice-stt/src/moonshine/model.rs +++ b/crates/voice-stt/src/moonshine/model.rs @@ -314,7 +314,7 @@ impl MoonshineAttention { // Scaled dot-product attention let o = if let Some(ref m) = mask { - quill_mlx::fast::scaled_dot_product_attention(&q, &k_exp, &v_exp, self.scale, m, None)? + quill_mlx::fast::scaled_dot_product_attention(&q, &k_exp, &v_exp, self.scale, m)? } else { quill_mlx::fast::scaled_dot_product_attention( &q, @@ -322,7 +322,6 @@ impl MoonshineAttention { &v_exp, self.scale, None::, - None, )? }; diff --git a/crates/voice-tts/Cargo.toml b/crates/voice-tts/Cargo.toml index cb88990..69e4bb0 100644 --- a/crates/voice-tts/Cargo.toml +++ b/crates/voice-tts/Cargo.toml @@ -8,12 +8,12 @@ license = "MIT" repository = "https://github.com/rgbkrk/voicers" [dependencies] -quill-mlx = { version = "0.25.4", features = ["safetensors"] } +quill-mlx = { version = "0.25.5", features = ["safetensors"] } serde = { version = "1", features = ["derive"] } serde_json = "1" hf-hub = "0.5" hound = "3" thiserror = "2" -quill-mlx-macros = "0.25.4" +quill-mlx-macros = "0.25.5" voice-dsp = { version = "0.2.0", path = "../voice-dsp" } voice-nn = { version = "0.2.0", path = "../voice-nn" }