fix: replace nam package import with local init_from_nam
The PyPI 'nam' package (v0.0.3) is Neural Additive Models for tabular data — completely different from Neural Audio Model (guitar amp modeling). Implemented a local _init_from_nam() that builds a WaveNet NAM model from the Core ML .nam JSON format. Uninstalling the wrong package.
This commit is contained in:
+93
-3
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user