fix: NAM SlimmableContainer (A2) model loading with correct weight import

Complete rewrite of _build_nam_model and _build_wavenet to handle:
- SlimmableContainer (A2): pick submodel by quality, recurse into WaveNet
- WaveNet: proper dilated convs with causal padding, 1x1 layer1x1 post-conv
- Linear: simple pass-through with gain/bias

Weight import matches NAM export order:
  rechannel -> per-layer(conv w+b, input_mixer w, layer1x1 w+b) -> head_rechannel w+b -> optional head -> head_scale

Also fixes MOCK_PRESETS in frontend to match actual factory preset names.

src/dsp/nam_host.py

@@ -1,11 +1,14 @@
 """NAM A2 model host — load, configure, and run inference on RPi 4B.
 
-Leverages the `neural-amp-modeler` (nam) Python package for model loading
+Leverages the neural-amp-modeler (NAM) .nam file format for model loading
-and inference. Supports Linear, WaveNet, and LSTM architectures.
+and inference via a lightweight PyTorch reimplementation. Supports:
+  - SlimmableContainer (A2 format): version 0.12+ with quality tier submodels
+  - WaveNet: standard dilated-convolution amp model
+  - Linear: simple pass-through/bias model
 
 On RPi 4B:
 - Python + PyTorch: good for Feather/Nano models only (~0.5-3ms at 256-block)
-- For Standard/Lite models, use `neural-amp-modeler-lv2` compiled natively
+- For Standard/Lite models, use neural-amp-modeler-lv2 compiled natively
   (NeuralAudio engine, LV2 plugin, ~1-2ms at 256-block)
 - For A2 Slimmable runtime quality dialing, port to OpenSauce/nam-rs
 """
@@ -24,23 +27,23 @@ from typing import Optional
 import numpy as np
 
 
-def _init_from_nam(data: dict) -> torch.nn.Module:
+# ── Architecture-aware NAM model builder ──────────────────────────────
-    """Build a PyTorch model from a Core ML NAM (.nam) JSON dictionary.
 
-    Supports WaveNet and Linear architectures. Falls back gracefully
+def _build_nam_model(data: dict) -> "torch.nn.Module":
-    if the architecture is unsupported.
+    """Build a PyTorch model from a parsed .nam file dict.
+
+    Dispatches to the appropriate builder based on ``architecture``:
+
+    * ``SlimmableContainer`` — A2 format: pick the highest-quality
+      submodel (``max_value`` closest to 1.0) and build its WaveNet.
+    * ``WaveNet`` — standard dilated-convolution architecture.
+    * ``Linear`` — simple single-weight pass-through.
 
     Args:
-        data: Parsed .nam file content with keys:
+        data: Parsed .nam JSON content.
-            architecture (str): ``\"WaveNet\"`` or ``\"Linear\"``.
-            config (dict): Architecture-specific config.
-            weights (list): Flat list of weight arrays.
 
     Returns:
-        A PyTorch ``nn.Module`` ready for inference.
+        A ``torch.nn.Module`` ready for inference.
-
-    Raises:
-        ValueError: If the architecture is not supported.
     """
     import torch.nn as nn
 
@@ -48,71 +51,308 @@ def _init_from_nam(data: dict) -> torch.nn.Module:
     config = data.get("config", {})
     weights = data.get("weights", [])
 
-    if arch == "WaveNet":
+    if arch == "SlimmableContainer":
+        # Pick the highest-quality submodel
+        submodels = config.get("submodels", [])
+        if not submodels:
+            raise ValueError("SlimmableContainer has no submodels")
+        # Sort by max_value descending — pick highest quality
+        best = max(submodels, key=lambda sm: sm.get("max_value", 0))
+        model_data = best.get("model", {})
+        if not model_data:
+            raise ValueError("SlimmableContainer submodel has no 'model' data")
+        # Recurse — the submodel's model field is a complete .nam export dict
+        return _build_nam_model(model_data)
+
+    elif arch == "WaveNet":
         return _build_wavenet(config, weights)
-    raise ValueError(f"Unsupported NAM architecture: {arch}")
+
+    elif arch == "Linear":
+        return _build_linear(config, weights)
+
+    else:
+        raise ValueError(f"Unsupported NAM architecture: {arch!r}")
 
 
-def _build_wavenet(config: dict, weights: list) -> torch.nn.Module:
+def _build_linear(config: dict, weights: list) -> "torch.nn.Module":
-    """Build a WaveNet NAM model from config and weights."""
+    """Build a Linear pass-through model.
+
+    NAM Linear model: just a scalar gain and optional bias.
+    """
+    import torch
+    import torch.nn as nn
+
+    has_bias = config.get("bias", False)
+
+    class LinearNAM(nn.Module):
+        def __init__(self, gain: float, bias: float = 0.0):
+            super().__init__()
+            self.gain = nn.Parameter(torch.tensor(gain, dtype=torch.float32))
+            self.bias = nn.Parameter(torch.tensor(bias, dtype=torch.float32)) if has_bias else 0.0
+
+        def forward(self, x):
+            # x: (B, T) or (T,)
+            if isinstance(self.bias, nn.Parameter):
+                return x * self.gain + self.bias
+            return x * self.gain
+
+    gain = float(weights[0]) if weights else 1.0
+    bias = float(weights[1]) if has_bias and len(weights) > 1 else 0.0
+    return LinearNAM(gain, bias)
+
+
+def _build_wavenet(config: dict, weights: list) -> "torch.nn.Module":
+    """Build a WaveNet NAM model from config and flat weight array.
+
+    Implements the standard NAM WaveNet architecture:
+    - RechannelIn: 1x1 conv (1 → C)
+    - N dilated conv layers with residual connections
+    - HeadRechannel: temporal conv (C → 1) + optional bias
+    - Optional WaveNet-level Head (activation + conv blocks)
+    - head_scale multiplier
+
+    Weights are imported in the exact order matching NAM's export format:
+
+    1. rechannel (weight only, no bias)
+    2. Per layer: conv(w+b), input_mixer(w, no bias)
+       (layer1x1, head1x1, FiLM omitted — not present in A2 models)
+    3. head_rechannel (w+b)
+    4. Optional head (conv w+b per block)
+    5. head_scale (single float)
+    """
     import torch
     import torch.nn as nn
 
     layers_cfg = config.get("layers", [])
-    head = config.get("head", 128)
+    head_cfg = config.get("head", None)  # WaveNet-level head
     head_scale = config.get("head_scale", 1.0)
 
+    if not layers_cfg:
+        raise ValueError("WaveNet config has no layers")
+
+    # ── Extract layer array config (we use first/only layer array) ─────
+    la = layers_cfg[0]  # A2 models have single layer array
+    input_size = la.get("input_size", 1)
+    condition_size = la.get("condition_size", 1)
+    channels = la.get("channels", 8)
+    kernel_sizes = la.get("kernel_sizes", la.get("kernel_size", [3]))
+    dilations = la.get("dilations", [1])
+    la_head_cfg = la.get("head", {"out_channels": 1, "kernel_size": 1, "bias": True})
+
+    # Normalize scalar kernel_size to list
+    if isinstance(kernel_sizes, int):
+        kernel_sizes = [kernel_sizes] * len(dilations)
+    if isinstance(dilations, int):
+        dilations = [dilations] * len(kernel_sizes)
+
+    num_layers = len(dilations)
+    out_channels = la_head_cfg.get("out_channels", 1)
+    head_kernel_size = la_head_cfg.get("kernel_size", 1)
+    head_bias = la_head_cfg.get("bias", True)
+
+    # ── Build modules ──────────────────────────────────────────────────
+
+    modules = []
+
+    # 1. RechannelIn: 1x1 conv, no bias
+    rechannel = nn.Conv1d(input_size, channels, 1, bias=False)
+    modules.append(("rechannel", rechannel))
+
+    # 2. Dilated conv layers + optional layer1x1
+    layer_convs = nn.ModuleList()
+    input_mixers = nn.ModuleList()
+    layer1x1s = nn.ModuleList()
+    for i in range(num_layers):
+        ks = kernel_sizes[i] if i < len(kernel_sizes) else kernel_sizes[-1]
+        d = dilations[i] if i < len(dilations) else 1
+        conv = nn.Conv1d(channels, channels, ks, dilation=d, padding=0)
+        im = nn.Conv1d(condition_size, channels, 1, bias=False)
+        layer_convs.append(conv)
+        input_mixers.append(im)
+        # Many A2 models include a 1x1 post-conv (layer1x1)
+        l1 = nn.Conv1d(channels, channels, 1, bias=True)
+        layer1x1s.append(l1)
+
+    modules.append(("layer_convs", layer_convs))
+    modules.append(("input_mixers", input_mixers))
+    modules.append(("layer1x1s", layer1x1s))
+
+    # 3. HeadRechannel: temporal conv
+    head_rechannel = nn.Conv1d(channels, out_channels, head_kernel_size, bias=head_bias)
+    modules.append(("head_rechannel", head_rechannel))
+
+    # 4. Optional WaveNet-level head
+    head_module = None
+    if head_cfg is not None and head_cfg.get("kernel_size", 0) > 0:
+        h_ks = head_cfg.get("kernel_size", 1)
+        h_bias = head_cfg.get("bias", True)
+        head_module = nn.Conv1d(out_channels, out_channels, h_ks, bias=h_bias)
+        modules.append(("head", head_module))
+
+    # ── Assemble WaveNetNAM module ─────────────────────────────────────
     class WaveNetNAM(nn.Module):
-        def __init__(self, layers_cfg, head, head_scale, weights):
+        def __init__(self, rechannel, layer_convs, input_mixers, layer1x1s,
+                     head_rechannel, head, head_scale, receptive_field):
             super().__init__()
+            self.rechannel = rechannel
+            self.layer_convs = layer_convs
+            self.input_mixers = input_mixers
+            self.layer1x1s = layer1x1s
+            self.head_rechannel = head_rechannel
             self.head = head
             self.head_scale = head_scale
-            self.layers = nn.ModuleList()
+            self.receptive_field = receptive_field
-            w_idx = 0
-            input_channels = 1
-            for lc in layers_cfg:
-                out_ch = lc.get("channels", 16)
-                kernel = lc.get("kernel_size", 3)
-                dilation = lc.get("dilation", 1)
-                conv = nn.Conv1d(input_channels, out_ch, kernel,
-                                padding=(kernel - 1) * dilation,
-                                dilation=dilation)
-                # Load weights
-                if w_idx < len(weights):
-                    conv.weight.data = torch.tensor(
-                        np.array(weights[w_idx]).reshape(conv.weight.shape),
-                        dtype=torch.float32)
-                    w_idx += 1
-                if w_idx < len(weights):
-                    conv.bias.data = torch.tensor(
-                        np.array(weights[w_idx]).reshape(conv.bias.shape),
-                        dtype=torch.float32)
-                    w_idx += 1
-                self.layers.append(conv)
-                out_ch, input_channels = input_channels, out_ch
-            self.out = nn.Conv1d(input_channels, 1, 1)
-            self.out.weight.data = torch.tensor(
-                np.array(weights[w_idx]).reshape(self.out.weight.shape),
-                dtype=torch.float32) if w_idx < len(weights) else torch.zeros(1, input_channels, 1)
-            w_idx += 1
-            if w_idx < len(weights):
-                self.out.bias.data = torch.tensor(
-                    np.array(weights[w_idx]).reshape(self.out.bias.shape),
-                    dtype=torch.float32)
-            self._receptive_field = sum(
-                (lc.get("kernel_size", 3) - 1) * lc.get("dilation", 1)
-                for lc in layers_cfg
-            )
 
         def forward(self, x):
-            x = x.unsqueeze(1)  # [B, 1, T]
+            # x: (B, T) or (B, 1, T)
-            for conv in self.layers:
+            if x.dim() == 2:
-                x = torch.tanh(conv(x))
+                x = x.unsqueeze(1)  # (B, 1, T)
-            x = self.out(x)
-            x = x.squeeze(1)  # [B, T]
-            return x * self.head_scale
 
-    return WaveNetNAM(layers_cfg, head, head_scale, weights)
+            # Rechannel
+            x = self.rechannel(x)  # (B, C, T)
+
+            # Dilated conv layers with residual
+            for conv, im, l1 in zip(self.layer_convs, self.input_mixers, self.layer1x1s):
+                # Causal padding (left-only) for dilated conv
+                pad = (conv.kernel_size[0] - 1) * conv.dilation[0]
+                if pad > 0:
+                    x_pad = nn.functional.pad(x, (pad, 0))
+                else:
+                    x_pad = x
+                # Dilated conv
+                out = conv(x_pad)
+                # Input mixer (condition path)
+                cond = torch.ones_like(x[:, :1, :])
+                c = im(cond)
+                # Ensure same length
+                if out.shape[-1] < x.shape[-1]:
+                    x = x[:, :, -out.shape[-1]:]
+                if c.shape[-1] < x.shape[-1]:
+                    c = c[:, :, -x.shape[-1]:]
+                # Gated activation
+                out = torch.tanh(out + c)
+                # 1x1 post-conv (layer1x1)
+                out = l1(out)
+                # Truncate to out length if needed
+                if out.shape[-1] < x.shape[-1]:
+                    x = x[:, :, -out.shape[-1]:]
+                # Residual
+                x = x + out
+
+            # Head rechannel
+            x = self.head_rechannel(x)  # (B, 1, T')
+            if x.shape[-1] > 0:
+                # Trim to receptive field boundary
+                pass
+
+            # Optional head
+            if self.head is not None:
+                pad = (self.head.kernel_size[0] - 1)
+                if pad > 0:
+                    x = nn.functional.pad(x, (pad, 0))
+                x = self.head(x)
+
+            # Scale
+            x = x * self.head_scale
+
+            return x.squeeze(1)  # (B, T')
+
+    # Compute receptive field
+    rf = 1
+    for i in range(num_layers):
+        ks = kernel_sizes[i] if i < len(kernel_sizes) else kernel_sizes[-1]
+        d = dilations[i] if i < len(dilations) else 1
+        rf += (ks - 1) * d
+    rf += head_kernel_size - 1
+    if head_module is not None:
+        rf += head_cfg.get("kernel_size", 1) - 1
+
+    model = WaveNetNAM(
+        rechannel, layer_convs, input_mixers, layer1x1s,
+        head_rechannel, head_module, head_scale, rf,
+    )
+
+    # ── Import weights from flat array ─────────────────────────────────
+    _import_wavenet_weights(model, weights)
+    return model
+
+
+def _import_wavenet_weights(model: "torch.nn.Module", weights: list) -> None:
+    """Import flat weight array into a WaveNetNAM model.
+
+    Weight order (matching NAM's export_weights):
+    1. rechannel.weight
+    2. For each layer: conv.weight, conv.bias, input_mixer.weight
+    3. head_rechannel.weight, head_rechannel.bias
+    4. head.weight, head.bias (if head exists)
+    5. head_scale
+    """
+    import torch as _torch
+    import numpy as _np
+
+    w = _torch.tensor(weights, dtype=_torch.float32)
+    i = 0
+
+    # rechannel.weight: [C, 1, 1]
+    n = model.rechannel.weight.numel()
+    model.rechannel.weight.data = w[i:i+n].reshape(model.rechannel.weight.shape)
+    i += n
+
+    # Per-layer weights
+    for conv, im, l1 in zip(model.layer_convs, model.input_mixers, model.layer1x1s):
+        # conv.weight: [C, C, ks]
+        n = conv.weight.numel()
+        conv.weight.data = w[i:i+n].reshape(conv.weight.shape)
+        i += n
+        # conv.bias: [C]
+        n = conv.bias.numel()
+        conv.bias.data = w[i:i+n].reshape(conv.bias.shape)
+        i += n
+        # input_mixer.weight: [C, condition_size, 1]
+        n = im.weight.numel()
+        im.weight.data = w[i:i+n].reshape(im.weight.shape)
+        i += n
+        # layer1x1.weight: [C, C, 1]
+        n = l1.weight.numel()
+        l1.weight.data = w[i:i+n].reshape(l1.weight.shape)
+        i += n
+        # layer1x1.bias: [C]
+        n = l1.bias.numel()
+        l1.bias.data = w[i:i+n].reshape(l1.bias.shape)
+        i += n
+
+    # head_rechannel.weight
+    n = model.head_rechannel.weight.numel()
+    model.head_rechannel.weight.data = w[i:i+n].reshape(model.head_rechannel.weight.shape)
+    i += n
+    # head_rechannel.bias
+    if model.head_rechannel.bias is not None:
+        n = model.head_rechannel.bias.numel()
+        model.head_rechannel.bias.data = w[i:i+n].reshape(model.head_rechannel.bias.shape)
+        i += n
+
+    # head weight + bias (if present)
+    if model.head is not None:
+        n = model.head.weight.numel()
+        model.head.weight.data = w[i:i+n].reshape(model.head.weight.shape)
+        i += n
+        if model.head.bias is not None:
+            n = model.head.bias.numel()
+            model.head.bias.data = w[i:i+n].reshape(model.head.bias.shape)
+            i += n
+
+    # head_scale (final float)
+    if i < len(w):
+        model.head_scale = float(w[i])
+        i += 1
+
+    # Verify we consumed all weights
+    if i != len(w):
+        warnings.warn(
+            f"NAM weight import consumed {i}/{len(w)} weights. "
+            f"Model has {sum(p.numel() for p in model.parameters())} params. "
+            f"Check weight order matches architecture."
+        )
 
 
 logger = logging.getLogger(__name__)
@@ -312,9 +552,8 @@ class NAMHost:
     def _build_inference(self, data: dict) -> bool:
         """Instantiate a PyTorch model from a NAM config dict.
 
-        Uses the ``nam`` package's ``init_from_nam()`` factory.
+        Uses our lightweight in-process builder that handles
-        Falls back gracefully if the package is unavailable or the
+        SlimmableContainer (A2), WaveNet, and Linear architectures.
-        architecture is unsupported.
         """
         self._import_torch()
 
@@ -328,7 +567,7 @@ class NAMHost:
             return True
 
         try:
-            model = _init_from_nam(data)
+            model = _build_nam_model(data)
             model.eval()
 
             # Move to target device
@@ -342,8 +581,8 @@ class NAMHost:
             self._inference_model = model
             return True
 
-        except ImportError:
+        except ImportError as exc:
-            logger.warning("nam package not installed; inference unavailable")
+            logger.warning("Required package not installed; inference unavailable: %s", exc)
             self._inference_model = None
             return False
         except Exception as exc:
@@ -561,4 +800,4 @@ def process_with_model(
     """Convenience: load a model, process one block, return audio."""
     host = NAMHost(device=device)
     host.load_model(model_path)
-    return host.process(audio)
+    return host.process(audio)