diff --git a/src/dsp/nam_host.py b/src/dsp/nam_host.py index f8cc3e3..5efad77 100644 --- a/src/dsp/nam_host.py +++ b/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 -and inference. Supports Linear, WaveNet, and LSTM architectures. +Leverages the neural-amp-modeler (NAM) .nam file format for model loading +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: - """Build a PyTorch model from a Core ML NAM (.nam) JSON dictionary. +# ── Architecture-aware NAM model builder ────────────────────────────── - Supports WaveNet and Linear architectures. Falls back gracefully - if the architecture is unsupported. +def _build_nam_model(data: dict) -> "torch.nn.Module": + """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: - architecture (str): ``\"WaveNet\"`` or ``\"Linear\"``. - config (dict): Architecture-specific config. - weights (list): Flat list of weight arrays. + data: Parsed .nam JSON content. Returns: - A PyTorch ``nn.Module`` ready for inference. - - Raises: - ValueError: If the architecture is not supported. + A ``torch.nn.Module`` ready for inference. """ 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: - """Build a WaveNet NAM model from config and weights.""" +def _build_linear(config: dict, weights: list) -> "torch.nn.Module": + """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() - 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 - ) + self.receptive_field = receptive_field def forward(self, x): - x = x.unsqueeze(1) # [B, 1, T] - for conv in self.layers: - x = torch.tanh(conv(x)) - x = self.out(x) - x = x.squeeze(1) # [B, T] - return x * self.head_scale + # x: (B, T) or (B, 1, T) + if x.dim() == 2: + x = x.unsqueeze(1) # (B, 1, T) - 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. - Falls back gracefully if the package is unavailable or the - architecture is unsupported. + Uses our lightweight in-process builder that handles + SlimmableContainer (A2), WaveNet, and Linear architectures. """ 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: - logger.warning("nam package not installed; inference unavailable") + except ImportError as exc: + 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) \ No newline at end of file + return host.process(audio)