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;