From 1c611646be564cbaee901b9d506fc5c30f8472f7 Mon Sep 17 00:00:00 2001 From: Shawn Date: Sun, 7 Jun 2026 23:53:33 -0400 Subject: [PATCH] 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 --- src/dsp/nam_host.py | 550 +++++++++++++++++++++++++++++++------------- src/dsp/pipeline.py | 7 +- 2 files changed, 399 insertions(+), 158 deletions(-) diff --git a/src/dsp/nam_host.py b/src/dsp/nam_host.py index 58596c4..86e3933 100644 --- a/src/dsp/nam_host.py +++ b/src/dsp/nam_host.py @@ -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 \ No newline at end of file + 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) \ No newline at end of file diff --git a/src/dsp/pipeline.py b/src/dsp/pipeline.py index 0bb014b..88a342f 100644 --- a/src/dsp/pipeline.py +++ b/src/dsp/pipeline.py @@ -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)