taconite_reuse/
lib.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! NVIDIA RE-USE (`nvidia/RE-USE`, universal speech enhancement: SEMamba,
5//! 30 bidirectional time-frequency Mamba blocks between a dense
6//! convolutional encoder and two decoders) on an AMD XDNA NPU: noisy,
7//! reverberant, clipped, band-limited or codec-damaged speech at 8 to 48
8//! kHz in, enhanced speech at the same rate out.
9//!
10//! The bundle `iron/applications/reuse/export_reuse.py` writes holds every
11//! compiled IRON kernel and the weights (the NPU ones pre-packed); this
12//! crate replays the forward the Python app runs (`reuse_common.py` /
13//! `reuse_npu.py`):
14//!
15//! | | NPU | host (here) |
16//! |---|---|---|
17//! | STFT (n_fft 320 at 8 kHz, scaled with the rate; hop 1/8 of it) | | `stft.rs` (and its inverse) |
18//! | encoder (dense block of 4 dilated 3 x 3 convs, a strided conv) | every conv as an `flm.GEMM` | the 1 x 1 input conv, the map layouts, the dilated taps' sums, InstanceNorm + PReLU |
19//! | 30 x (time Mamba, frequency Mamba), both directions | in_proj, causal conv1d + x_proj, the selective scan, the output projection (the `MambaScan` operator's phases) | token layouts, both directions summed, LayerNorm, residual |
20//! | decoders (dense block, two sub-pixel upsampling convs) | every conv as an `flm.GEMM` | layouts, InstanceNorm + PReLU, the 1 x 1 heads, atan2 |
21//!
22//! [`Reuse::enhance`] takes any length: up to the bundle's longest Mamba
23//! grid (`t_max`, ~5 s) in one pass, as the checkpoint's `inference.py`;
24//! longer audio in 5 s chunks at a 2.5 s hop, overlap-added under a Hann
25//! window, as its `inference_chunk.py`.
26
27use std::collections::HashMap;
28use std::fmt;
29use std::path::Path;
30use std::time::Instant;
31
32use taconite_bundle::{Manifest, Store};
33
34pub use taconite::Timing;
35
36#[cfg(feature = "embed")]
37pub mod embed;
38#[cfg(feature = "encode")]
39pub mod encode;
40#[cfg(feature = "decode")]
41pub mod media;
42pub mod model;
43pub mod npu;
44pub mod stft;
45pub mod wav;
46
47use model::Model;
48use npu::Npu;
49use stft::Stft;
50
51pub const VERSION: u32 = 1;
52
53#[derive(Debug)]
54pub enum Error {
55    Bundle(String),
56    Npu(String),
57    Input(String),
58}
59
60impl fmt::Display for Error {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        match self {
63            Error::Bundle(m) => write!(f, "bundle: {m}"),
64            Error::Npu(m) => write!(f, "NPU: {m}"),
65            Error::Input(m) => write!(f, "input: {m}"),
66        }
67    }
68}
69
70impl std::error::Error for Error {}
71
72impl From<taconite_bundle::Error> for Error {
73    fn from(e: taconite_bundle::Error) -> Self {
74        Error::Bundle(e.to_string())
75    }
76}
77
78impl From<taconite::Error> for Error {
79    fn from(e: taconite::Error) -> Self {
80        Error::Npu(e.to_string())
81    }
82}
83
84/// The model's constants and the bundle's layout (the manifest's params).
85#[derive(Debug, Clone)]
86pub struct Config {
87    pub blocks: usize,
88    /// n_fft, hop, window at `base_sr` (scaled to the input's rate)
89    pub n_fft: usize,
90    pub hop: usize,
91    pub win: usize,
92    pub base_sr: usize,
93    /// the LayerNorms' and the InstanceNorms' epsilon
94    pub eps: f32,
95    pub in_eps: f32,
96    pub pix: usize,
97    pub padt: usize,
98    pub skip: usize,
99    pub rec: usize,
100    pub yw: usize,
101    /// the sample rates the bundle has kernels for
102    pub rates: Vec<u32>,
103    /// the Mamba grid's frame buckets: t_step, 2 t_step, ... t_max
104    pub t_step: usize,
105    pub t_max: usize,
106    pub chunk_secs: f32,
107    pub chunk_hop: f32,
108}
109
110impl Config {
111    fn load(m: &Manifest) -> Result<Self, Error> {
112        let p = |k: &str| m.param_as::<usize>(k);
113        let f = |k: &str| m.param_as::<f32>(k);
114        if p("hid")? != 64 || p("d_state")? != 16 {
115            return Err(Error::Bundle("not a RE-USE (hid 64, d_state 16) bundle".into()));
116        }
117        Ok(Config {
118            blocks: p("blocks")?,
119            n_fft: p("n_fft")?,
120            hop: p("hop")?,
121            win: p("win")?,
122            base_sr: p("base_sr")?,
123            eps: f("eps")?,
124            in_eps: f("in_eps")?,
125            pix: p("pix")?,
126            padt: p("padt")?,
127            skip: p("skip")?,
128            rec: p("rec")?,
129            yw: p("yw")?,
130            rates: m.list("rates")?,
131            t_step: p("t_step")?,
132            t_max: p("t_max")?,
133            chunk_secs: f("chunk_secs")?,
134            chunk_hop: f("chunk_hop")?,
135        })
136    }
137
138    /// n_fft, hop, window at `sr` (inference.py: v sr // base_sr, rounded
139    /// up to even).
140    pub fn stft_sizes(&self, sr: u32) -> (usize, usize, usize) {
141        let even = |v: usize| v + v % 2;
142        let s = |v: usize| even(v * sr as usize / self.base_sr);
143        (s(self.n_fft), s(self.hop), s(self.win))
144    }
145
146    /// The Mamba grid's frames for `samples` samples at `sr`.
147    pub fn grid_frames(&self, sr: u32, samples: usize) -> usize {
148        let (_, hop, _) = self.stft_sizes(sr);
149        (1 + samples / hop + 2 - 1) / 4 + 1
150    }
151}
152
153pub struct Reuse {
154    pub cfg: Config,
155    pub npu: Npu,
156    pub model: Model,
157    pub manifest: Manifest,
158    pub store: Store,
159    stfts: HashMap<u32, Stft>,
160    /// where the last call spent its time (`npu:<kernel>`, host stages)
161    pub timing: Timing,
162}
163
164impl Reuse {
165    /// Loads a bundle: opens the NPU and uploads the packed weights
166    /// (kernels load on first use).
167    pub fn load(dir: &Path) -> Result<Self, Error> {
168        let manifest = Manifest::load(dir, VERSION)?;
169        let store = Store::load(dir)?;
170        let cfg = Config::load(&manifest)?;
171        let npu = Npu::open(&manifest)?;
172        let model = Model::load(&store, &npu)?;
173        Ok(Reuse { cfg, npu, model, manifest, store, stfts: HashMap::new(), timing: Timing::default() })
174    }
175
176    /// Hardware contexts the bundle's kernels use.
177    pub fn contexts(&self) -> usize {
178        self.npu.contexts
179    }
180
181    /// Frees every NPU hardware context the model holds, for other models
182    /// (NPU2 has 16 across every process); the kernels load again on the
183    /// next call. The number of contexts freed.
184    pub fn release_contexts(&mut self) -> usize {
185        self.npu.release()
186    }
187
188    /// Enhances one channel `wav` at `sr` (one of the bundle's rates):
189    /// in one pass when its Mamba grid fits the bundle (`t_max`), else in
190    /// chunks (inference_chunk.py's). Same length out.
191    pub fn enhance(&mut self, wav: &[f32], sr: u32) -> Result<Vec<f32>, Error> {
192        self.enhance_with(wav, sr, &mut |_| {})
193    }
194
195    /// `enhance`, calling `progress` with the fraction done (0 to 1, over
196    /// every chunk) as it goes: a few dozen calls a pass, on this thread.
197    pub fn enhance_with(&mut self, wav: &[f32], sr: u32, progress: &mut dyn FnMut(f64)) -> Result<Vec<f32>, Error> {
198        if !self.cfg.rates.contains(&sr) {
199            return Err(Error::Input(format!(
200                "{sr} Hz: this bundle has kernels for {:?} Hz (export_reuse.py --rates; or resample)",
201                self.cfg.rates
202            )));
203        }
204        self.timing.clear();
205        if self.cfg.grid_frames(sr, wav.len()) <= self.cfg.t_max {
206            let y = self.enhance_whole_with(wav, sr, progress)?;
207            progress(1.0);
208            return Ok(y);
209        }
210        let size = (self.cfg.chunk_secs * sr as f32) as usize;
211        let hop = (self.cfg.chunk_hop * size as f32) as usize;
212        let s = wav.len();
213        let n = (s - size).div_ceil(hop) + 1;
214        let win: Vec<f32> =
215            (0..size).map(|i| (0.5 - 0.5 * (2.0 * std::f64::consts::PI * i as f64 / size as f64).cos()).max(1e-7) as f32).collect();
216        let mut out = vec![0f32; s];
217        let mut wsum = vec![0f32; s];
218        for i in 0..n {
219            let lo = if i < n - 1 { i * hop } else { s - size };
220            let y = self.enhance_whole_with(&wav[lo..lo + size], sr, &mut |f| progress((i as f64 + f) / n as f64))?;
221            for (k, &v) in y.iter().enumerate() {
222                out[lo + k] += v * win[k];
223                wsum[lo + k] += win[k];
224            }
225        }
226        progress(1.0);
227        Ok(out.iter().zip(&wsum).map(|(a, b)| a / b).collect())
228    }
229
230    /// One pass of the network over the whole signal (inference.py): STFT,
231    /// the forward, frames mostly at zero magnitude muted, inverse STFT,
232    /// the length matched to the input (padded with 1e-8).
233    pub fn enhance_whole(&mut self, wav: &[f32], sr: u32) -> Result<Vec<f32>, Error> {
234        self.enhance_whole_with(wav, sr, &mut |_| {})
235    }
236
237    /// `enhance_whole`, calling `tick` with the fraction of the pass done.
238    pub fn enhance_whole_with(&mut self, wav: &[f32], sr: u32, tick: &mut dyn FnMut(f64)) -> Result<Vec<f32>, Error> {
239        let (n_fft, hop, win) = self.cfg.stft_sizes(sr);
240        if wav.len() <= n_fft / 2 {
241            return Err(Error::Input(format!("{} samples: shorter than half an STFT window ({})", wav.len(), n_fft / 2)));
242        }
243        let st = self.stfts.entry(sr).or_insert_with(|| Stft::new(n_fft, hop, win));
244        let t0 = Instant::now();
245        let (mag, pha) = st.forward(wav);
246        let (fq, tf) = (st.bins(), st.frames(wav.len()));
247        self.timing.add("host.stft", t0.elapsed());
248        let (em, ep) = self.model.forward_with(&self.cfg, &self.npu, &mag, &pha, fq, tf, &mut self.timing, tick)?;
249        let t0 = Instant::now();
250        let mut lin = vec![0f32; fq * tf];
251        for t in 0..tf {
252            let zero = (0..fq).filter(|&f| em[f * tf + t] <= 0.0).count();
253            if zero * 2 > fq {
254                continue; // muted: inference.py's "strange sweep artifact" guard
255            }
256            for f in 0..fq {
257                lin[f * tf + t] = em[f * tf + t].max(0.0).exp_m1();
258            }
259        }
260        let st = &self.stfts[&sr];
261        let mut y = st.inverse(&lin, &ep, tf);
262        y.resize(wav.len(), 1e-8);
263        self.timing.add("host.istft", t0.elapsed());
264        Ok(y)
265    }
266}
267
268/// SNR of `x` against `reference`, in dB.
269pub fn snr(x: &[f32], reference: &[f32]) -> f64 {
270    let (mut s, mut e) = (0f64, 0f64);
271    for (a, b) in x.iter().zip(reference) {
272        s += (*b as f64).powi(2);
273        e += (*a as f64 - *b as f64).powi(2);
274    }
275    10.0 * (s / e).log10()
276}
taconite_reuse/lib.rs

taconite_reuse/
lib.rs
