feat: add NAM neural amp modeler stage (#243)

Author: OpenSauce

Cargo.lock

@@ -17,6 +17,7 @@ rustfft = "6.4"
 realfft = "3.5"
 arc-swap = "1.8"
 assert_no_alloc = { version = "1.1", features = ["warn_debug"] }
+nam-rs = "0.1.0"
 
 [dev-dependencies]
 criterion = { version = "0.8", features = ["html_reports"] }

nam/.gitkeep

@@ -6,6 +6,7 @@ pub mod eq;
 pub mod filter;
 pub mod level;
 pub mod multiband_saturator;
+pub mod nam;
 pub mod noise_gate;
 pub mod poweramp;
 pub mod preamp;

rustortion-core/Cargo.toml

@@ -0,0 +1,193 @@
+use log::warn;
+use nam_rs::WaveNet;
+use serde::{Deserialize, Serialize};
+
+use crate::amp::stages::Stage;
+use crate::amp::stages::common::db_to_lin;
+use crate::nam::registry;
+
+/// Valid range for the input/output gain knobs, matching the UI and plugin params.
+const GAIN_DB_MIN: f32 = -24.0;
+const GAIN_DB_MAX: f32 = 24.0;
+
+/// A Neural Amp Modeler stage running a WaveNet `.nam` model.
+///
+/// With no model loaded the stage is a passthrough. Input/output gain are applied
+/// around the model and the wet output is blended with the dry signal via `mix`.
+pub struct NamStage {
+    wavenet: Option<WaveNet>,
+    input_gain: f32,
+    output_gain: f32,
+    mix: f32,
+    /// Native sample rate of the loaded model (0.0 if none), for UI display.
+    native_sample_rate: f32,
+    /// True if the model's native rate differs from the engine rate.
+    sample_rate_mismatch: bool,
+}
+
+impl NamStage {
+    const fn passthrough(input_gain: f32, output_gain: f32, mix: f32) -> Self {
+        Self {
+            wavenet: None,
+            input_gain,
+            output_gain,
+            mix,
+            native_sample_rate: 0.0,
+            sample_rate_mismatch: false,
+        }
+    }
+}
+
+impl Stage for NamStage {
+    fn process(&mut self, input: f32) -> f32 {
+        let Some(wavenet) = self.wavenet.as_mut() else {
+            return input;
+        };
+        let wet = wavenet.process_sample(input * self.input_gain) * self.output_gain;
+        self.mix.mul_add(wet - input, input)
+    }
+
+    fn set_parameter(&mut self, name: &str, value: f32) -> Result<(), &'static str> {
+        match name {
+            "input_gain_db" => {
+                if (GAIN_DB_MIN..=GAIN_DB_MAX).contains(&value) {
+                    self.input_gain = db_to_lin(value);
+                    Ok(())
+                } else {
+                    Err("Input gain must be between -24 and 24 dB")
+                }
+            }
+            "output_gain_db" => {
+                if (GAIN_DB_MIN..=GAIN_DB_MAX).contains(&value) {
+                    self.output_gain = db_to_lin(value);
+                    Ok(())
+                } else {
+                    Err("Output gain must be between -24 and 24 dB")
+                }
+            }
+            "mix" => {
+                if (0.0..=1.0).contains(&value) {
+                    self.mix = value;
+                    Ok(())
+                } else {
+                    Err("Mix must be between 0.0 and 1.0")
+                }
+            }
+            "native_sample_rate" | "sample_rate_mismatch" => Err("Parameter is read-only"),
+            _ => Err("Unknown parameter"),
+        }
+    }
+
+    fn get_parameter(&self, name: &str) -> Result<f32, &'static str> {
+        match name {
+            "input_gain_db" => Ok(20.0 * self.input_gain.log10()),
+            "output_gain_db" => Ok(20.0 * self.output_gain.log10()),
+            "mix" => Ok(self.mix),
+            "native_sample_rate" => Ok(self.native_sample_rate),
+            "sample_rate_mismatch" => Ok(f32::from(u8::from(self.sample_rate_mismatch))),
+            _ => Err("Unknown parameter name"),
+        }
+    }
+}
+
+// --- Config ---
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct NamConfig {
+    /// Display name of the selected model, or `None` for passthrough.
+    #[serde(default)]
+    pub model_name: Option<String>,
+    pub input_gain_db: f32,
+    pub output_gain_db: f32,
+    pub mix: f32,
+    #[serde(default)]
+    pub bypassed: bool,
+}
+
+impl Default for NamConfig {
+    fn default() -> Self {
+        Self {
+            model_name: None,
+            input_gain_db: 0.0,
+            output_gain_db: 0.0,
+            mix: 1.0,
+            bypassed: false,
+        }
+    }
+}
+
+impl NamConfig {
+    /// Build a runnable stage. Resolves the model from the global registry and
+    /// allocates the `WaveNet` here (off the real-time thread). On any failure the
+    /// stage falls back to passthrough with a warning.
+    pub fn to_stage(&self, sample_rate: f32) -> NamStage {
+        let input_gain = db_to_lin(self.input_gain_db.clamp(GAIN_DB_MIN, GAIN_DB_MAX));
+        let output_gain = db_to_lin(self.output_gain_db.clamp(GAIN_DB_MIN, GAIN_DB_MAX));
+        let mix = self.mix.clamp(0.0, 1.0);
+
+        let Some(name) = self.model_name.as_deref() else {
+            return NamStage::passthrough(input_gain, output_gain, mix);
+        };
+
+        let Some(model) = registry::get(name) else {
+            warn!("NAM model '{name}' not found in registry; using passthrough");
+            return NamStage::passthrough(input_gain, output_gain, mix);
+        };
+
+        let native_sample_rate = model.sample_rate() as f32;
+        let sample_rate_mismatch = (native_sample_rate - sample_rate).abs() > 1.0;
+        if sample_rate_mismatch {
+            warn!(
+                "NAM model '{name}' native rate {native_sample_rate} Hz differs from engine \
+                 rate {sample_rate} Hz; tone may be affected"
+            );
+        }
+
+        match WaveNet::new(&model) {
+            Ok(wavenet) => NamStage {
+                wavenet: Some(wavenet),
+                input_gain,
+                output_gain,
+                mix,
+                native_sample_rate,
+                sample_rate_mismatch,
+            },
+            Err(e) => {
+                warn!("Failed to build NAM model '{name}': {e}; using passthrough");
+                NamStage::passthrough(input_gain, output_gain, mix)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn passthrough_when_no_model() {
+        let stage = NamConfig::default().to_stage(48_000.0);
+        let mut stage = stage;
+        for x in [-1.0, 0.0, 0.25, 0.9] {
+            assert_eq!(stage.process(x), x);
+        }
+    }
+
+    #[test]
+    fn gain_and_mix_round_trip() {
+        let mut stage = NamConfig::default().to_stage(48_000.0);
+        stage.set_parameter("mix", 0.5).unwrap();
+        assert!((stage.get_parameter("mix").unwrap() - 0.5).abs() < 1e-6);
+
+        stage.set_parameter("input_gain_db", 6.0).unwrap();
+        assert!((stage.get_parameter("input_gain_db").unwrap() - 6.0).abs() < 1e-3);
+
+        assert!(stage.set_parameter("mix", 2.0).is_err());
+        assert!(stage.set_parameter("native_sample_rate", 1.0).is_err());
+
+        // Gains outside ±24 dB (and NaN) are rejected.
+        assert!(stage.set_parameter("input_gain_db", 30.0).is_err());
+        assert!(stage.set_parameter("output_gain_db", -30.0).is_err());
+        assert!(stage.set_parameter("input_gain_db", f32::NAN).is_err());
+    }
+}

rustortion-core/src/amp/stages/mod.rs

@@ -33,5 +33,6 @@ pub mod amp;
 pub mod audio;
 pub mod ir;
 pub mod metronome;
+pub mod nam;
 pub mod preset;
 pub mod tuner;

rustortion-core/src/amp/stages/nam.rs

@@ -0,0 +1,88 @@
+use std::collections::BTreeMap;
+use std::path::Path;
+use std::sync::Arc;
+
+use anyhow::{Context, Result};
+use log::{info, warn};
+use nam_rs::NamModel;
+
+/// Scans a directory for `*.nam` files and parses each into memory.
+///
+/// Parsing happens once, at construction (off the real-time thread). Models are
+/// keyed by display name (the file stem). Unparseable files are skipped with a
+/// warning rather than failing the whole scan, matching the IR loader's tolerant
+/// behaviour.
+pub struct NamLoader {
+    models: BTreeMap<String, Arc<NamModel>>,
+}
+
+impl NamLoader {
+    /// Scan `directory` and parse every `*.nam` file found.
+    ///
+    /// A missing directory yields an empty loader (warn, not error) so the app can
+    /// run without a nam folder present.
+    pub fn new(directory: &Path) -> Result<Self> {
+        let mut models = BTreeMap::new();
+
+        if !directory.is_dir() {
+            warn!(
+                "NAM directory '{}' does not exist; no models loaded",
+                directory.display()
+            );
+            return Ok(Self { models });
+        }
+
+        let entries = std::fs::read_dir(directory)
+            .with_context(|| format!("Failed to read NAM directory '{}'", directory.display()))?;
+
+        for entry in entries {
+            let entry = match entry {
+                Ok(entry) => entry,
+                Err(e) => {
+                    warn!("Skipping unreadable entry in NAM directory: {e}");
+                    continue;
+                }
+            };
+            let path = entry.path();
+            if path.extension().and_then(|e| e.to_str()) != Some("nam") {
+                continue;
+            }
+            let Some(name) = path.file_stem().and_then(|s| s.to_str()).map(str::to_owned) else {
+                continue;
+            };
+
+            match std::fs::read_to_string(&path)
+                .map_err(anyhow::Error::from)
+                .and_then(|json| NamModel::from_json_str(&json).map_err(anyhow::Error::from))
+            {
+                Ok(model) => {
+                    info!(
+                        "Loaded NAM model '{name}' ({} Hz)",
+                        model.sample_rate() as u32
+                    );
+                    models.insert(name, Arc::new(model));
+                }
+                Err(e) => warn!("Skipping NAM file '{}': {e}", path.display()),
+            }
+        }
+
+        Ok(Self { models })
+    }
+
+    /// Sorted list of available model display names.
+    #[must_use]
+    pub fn available_names(&self) -> Vec<String> {
+        self.models.keys().cloned().collect()
+    }
+
+    /// Look up a parsed model by display name.
+    #[must_use]
+    pub fn get(&self, name: &str) -> Option<Arc<NamModel>> {
+        self.models.get(name).cloned()
+    }
+
+    /// All parsed models, for populating the global registry.
+    pub fn models(&self) -> impl Iterator<Item = (&String, &Arc<NamModel>)> {
+        self.models.iter()
+    }
+}

rustortion-core/src/lib.rs

@@ -0,0 +1,12 @@
+//! NAM (Neural Amp Modeler) model loading and a process-global parsed-model
+//! registry.
+//!
+//! `.nam` models are parsed (and the `WaveNet` allocated) off the real-time thread.
+//! The [`loader`] scans a directory and parses every `*.nam` file into memory at
+//! startup; the [`registry`] makes those parsed models reachable from
+//! `StageConfig::to_runtime`, which has no other handle to the loader.
+
+pub mod loader;
+pub mod registry;
+
+pub use loader::NamLoader;

rustortion-core/src/nam/loader.rs

@@ -0,0 +1,53 @@
+//! Process-global registry of parsed NAM models.
+//!
+//! `StageConfig::to_runtime(sample_rate)` builds stages off the real-time thread but
+//! has no handle to the [`NamLoader`](super::loader::NamLoader). Since the nam folder
+//! is a singleton resource, a process-global registry lets `NamConfig::to_runtime`
+//! resolve a model by name without threading a loader through every `to_runtime` call
+//! site.
+
+use std::collections::HashMap;
+use std::sync::{Arc, OnceLock, RwLock};
+
+use nam_rs::NamModel;
+
+use super::loader::NamLoader;
+
+type Store = RwLock<HashMap<String, Arc<NamModel>>>;
+
+static NAM_REGISTRY: OnceLock<Store> = OnceLock::new();
+
+fn store() -> &'static Store {
+    NAM_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
+}
+
+/// Populate (or replace) the global registry from a loader's parsed models.
+pub fn init_from_loader(loader: &NamLoader) {
+    let mut map = store()
+        .write()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    map.clear();
+    for (name, model) in loader.models() {
+        map.insert(name.clone(), Arc::clone(model));
+    }
+}
+
+/// Look up a parsed model by display name.
+#[must_use]
+pub fn get(name: &str) -> Option<Arc<NamModel>> {
+    let map = store()
+        .read()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    map.get(name).cloned()
+}
+
+/// Sorted list of available model display names.
+#[must_use]
+pub fn available_names() -> Vec<String> {
+    let map = store()
+        .read()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    let mut names: Vec<String> = map.keys().cloned().collect();
+    names.sort();
+    names
+}