open_stomp/nam/
wavenet.rs

1//! NAM WaveNet inference, ported from NeuralAmpModelerCore (`NAM/wavenet/model.cpp`, MIT).
2//!
3//! Supported: any number of layer arrays; per-layer kernel sizes, dilations and activations;
4//! gating `none` and `gated` (including the legacy `"gated": true`); bottleneck; `layer1x1` and
5//! `head1x1`; heads with a kernel (`"head": {...}`) or the legacy `head_size`/`head_bias`.
6//! Not supported (refused with an error): FiLM, grouped convolutions, blended gating, a
7//! post-stack head and a condition DSP.
8//!
9//! Data is frame-major: `x[t * channels + c]`.
10
11use serde_json::Value;
12
13use super::NamError;
14use super::activation::Activation;
15
16/// Pull `n` weights from the flat weight list.
17struct Weights<'a> {
18    w: &'a [f32],
19    pos: usize,
20}
21
22impl Weights<'_> {
23    fn take(&mut self, n: usize) -> Result<&[f32], NamError> {
24        let s = self.w.get(self.pos..self.pos + n).ok_or(NamError::Weights { expected: self.pos + n, got: self.w.len() })?;
25        self.pos += n;
26        Ok(s)
27    }
28}
29
30/// A causal dilated 1-D convolution with its own input history.
31#[derive(Clone, Debug)]
32struct Conv1d {
33    cin: usize,
34    cout: usize,
35    kernel: usize,
36    dilation: usize,
37    /// `w[(k * cout + o) * cin + i]`; tap `k = kernel - 1` is the current sample.
38    w: Vec<f32>,
39    bias: Option<Vec<f32>>,
40    /// The last `(kernel - 1) * dilation` input frames.
41    history: Vec<f32>,
42    ext: Vec<f32>,
43}
44
45impl Conv1d {
46    fn new(cin: usize, cout: usize, kernel: usize, dilation: usize, bias: bool) -> Self {
47        Conv1d {
48            cin,
49            cout,
50            kernel,
51            dilation,
52            w: vec![0.0; kernel * cout * cin],
53            bias: bias.then(|| vec![0.0; cout]),
54            history: vec![0.0; (kernel - 1) * dilation * cin],
55            ext: vec![],
56        }
57    }
58
59    /// NeuralAmpModelerCore's flattening: for out, for in, for tap; then the bias.
60    fn load(&mut self, w: &mut Weights) -> Result<(), NamError> {
61        let src = w.take(self.cout * self.cin * self.kernel)?;
62        let mut it = src.iter();
63        for o in 0..self.cout {
64            for i in 0..self.cin {
65                for k in 0..self.kernel {
66                    self.w[(k * self.cout + o) * self.cin + i] = *it.next().expect("sized");
67                }
68            }
69        }
70        if let Some(b) = &mut self.bias {
71            let n = b.len();
72            b.copy_from_slice(w.take(n)?);
73        }
74        Ok(())
75    }
76
77    fn receptive_field(&self) -> usize {
78        (self.kernel - 1) * self.dilation
79    }
80
81    /// `out` gets `n` frames of `cout` values.
82    fn process(&mut self, input: &[f32], n: usize, out: &mut Vec<f32>) {
83        let (cin, cout, h) = (self.cin, self.cout, self.history.len() / self.cin.max(1));
84        self.ext.clear();
85        self.ext.extend_from_slice(&self.history);
86        self.ext.extend_from_slice(&input[..n * cin]);
87        out.clear();
88        out.resize(n * cout, 0.0);
89        for t in 0..n {
90            let o_row = &mut out[t * cout..(t + 1) * cout];
91            if let Some(b) = &self.bias {
92                o_row.copy_from_slice(b);
93            }
94            for k in 0..self.kernel {
95                // frame in ext: h + t - (kernel - 1 - k) * dilation
96                let src_t = h + t - (self.kernel - 1 - k) * self.dilation;
97                let x = &self.ext[src_t * cin..(src_t + 1) * cin];
98                let wk = &self.w[k * cout * cin..(k + 1) * cout * cin];
99                for (o, acc) in o_row.iter_mut().enumerate() {
100                    let wr = &wk[o * cin..(o + 1) * cin];
101                    *acc += wr.iter().zip(x).map(|(a, b)| a * b).sum::<f32>();
102                }
103            }
104        }
105        if h > 0 {
106            let total = self.ext.len();
107            self.history.copy_from_slice(&self.ext[total - h * cin..]);
108        }
109    }
110
111    fn reset(&mut self) {
112        self.history.iter_mut().for_each(|x| *x = 0.0);
113    }
114}
115
116#[derive(Clone, Copy, Debug, PartialEq, Eq)]
117enum Gating {
118    None,
119    Gated,
120}
121
122#[derive(Clone, Debug)]
123struct Layer {
124    conv: Conv1d,
125    mixin: Conv1d,
126    activation: Activation,
127    gating: Gating,
128    secondary: Activation,
129    bottleneck: usize,
130    layer1x1: Option<Conv1d>,
131    head1x1: Option<Conv1d>,
132    // scratch
133    z: Vec<f32>,
134    mix: Vec<f32>,
135    tmp: Vec<f32>,
136}
137
138impl Layer {
139    /// `input`: n × channels; `cond`: n × condition_size. Writes the head contribution and the
140    /// next layer's input.
141    fn process(&mut self, input: &[f32], cond: &[f32], n: usize, head: &mut Vec<f32>, next: &mut Vec<f32>) {
142        let zc = self.conv.cout;
143        let b = self.bottleneck;
144        self.conv.process(input, n, &mut self.z);
145        self.mixin.process(cond, n, &mut self.mix);
146        for (z, m) in self.z.iter_mut().zip(&self.mix) {
147            *z += m;
148        }
149        // activated layer output: n × bottleneck
150        let act: &mut Vec<f32> = &mut self.tmp;
151        match self.gating {
152            Gating::None => {
153                self.activation.apply(&mut self.z, zc);
154                act.clear();
155                act.extend_from_slice(&self.z);
156            }
157            Gating::Gated => {
158                act.clear();
159                act.resize(n * b, 0.0);
160                for t in 0..n {
161                    for c in 0..b {
162                        let p = self.activation.apply_one(self.z[t * zc + c], c);
163                        let g = self.secondary.apply_one(self.z[t * zc + b + c], c);
164                        act[t * b + c] = p * g;
165                    }
166                }
167            }
168        }
169        match &mut self.head1x1 {
170            Some(h) => h.process(act, n, head),
171            None => {
172                head.clear();
173                head.extend_from_slice(act);
174            }
175        }
176        let channels = self.conv.cin;
177        next.clear();
178        next.extend_from_slice(&input[..n * channels]);
179        if let Some(l) = &mut self.layer1x1 {
180            l.process(act, n, &mut self.mix);
181            for (x, y) in next.iter_mut().zip(&self.mix) {
182                *x += y;
183            }
184        }
185    }
186
187    fn reset(&mut self) {
188        self.conv.reset();
189    }
190}
191
192#[derive(Clone, Debug)]
193struct LayerArray {
194    rechannel: Conv1d,
195    layers: Vec<Layer>,
196    head_rechannel: Conv1d,
197    head_in: usize,
198    // scratch
199    a: Vec<f32>,
200    b: Vec<f32>,
201    head_sum: Vec<f32>,
202    head_part: Vec<f32>,
203}
204
205impl LayerArray {
206    fn from_json(cfg: &Value) -> Result<LayerArray, NamError> {
207        let int = |k: &str| cfg.get(k).and_then(Value::as_u64).map(|x| x as usize);
208        let req = |k: &str| int(k).ok_or_else(|| NamError::Format(format!("layer array: missing {k}")));
209        let input_size = req("input_size")?;
210        let condition_size = req("condition_size")?;
211        let channels = req("channels")?;
212        let bottleneck = int("bottleneck").unwrap_or(channels);
213        for key in ["groups_input", "groups_input_mixin"] {
214            if int(key).is_some_and(|g| g != 1) {
215                return Err(NamError::Unsupported(format!("{key} != 1")));
216            }
217        }
218        for key in [
219            "conv_pre_film",
220            "conv_post_film",
221            "input_mixin_pre_film",
222            "input_mixin_post_film",
223            "activation_pre_film",
224            "activation_post_film",
225            "layer1x1_post_film",
226            "head1x1_post_film",
227        ] {
228            if cfg.get(key).and_then(|f| f.get("active")).and_then(Value::as_bool) == Some(true) {
229                return Err(NamError::Unsupported(format!("FiLM ({key})")));
230            }
231        }
232        let sub_active = |key: &str, default: bool| -> Result<(bool, usize), NamError> {
233            match cfg.get(key) {
234                Some(v) if !v.is_null() => {
235                    if v.get("groups").and_then(Value::as_u64).is_some_and(|g| g != 1) {
236                        return Err(NamError::Unsupported(format!("{key} groups != 1")));
237                    }
238                    let active = v.get("active").and_then(Value::as_bool).unwrap_or(default);
239                    let out = v.get("out_channels").and_then(Value::as_u64).map_or(0, |x| x as usize);
240                    Ok((active, out))
241                }
242                _ => Ok((default, 0)),
243            }
244        };
245        let (layer1x1, _) = sub_active("layer1x1", true)?;
246        let (head1x1, head1x1_out) = sub_active("head1x1", false)?;
247        if !layer1x1 && bottleneck != channels {
248            return Err(NamError::Format("bottleneck must equal channels without layer1x1".into()));
249        }
250
251        let (head_size, head_kernel, head_bias, head_dilation) = match cfg.get("head") {
252            Some(h) if !h.is_null() => (
253                h.get("out_channels").and_then(Value::as_u64).ok_or_else(|| NamError::Format("head.out_channels".into()))?
254                    as usize,
255                h.get("kernel_size").and_then(Value::as_u64).unwrap_or(1) as usize,
256                h.get("bias").and_then(Value::as_bool).unwrap_or(false),
257                h.get("head_dilation").and_then(Value::as_u64).unwrap_or(1) as usize,
258            ),
259            _ => (req("head_size")?, 1, cfg.get("head_bias").and_then(Value::as_bool).unwrap_or(false), 1),
260        };
261
262        let dilations: Vec<usize> = cfg
263            .get("dilations")
264            .and_then(Value::as_array)
265            .ok_or_else(|| NamError::Format("layer array: missing dilations".into()))?
266            .iter()
267            .map(|d| d.as_u64().map(|x| x as usize).ok_or_else(|| NamError::Format("dilation".into())))
268            .collect::<Result<_, _>>()?;
269        let n = dilations.len();
270        let kernels: Vec<usize> = match (cfg.get("kernel_sizes").and_then(Value::as_array), int("kernel_size")) {
271            (Some(a), _) => a.iter().map(|k| k.as_u64().unwrap_or(1) as usize).collect(),
272            (None, Some(k)) => vec![k; n],
273            _ => return Err(NamError::Format("layer array: missing kernel_size(s)".into())),
274        };
275        let per_layer = |key: &str| -> Option<Vec<Value>> {
276            match cfg.get(key) {
277                Some(Value::Array(a)) => Some(a.clone()),
278                Some(v) if !v.is_null() => Some(vec![v.clone(); n]),
279                _ => None,
280            }
281        };
282        let activations: Vec<Activation> = per_layer("activation")
283            .ok_or_else(|| NamError::Format("layer array: missing activation".into()))?
284            .iter()
285            .map(Activation::from_json)
286            .collect::<Result<_, _>>()?;
287        let gatings: Vec<Gating> = match per_layer("gating_mode") {
288            Some(v) => v
289                .iter()
290                .map(|g| match g.as_str() {
291                    Some("none") => Ok(Gating::None),
292                    Some("gated") => Ok(Gating::Gated),
293                    other => Err(NamError::Unsupported(format!("gating_mode {other:?}"))),
294                })
295                .collect::<Result<_, _>>()?,
296            None => vec![if cfg.get("gated").and_then(Value::as_bool) == Some(true) { Gating::Gated } else { Gating::None }; n],
297        };
298        let secondaries: Vec<Activation> = match per_layer("secondary_activation") {
299            Some(v) => v
300                .iter()
301                .map(|a| if a.is_null() { Ok(Activation::Sigmoid) } else { Activation::from_json(a) })
302                .collect::<Result<_, _>>()?,
303            None => vec![Activation::Sigmoid; n],
304        };
305        if [kernels.len(), activations.len(), gatings.len(), secondaries.len()].iter().any(|&l| l != n) {
306            return Err(NamError::Format("layer array: per-layer lists differ in length".into()));
307        }
308
309        let layers = (0..n)
310            .map(|i| {
311                let zc = if gatings[i] == Gating::None { bottleneck } else { 2 * bottleneck };
312                Layer {
313                    conv: Conv1d::new(channels, zc, kernels[i], dilations[i], true),
314                    mixin: Conv1d::new(condition_size, zc, 1, 1, false),
315                    activation: activations[i].clone(),
316                    gating: gatings[i],
317                    secondary: secondaries[i].clone(),
318                    bottleneck,
319                    layer1x1: layer1x1.then(|| Conv1d::new(bottleneck, channels, 1, 1, true)),
320                    head1x1: head1x1.then(|| Conv1d::new(bottleneck, head1x1_out, 1, 1, true)),
321                    z: vec![],
322                    mix: vec![],
323                    tmp: vec![],
324                }
325            })
326            .collect();
327        let head_in = if head1x1 { head1x1_out } else { bottleneck };
328        Ok(LayerArray {
329            rechannel: Conv1d::new(input_size, channels, 1, 1, false),
330            layers,
331            head_rechannel: Conv1d::new(head_in, head_size, head_kernel, head_dilation, head_bias),
332            head_in,
333            a: vec![],
334            b: vec![],
335            head_sum: vec![],
336            head_part: vec![],
337        })
338    }
339
340    fn load(&mut self, w: &mut Weights) -> Result<(), NamError> {
341        self.rechannel.load(w)?;
342        for l in &mut self.layers {
343            l.conv.load(w)?;
344            l.mixin.load(w)?;
345            if let Some(c) = &mut l.layer1x1 {
346                c.load(w)?;
347            }
348            if let Some(c) = &mut l.head1x1 {
349                c.load(w)?;
350            }
351        }
352        self.head_rechannel.load(w)
353    }
354
355    fn receptive_field(&self) -> usize {
356        self.layers.iter().map(|l| l.conv.receptive_field()).sum::<usize>() + self.head_rechannel.receptive_field()
357    }
358
359    /// Returns (layer outputs, head outputs), both n frames.
360    fn process(&mut self, input: &[f32], cond: &[f32], head_in: Option<&[f32]>, n: usize) -> (&[f32], &[f32]) {
361        self.rechannel.process(input, n, &mut self.a);
362        self.head_sum.clear();
363        match head_in {
364            Some(h) => self.head_sum.extend_from_slice(&h[..n * self.head_in]),
365            None => self.head_sum.resize(n * self.head_in, 0.0),
366        }
367        for l in &mut self.layers {
368            l.process(&self.a, cond, n, &mut self.head_part, &mut self.b);
369            for (s, p) in self.head_sum.iter_mut().zip(&self.head_part) {
370                *s += p;
371            }
372            std::mem::swap(&mut self.a, &mut self.b);
373        }
374        self.head_rechannel.process(&self.head_sum, n, &mut self.b);
375        (&self.a, &self.b)
376    }
377
378    fn reset(&mut self) {
379        for l in &mut self.layers {
380            l.reset();
381        }
382        self.head_rechannel.reset();
383    }
384}
385
386/// A NAM WaveNet (mono in, mono out).
387#[derive(Clone, Debug)]
388pub struct WaveNet {
389    arrays: Vec<LayerArray>,
390    head_scale: f32,
391    sample_rate: Option<f64>,
392    // scratch: previous array's layer outputs and head outputs
393    layer_out: Vec<f32>,
394    head_out: Vec<f32>,
395}
396
397impl WaveNet {
398    pub fn from_json(config: &Value, weights: &[f32], sample_rate: Option<f64>) -> Result<WaveNet, NamError> {
399        if config.get("condition_dsp").is_some_and(|c| !c.is_null()) {
400            return Err(NamError::Unsupported("condition_dsp".into()));
401        }
402        if config.get("head").is_some_and(|h| !h.is_null()) {
403            return Err(NamError::Unsupported("post-stack head".into()));
404        }
405        let mut arrays: Vec<LayerArray> = config
406            .get("layers")
407            .and_then(Value::as_array)
408            .ok_or_else(|| NamError::Format("WaveNet: missing layers".into()))?
409            .iter()
410            .map(LayerArray::from_json)
411            .collect::<Result<_, _>>()?;
412        if arrays.is_empty() {
413            return Err(NamError::Format("WaveNet: no layer arrays".into()));
414        }
415        if arrays[0].rechannel.cin != 1 || arrays[0].layers.first().is_some_and(|l| l.mixin.cin != 1) {
416            return Err(NamError::Unsupported("multi-channel input".into()));
417        }
418        if arrays.last().is_some_and(|a| a.head_rechannel.cout != 1) {
419            return Err(NamError::Unsupported("multi-channel output".into()));
420        }
421        let mut w = Weights { w: weights, pos: 0 };
422        for a in &mut arrays {
423            a.load(&mut w)?;
424        }
425        let head_scale = *w.take(1)?.first().expect("1");
426        if w.pos != weights.len() {
427            return Err(NamError::Weights { expected: w.pos, got: weights.len() });
428        }
429        let mut net = WaveNet { arrays, head_scale, sample_rate, layer_out: vec![], head_out: vec![] };
430        net.reset();
431        Ok(net)
432    }
433
434    /// Samples of silence needed to settle every buffer (NeuralAmpModelerCore's prewarm length).
435    pub fn prewarm_samples(&self) -> usize {
436        1 + self.arrays.iter().map(LayerArray::receptive_field).sum::<usize>()
437    }
438
439    pub fn sample_rate(&self) -> Option<f64> {
440        self.sample_rate
441    }
442
443    /// Clear the history, then run silence through it, like NeuralAmpModelerCore's `Reset` + `prewarm`.
444    pub fn reset(&mut self) {
445        for a in &mut self.arrays {
446            a.reset();
447        }
448        let zeros = vec![0.0; 4096];
449        let mut out = vec![0.0; 4096];
450        let mut left = self.prewarm_samples();
451        while left > 0 {
452            let n = left.min(4096);
453            self.process(&zeros[..n], &mut out[..n]);
454            left -= n;
455        }
456    }
457
458    pub fn process(&mut self, input: &[f32], output: &mut [f32]) {
459        let n = input.len().min(output.len());
460        for i in 0..self.arrays.len() {
461            let (prev, rest) = self.arrays.split_at_mut(i);
462            let arr = &mut rest[0];
463            let (lo, ho) = match prev.last() {
464                None => arr.process(input, input, None, n),
465                Some(_) => arr.process(&self.layer_out, input, Some(&self.head_out), n),
466            };
467            let (lo, ho) = (lo.to_vec(), ho.to_vec());
468            self.layer_out = lo;
469            self.head_out = ho;
470        }
471        for (o, h) in output[..n].iter_mut().zip(&self.head_out) {
472            *o = self.head_scale * h;
473        }
474    }
475}
open_stomp/nam/wavenet.rs

open_stomp/nam/
wavenet.rs
