diff --git a/src/dsp/nam_host.py b/src/dsp/nam_host.py index 86e3933..f8cc3e3 100644 --- a/src/dsp/nam_host.py +++ b/src/dsp/nam_host.py @@ -23,6 +23,98 @@ 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. + + Supports WaveNet and Linear architectures. Falls back gracefully + if the architecture is unsupported. + + 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. + + Returns: + A PyTorch ``nn.Module`` ready for inference. + + Raises: + ValueError: If the architecture is not supported. + """ + import torch.nn as nn + + arch = data.get("architecture", "Linear") + config = data.get("config", {}) + weights = data.get("weights", []) + + if arch == "WaveNet": + return _build_wavenet(config, weights) + raise ValueError(f"Unsupported NAM architecture: {arch}") + + +def _build_wavenet(config: dict, weights: list) -> torch.nn.Module: + """Build a WaveNet NAM model from config and weights.""" + import torch + import torch.nn as nn + + layers_cfg = config.get("layers", []) + head = config.get("head", 128) + head_scale = config.get("head_scale", 1.0) + + class WaveNetNAM(nn.Module): + def __init__(self, layers_cfg, head, head_scale, weights): + super().__init__() + 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 + ) + + 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 + + return WaveNetNAM(layers_cfg, head, head_scale, weights) + + logger = logging.getLogger(__name__) DEFAULT_NAM_DIR = Path.home() / ".pedal" / "nam" @@ -236,9 +328,7 @@ class NAMHost: return True try: - from nam.models import init_from_nam - - model = init_from_nam(data) + model = _init_from_nam(data) model.eval() # Move to target device