"""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. 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 (NeuralAudio engine, LV2 plugin, ~1-2ms at 256-block) - For A2 Slimmable runtime quality dialing, port to OpenSauce/nam-rs """ from __future__ import annotations import json import logging from dataclasses import dataclass, field from pathlib import Path from typing import Optional, Any 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" @dataclass class NAMModel: """Metadata for a loaded NAM model.""" name: str path: str size_mb: float architecture: str channels: int sample_rate: int = 48000 latency_samples: int = 0 compatible: bool = True @property def family(self) -> str: """Categorize the model by size.""" if self.size_mb < 0.1: return "nano" elif self.size_mb < 1.0: return "feather" elif self.size_mb < 4.0: return "lite" elif self.size_mb < 10.0: return "standard" else: return "heavy" @property def estimated_latency_ms(self) -> str: """Return estimated per-block latency on RPi 4B at 256-block / 48kHz.""" 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)", } return estimates.get(self.family, "unknown") 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. """ def __init__( self, models_dir: str | Path = DEFAULT_NAM_DIR, lv2_dir: str | Path = DEFAULT_LV2_MODEL_DIR, use_lv2: bool = True, ): self._models_dir = Path(models_dir) self._lv2_dir = Path(lv2_dir) self._use_lv2 = use_lv2 self._loaded_model: Optional[NAMModel] = None self._inference_model: Any = None # PyTorch model instance self._torch = None # lazy import self._models_dir.mkdir(parents=True, exist_ok=True) def _import_torch(self): """Lazy-import torch to avoid startup cost when using LV2.""" if self._torch is None: import torch self._torch = torch 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") @property def is_loaded(self) -> bool: return self._loaded_model is not None @property def current_model(self) -> Optional[NAMModel]: return self._loaded_model