open_stomp/nam/
mod.rs

1//! Neural Amp Modeler (`.nam`) models (feature `nam`): loading and inference.
2//!
3//! A pure-Rust port of the parts of [NeuralAmpModelerCore](https://github.com/sdatkinson/NeuralAmpModelerCore)
4//! (MIT) that the pedal's converter needs: WaveNet models, plain or inside a `SlimmableContainer`
5//! (NAM 0.5 "classic" and 0.7 files). The converter (docs/05) runs test signals through the model,
6//! so this must match NeuralAmpModelerCore numerically; `tests/nam_inference.rs` checks it against
7//! the vendor converter's recorded model output where available.
8//!
9//! ```no_run
10//! # fn main() -> Result<(), open_stomp::nam::NamError> {
11//! let mut model = open_stomp::nam::Model::load("amp.nam")?;
12//! let input = vec![0.0f32; 512];
13//! let mut output = vec![0.0f32; 512];
14//! model.process(&input, &mut output);
15//! # Ok(()) }
16//! ```
17
18mod activation;
19pub mod convert;
20pub mod optimise;
21mod wavenet;
22
23use std::path::Path;
24
25use serde_json::Value;
26
27pub use activation::Activation;
28pub use convert::{Converted, convert};
29pub use wavenet::WaveNet;
30
31#[derive(Debug)]
32pub enum NamError {
33    Io(std::io::Error),
34    Json(serde_json::Error),
35    /// The file is not a valid model.
36    Format(String),
37    /// A valid model using a feature this port doesn't implement.
38    Unsupported(String),
39    /// The weight list doesn't fit the architecture.
40    Weights { expected: usize, got: usize },
41}
42
43impl std::fmt::Display for NamError {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            NamError::Io(e) => write!(f, "{e}"),
47            NamError::Json(e) => write!(f, "invalid JSON: {e}"),
48            NamError::Format(s) => write!(f, "invalid model: {s}"),
49            NamError::Unsupported(s) => write!(f, "unsupported model feature: {s}"),
50            NamError::Weights { expected, got } => write!(f, "weight count mismatch: model needs {expected}, file has {got}"),
51        }
52    }
53}
54
55impl std::error::Error for NamError {
56    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
57        match self {
58            NamError::Io(e) => Some(e),
59            NamError::Json(e) => Some(e),
60            _ => None,
61        }
62    }
63}
64
65impl From<std::io::Error> for NamError {
66    fn from(e: std::io::Error) -> Self {
67        NamError::Io(e)
68    }
69}
70
71impl From<serde_json::Error> for NamError {
72    fn from(e: serde_json::Error) -> Self {
73        NamError::Json(e)
74    }
75}
76
77/// A loaded model, ready to process audio (mono, at its sample rate, normally 48 kHz).
78#[derive(Clone, Debug)]
79pub struct Model {
80    net: WaveNet,
81    /// Display name from the file's metadata, if any.
82    pub name: Option<String>,
83    /// The architecture as written in the file (`"WaveNet"`, `"SlimmableContainer"`).
84    pub architecture: String,
85    /// The model's loudness in dB, from the file's metadata (written by the NAM trainer: the
86    /// model's output level on its standard input). NAM plugins normalise the output with it.
87    pub loudness: Option<f64>,
88}
89
90fn weights_of(v: &Value) -> Result<Vec<f32>, NamError> {
91    match v.get("weights") {
92        Some(Value::Array(a)) => a
93            .iter()
94            .map(|x| x.as_f64().map(|f| f as f32).ok_or_else(|| NamError::Format("non-numeric weight".into())))
95            .collect(),
96        None => Ok(vec![]),
97        _ => Err(NamError::Format("weights must be an array".into())),
98    }
99}
100
101fn build(v: &Value) -> Result<WaveNet, NamError> {
102    let arch = v.get("architecture").and_then(Value::as_str).unwrap_or("");
103    let config = v.get("config").ok_or_else(|| NamError::Format("missing config".into()))?;
104    let sample_rate = v.get("sample_rate").and_then(Value::as_f64);
105    match arch {
106        "WaveNet" => WaveNet::from_json(config, &weights_of(v)?, sample_rate),
107        "SlimmableContainer" => {
108            // NeuralAmpModelerCore's ContainerModel defaults to the last (full-size) submodel.
109            let subs = config
110                .get("submodels")
111                .and_then(Value::as_array)
112                .filter(|s| !s.is_empty())
113                .ok_or_else(|| NamError::Format("SlimmableContainer without submodels".into()))?;
114            let last = subs.last().and_then(|s| s.get("model")).ok_or_else(|| NamError::Format("submodel".into()))?;
115            build(last)
116        }
117        other => Err(NamError::Unsupported(format!("architecture {other:?}"))),
118    }
119}
120
121impl Model {
122    pub fn load(path: impl AsRef<Path>) -> Result<Model, NamError> {
123        Model::from_json_str(&std::fs::read_to_string(path)?)
124    }
125
126    pub fn from_json_str(s: &str) -> Result<Model, NamError> {
127        let v: Value = serde_json::from_str(s)?;
128        let net = build(&v)?;
129        let name = v.get("metadata").and_then(|m| m.get("name")).and_then(Value::as_str).map(str::to_owned);
130        let architecture = v.get("architecture").and_then(Value::as_str).unwrap_or("").to_owned();
131        let loudness = v.get("metadata").and_then(|m| m.get("loudness")).and_then(Value::as_f64);
132        Ok(Model { net, name, architecture, loudness })
133    }
134
135    /// Process a block (any length; state carries over between calls).
136    pub fn process(&mut self, input: &[f32], output: &mut [f32]) {
137        for (i, o) in input.chunks(4096).zip(output.chunks_mut(4096)) {
138            self.net.process(i, o);
139        }
140    }
141
142    /// Back to the state after loading (history cleared, then prewarmed with silence).
143    pub fn reset(&mut self) {
144        self.net.reset();
145    }
146
147    /// The sample rate the model was trained at, if the file says.
148    pub fn sample_rate(&self) -> Option<f64> {
149        self.net.sample_rate()
150    }
151}
open_stomp/nam/mod.rs

open_stomp/nam/
mod.rs
