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.
This commit is contained in:
+310
-71
@@ -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)
|
||||
return host.process(audio)
|
||||
|
||||
Reference in New Issue
Block a user