nam_rs/models/
static_model.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
4use super::slimmable::SlimmableModel;
5use super::{NamModel, StaticModel};
6use std::sync::Arc;
7
8impl StaticModel {
9    /// Injects `RtStatusFlags` into the model so it can signal its state
10    /// to the UI via atomic flags.
11    pub fn inject_rt_status(&mut self, rt_status: Arc<crate::common::spsc::RtStatusFlags>) {
12        match self {
13            Self::WavenetA2Full(m) => m.inject_rt_status(rt_status),
14            Self::WavenetA2Lite(m) => m.inject_rt_status(rt_status),
15            Self::WavenetA2Dyn(_) => {}
16            Self::WavenetA2Cascade(_) => {}
17            _ => {}
18        }
19    }
20
21    /// Scalar processing path for LSTM models (exact tanh/sigmoid via libm).
22    ///
23    /// Only available under `#[cfg(test)]` or `#[cfg(feature = "testing")]`.
24    /// For non-LSTM models, delegates to `process()`.
25    #[cfg(any(test, feature = "testing"))]
26    pub fn process_scalar(&mut self, input: &[f32], output: &mut [f32]) {
27        match self {
28            Self::Lstm1x3(m) => m.process_scalar(input, output),
29            Self::Lstm1x8(m) => m.process_scalar(input, output),
30            Self::Lstm1x12(m) => m.process_scalar(input, output),
31            Self::Lstm1x16(m) => m.process_scalar(input, output),
32            Self::Lstm1x24(m) => m.process_scalar(input, output),
33            Self::Lstm2x8(m) => m.process_scalar(input, output),
34            Self::Lstm2x12(m) => m.process_scalar(input, output),
35            Self::Lstm2x16(m) => m.process_scalar(input, output),
36            Self::Lstm1x40(m) => m.process_scalar(input, output),
37            Self::Lstm2x24(m) => m.process_scalar(input, output),
38            Self::LstmDyn(m) => m.process_scalar(input, output),
39            Self::ConvNet(m) => m.process(input, output),
40            Self::WavenetA2Cascade(m) => m.process(input, output),
41            other => other.process(input, output),
42        }
43    }
44
45    /// Sets the effective number of layers for soft-degrade.
46    /// Only applies to WaveNet variants. LSTM handles reduction at the pipeline level.
47    #[inline(always)]
48    pub fn set_effective_layers(&mut self, n: usize) {
49        match self {
50            Self::WavenetStandard(m) => m.set_effective_layers(n),
51            Self::WavenetLite(m) => m.set_effective_layers(n),
52            Self::WavenetFeather(m) => m.set_effective_layers(n),
53            Self::WavenetNano(m) => m.set_effective_layers(n),
54            Self::WavenetDyn(m) => m.set_effective_layers(n),
55            // LSTM, A2, Container, and Linear: no-op — reduction handled at pipeline level
56            Self::WavenetA2Cascade(_)
57            | Self::WavenetA2Full(_)
58            | Self::WavenetA2Lite(_)
59            | Self::WavenetA2Dyn(_)
60            | Self::Container(_)
61            | Self::Lstm1x3(_)
62            | Self::Lstm1x8(_)
63            | Self::Lstm1x12(_)
64            | Self::Lstm1x16(_)
65            | Self::Lstm1x24(_)
66            | Self::Lstm2x8(_)
67            | Self::Lstm2x12(_)
68            | Self::Lstm2x16(_)
69            | Self::Lstm1x40(_)
70            | Self::Lstm2x24(_)
71            | Self::LstmDyn(_) => {}
72            Self::Linear(_) => {}
73            Self::ConvNet(_) => {}
74        }
75    }
76
77    /// Backs up WaveNet buffer starts. No-op for non-WaveNet models.
78    #[inline(always)]
79    pub fn backup_buffer_starts(&self, starts: &mut [usize], offset: &mut usize) {
80        match self {
81            Self::WavenetStandard(m) => m.backup_buffer_starts(starts, offset),
82            Self::WavenetLite(m) => m.backup_buffer_starts(starts, offset),
83            Self::WavenetFeather(m) => m.backup_buffer_starts(starts, offset),
84            Self::WavenetNano(m) => m.backup_buffer_starts(starts, offset),
85            Self::WavenetDyn(m) => m.backup_buffer_starts(starts, offset),
86            Self::WavenetA2Cascade(_) => {}
87            _ => {}
88        }
89    }
90
91    /// Restores WaveNet buffer starts. No-op for non-WaveNet models.
92    #[inline(always)]
93    pub fn restore_buffer_starts(&mut self, starts: &[usize], offset: &mut usize) {
94        match self {
95            Self::WavenetStandard(m) => m.restore_buffer_starts(starts, offset),
96            Self::WavenetLite(m) => m.restore_buffer_starts(starts, offset),
97            Self::WavenetFeather(m) => m.restore_buffer_starts(starts, offset),
98            Self::WavenetNano(m) => m.restore_buffer_starts(starts, offset),
99            Self::WavenetDyn(m) => m.restore_buffer_starts(starts, offset),
100            Self::WavenetA2Cascade(_) => {}
101            _ => {}
102        }
103    }
104
105    /// Sets the slimmable quality level for `SlimmableModel` variants.
106    ///
107    /// Only applies to `Container`. Other variants are a no-op.
108    ///
109    /// `rt_status` — when on the RT hot-path, pass `Some(&rt_status)` so errors
110    /// can be signaled atomically. Pass `None` in tests or off-RT contexts.
111    #[inline(always)]
112    pub fn set_slimmable_size(
113        &mut self,
114        val: f32,
115        rt_status: Option<&crate::common::spsc::RtStatusFlags>,
116    ) {
117        if let Self::Container(c) = self {
118            c.set_slimmable_size(val, rt_status);
119        }
120    }
121
122    /// Returns the slimmable quality breakpoints for this model.
123    ///
124    /// Delegates to `ContainerModel::slimmable_breakpoints()` when this is a
125    /// `Container` variant. Returns an empty vector for all other variants.
126    pub fn slimmable_breakpoints(&self) -> Vec<f64> {
127        if let Self::Container(c) = self {
128            SlimmableModel::slimmable_breakpoints(c.as_ref())
129        } else {
130            vec![]
131        }
132    }
133
134    /// Returns the total number of layers for the model (0 for non-WaveNet).
135    /// Used by the adaptive FSM to compute how many to keep.
136    #[inline(always)]
137    pub fn layer_count(&self) -> usize {
138        match self {
139            Self::WavenetStandard(m) => m.array1.layers.len(),
140            Self::WavenetLite(m) => m.array1.layers.len(),
141            Self::WavenetFeather(m) => m.array1.layers.len(),
142            Self::WavenetNano(m) => m.array1.layers.len(),
143            Self::WavenetDyn(m) => m.arrays[0].layers.len(),
144            Self::WavenetA2Full(_) | Self::WavenetA2Lite(_) | Self::WavenetA2Dyn(_) => {
145                crate::models::a2::A2_NUM_LAYERS
146            }
147            Self::WavenetA2Cascade(m) => m.arrays.iter().map(|a| a.num_layers).sum(),
148            Self::Container(c) => c.active().layer_count(),
149            Self::Lstm2x8(_) | Self::Lstm2x12(_) | Self::Lstm2x16(_) | Self::Lstm2x24(_) => 2,
150            Self::LstmDyn(m) => m.layers.len(),
151            Self::Lstm1x3(_)
152            | Self::Lstm1x8(_)
153            | Self::Lstm1x12(_)
154            | Self::Lstm1x16(_)
155            | Self::Lstm1x24(_)
156            | Self::Lstm1x40(_) => 1,
157            Self::Linear(_) => 0,
158            Self::ConvNet(m) => m.blocks.len(),
159        }
160    }
161
162    /// Returns a human-readable classification label for this model variant.
163    ///
164    /// Used by `nondist_validation.rs` to cross-check against the manifest's
165    /// `expected_class` field and detect dispatcher routing errors.
166    #[cold]
167    pub fn class_label(&self) -> String {
168        match self {
169            Self::WavenetStandard(_) => "WaveNet A1 Standard (CH=16)".into(),
170            Self::WavenetLite(_) => "WaveNet A1 Lite (CH=12)".into(),
171            Self::WavenetFeather(_) => "WaveNet A1 Feather (CH=8)".into(),
172            Self::WavenetNano(_) => "WaveNet A1 Nano (CH=4)".into(),
173            Self::WavenetA2Full(_) => "WaveNet A2 (CH=8)".into(),
174            Self::WavenetA2Lite(_) => "WaveNet A2 Lite (CH=3)".into(),
175            Self::WavenetA2Dyn(m) => format!("WaveNet A2 (CH={})", m.channels),
176            Self::WavenetA2Cascade(m) => format!("WaveNet A2 Cascade ({} arrays)", m.arrays.len()),
177            Self::WavenetDyn(m) => {
178                if m.arrays.len() == 2 {
179                    "WaveNet A1 (Custom)".into()
180                } else {
181                    "WaveNet (Custom Layers)".into()
182                }
183            }
184            Self::Container(c) => {
185                let active_label = c.active().class_label();
186                if active_label.starts_with("WaveNet") {
187                    active_label
188                } else {
189                    "SlimmableContainer".into()
190                }
191            }
192            Self::Lstm1x3(_) => "LSTM 1x3".into(),
193            Self::Lstm1x8(_) => "LSTM 1x8".into(),
194            Self::Lstm1x12(_) => "LSTM 1x12".into(),
195            Self::Lstm1x16(_) => "LSTM 1x16".into(),
196            Self::Lstm1x24(_) => "LSTM 1x24".into(),
197            Self::Lstm1x40(_) => "LSTM 1x40".into(),
198            Self::Lstm2x8(_) => "LSTM 2x8".into(),
199            Self::Lstm2x12(_) => "LSTM 2x12".into(),
200            Self::Lstm2x16(_) => "LSTM 2x16".into(),
201            Self::Lstm2x24(_) => "LSTM 2x24".into(),
202            Self::LstmDyn(m) => format!("LSTM {}x{}", m.layers.len(), m.head_weights.len()),
203            Self::Linear(_) => "Linear".into(),
204            Self::ConvNet(m) => format!("ConvNet (CH={})", m.in_channels()),
205        }
206    }
207
208    /// Returns `true` if this is an LSTM model.
209    #[inline(always)]
210    pub fn is_lstm(&self) -> bool {
211        matches!(
212            self,
213            Self::Lstm1x3(_)
214                | Self::Lstm1x8(_)
215                | Self::Lstm1x12(_)
216                | Self::Lstm1x16(_)
217                | Self::Lstm1x24(_)
218                | Self::Lstm2x8(_)
219                | Self::Lstm2x12(_)
220                | Self::Lstm2x16(_)
221                | Self::Lstm1x40(_)
222                | Self::Lstm2x24(_)
223                | Self::LstmDyn(_)
224        )
225    }
226
227    /// Returns `true` if this is a WaveNet model.
228    #[inline(always)]
229    pub fn is_wavenet(&self) -> bool {
230        matches!(
231            self,
232            Self::WavenetStandard(_)
233                | Self::WavenetLite(_)
234                | Self::WavenetFeather(_)
235                | Self::WavenetNano(_)
236                | Self::WavenetA2Full(_)
237                | Self::WavenetA2Lite(_)
238                | Self::WavenetA2Dyn(_)
239                | Self::WavenetA2Cascade(_)
240                | Self::WavenetDyn(_)
241        )
242    }
243
244    /// Returns `true` if this model supports layer skipping (WaveNet A1 variants only).
245    /// A2 architectures use `MirroredBuffer` + `layer_buffer_starts`/`head_write_pos`
246    /// and cannot safely participate in the double-pass crossfade mechanism.
247    #[inline(always)]
248    pub fn supports_layer_skip(&self) -> bool {
249        matches!(
250            self,
251            Self::WavenetStandard(_)
252                | Self::WavenetLite(_)
253                | Self::WavenetFeather(_)
254                | Self::WavenetNano(_)
255                | Self::WavenetDyn(_)
256        )
257    }
258
259    /// Returns `true` if this is a SlimmableContainer model.
260    #[inline(always)]
261    pub fn is_container(&self) -> bool {
262        matches!(self, Self::Container(_))
263    }
264
265    /// Returns the number of channels inside the model.
266    pub fn channels(&self) -> usize {
267        match self {
268            Self::WavenetStandard(_) => 16,
269            Self::WavenetLite(_) => 12,
270            Self::WavenetFeather(_) => 8,
271            Self::WavenetNano(_) => 4,
272            Self::WavenetA2Full(_) => 8,
273            Self::WavenetA2Lite(_) => 3,
274            Self::WavenetA2Dyn(m) => m.channels,
275            Self::WavenetA2Cascade(m) => m.channels(),
276            Self::WavenetDyn(m) => m.ch,
277            Self::Container(c) => c.active().channels(),
278            Self::Lstm1x3(_) => 3,
279            Self::Lstm1x8(_) | Self::Lstm2x8(_) => 8,
280            Self::Lstm1x12(_) | Self::Lstm2x12(_) => 12,
281            Self::Lstm1x16(_) | Self::Lstm2x16(_) => 16,
282            Self::Lstm1x24(_) | Self::Lstm2x24(_) => 24,
283            Self::Lstm1x40(_) => 40,
284            Self::LstmDyn(m) => m.head_weights.len(),
285            Self::Linear(_) => 1,
286            Self::ConvNet(m) => m.in_channels(),
287        }
288    }
289
290    /// Returns the number of input channels consumed by the model's `process()`.
291    ///
292    /// - WaveNet variants and LSTM: always 1 (single sample per frame).
293    /// - ConvNet: `in_channels` from configuration (defaults to 1).
294    /// - Container: delegates to active sub-model.
295    pub fn in_channels(&self) -> usize {
296        match self {
297            Self::WavenetStandard(_)
298            | Self::WavenetLite(_)
299            | Self::WavenetFeather(_)
300            | Self::WavenetNano(_) => 1,
301            Self::WavenetA2Full(_) => 1,
302            Self::WavenetA2Lite(_) => 1,
303            Self::WavenetA2Dyn(_) => 1,
304            Self::WavenetA2Cascade(_) => 1,
305            Self::WavenetDyn(_) => 1,
306            Self::Container(c) => c.active().in_channels(),
307            Self::Lstm1x3(_)
308            | Self::Lstm1x8(_)
309            | Self::Lstm1x12(_)
310            | Self::Lstm1x16(_)
311            | Self::Lstm1x24(_)
312            | Self::Lstm2x8(_)
313            | Self::Lstm2x12(_)
314            | Self::Lstm2x16(_)
315            | Self::Lstm1x40(_)
316            | Self::Lstm2x24(_)
317            | Self::LstmDyn(_) => 1,
318            Self::Linear(_) => 1,
319            Self::ConvNet(m) => m.in_channels(),
320        }
321    }
322
323    /// Returns the receptive field size of the model (or 0 for LSTM/Container).
324    pub fn receptive_field(&self) -> usize {
325        match self {
326            Self::Container(_) => 0,
327            Self::Linear(m) => m.receptive_field,
328            _ if self.is_lstm() => 0,
329            _ => self.prewarm_samples(),
330        }
331    }
332
333    /// Returns the number of output channels produced by the model's `process()`.
334    ///
335    /// Mirrors C++ `DSP::NumOutputChannels()`:
336    /// - WaveNet A1 without post-stack head: last array's `head_size` (`wave_net_output_channels`)
337    /// - LSTM: `hidden_size`
338    /// - Linear: 1
339    /// - Container: delegates to active sub-model
340    pub fn num_output_channels(&self) -> usize {
341        match self {
342            Self::WavenetStandard(_)
343            | Self::WavenetLite(_)
344            | Self::WavenetFeather(_)
345            | Self::WavenetNano(_) => 1,
346            Self::WavenetA2Full(_) => 1,
347            Self::WavenetA2Lite(_) => 1,
348            Self::WavenetA2Dyn(_) => 1,
349            Self::WavenetA2Cascade(m) => m.arrays.last().map(|a| a.head_size).unwrap_or(1),
350            Self::WavenetDyn(m) => m.arrays.last().map(|a| a.head).unwrap_or(0),
351            Self::Container(c) => c.active().num_output_channels(),
352            Self::Lstm1x3(_)
353            | Self::Lstm1x8(_)
354            | Self::Lstm2x8(_)
355            | Self::Lstm1x12(_)
356            | Self::Lstm2x12(_)
357            | Self::Lstm1x16(_)
358            | Self::Lstm2x16(_)
359            | Self::Lstm1x24(_)
360            | Self::Lstm2x24(_)
361            | Self::Lstm1x40(_)
362            | Self::LstmDyn(_) => 1,
363            Self::Linear(_) => 1,
364            Self::ConvNet(m) => m.out_channels(),
365        }
366    }
367}
368
369/// Clones a `StaticModel` for use in `condition_dsp` preservation during
370/// channel slimming (`slice_wavenet_model`).
371///
372/// Only `WavenetDyn` is currently supported (the only variant that can appear
373/// as a recursive `condition_dsp` sub-model in the current dispatcher).
374/// All other variants return `None`, matching the pre-Task-2.1.2 behavior.
375pub(crate) fn clone_condition_dsp(model: &Option<Box<StaticModel>>) -> Option<Box<StaticModel>> {
376    model.as_ref().and_then(|m| match m.as_ref() {
377        StaticModel::WavenetDyn(w) => Some(Box::new(StaticModel::WavenetDyn(w.clone()))),
378        StaticModel::WavenetA2Cascade(_) => None,
379        _ => None,
380    })
381}
nam_rs/models/static_model.rs

nam_rs/models/
static_model.rs
