Skip to main content

nam_rs/models/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.
3
4//! Neural Inference Architectures (Brain Engines) module for NAM-rs.
5//!
6//! This module contains the acoustic brains of the program: neural networks that have learned how,
7//! for example, a real amplifier or pedal distorts and colors a guitar sound.
8
9pub mod a2;
10pub mod container;
11/// ConvNet feed-forward architecture (F4).
12pub mod convnet;
13pub mod linear;
14pub mod linear_fft;
15pub mod lstm;
16pub mod slimmable;
17pub mod wavenet;
18
19/// NamModel trait implementation for StaticModel (dispatch methods).
20mod nam_model;
21mod static_model;
22
23// =============================================================================
24// Sealed Pattern — Prevents external implementations of NamModel
25// =============================================================================
26
27mod sealed {
28    pub trait Sealed {}
29}
30
31// =============================================================================
32// Trait NamModel — Public Contract
33// =============================================================================
34
35/// The interface (standard connector) for any neural model (amplifiers, pedals, etc.).
36///
37/// Sealed via private supertrait — only types within this crate can implement `NamModel`.
38pub trait NamModel: Send + Sync + sealed::Sealed {
39    /// Invoked by the DSP RT-Thread to process acoustic sample blocks (Float32).
40    fn process(&mut self, input: &[f32], output: &mut [f32]);
41
42    /// "Heats up" the virtual tubes of the neural engine (`prewarm`).
43    fn prewarm(&mut self, num_samples: usize);
44
45    /// Returns whether prewarm should be executed on `reset()`.
46    ///
47    /// Default: `true` (prewarm on every reset).
48    fn prewarm_on_reset(&self) -> bool {
49        true
50    }
51
52    /// Sets whether prewarm should be executed on `reset()`.
53    ///
54    /// Default: no-op (fixed-size models ignore this flag).
55    fn set_prewarm_on_reset(&mut self, _val: bool) {}
56
57    /// Resets the model's internal state with a new sample rate and max buffer size.
58    ///
59    /// The default implementation calls `prewarm(max_buffer_size)` if `prewarm_on_reset()`
60    /// returns `true`. Architectures with recurrent state (LSTM) may override this
61    /// for a lighter reset (only zero the internal states without reprocessing a full prewarm).
62    fn reset(&mut self, _sample_rate: u32, max_buffer_size: usize) -> anyhow::Result<()> {
63        if self.prewarm_on_reset() {
64            self.prewarm(max_buffer_size);
65        }
66        Ok(())
67    }
68
69    /// Reallocates internal buffers to support the given maximum block size.
70    ///
71    /// Models with fixed (const-generic) buffer sizes can use the default no-op.
72    ///
73    /// Default: no-op (suitable for static models and LSTM).
74    fn set_max_buffer_size(&mut self, _max_buf: usize) -> anyhow::Result<()> {
75        Ok(())
76    }
77
78    /// Returns the number of samples needed to fully stabilize the model's internal
79    /// state (receptive field / recurrent memory depth).
80    ///
81    /// Default: `0` (suitable for LSTM, which stabilizes via recurrence).
82    /// WaveNet variants override this to return the sum of all array receptive
83    /// fields plus any condition_dsp prewarm samples.
84    fn prewarm_samples(&self) -> usize {
85        0
86    }
87
88    /// Returns the breakpoints at which slimmable quality transitions occur.
89    ///
90    /// Each breakpoint represents a normalized value in `[0.0, 1.0]` where the
91    /// model switches to a different submodel or internal quality tier. Hosts
92    /// and plugins can use these to map and snap discrete quality parameters.
93    ///
94    /// Delegates to the inner model when applicable (e.g., `ContainerModel`).
95    /// Defaults to an empty vector for models without discrete breakpoints.
96    fn slimmable_breakpoints(&self) -> Vec<f64> {
97        vec![]
98    }
99}
100
101/// Wrapper enum for trained model variants.
102/// Enables static dispatch of DSP calls to the concrete variant, avoiding vtable overhead.
103///
104/// Named `StaticModel` because all variants are compile-time-fixed geometries.
105/// The legacy "Dynamic" mode (arbitrary geometry at runtime) has been retired.
106pub enum StaticModel {
107    /// WaveNet Standard (16 channels, kernel 3, dilation 8).
108    WavenetStandard(Box<wavenet::WaveNetModel<16, 3, 8>>),
109    /// WaveNet Lite (12 channels, kernel 3, dilation 6).
110    WavenetLite(Box<wavenet::WaveNetModel<12, 3, 6>>),
111    /// WaveNet Feather (8 channels, kernel 3, dilation 4).
112    WavenetFeather(Box<wavenet::WaveNetModel<8, 3, 4>>),
113    /// WaveNet Nano (4 channels, kernel 3, dilation 2).
114    WavenetNano(Box<wavenet::WaveNetModel<4, 3, 2>>),
115    /// WaveNet A2 Full (8 channels, real inference).
116    WavenetA2Full(Box<a2::WaveNetA2<8>>),
117    /// WaveNet A2 Lite (3 channels, real inference).
118    WavenetA2Lite(Box<a2::WaveNetA2<3>>),
119    /// WaveNet A2 Dynamic (runtime-dimensioned, full topology spectrum).
120    WavenetA2Dyn(Box<a2::WaveNetA2Dyn>),
121    /// WaveNet A2 Cascade (multi-array chain of Dynamic engines).
122    WavenetA2Cascade(Box<a2::WaveNetA2Cascade>),
123    /// WaveNet Dynamic (runtime-dimensioned, free geometry).
124    WavenetDyn(Box<wavenet::WaveNetModelDyn>),
125    /// LSTM 1 Layer × 3 hidden units.
126    Lstm1x3(Box<lstm::Lstm1x3>),
127    /// LSTM 1 Layer × 8 hidden units.
128    Lstm1x8(Box<lstm::Lstm1x8>),
129    /// LSTM 1 Layer × 12 hidden units.
130    Lstm1x12(Box<lstm::Lstm1x12>),
131    /// LSTM 1 Layer × 16 hidden units.
132    Lstm1x16(Box<lstm::Lstm1x16>),
133    /// LSTM 1 Layer × 24 hidden units.
134    Lstm1x24(Box<lstm::Lstm1x24>),
135    /// LSTM 2 Layers × 8 hidden units.
136    Lstm2x8(Box<lstm::Lstm2x8>),
137    /// LSTM 2 Layers × 12 hidden units.
138    Lstm2x12(Box<lstm::Lstm2x12>),
139    /// LSTM 2 Layers × 16 hidden units.
140    Lstm2x16(Box<lstm::Lstm2x16>),
141    /// LSTM 1 Layer × 40 hidden units.
142    Lstm1x40(Box<lstm::Lstm1x40>),
143    /// LSTM 2 Layers × 24 hidden units.
144    Lstm2x24(Box<lstm::Lstm2x24>),
145    /// LSTM Dynamic — runtime-dimensioned, free geometry (F7 fallback).
146    LstmDyn(Box<lstm::LstmModelDyn>),
147    /// SlimmableContainer — bundle of submodels selected by quality threshold.
148    Container(Box<container::ContainerModel>),
149    /// Linear — FIR-based model (dot product of input history with weights + bias).
150    Linear(Box<linear::LinearModel>),
151    /// ConvNet feed-forward model (F4).
152    ConvNet(Box<convnet::ConvNetModel>),
153}
154
155impl sealed::Sealed for StaticModel {}
156
157pub(crate) use static_model::clone_condition_dsp;