nam_rs/models/wavenet/
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//! WaveNet Module — Dilated Causal Neural Architecture for amplifier and pedal emulation.
5//!
6//! This module provides WaveNet inference engines optimized for real-time DSP execution
7//! using Const Generics for fixed dimensions, eliminating bounds checks and maximizing throughput.
8//!
9//! ## Sub-modules
10//!
11//! | Module        | Description                                                             |
12//! | ------------- | ----------------------------------------------------------------------- |
13//! | `common`      | Fundamental constants and types (`WaveNetLayerState`, `WavenetProcessContext`) |
14//! | `conv1d`      | Static causal 1D convolution (`Conv1d`)                                |
15//! | `conv1d_dyn`  | Runtime-dimensional causal 1D convolution (`Conv1dDyn`)                 |
16//! | `dense`       | Static 1x1 dense layer (`DenseLayer`)                                   |
17//! | `dense_dyn`   | Runtime-dimensional 1x1 dense layer (`DenseLayerDyn`)                   |
18//! | `layer`       | Static WaveNet layer (`WaveNetLayer`)                                   |
19//! | `layer_dyn`   | Runtime-dimensional WaveNet layer (`WaveNetLayerDyn`)                   |
20//! | `layer_array` | Static WaveNet layer array (`WaveNetLayerArray`)                        |
21//! | `layer_array_dyn` | Runtime-dimensional WaveNet layer array (`WaveNetLayerArrayDyn`)    |
22//! | `model`           | Complete static model (`WaveNetModel`)                                  |
23//! | `model_dyn`       | Runtime-dimensional model (`WaveNetModelDyn`)                           |
24//! | `post_stack_head` | Post-stack head sub-object (Conv1D + activation)                        |
25
26pub mod common;
27pub mod conv1d;
28pub mod conv1d_dual;
29pub mod conv1d_dyn;
30/// Dual-frame kernel for Conv1dDyn (Temporal-Tiling dual-frame processing).
31pub mod conv1d_dyn_dual;
32/// F32-native dot-product helpers (standalone functions).
33pub mod conv_input;
34/// Static 1x1 dense layer (`DenseLayer<IN, OUT>`).
35pub mod dense;
36/// Runtime-dimensional 1x1 dense layer (`DenseLayerDyn`).
37pub mod dense_dyn;
38/// Static WaveNet layer (`WaveNetLayer<COND, CH, K>`).
39pub mod layer;
40/// Static WaveNet layer array (`WaveNetLayerArray<IN, COND, CH, K, HEAD>`).
41pub mod layer_array;
42/// Runtime-dimensional WaveNet layer array (`WaveNetLayerArrayDyn`).
43pub mod layer_array_dyn;
44/// Runtime-dimensional WaveNet layer (`WaveNetLayerDyn`).
45pub mod layer_dyn;
46/// Static model (`WaveNetModel<CH, K, HEAD>`).
47pub mod model;
48/// Runtime-dimensional model (`WaveNetModelDyn`).
49pub mod model_dyn;
50/// Post-stack head sub-object (Conv1D + activation).
51pub mod post_stack_head;
52
53use super::NamModel;
54use super::sealed;
55
56// =============================================================================
57// sealed::Sealed for WaveNet
58// =============================================================================
59
60impl<const CH: usize, const K: usize, const HEAD: usize> sealed::Sealed
61    for model::WaveNetModel<CH, K, HEAD>
62{
63}
64
65impl sealed::Sealed for model_dyn::WaveNetModelDyn {}
66
67// =============================================================================
68// NamModel for WaveNet (Const Generics)
69// =============================================================================
70
71impl<const CH: usize, const K: usize, const HEAD: usize> NamModel
72    for model::WaveNetModel<CH, K, HEAD>
73{
74    fn process(&mut self, input: &[f32], output: &mut [f32]) {
75        // Delegates to the inherent WaveNetModel::process method (inherent methods have priority)
76        self.process(input, output);
77    }
78
79    fn prewarm(&mut self, _num_samples: usize) {
80        // WaveNet prewarm is one-shot: fills the receptive field via copy_buffer.
81        // C++ runs `model->Prewarm()` without a parameter (unlike LSTM).
82        self.prewarm();
83    }
84
85    fn prewarm_samples(&self) -> usize {
86        self.array1.receptive_field_size + self.array2.receptive_field_size
87    }
88
89    fn prewarm_on_reset(&self) -> bool {
90        self.prewarm_on_reset
91    }
92
93    fn set_prewarm_on_reset(&mut self, val: bool) {
94        self.prewarm_on_reset = val;
95    }
96}
97
98// =============================================================================
99// NamModel for WaveNet (Dynamic / Runtime-Dimensional)
100// =============================================================================
101
102impl NamModel for model_dyn::WaveNetModelDyn {
103    fn process(&mut self, input: &[f32], output: &mut [f32]) {
104        self.process(input, output);
105    }
106
107    fn prewarm(&mut self, _num_samples: usize) {
108        self.prewarm();
109    }
110
111    fn prewarm_samples(&self) -> usize {
112        let mut rf: usize = self.arrays.iter().map(|a| a.receptive_field_size).sum();
113        if let Some(ref cond_dsp) = self.condition_dsp {
114            rf += cond_dsp.prewarm_samples();
115        }
116        if let Some(ref head_proc) = self.post_stack_head {
117            rf += head_proc.receptive_field() - 1;
118        }
119        rf
120    }
121
122    fn set_max_buffer_size(&mut self, max_buf: usize) -> anyhow::Result<()> {
123        if let Some(ref mut cond_dsp) = self.condition_dsp {
124            cond_dsp.set_max_buffer_size(max_buf)?;
125        }
126        Ok(())
127    }
128
129    fn prewarm_on_reset(&self) -> bool {
130        self.prewarm_on_reset
131    }
132
133    fn set_prewarm_on_reset(&mut self, val: bool) {
134        self.prewarm_on_reset = val;
135        if let Some(ref mut cond_dsp) = self.condition_dsp {
136            cond_dsp.set_prewarm_on_reset(val);
137        }
138    }
139}
140
141// =============================================================================
142// Public re-exports
143// =============================================================================
144
145pub use common::{
146    LAYER_ARRAY_BUFFER_PADDING, MAX_KERNEL, WAVENET_MAX_NUM_FRAMES, WaveNetLayerState,
147    WavenetProcessContext,
148};
149pub use conv1d::Conv1d;
150pub use conv1d_dyn::Conv1dDyn;
151pub use dense::DenseLayer;
152pub use dense_dyn::DenseLayerDyn;
153pub use layer::WaveNetLayer;
154pub use layer_array::WaveNetLayerArray;
155pub use layer_array_dyn::WaveNetLayerArrayDyn;
156pub use layer_dyn::WaveNetLayerDyn;
157pub use model::WaveNetModel;
158pub use model_dyn::WaveNetModelDyn;
159pub use post_stack_head::PostStackHead;
160
161#[cfg(test)]
162mod wavenet_test;
nam_rs/models/wavenet/mod.rs

nam_rs/models/wavenet/
mod.rs
