ferraus_denoise/
ml.rs

1//! Model-based denoising: RNNoise (pure-Rust `nnnoiseless`, weights embedded) and the
2//! model slots for DeepFilterNet 3, music-aware denoise, and de-reverb checkpoints.
3
4use ferraus_core::{AudioBuffer, Error, ErrorKind, OpCtx, Result};
5use serde_json::{json, Value};
6
7/// A registry slot for a denoise-family model. The models crate can list these to know
8/// which names `ferraus models pull` should offer for denoising.
9#[derive(Clone, Debug)]
10pub struct ModelSlot {
11    /// Registry name pattern (`denoise/dfn3`, `denoise/music-*`, `denoise/dereverb-*`).
12    pub pattern: &'static str,
13    /// The concrete name used in hints.
14    pub default_name: &'static str,
15    pub purpose: &'static str,
16    pub domain: &'static str,
17    pub license: &'static str,
18    /// Whether this build can run the model once it is cached.
19    pub runnable: bool,
20}
21
22/// Every denoise model slot this crate knows about.
23pub fn model_slots() -> Vec<ModelSlot> {
24    vec![
25        ModelSlot {
26            pattern: "denoise/rnnoise",
27            default_name: "denoise/rnnoise",
28            purpose: "RNNoise speech denoiser (weights embedded in the binary; no download)",
29            domain: "speech",
30            license: "BSD-3-Clause",
31            runnable: cfg!(feature = "rnnoise"),
32        },
33        ModelSlot {
34            pattern: "denoise/dfn3",
35            default_name: "denoise/dfn3",
36            purpose: "DeepFilterNet 3 speech enhancement at 48 kHz",
37            domain: "speech",
38            license: "MIT OR Apache-2.0",
39            runnable: cfg!(feature = "dfn3"),
40        },
41        ModelSlot {
42            pattern: "denoise/music-*",
43            default_name: "denoise/music-roformer",
44            purpose: "Music-aware denoise checkpoints (community RoFormer); shipped only with an acceptable license",
45            domain: "music",
46            license: "per checkpoint",
47            runnable: false,
48        },
49        ModelSlot {
50            pattern: "denoise/dereverb-*",
51            default_name: "denoise/dereverb-roformer",
52            purpose: "De-reverb checkpoints (best-effort)",
53            domain: "speech or music, per checkpoint",
54            license: "per checkpoint",
55            runnable: false,
56        },
57    ]
58}
59
60/// Look for a cached model whose registry name matches `prefix` (e.g. `denoise/dereverb-`)
61/// under the models directory. Returns the first match.
62pub fn find_cached(prefix: &str) -> Option<std::path::PathBuf> {
63    let dir = ferraus_core::paths::models_dir();
64    let (family, stem) = prefix.split_once('/')?;
65    let fam = dir.join(family);
66    let mut entries: Vec<_> = std::fs::read_dir(&fam).ok()?.filter_map(|e| e.ok()).collect();
67    entries.sort_by_key(|e| e.file_name());
68    entries.into_iter().find(|e| e.file_name().to_string_lossy().starts_with(stem)).map(|e| e.path())
69}
70
71fn not_cached(name: &str, what: &str) -> Error {
72    Error::new(
73        ErrorKind::ModelNotCached,
74        "model_not_cached",
75        format!("{what} needs the model '{name}', which is not in the model cache"),
76    )
77    .with_hint(format!("ferraus models pull {name}"))
78}
79
80/// Resolve a model slot that this build cannot run itself. Always an error: either the
81/// model is missing (exit 5 with a pull hint) or it is cached but no runtime for it is
82/// compiled into this build. Never falls back silently.
83pub fn unavailable_slot(prefix: &str, default_name: &str, what: &str) -> Error {
84    match find_cached(prefix) {
85        None => not_cached(default_name, what),
86        Some(p) => Error::unsupported(format!(
87            "{what}: found {} in the model cache, but this build has no runtime for it yet",
88            p.display()
89        ))
90        .with_hint("ferraus doctor"),
91    }
92}
93
94/// Run DeepFilterNet 3 (`denoise/dfn3`, through the model cache and its download policy)
95/// on every channel.
96#[cfg(feature = "dfn3")]
97pub fn dfn3(buf: &AudioBuffer, ctx: &mut OpCtx) -> Result<AudioBuffer> {
98    let model = crate::dfn3::Dfn3::load(ctx)?;
99    let c: &OpCtx = ctx;
100    model.enhance(buf, &mut |f| {
101        c.set_progress(f);
102        c.check_cancel()
103    })
104}
105
106/// DeepFilterNet 3 in a build without the `dfn3` feature: always an explicit error.
107#[cfg(not(feature = "dfn3"))]
108pub fn dfn3(_buf: &AudioBuffer, _ctx: &mut OpCtx) -> Result<AudioBuffer> {
109    Err(dfn3_unavailable())
110}
111
112/// The error for `--model dfn3` in a build without the `dfn3` feature.
113pub fn dfn3_unavailable() -> Error {
114    match find_cached("denoise/dfn3") {
115        None => not_cached("denoise/dfn3", "--model dfn3 (DeepFilterNet 3)"),
116        Some(p) => Error::unsupported(format!(
117            "--model dfn3: found {} in the model cache, but this build was compiled without the DeepFilterNet runtime; use --model rnnoise or the profile path",
118            p.display()
119        ))
120        .with_hint("ferraus denoise <input> --mode ml --model rnnoise -o clean.wav"),
121    }
122}
123
124/// Model identity recorded in provenance for RNNoise.
125pub const RNNOISE_NAME: &str = "denoise/rnnoise";
126pub const RNNOISE_VERSION: &str = "nnnoiseless-0.5.2";
127/// SHA-256 of the `weights.rnn` embedded by nnnoiseless 0.5.2.
128pub const RNNOISE_SHA256: &str = "e6de5fbfadf7ec91d1b24d6a6ccfd0290cb4d8bf555c5eab3ce41506f67a58b1";
129/// RNNoise's internal rate and its algorithmic delay (one 10 ms frame).
130pub const RNNOISE_RATE: u32 = 48_000;
131
132/// Run RNNoise on every channel independently (48 kHz internally, resampled in and out).
133#[cfg(feature = "rnnoise")]
134pub fn rnnoise(buf: &AudioBuffer) -> Result<AudioBuffer> {
135    use ferraus_core::resample::{resample, Quality};
136    use nnnoiseless::DenoiseState;
137    let frame = DenoiseState::FRAME_SIZE;
138    let x = resample(buf, RNNOISE_RATE, Quality::High);
139    let len = x.frames();
140    let mut data = Vec::with_capacity(x.channels());
141    for ch in &x.data {
142        let mut st = DenoiseState::new();
143        let mut out = Vec::with_capacity(len + 2 * frame);
144        let mut inb = vec![0f32; frame];
145        let mut outb = vec![0f32; frame];
146        let mut pos = 0;
147        // one extra frame flushes the algorithmic delay
148        while pos < len + frame {
149            for (i, v) in inb.iter_mut().enumerate() {
150                *v = ch.get(pos + i).copied().unwrap_or(0.0) * 32768.0;
151            }
152            st.process_frame(&mut outb, &inb);
153            out.extend(outb.iter().map(|v| v / 32768.0));
154            pos += frame;
155        }
156        // compensate the one-frame delay
157        out.drain(..frame.min(out.len()));
158        out.truncate(len);
159        data.push(out);
160    }
161    let y = AudioBuffer::from_planar(RNNOISE_RATE, data);
162    let mut y = resample(&y, buf.sample_rate, Quality::High);
163    for c in &mut y.data {
164        c.resize(buf.frames(), 0.0);
165    }
166    Ok(y)
167}
168
169#[cfg(not(feature = "rnnoise"))]
170pub fn rnnoise(_buf: &AudioBuffer) -> Result<AudioBuffer> {
171    Err(Error::unsupported("--model rnnoise: this build was compiled without the 'rnnoise' feature"))
172}
173
174/// The JSON note attached to every model-path run.
175pub fn speech_note(model: &str) -> Value {
176    json!({
177        "model": model,
178        "domain": "speech",
179        "note": "speech-tuned model; on music the profile path (--mode profile) usually sounds better",
180    })
181}
ferraus_denoise/ml.rs

ferraus_denoise/
ml.rs
