nam_rs/models/a2/
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//! A2 Architecture.
5//!
6//! Este módulo isola componentes da arquitetura A2 (v0.6+), incluindo
7//! ativações, FiLM, gating, parâmetros e os modelos A2.
8//!
9//! ## Status
10//!
11//! O scaffolding existente inclui a definição completa de parâmetros e
12//! constantes arquiteturais (espelhando `a2_fast.h`). O fast-path
13//! A2-Full/Lite (fast-path) utiliza um subconjunto destes structs;
14//! os campos de FiLM, gating (`Gated`/`Blended`), `head1x1`, `bottleneck ≠ channels`
15//! e ativações heterogêneas estão reservados para o motor A2 geral (futuro).
16
17pub mod activations;
18pub mod conv1d;
19pub mod conv1d_ch;
20pub mod conv1d_ch3;
21pub mod conv1d_ch8;
22pub mod conv1d_fallback;
23pub mod film;
24pub mod gating;
25pub mod grouped_conv1d;
26pub mod head;
27pub mod layer;
28pub mod model;
29pub mod params;
30pub mod weights_layout;
31/// Public re-exports for easy access.
32pub use activations::{ActivationFn, ActivationType};
33pub use conv1d::A2Conv1d;
34pub use film::{FiLMConfig, FiLMLayer, FilmBlock};
35pub use gating::GatingMode;
36pub use grouped_conv1d::{
37    A2GroupedConv1d, grouped_conv1d_block_ref, grouped_conv1d_single_frame_ref,
38};
39pub use head::{
40    A2HeadConv, a2_head_block_scalar_ref, a2_head_single_frame_scalar_ref, head_process_ch3_sse,
41    head_process_ch8_avx2,
42};
43pub use layer::{A2ConvCh, A2Layer, a2_layer_single_frame_scalar_ref};
44pub use model::WaveNetA2;
45pub use model::cascade::WaveNetA2Cascade;
46pub use model::dynamic::WaveNetA2Dyn;
47pub use params::{
48    A2_DILATIONS, A2_HEAD_KERNEL_SIZE, A2_KERNEL_SIZES, A2_LEAKY_SLOPE, A2_NUM_LAYERS,
49    A2_VALID_CHANNELS, HeadParams, LayerArrayParamsA2, LayerParamsA2,
50};
51
52/// Total weight count for `WaveNetA2<CH>`.
53///
54/// Formula: `rechannel(ch) + Σ(kernel sizes: ch²·k + ch + ch + ch² + ch) + head(16·ch + 2)`.
55/// This is the single canonical source — topology changes require updates here only.
56pub const fn a2_weight_count<const CH: usize>() -> usize {
57    let mut total = CH; // rechannel_w
58    let mut i = 0;
59    while i < A2_NUM_LAYERS {
60        let k = params::A2_KERNEL_SIZES[i];
61        total += CH * CH * k; // conv_w
62        total += CH; // conv_b
63        total += CH; // mixin_w
64        total += CH * CH; // l1x1_w
65        total += CH; // l1x1_b
66        i += 1;
67    }
68    total += A2_HEAD_KERNEL_SIZE * CH; // head_w
69    total += 1; // head_b
70    total += 1; // head_scale
71    total
72}
73
74use crate::models::NamModel;
75use crate::models::sealed;
76
77impl<const CH: usize> sealed::Sealed for model::WaveNetA2<CH> {}
78
79impl<const CH: usize> NamModel for model::WaveNetA2<CH> {
80    fn process(&mut self, input: &[f32], output: &mut [f32]) {
81        self.process(input, output);
82    }
83
84    fn prewarm(&mut self, _num_samples: usize) {
85        self.prewarm();
86    }
87
88    fn prewarm_samples(&self) -> usize {
89        self.receptive_field_size
90    }
91
92    fn set_max_buffer_size(&mut self, max_buf: usize) -> anyhow::Result<()> {
93        self.set_max_buffer_size(max_buf)
94    }
95
96    fn reset(&mut self, _sample_rate: u32, max_buffer_size: usize) -> anyhow::Result<()> {
97        self.reset(_sample_rate, max_buffer_size)
98    }
99
100    fn prewarm_on_reset(&self) -> bool {
101        self.prewarm_on_reset
102    }
103
104    fn set_prewarm_on_reset(&mut self, val: bool) {
105        self.prewarm_on_reset = val;
106    }
107}
108
109impl sealed::Sealed for model::dynamic::WaveNetA2Dyn {}
110
111impl NamModel for model::dynamic::WaveNetA2Dyn {
112    fn process(&mut self, input: &[f32], output: &mut [f32]) {
113        self.process(input, output);
114    }
115
116    fn prewarm(&mut self, _num_samples: usize) {
117        self.prewarm();
118    }
119
120    fn prewarm_samples(&self) -> usize {
121        self.receptive_field_size
122    }
123
124    fn set_max_buffer_size(&mut self, max_buf: usize) -> anyhow::Result<()> {
125        self.set_max_buffer_size(max_buf)
126    }
127
128    fn reset(&mut self, _sample_rate: u32, max_buffer_size: usize) -> anyhow::Result<()> {
129        self.reset(_sample_rate, max_buffer_size)
130    }
131
132    fn prewarm_on_reset(&self) -> bool {
133        self.prewarm_on_reset
134    }
135
136    fn set_prewarm_on_reset(&mut self, val: bool) {
137        self.prewarm_on_reset = val;
138    }
139}
140
141impl sealed::Sealed for model::cascade::WaveNetA2Cascade {}
142
143impl NamModel for model::cascade::WaveNetA2Cascade {
144    fn process(&mut self, input: &[f32], output: &mut [f32]) {
145        self.process(input, output);
146    }
147
148    fn prewarm(&mut self, _num_samples: usize) {
149        self.prewarm();
150    }
151
152    fn prewarm_samples(&self) -> usize {
153        self.receptive_field_size
154    }
155
156    fn set_max_buffer_size(&mut self, max_buf: usize) -> anyhow::Result<()> {
157        self.set_max_buffer_size(max_buf)
158    }
159
160    fn reset(&mut self, _sample_rate: u32, max_buffer_size: usize) -> anyhow::Result<()> {
161        self.reset(_sample_rate, max_buffer_size)
162    }
163
164    fn prewarm_on_reset(&self) -> bool {
165        self.prewarm_on_reset
166    }
167
168    fn set_prewarm_on_reset(&mut self, val: bool) {
169        self.prewarm_on_reset = val;
170    }
171}
172
173#[cfg(test)]
174mod a2_test;
nam_rs/models/a2/mod.rs

nam_rs/models/a2/
mod.rs
