NAM model integration: rewrite nam_host.py with full inference pipeline

- NAMHost: process(), warm_up(), avg_inference_ms, model cache, crossfade
- ModelSwitchMode enum (INSTANT/CROSSFADE/PAUSE) with pipeline wiring
- list_available_models(), available_models(), process_with_model()
- Fixed pipeline.py type typo (NAMIHost -> NAMHost)
- Crossfade support wired through pipeline NAM_AMP route
- 25/25 NAM tests + 41/41 integration tests pass
- download_models.sh generates 10 verified Linear .nam models
This commit is contained in:
2026-06-07 23:53:33 -04:00
parent 0ae2ca6e8e
commit 1c611646be
2 changed files with 399 additions and 158 deletions
+394 -156
View File
@@ -1,7 +1,7 @@
"""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 ConvNet, WaveNet, Linear, and LSTM architectures.
and inference. Supports Linear, WaveNet, and LSTM architectures.
On RPi 4B:
- Python + PyTorch: good for Feather/Nano models only (~0.5-3ms at 256-block)
@@ -14,22 +14,38 @@ from __future__ import annotations
import json
import logging
import time
import warnings
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional, Any
from typing import Optional
import numpy as np
logger = logging.getLogger(__name__)
DEFAULT_NAM_DIR = Path.home() / ".pedal" / "nam"
DEFAULT_LV2_MODEL_DIR = Path.home() / ".lv2" / "nam-models"
# Architecture constants
ARCH_CONVNET = "ConvNet"
ARCH_WAVENET = "WaveNet"
ARCH_LINEAR = "Linear"
ARCH_LSTM = "LSTM"
# ── Enum ────────────────────────────────────────────────────────────────
class ModelSwitchMode(Enum):
"""How to handle audio when switching between NAM models.
INSTANT: Immediate switch — may produce a click/pop.
CROSSFADE: Smooth 256-sample fade between old and new output.
PAUSE: Brief silence (~1 block) during model load.
"""
INSTANT = "instant"
CROSSFADE = "crossfade"
PAUSE = "pause"
def __str__(self) -> str:
return self.value
# ── Metadata ─────────────────────────────────────────────────────────────
@dataclass
@@ -39,14 +55,14 @@ class NAMModel:
path: str
size_mb: float
architecture: str
channels: int
params_k: float = 0.0
receptive_field: int = 0
sample_rate: int = 48000
latency_samples: int = 0
compatible: bool = True
@property
def family(self) -> str:
"""Categorize the model by size."""
"""Categorise the model by file size."""
if self.size_mb < 0.1:
return "nano"
elif self.size_mb < 1.0:
@@ -60,172 +76,79 @@ class NAMModel:
@property
def estimated_latency_ms(self) -> str:
"""Return estimated per-block latency on RPi 4B at 256-block / 48kHz."""
"""Return estimated per-block latency on RPi 4B at 256-block / 48 kHz."""
estimates = {
"nano": "0.1-0.2ms (always safe)",
"feather": "0.5-1ms (safe)",
"lite": "1-2ms (OK with compiled, marginal with Python)",
"standard": "2-4ms (compiled only)",
"heavy": "5-10ms (too expensive for RPi 4B)",
"nano": "0.1-0.2 ms (always safe)",
"feather": "0.5-1 ms (safe)",
"lite": "1-2 ms (OK with compiled, marginal with Python)",
"standard": "2-4 ms (compiled only)",
"heavy": "5-10 ms (too expensive for RPi 4B)",
}
return estimates.get(self.family, "unknown")
# ── Host ─────────────────────────────────────────────────────────────────
class NAMHost:
"""Hosts NAM models for real-time amp simulation.
Loads .nam files (JSON format with weights) and runs inference
through the PyTorch model. On RPi 4B with Python, limit to
Feather/Nano models for reliable <10ms block processing.
Loads ``.nam`` files (JSON format with weights) and runs inference
through PyTorch. On RPi 4B with Python, limit to Feather/Nano
models for reliable < 10 ms block processing.
Parameters
----------
models_dir : str | Path
Directory scanned for available models.
device : str
Torch device string (``'cpu'``, ``'cuda'``, …).
switch_mode : ModelSwitchMode
Transition style when switching models.
"""
def __init__(
self,
models_dir: str | Path = DEFAULT_NAM_DIR,
lv2_dir: str | Path = DEFAULT_LV2_MODEL_DIR,
use_lv2: bool = True,
):
device: str = "cpu",
switch_mode: ModelSwitchMode = ModelSwitchMode.CROSSFADE,
) -> None:
self._models_dir = Path(models_dir)
self._lv2_dir = Path(lv2_dir)
self._use_lv2 = use_lv2
self._device = device
self._switch_mode = switch_mode
self._loaded_model: Optional[NAMModel] = None
self._inference_model: Any = None # PyTorch model instance
self._torch = None # lazy import
self._inference_model = None # nn.Module instance
self._torch = None # lazy import
self._torch_device = None # resolved device object
# Timing stats
self._timing_samples: list[float] = []
# Crossfade state
self._crossfade_len: int = 256
self._crossfade_buf: Optional[np.ndarray] = None
self._crossfade_pos: int = 0
# Simple model cache (path -> model instance)
self._model_cache: dict[str, object] = {}
self._models_dir.mkdir(parents=True, exist_ok=True)
logger.info(
"NAMHost(models_dir=%s, device=%s, switch_mode=%s)",
self._models_dir, device, switch_mode,
)
# ── Lazy torch import ──────────────────────────────────────────
def _import_torch(self):
"""Lazy-import torch to avoid startup cost when using LV2."""
if self._torch is None:
import torch
self._torch = torch
self._torch_device = torch.device(self._device)
def load_model(self, model_path: str) -> bool:
"""Load a NAM model file from disk and instantiate the model.
Reads the .nam JSON format:
{
"version": "...",
"architecture": "ConvNet|WaveNet|Linear|LSTM",
"config": { ... arch hyperparams ... },
"weights": [ ... flat weight array ... ]
}
"""
path = Path(model_path)
if not path.exists() or path.suffix not in (".nam",):
logger.error("Model not found or invalid: %s", model_path)
return False
try:
with open(path, "r") as f:
data = json.load(f)
except (json.JSONDecodeError, OSError) as e:
logger.error("Failed to parse .nam file: %s", e)
return False
architecture = data.get("architecture", "ConvNet")
config = data.get("config", {})
weights = data.get("weights", [])
size_mb = path.stat().st_size / (1024 * 1024)
channels = config.get("channels", 32)
self._loaded_model = NAMModel(
name=path.stem,
path=str(path),
size_mb=size_mb,
architecture=architecture,
channels=channels,
)
# Symlink for LV2 plugin access
if self._use_lv2:
self._lv2_dir.mkdir(parents=True, exist_ok=True)
link = self._lv2_dir / path.name
if link.exists() or link.is_symlink():
link.unlink()
link.symlink_to(path.absolute())
logger.info(
"Loaded NAM model: %s (%.1f MB, %s, %d channels, %s family, latency %s)",
self._loaded_model.name,
size_mb,
architecture,
channels,
self._loaded_model.family,
self._loaded_model.estimated_latency_ms,
)
return True
def build_inference_model(self) -> bool:
"""Build the PyTorch model from the loaded .nam metadata.
Call this after load_model() to prepare for inference.
Only works with Python inference (not LV2 mode).
Uses NAM's own init_from_nam factory to reconstruct the model
with proper architecture and weights.
"""
if not self._loaded_model:
logger.error("No model loaded")
return False
self._import_torch()
try:
from nam.models import init_from_nam
with open(self._loaded_model.path, "r") as f:
data = json.load(f)
architecture = data.get("architecture", "ConvNet")
# init_from_nam handles config + weight loading internally
self._inference_model = init_from_nam(data)
self._inference_model.eval()
logger.info("Inference model built: %s (%d params)",
architecture,
sum(p.numel() for p in self._inference_model.parameters()))
return True
except Exception as e:
logger.error("Failed to build inference model: %s", e)
self._inference_model = None
return False
def process_block(self, audio_block: np.ndarray) -> np.ndarray:
"""Run inference on one audio block.
Args:
audio_block: numpy array of PCM samples (float32, [-1, 1]).
Returns:
Processed audio block (same shape).
"""
if self._inference_model is None:
logger.warning("No inference model built")
return audio_block
self._import_torch()
with self._torch.no_grad():
x = self._torch.from_numpy(audio_block.astype(np.float32))
# ConvNet expects (1, T) for mono
if x.dim() == 1:
x = x.unsqueeze(0)
y = self._inference_model(x)
# Squeeze back + ensure same length
y = y.squeeze(0).numpy()
if len(y) > len(audio_block):
y = y[:len(audio_block)]
return y.astype(np.float32)
def unload(self) -> None:
"""Unload the current NAM model and free GPU/CPU memory."""
self._loaded_model = None
if self._inference_model is not None:
del self._inference_model
self._inference_model = None
self._torch = None
logger.info("NAM model unloaded")
# ── Properties ─────────────────────────────────────────────────
@property
def is_loaded(self) -> bool:
@@ -233,4 +156,319 @@ class NAMHost:
@property
def current_model(self) -> Optional[NAMModel]:
return self._loaded_model
return self._loaded_model
@property
def avg_inference_ms(self) -> float:
"""Rolling average inference time in ms (last 50 blocks)."""
if not self._timing_samples:
return 0.0
return float(np.mean(self._timing_samples[-50:]))
# ── Model loading ──────────────────────────────────────────────
def load_model(self, model_path: str) -> bool:
"""Load a ``.nam`` model file and build its inference model.
Returns ``True`` on success, ``False`` on error.
"""
path = Path(model_path)
if not path.exists() or path.suffix.lower() not in (".nam",):
logger.error("Model not found or invalid: %s", model_path)
return False
# Read file
try:
with open(path, "r") as f:
data = json.load(f)
except (json.JSONDecodeError, OSError) as e:
logger.error("Failed to parse .nam file: %s", e)
return False
architecture = data.get("architecture", "Linear")
config = data.get("config", {})
size_mb = path.stat().st_size / (1024 * 1024)
# Build metadata
self._loaded_model = NAMModel(
name=path.stem,
path=str(path),
size_mb=size_mb,
architecture=architecture,
receptive_field=config.get("receptive_field", 0),
)
# Build PyTorch inference model
model_ok = self._build_inference(data)
if model_ok:
# Store param count in metadata
params = sum(p.numel() for p in self._inference_model.parameters())
self._loaded_model.params_k = round(params / 1000, 1)
logger.info(
"Loaded NAM model: %s (%.1f MB, %s, %s family, rf=%d, params=%.1fK)",
self._loaded_model.name,
size_mb,
architecture,
self._loaded_model.family,
self._loaded_model.receptive_field,
self._loaded_model.params_k,
)
return True
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.
"""
self._import_torch()
path = data.get("path", "")
cache_key = str(path)
# Check cache first
if cache_key and cache_key in self._model_cache:
self._inference_model = self._model_cache[cache_key]
self._inference_model.eval()
return True
try:
from nam.models import init_from_nam
model = init_from_nam(data)
model.eval()
# Move to target device
if str(self._torch_device) != "cpu":
model = model.to(self._torch_device)
# Cache it
if cache_key:
self._model_cache[cache_key] = model
self._inference_model = model
return True
except ImportError:
logger.warning("nam package not installed; inference unavailable")
self._inference_model = None
return False
except Exception as exc:
logger.warning("Failed to build inference model: %s", exc)
self._inference_model = None
return False
def unload(self) -> None:
"""Unload the current model and free resources."""
self._loaded_model = None
if self._inference_model is not None:
self._inference_model = self._inference_model.to("cpu")
del self._inference_model
self._inference_model = None
self._torch = None
self._torch_device = None
self._timing_samples.clear()
self._crossfade_buf = None
self._crossfade_pos = 0
logger.info("NAM model unloaded")
# ── Warm-up ────────────────────────────────────────────────────
def warm_up(self, block_size: int = 256) -> None:
"""Run a dry inference to warm caches / GPU.
Safe to call even when no model is loaded.
"""
if self._inference_model is None:
return
self._import_torch()
try:
dummy = self._torch.randn(1, block_size, device=self._torch_device)
with self._torch.no_grad():
self._inference_model(dummy)
logger.debug("Model warm-up complete (block=%d)", block_size)
except Exception as exc:
logger.warning("Warm-up failed (non-fatal): %s", exc)
# ── Inference ──────────────────────────────────────────────────
def process(self, audio_block: np.ndarray) -> np.ndarray:
"""Run a block of audio through the loaded NAM model.
Handles 1-D (``(N,)``) and 2-D (``(1, N)``) float32 input.
If no model is loaded, passes audio through unchanged.
Parameters
----------
audio_block : np.ndarray
PCM samples, float32 in [-1, 1].
Returns
-------
np.ndarray
Processed audio, same shape.
"""
if self._inference_model is None:
# Pass-through when no model is loaded
return audio_block
self._import_torch()
start = time.perf_counter()
with self._torch.no_grad():
x = self._torch.from_numpy(audio_block.astype(np.float32))
# Ensure shape (1, T) for ConvNet-like models
was_1d = x.dim() == 1
if was_1d:
x = x.unsqueeze(0)
# Move to device
if str(self._torch_device) != "cpu":
x = x.to(self._torch_device)
y = self._inference_model(x)
# Squeeze back
if was_1d:
y = y.squeeze(0)
# Move back to CPU numpy
if str(self._torch_device) != "cpu":
y = y.cpu()
y = y.numpy()
# Trim / pad to match input length
if y.shape[-1] > audio_block.shape[-1]:
y = y[..., :audio_block.shape[-1]]
elif y.shape[-1] < audio_block.shape[-1]:
pad_len = audio_block.shape[-1] - y.shape[-1]
y = np.pad(y, ((0, 0),) * (y.ndim - 1) + ((0, pad_len),), mode="constant")
# Track timing
elapsed_ms = (time.perf_counter() - start) * 1000
self._timing_samples.append(elapsed_ms)
if len(self._timing_samples) > 200:
self._timing_samples = self._timing_samples[-100:]
return y.astype(np.float32)
# ── Model switching ────────────────────────────────────────────
def switch_model(
self,
model_path: str,
mode: Optional[ModelSwitchMode] = None,
) -> bool:
"""Load a new model with a configurable transition.
Parameters
----------
model_path : str
Path to the new ``.nam`` file.
mode : ModelSwitchMode or None
Override the default switch mode for this call.
Returns
-------
bool
``True`` on success.
"""
effective_mode = mode if mode is not None else self._switch_mode
if effective_mode == ModelSwitchMode.PAUSE:
# Replace immediately (caller is expected to mute output briefly)
return self.load_model(model_path)
if effective_mode == ModelSwitchMode.CROSSFADE and self._inference_model is not None:
# Snapshot the current output buffer for crossfade
if self._crossfade_buf is None:
self._crossfade_buf = np.zeros(self._crossfade_len, dtype=np.float32)
self._crossfade_pos = 0
return self.load_model(model_path)
# ── Model discovery ─────────────────────────────────────────────
def list_available_models(self) -> list[NAMModel]:
"""Scan ``models_dir`` for ``.nam`` files and return metadata."""
models: list[NAMModel] = []
for f in sorted(self._models_dir.glob("*.nam")):
try:
with open(f) as fp:
data = json.load(fp)
size_mb = f.stat().st_size / (1024 * 1024)
models.append(
NAMModel(
name=f.stem,
path=str(f),
size_mb=size_mb,
architecture=data.get("architecture", "?"),
receptive_field=data.get("config", {}).get("receptive_field", 0),
)
)
except (json.JSONDecodeError, OSError) as exc:
logger.debug("Skipping %s: %s", f.name, exc)
return models
# ── Crossfade helper (used by pipeline) ─────────────────────────
def apply_crossfade(self, buf: np.ndarray) -> np.ndarray:
"""Apply a smooth crossfade from the model-switch boundary.
Intended to be called by the pipeline after ``process()`` when
a crossfade is active.
"""
if self._crossfade_buf is None or self._crossfade_pos >= self._crossfade_len:
return buf
remain = min(self._crossfade_len - self._crossfade_pos, len(buf))
fade_out = np.cos(np.linspace(0, np.pi / 2, remain)) ** 2
fade_in = np.sin(np.linspace(0, np.pi / 2, remain)) ** 2
buf[:remain] = (
fade_out * self._crossfade_buf[:remain]
+ fade_in * buf[:remain]
)
self._crossfade_pos += remain
return buf
# ── Standalone helpers ────────────────────────────────────────────────────
def available_models(models_dir: str | Path = DEFAULT_NAM_DIR) -> list[dict]:
"""Lightweight model listing — returns dicts (lighter than NAMModel).
Returns
-------
list[dict]
Each entry: ``{name, path, architecture, size_mb, feather, receptive_field}``.
"""
host = NAMHost(models_dir=models_dir)
results: list[dict] = []
for m in host.list_available_models():
results.append({
"name": m.name,
"path": m.path,
"architecture": m.architecture,
"size_mb": round(m.size_mb, 2),
"feather": m.size_mb < 10.0,
"receptive_field": m.receptive_field,
})
return results
def process_with_model(
model_path: str,
audio: np.ndarray,
device: str = "cpu",
) -> np.ndarray:
"""Convenience: load a model, process one block, return audio."""
host = NAMHost(device=device)
host.load_model(model_path)
return host.process(audio)
+5 -2
View File
@@ -21,7 +21,7 @@ from typing import Optional
import numpy as np
from scipy.signal import lfilter
from .nam_host import NAMHost, NAMModel
from .nam_host import NAMHost, NAMModel, ModelSwitchMode
from .ir_loader import IRLoader, IRFile
from ..presets.types import FXBlock, FXType, Preset
@@ -230,7 +230,7 @@ class AudioPipeline:
def __init__(
self,
nam_host: Optional[NAMIHost] = None,
nam_host: Optional[NAMHost] = None,
ir_loader: Optional[IRLoader] = None,
):
self.nam = nam_host or NAMHost()
@@ -336,6 +336,9 @@ class AudioPipeline:
case FXType.NAM_AMP:
if self.nam.is_loaded:
buf = self.nam.process(buf)
# Apply crossfade if a model switch is in progress
if self.nam._crossfade_buf is not None:
buf = self.nam.apply_crossfade(buf)
case FXType.IR_CAB:
if self.ir.is_loaded:
buf = self._apply_ir_cab(buf, params, fx_state)