shawn
0e77adb4c3
- Full FFT overlap-add IR convolution in IRLoader (process(), set_mix(), toggle) - Lazy FFT computation — IR FFT padded to correct block+ir size on first process() - Wet/dry mix control, enabled/disabled toggle with tail clearing - Fixed pipeline._apply_ir_cab() to delegate to IRLoader.process() instead of poking internals (old code had array-size mismatch bug: IR FFT at ir_len vs block FFT at conv_size) - 46 tests: loading, convolution correctness, overlap-add state, mix, toggle, directory listing, performance budget (all <5ms even at 8192 taps), edge cases - scripts/download_irs.sh: free IR pack downloader (God's Cab, Seacow)
236 lines
7.4 KiB
Python
236 lines
7.4 KiB
Python
"""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 |