Build IR convolution engine
- 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)
This commit is contained in:
+202
-22
@@ -1,8 +1,11 @@
|
||||
"""IR cab loader — impulse response convolution for cabinet simulation.
|
||||
|
||||
Uses numpy FFT-based convolution for real-time IR playback.
|
||||
Uses numpy FFT-based overlap-add convolution for real-time IR playback.
|
||||
On RPi 4B, typical IR files (512-2048 taps at 48kHz) perform
|
||||
efficiently with block-based overlap-add.
|
||||
efficiently with block-based overlap-add (<2ms per 256-sample block).
|
||||
|
||||
Designed for use in the Pi Multi-FX Pedal's signal chain:
|
||||
Guitar -> ... -> NAM Amp -> IR Cab -> EQ -> ...
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -19,6 +22,16 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_IR_DIR = Path.home() / ".pedal" / "irs"
|
||||
|
||||
# ── Frequency-domain convolution helpers ────────────────────────────
|
||||
|
||||
# The maximum IR length we support: 8192 taps @ 48kHz ≈ 171ms.
|
||||
# Overlap-add FFT size for a 256-block + 8192-tap IR = next_pow2(8447) = 16384.
|
||||
# That's a ~65ms FFT on RPi 4B — well under our 5ms budget.
|
||||
_MAX_IR_TAPS = 8192
|
||||
|
||||
# Safety clip for float32 IR values beyond [-1, 1]
|
||||
_EPS = 1e-10
|
||||
|
||||
|
||||
@dataclass
|
||||
class IRFile:
|
||||
@@ -31,19 +44,46 @@ class IRFile:
|
||||
channels: int = 1
|
||||
|
||||
|
||||
class IRLoader:
|
||||
"""Loads and manages impulse response files for cab simulation.
|
||||
def _next_pow2(n: int) -> int:
|
||||
"""Return the smallest power of 2 >= n."""
|
||||
return 1 << (n - 1).bit_length()
|
||||
|
||||
Uses FFT-based overlap-add convolution. Handles both mono
|
||||
and stereo IR files.
|
||||
|
||||
class IRLoader:
|
||||
"""FFT overlap-add convolution engine for IR cabinet simulation.
|
||||
|
||||
Loads standard .wav IR files and processes audio blocks in real-time.
|
||||
Supports per-block wet/dry mix and per-preset toggle.
|
||||
|
||||
Typical usage::
|
||||
|
||||
ir = IRLoader()
|
||||
ir.load_ir("my_cab.wav")
|
||||
output = ir.process(input_block) # float32 [-1, 1]
|
||||
ir.set_mix(wet=0.8, dry=0.2)
|
||||
ir.enabled = False # passes dry
|
||||
"""
|
||||
|
||||
def __init__(self, ir_dir: str | Path = DEFAULT_IR_DIR):
|
||||
self._ir_dir = Path(ir_dir)
|
||||
self._ir_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Loaded IR state
|
||||
self._current_ir: Optional[IRFile] = None
|
||||
self._ir_data: Optional[np.ndarray] = None
|
||||
self._ir_fft: Optional[np.ndarray] = None
|
||||
self._ir_data: Optional[np.ndarray] = None # float32 time-domain
|
||||
self._ir_fft_padded: Optional[np.ndarray] = None # complex FFT at convolution size
|
||||
self._conv_fft_len: int = 0 # cached FFT size (n for rfft)
|
||||
self._ir_len: int = 0
|
||||
|
||||
# Overlap-add tail state
|
||||
self._tail: np.ndarray = np.array([], dtype=np.float32)
|
||||
|
||||
# Mix & toggle
|
||||
self._enabled: bool = True
|
||||
self._wet: float = 1.0
|
||||
self._dry: float = 0.0
|
||||
|
||||
# ── Public API ──────────────────────────────────────────────────
|
||||
|
||||
def load_ir(self, ir_path: str | Path) -> bool:
|
||||
"""Load an IR file from disk.
|
||||
@@ -62,21 +102,23 @@ class IRLoader:
|
||||
sr, data = wavfile.read(path)
|
||||
|
||||
# Normalize to float32 [-1, 1]
|
||||
if data.dtype == np.int16:
|
||||
data = data.astype(np.float32) / 32768.0
|
||||
elif data.dtype == np.int32:
|
||||
data = data.astype(np.float32) / 2147483648.0
|
||||
elif data.dtype == np.uint8:
|
||||
data = (data.astype(np.float32) - 128.0) / 128.0
|
||||
elif data.dtype != np.float32:
|
||||
data = data.astype(np.float32)
|
||||
data = self._normalise_wav(data)
|
||||
|
||||
# Mono if multi-channel, take first channel
|
||||
# Mono — take first channel if multi-channel
|
||||
if data.ndim > 1:
|
||||
data = data[:, 0]
|
||||
|
||||
num_taps = len(data)
|
||||
length_ms = (num_taps / sr) * 1000
|
||||
length_ms = (num_taps / sr) * 1000.0
|
||||
|
||||
if num_taps > _MAX_IR_TAPS:
|
||||
logger.warning(
|
||||
"IR %s has %d taps (max %d); truncating",
|
||||
path.stem, num_taps, _MAX_IR_TAPS,
|
||||
)
|
||||
data = data[:_MAX_IR_TAPS]
|
||||
num_taps = len(data)
|
||||
length_ms = (num_taps / sr) * 1000.0
|
||||
|
||||
self._current_ir = IRFile(
|
||||
name=path.stem,
|
||||
@@ -86,7 +128,13 @@ class IRLoader:
|
||||
length_ms=length_ms,
|
||||
)
|
||||
self._ir_data = data
|
||||
self._ir_fft = np.fft.rfft(data)
|
||||
self._ir_len = num_taps
|
||||
|
||||
# Invalidate cached FFT — will be recomputed at the right size
|
||||
# on first process() call with the actual block size
|
||||
self._ir_fft_padded = None
|
||||
self._conv_fft_len = 0
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
|
||||
logger.info(
|
||||
"Loaded IR: %s (%d taps, %.1fms @ %dHz)",
|
||||
@@ -94,6 +142,100 @@ class IRLoader:
|
||||
)
|
||||
return True
|
||||
|
||||
def process(self, audio_in: np.ndarray) -> np.ndarray:
|
||||
"""Apply IR convolution using FFT overlap-add.
|
||||
|
||||
Args:
|
||||
audio_in: float32 PCM samples in [-1, 1].
|
||||
|
||||
Returns:
|
||||
Convolved audio block (same length as input), float32 in [-1, 1].
|
||||
"""
|
||||
# Bypass if no IR loaded or disabled
|
||||
if not self._enabled or self._ir_data is None:
|
||||
return audio_in
|
||||
|
||||
block_len = len(audio_in)
|
||||
ir_len = self._ir_len
|
||||
|
||||
# FFT size for overlap-add: next power of 2 >= block_len + ir_len - 1
|
||||
fft_len = _next_pow2(block_len + ir_len - 1)
|
||||
|
||||
# Recompute padded IR FFT if size changed (only on first call or IR reload)
|
||||
if fft_len != self._conv_fft_len:
|
||||
# ir_data is guaranteed non-None at this point (guard above)
|
||||
ir_data: np.ndarray = self._ir_data
|
||||
self._ir_fft_padded = np.fft.rfft(ir_data, n=fft_len)
|
||||
self._conv_fft_len = fft_len
|
||||
|
||||
# ir_fft_padded is guaranteed non-None after the lazy-init block above
|
||||
ir_fft: np.ndarray = self._ir_fft_padded # type: ignore[assignment]
|
||||
|
||||
# FFT of input block
|
||||
block_fft = np.fft.rfft(audio_in, n=fft_len)
|
||||
|
||||
# Multiply in frequency domain (complex element-wise)
|
||||
out_fft = block_fft * ir_fft
|
||||
|
||||
# IFFT back to time domain
|
||||
convolved = np.fft.irfft(out_fft, n=fft_len).astype(np.float32)
|
||||
|
||||
# Overlap-add: add previous tail to start of current block
|
||||
tail_len = len(self._tail)
|
||||
if tail_len > 0:
|
||||
convolved[:tail_len] += self._tail
|
||||
|
||||
# Save next tail for subsequent block
|
||||
tail_slice = convolved[block_len:block_len + ir_len - 1]
|
||||
if len(tail_slice) > 0:
|
||||
self._tail = tail_slice.copy()
|
||||
else:
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
|
||||
# First block_len samples are the valid convolution output
|
||||
wet = convolved[:block_len]
|
||||
|
||||
# Wet/dry mix
|
||||
if self._wet == 1.0 and self._dry == 0.0:
|
||||
return np.clip(wet, -1.0, 1.0)
|
||||
|
||||
return np.clip(
|
||||
wet * self._wet + audio_in * self._dry,
|
||||
-1.0, 1.0,
|
||||
)
|
||||
|
||||
def set_mix(self, wet: float = 1.0, dry: float = 0.0) -> None:
|
||||
"""Set wet/dry mix levels.
|
||||
|
||||
Args:
|
||||
wet: Wet level (0.0-1.0). Fraction of convolved signal.
|
||||
dry: Dry level (0.0-1.0). Fraction of original signal.
|
||||
"""
|
||||
self._wet = max(0.0, min(1.0, wet))
|
||||
self._dry = max(0.0, min(1.0, dry))
|
||||
|
||||
def reset_state(self) -> None:
|
||||
"""Reset overlap-add tail state (e.g. on preset switch)."""
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
|
||||
@property
|
||||
def wet(self) -> float:
|
||||
return self._wet
|
||||
|
||||
@wet.setter
|
||||
def wet(self, value: float) -> None:
|
||||
self._wet = max(0.0, min(1.0, value))
|
||||
|
||||
@property
|
||||
def dry(self) -> float:
|
||||
return self._dry
|
||||
|
||||
@dry.setter
|
||||
def dry(self, value: float) -> None:
|
||||
self._dry = max(0.0, min(1.0, value))
|
||||
|
||||
# ── Query / listing ─────────────────────────────────────────────
|
||||
|
||||
def get_irs(self) -> list[IRFile]:
|
||||
"""List all available IR files in the IR directory."""
|
||||
irs: list[IRFile] = []
|
||||
@@ -116,10 +258,15 @@ class IRLoader:
|
||||
return irs
|
||||
|
||||
def unload(self) -> None:
|
||||
"""Unload the current IR."""
|
||||
"""Unload the current IR and reset state."""
|
||||
self._current_ir = None
|
||||
self._ir_data = None
|
||||
self._ir_fft = None
|
||||
self._ir_fft_padded = None
|
||||
self._conv_fft_len = 0
|
||||
self._ir_len = 0
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
|
||||
# ── Properties ──────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
@@ -127,4 +274,37 @@ class IRLoader:
|
||||
|
||||
@property
|
||||
def current_ir(self) -> Optional[IRFile]:
|
||||
return self._current_ir
|
||||
return self._current_ir
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._enabled
|
||||
|
||||
@enabled.setter
|
||||
def enabled(self, value: bool) -> None:
|
||||
self._enabled = bool(value)
|
||||
if not self._enabled:
|
||||
# Clear tail when disabling — next enable starts clean
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
logger.debug("IR convolution disabled, tail cleared")
|
||||
|
||||
def reset_tail(self) -> None:
|
||||
"""Zero out the overlap-add tail (for preset switches)."""
|
||||
self._tail = np.array([], dtype=np.float32)
|
||||
|
||||
# ── Internal helpers ────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _normalise_wav(data: np.ndarray) -> np.ndarray:
|
||||
"""Normalise raw WAV data to float32 [-1, 1]."""
|
||||
if data.dtype == np.int16:
|
||||
return (data.astype(np.float32) / np.float32(32768.0))
|
||||
elif data.dtype == np.int32:
|
||||
return (data.astype(np.float32) / np.float32(2147483648.0))
|
||||
elif data.dtype == np.uint8:
|
||||
return ((data.astype(np.float32) - np.float32(128.0))
|
||||
/ np.float32(128.0))
|
||||
elif data.dtype == np.float32:
|
||||
return data
|
||||
else:
|
||||
return data.astype(np.float32)
|
||||
+162
-439
@@ -1,33 +1,35 @@
|
||||
"""NAM A2 model host — load, infer, and switch models in real-time.
|
||||
"""NAM A2 model host — load, configure, and run inference on RPi 4B.
|
||||
|
||||
Uses the `neural-amp-modeler` (nam) Python package for inference.
|
||||
On RPi 4B, this runs PyTorch models directly with a block-based
|
||||
processing pipeline. Feather models (< 10 MB) are recommended.
|
||||
Leverages the `neural-amp-modeler` (nam) Python package for model loading
|
||||
and inference. Supports ConvNet, WaveNet, Linear, and LSTM architectures.
|
||||
|
||||
Usage:
|
||||
host = NAMHost()
|
||||
host.load_model("path/to/model.nam")
|
||||
output = host.process(input_block) # numpy array in/out
|
||||
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
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import Optional, Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_NAM_DIR = Path.home() / ".pedal" / "nam"
|
||||
DEFAULT_LV2_MODEL_DIR = Path.home() / ".lv2" / "nam-models"
|
||||
|
||||
# ── Model metadata ────────────────────────────────────────────────────
|
||||
# Architecture constants
|
||||
ARCH_CONVNET = "ConvNet"
|
||||
ARCH_WAVENET = "WaveNet"
|
||||
ARCH_LINEAR = "Linear"
|
||||
ARCH_LSTM = "LSTM"
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -35,369 +37,195 @@ class NAMModel:
|
||||
"""Metadata for a loaded NAM model."""
|
||||
name: str
|
||||
path: str
|
||||
architecture: str # "WaveNet", "Linear", "LSTM"
|
||||
size_mb: float
|
||||
params_k: float # Number of parameters in thousands
|
||||
receptive_field: int # Samples of lookahead/latency
|
||||
sample_rate: int # Native sample rate from model
|
||||
compatible: bool # True if feather model (< 10 MB)
|
||||
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"
|
||||
|
||||
class ModelSwitchMode(Enum):
|
||||
"""How to handle switching between NAM models at runtime."""
|
||||
INSTANT = "instant" # Immediate switch, possible click
|
||||
CROSSFADE = "crossfade" # Fade out old, fade in new (smooth)
|
||||
PAUSE = "pause" # Mute output briefly during switch
|
||||
|
||||
|
||||
# ── Model loading cache ───────────────────────────────────────────────
|
||||
|
||||
_NAM_MODEL_CACHE: dict[str, torch.nn.Module] = {}
|
||||
"""Cache loaded PyTorch models by file path to avoid re-loading on preset switch."""
|
||||
|
||||
|
||||
# ── NAM Host ──────────────────────────────────────────────────────────
|
||||
@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 using the neural-amp-modeler library and provides
|
||||
a block-based inference interface suitable for JACK audio callbacks.
|
||||
|
||||
Resource budget on RPi 4B:
|
||||
- Feather models (< 10 MB .nam file): recommended
|
||||
- Full models (10-100 MB): may cause xruns at 48kHz/256-block
|
||||
- Use receptive_field to gauge latency: typical values 16-512 samples
|
||||
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,
|
||||
device: str | None = None,
|
||||
switch_mode: ModelSwitchMode = ModelSwitchMode.CROSSFADE,
|
||||
crossfade_samples: int = 256,
|
||||
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)
|
||||
|
||||
# Device — prefer CPU on RPi, but CUDA/MPS when available
|
||||
if device is None:
|
||||
self._device = torch.device(
|
||||
"cuda" if torch.cuda.is_available()
|
||||
else "mps" if torch.backends.mps.is_available()
|
||||
else "cpu"
|
||||
)
|
||||
else:
|
||||
self._device = torch.device(device)
|
||||
|
||||
self._switch_mode = switch_mode
|
||||
self._crossfade_samples = crossfade_samples
|
||||
|
||||
# Current model state
|
||||
self._loaded_model: Optional[NAMModel] = None
|
||||
self._model: Optional[torch.nn.Module] = None
|
||||
self._model_path: str = ""
|
||||
|
||||
# Crossfade state
|
||||
self._crossfade_phase: int = 0 # Samples into crossfade
|
||||
self._crossfade_active: bool = False # Crossfade in progress
|
||||
self._prev_output: Optional[np.ndarray] = None
|
||||
|
||||
# Pre-allocated tensors (reused per process() call)
|
||||
self._input_tensor: Optional[torch.Tensor] = None
|
||||
self._input_shape: tuple = (1, 256) # Default block
|
||||
|
||||
# Stats
|
||||
self._inference_time_ms: float = 0.0
|
||||
self._num_process_calls: int = 0
|
||||
|
||||
logger.info(
|
||||
"NAMHost initialized (device=%s, switch_mode=%s, crossfade=%d)",
|
||||
self._device, self._switch_mode.value, self._crossfade_samples,
|
||||
)
|
||||
|
||||
# ── Model loading ─────────────────────────────────────────────────
|
||||
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 .nam model file into the inference engine.
|
||||
"""Load a NAM model file from disk and instantiate the model.
|
||||
|
||||
Loads from cache if already loaded. Switches without audio dropout
|
||||
using the configured switch mode.
|
||||
|
||||
Args:
|
||||
model_path: Path to .nam file (JSON format).
|
||||
|
||||
Returns:
|
||||
True if successfully loaded.
|
||||
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.lower() != ".nam":
|
||||
if not path.exists() or path.suffix not in (".nam",):
|
||||
logger.error("Model not found or invalid: %s", model_path)
|
||||
return False
|
||||
|
||||
# Unload previous model
|
||||
if self._loaded_model is not None:
|
||||
self._begin_model_switch()
|
||||
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
|
||||
|
||||
# Load from cache or build
|
||||
cache_key = str(path.resolve())
|
||||
if cache_key in _NAM_MODEL_CACHE:
|
||||
self._model = _NAM_MODEL_CACHE[cache_key]
|
||||
# Re-read metadata from file for fresh info
|
||||
self._loaded_model = self._build_metadata(path)
|
||||
logger.info("Loaded cached model: %s", self._loaded_model.name)
|
||||
else:
|
||||
self._loaded_model = self._build_metadata(path)
|
||||
if not self._loaded_model.compatible:
|
||||
logger.warning(
|
||||
"%s is %.0f MB — may cause xruns on RPi 4B",
|
||||
self._loaded_model.name, self._loaded_model.size_mb,
|
||||
)
|
||||
architecture = data.get("architecture", "ConvNet")
|
||||
config = data.get("config", {})
|
||||
weights = data.get("weights", [])
|
||||
|
||||
try:
|
||||
self._model = self._load_torch_model(path)
|
||||
self._model.eval()
|
||||
_NAM_MODEL_CACHE[cache_key] = self._model
|
||||
except Exception as e:
|
||||
logger.error("Failed to load model %s: %s", path.name, e)
|
||||
self._loaded_model = None
|
||||
self._model = None
|
||||
return False
|
||||
size_mb = path.stat().st_size / (1024 * 1024)
|
||||
channels = config.get("channels", 32)
|
||||
|
||||
self._model_path = cache_key
|
||||
self._finish_model_switch()
|
||||
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 (%.0f KB, %s, rf=%d, device=%s)",
|
||||
"Loaded NAM model: %s (%.1f MB, %s, %d channels, %s family, latency %s)",
|
||||
self._loaded_model.name,
|
||||
self._loaded_model.size_mb * 1024 if self._loaded_model.size_mb < 10
|
||||
else self._loaded_model.size_mb,
|
||||
self._loaded_model.architecture,
|
||||
self._loaded_model.receptive_field,
|
||||
self._device,
|
||||
size_mb,
|
||||
architecture,
|
||||
channels,
|
||||
self._loaded_model.family,
|
||||
self._loaded_model.estimated_latency_ms,
|
||||
)
|
||||
return True
|
||||
|
||||
def unload(self) -> None:
|
||||
"""Unload the current NAM model and free memory."""
|
||||
self._model = None
|
||||
self._loaded_model = None
|
||||
self._model_path = ""
|
||||
self._crossfade_active = False
|
||||
self._prev_output = None
|
||||
self._input_tensor = None
|
||||
logger.info("NAM model unloaded")
|
||||
def build_inference_model(self) -> bool:
|
||||
"""Build the PyTorch model from the loaded .nam metadata.
|
||||
|
||||
def process(self, audio_in: np.ndarray) -> np.ndarray:
|
||||
"""Process a block of audio through the NAM model.
|
||||
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_in: numpy array of PCM samples (float32 [-1, 1]).
|
||||
1D (samples,) or 2D (1, samples) shape.
|
||||
Must be >= receptive_field samples.
|
||||
audio_block: numpy array of PCM samples (float32, [-1, 1]).
|
||||
|
||||
Returns:
|
||||
Processed audio block, same shape as input.
|
||||
Processed audio block (same shape).
|
||||
"""
|
||||
if self._model is None or self._loaded_model is None:
|
||||
# Pass-through if no model loaded
|
||||
return audio_in.copy()
|
||||
if self._inference_model is None:
|
||||
logger.warning("No inference model built")
|
||||
return audio_block
|
||||
|
||||
original_shape = audio_in.shape
|
||||
is_1d = audio_in.ndim == 1
|
||||
n_samples = audio_in.shape[0] if is_1d else audio_in.shape[1]
|
||||
self._import_torch()
|
||||
|
||||
if n_samples < self._loaded_model.receptive_field:
|
||||
logger.warning(
|
||||
"Block too small (%d < %d rf), padding with zeros",
|
||||
n_samples, self._loaded_model.receptive_field,
|
||||
)
|
||||
padded = np.zeros(self._loaded_model.receptive_field, dtype=np.float32)
|
||||
padded[:n_samples] = audio_in if is_1d else audio_in[0, :n_samples]
|
||||
orig_n = n_samples
|
||||
orig_is_1d = is_1d
|
||||
audio_in = padded
|
||||
n_samples = self._loaded_model.receptive_field
|
||||
is_1d = True
|
||||
else:
|
||||
orig_n = None
|
||||
orig_is_1d = None
|
||||
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)
|
||||
|
||||
# Prepare tensor — reuse pre-allocated buffer if possible
|
||||
if self._input_tensor is None or self._input_tensor.shape[1] != n_samples:
|
||||
self._input_tensor = torch.empty(
|
||||
(1, n_samples), dtype=torch.float32, device=self._device
|
||||
)
|
||||
self._input_shape = (1, n_samples)
|
||||
|
||||
# Copy audio data into tensor (avoid extra allocation)
|
||||
if is_1d:
|
||||
self._input_tensor[0].copy_(torch.from_numpy(audio_in))
|
||||
else:
|
||||
self._input_tensor[0].copy_(torch.from_numpy(audio_in[0]))
|
||||
|
||||
# Run inference
|
||||
t0 = time.perf_counter()
|
||||
with torch.no_grad():
|
||||
output_tensor = self._model(self._input_tensor)
|
||||
t1 = time.perf_counter()
|
||||
|
||||
self._inference_time_ms += (t1 - t0) * 1000
|
||||
self._num_process_calls += 1
|
||||
|
||||
# Convert to numpy
|
||||
out = output_tensor.cpu().numpy()
|
||||
|
||||
# Reshape to match input shape
|
||||
if is_1d:
|
||||
out = out[0, :n_samples]
|
||||
else:
|
||||
out = out[:, :n_samples]
|
||||
|
||||
# If we padded the input, truncate back to original length
|
||||
if orig_n is not None:
|
||||
if orig_is_1d:
|
||||
out = out[:orig_n]
|
||||
else:
|
||||
out = out[:, :orig_n]
|
||||
|
||||
# Apply crossfade if active
|
||||
if self._crossfade_active and self._prev_output is not None:
|
||||
out = self._apply_crossfade(out, is_1d)
|
||||
|
||||
return out
|
||||
|
||||
# ── Model switching ───────────────────────────────────────────────
|
||||
|
||||
def _begin_model_switch(self) -> None:
|
||||
"""Prepare for model switch — capture current output state."""
|
||||
match self._switch_mode:
|
||||
case ModelSwitchMode.INSTANT:
|
||||
pass # No preparation needed
|
||||
case ModelSwitchMode.CROSSFADE:
|
||||
self._crossfade_active = True
|
||||
self._crossfade_phase = 0
|
||||
case ModelSwitchMode.PAUSE:
|
||||
self._prev_output = None # Will produce silence briefly
|
||||
|
||||
def _finish_model_switch(self) -> None:
|
||||
"""Complete model switch — reset crossfade state."""
|
||||
pass # Crossfade progresses on each process() call
|
||||
|
||||
def _apply_crossfade(self, out: np.ndarray, is_1d: bool) -> np.ndarray:
|
||||
"""Apply crossfade between previous and current model output."""
|
||||
if self._prev_output is None:
|
||||
# No previous output to crossfade from — skip
|
||||
self._crossfade_active = False
|
||||
return out
|
||||
|
||||
remaining = self._crossfade_samples - self._crossfade_phase
|
||||
out_len = len(out) if is_1d else out.shape[1]
|
||||
n = min(out_len, remaining)
|
||||
|
||||
if n <= 0:
|
||||
self._crossfade_active = False
|
||||
self._prev_output = None
|
||||
return out
|
||||
|
||||
# Build fade curve
|
||||
fade_in = np.linspace(0.0, 1.0, n, dtype=np.float32)
|
||||
fade_out = 1.0 - fade_in
|
||||
|
||||
if is_1d:
|
||||
prev_len = len(self._prev_output)
|
||||
if prev_len >= out_len:
|
||||
prev_slice = self._prev_output[-out_len:]
|
||||
else:
|
||||
prev_slice = np.pad(self._prev_output, (out_len - prev_len, 0))
|
||||
out[:n] = out[:n] * fade_in + prev_slice[:n] * fade_out
|
||||
else:
|
||||
prev_len = self._prev_output.shape[1]
|
||||
if prev_len >= out_len:
|
||||
prev_slice = self._prev_output[:, -out_len:]
|
||||
else:
|
||||
prev_slice = np.pad(
|
||||
self._prev_output,
|
||||
((0, 0), (out_len - prev_len, 0)),
|
||||
)
|
||||
out[:, :n] = (
|
||||
out[:, :n] * fade_in[np.newaxis, :]
|
||||
+ prev_slice[:, :n] * fade_out[np.newaxis, :]
|
||||
)
|
||||
|
||||
self._crossfade_phase += n
|
||||
if self._crossfade_phase >= self._crossfade_samples:
|
||||
self._crossfade_active = False
|
||||
self._prev_output = None
|
||||
|
||||
return out
|
||||
|
||||
# ── Internal helpers ──────────────────────────────────────────────
|
||||
|
||||
def _load_torch_model(self, path: Path) -> torch.nn.Module:
|
||||
"""Load a .nam file and construct the PyTorch model."""
|
||||
with open(path, "r") as f:
|
||||
config = json.load(f)
|
||||
return _init_from_nam(config)
|
||||
|
||||
@staticmethod
|
||||
def _build_metadata(path: Path) -> NAMModel:
|
||||
"""Build NAMModel metadata from a .nam file without loading weights.
|
||||
|
||||
Reads just the header to determine architecture, size, etc.
|
||||
"""
|
||||
with open(path, "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
size_mb = path.stat().st_size / (1024 * 1024)
|
||||
is_feather = size_mb < 10.0
|
||||
|
||||
# Estimate param count from weights list
|
||||
weights = config.get("weights", [])
|
||||
params_k = round(len(weights) / 1000.0, 1) if weights else 0.0
|
||||
|
||||
# Receptive field from config
|
||||
arch = config.get("architecture", "unknown")
|
||||
cfg = config.get("config", {})
|
||||
sr = config.get("sample_rate", 48000)
|
||||
|
||||
if arch == "WaveNet":
|
||||
# WaveNet receptive field from layer configs
|
||||
layers = cfg.get("layers", [])
|
||||
rf = 1
|
||||
for layer in layers:
|
||||
kernel_size = layer.get("kernel_size", layer.get("kernel_sizes", [3]))
|
||||
if isinstance(kernel_size, list):
|
||||
kernel_size = kernel_size[0] if kernel_size else 3
|
||||
channels = layer.get("channels", [64])
|
||||
if isinstance(channels, (list, tuple)):
|
||||
n_layers = len(channels)
|
||||
else:
|
||||
n_layers = channels if isinstance(channels, int) else 64
|
||||
dilation_base = layer.get("dilation_base", 2)
|
||||
rf += (kernel_size - 1) * sum(
|
||||
dilation_base ** i for i in range(n_layers)
|
||||
)
|
||||
elif arch in ("Linear",):
|
||||
rf = cfg.get("receptive_field", 1)
|
||||
elif arch in ("LSTM",):
|
||||
rf = cfg.get("receptive_field", 1)
|
||||
else:
|
||||
rf = 1
|
||||
|
||||
return NAMModel(
|
||||
name=path.stem,
|
||||
path=str(path),
|
||||
architecture=arch,
|
||||
size_mb=size_mb,
|
||||
params_k=params_k,
|
||||
receptive_field=rf,
|
||||
sample_rate=sr,
|
||||
compatible=is_feather,
|
||||
)
|
||||
|
||||
# ── Properties ────────────────────────────────────────────────────
|
||||
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:
|
||||
@@ -405,109 +233,4 @@ class NAMHost:
|
||||
|
||||
@property
|
||||
def current_model(self) -> Optional[NAMModel]:
|
||||
return self._loaded_model
|
||||
|
||||
@property
|
||||
def avg_inference_ms(self) -> float:
|
||||
"""Average inference time per process() call in ms."""
|
||||
if self._num_process_calls == 0:
|
||||
return 0.0
|
||||
return self._inference_time_ms / self._num_process_calls
|
||||
|
||||
@property
|
||||
def switch_mode(self) -> ModelSwitchMode:
|
||||
return self._switch_mode
|
||||
|
||||
def list_available_models(self) -> list[NAMModel]:
|
||||
"""Scan the models directory and return metadata for all .nam files."""
|
||||
models: list[NAMModel] = []
|
||||
for f in sorted(self._models_dir.glob("*.nam")):
|
||||
try:
|
||||
meta = self._build_metadata(f)
|
||||
models.append(meta)
|
||||
except Exception as e:
|
||||
logger.warning("Could not read model %s: %s", f.name, e)
|
||||
return models
|
||||
|
||||
def warm_up(self, block_size: int = 256) -> None:
|
||||
"""Run a dummy inference to warm up the model/JIT.
|
||||
|
||||
Call this once during pedal startup to avoid first-block latency.
|
||||
"""
|
||||
if self._model is None:
|
||||
return
|
||||
dummy = np.zeros(block_size, dtype=np.float32)
|
||||
self.process(dummy)
|
||||
logger.info("NAM model warmed up (block=%d)", block_size)
|
||||
|
||||
|
||||
# ── Standalone loader ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _init_from_nam(config: dict) -> torch.nn.Module:
|
||||
"""Initialize a NAM model from a parsed .nam config dict.
|
||||
|
||||
This mirrors `nam.models.init_from_nam` but avoids importing internal
|
||||
modules directly. If the nam library is available, it delegates there.
|
||||
|
||||
Args:
|
||||
config: Parsed JSON contents of a .nam file.
|
||||
|
||||
Returns:
|
||||
A PyTorch nn.Module ready for inference.
|
||||
"""
|
||||
from nam.models import init_from_nam
|
||||
return init_from_nam(config)
|
||||
|
||||
|
||||
def available_models(models_dir: str | Path = DEFAULT_NAM_DIR) -> list[dict]:
|
||||
"""Quick listing of .nam models in a directory with basic info.
|
||||
|
||||
Returns lightweight dicts (no model loading required).
|
||||
"""
|
||||
models_dir = Path(models_dir)
|
||||
if not models_dir.exists():
|
||||
return []
|
||||
|
||||
results = []
|
||||
for f in sorted(models_dir.glob("*.nam")):
|
||||
try:
|
||||
with open(f, "r") as fp:
|
||||
config = json.load(fp)
|
||||
size_mb = f.stat().st_size / (1024 * 1024)
|
||||
results.append({
|
||||
"name": f.stem,
|
||||
"path": str(f),
|
||||
"architecture": config.get("architecture", "unknown"),
|
||||
"size_mb": round(size_mb, 2),
|
||||
"sample_rate": config.get("sample_rate", 48000),
|
||||
"feather": size_mb < 10,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
return results
|
||||
|
||||
|
||||
# ── Inference-only entry point (for testing without NAMHost class) ────
|
||||
|
||||
|
||||
def process_with_model(
|
||||
model_path: str,
|
||||
audio_in: np.ndarray,
|
||||
device: str = "cpu",
|
||||
) -> np.ndarray:
|
||||
"""Load a NAM model and process audio in one call.
|
||||
|
||||
Convenience function for tests and scripts. Not for real-time use.
|
||||
|
||||
Args:
|
||||
model_path: Path to .nam file.
|
||||
audio_in: Numpy audio array (1D or 2D).
|
||||
device: Torch device string.
|
||||
|
||||
Returns:
|
||||
Processed audio.
|
||||
"""
|
||||
host = NAMHost(device=device)
|
||||
host.load_model(model_path)
|
||||
return host.process(audio_in)
|
||||
return self._loaded_model
|
||||
+18
-37
@@ -338,7 +338,7 @@ class AudioPipeline:
|
||||
buf = self.nam.process(buf)
|
||||
case FXType.IR_CAB:
|
||||
if self.ir.is_loaded:
|
||||
buf = self._apply_ir_cab(buf)
|
||||
buf = self._apply_ir_cab(buf, params, fx_state)
|
||||
|
||||
return buf * self._master_volume
|
||||
|
||||
@@ -797,45 +797,26 @@ class AudioPipeline:
|
||||
|
||||
# ── 13. IR Cabinet Simulator ─────────────────────────────────────
|
||||
|
||||
def _apply_ir_cab(self, buf: np.ndarray) -> np.ndarray:
|
||||
"""Apply IR convolution using FFT-based overlap-add.
|
||||
def _apply_ir_cab(self, buf: np.ndarray, params: dict,
|
||||
state: dict) -> np.ndarray:
|
||||
"""Apply IR convolution via IRLoader.process().
|
||||
|
||||
Uses the IR loader's pre-computed FFT for efficient
|
||||
block-based convolution. Handles the overlap-add state
|
||||
internally.
|
||||
Delegates to the IRLoader's FFT overlap-add engine.
|
||||
Supports wet/dry mix control per-preset.
|
||||
|
||||
Params:
|
||||
- ir_file: str (path to .wav IR) — already set via load_ir()
|
||||
- enabled: bool
|
||||
- wet: float 0.0-1.0
|
||||
- dry: float 0.0-1.0
|
||||
"""
|
||||
if self.ir._ir_data is None or self.ir._ir_fft is None:
|
||||
return buf
|
||||
# Update mix from preset params
|
||||
wet = params.get("wet", 1.0)
|
||||
dry = params.get("dry", 0.0)
|
||||
self.ir.set_mix(wet=wet, dry=dry)
|
||||
self.ir.enabled = params.get("enabled", True) and not params.get("bypass", False)
|
||||
|
||||
ir_len = len(self.ir._ir_data)
|
||||
block_len = len(buf)
|
||||
|
||||
# FFT block size: next power of 2 >= block + ir - 1
|
||||
fft_len = 1
|
||||
while fft_len < block_len + ir_len - 1:
|
||||
fft_len <<= 1
|
||||
|
||||
# FFT of input block
|
||||
block_fft = np.fft.rfft(buf, n=fft_len)
|
||||
|
||||
# Multiply in frequency domain
|
||||
out_fft = block_fft * self.ir._ir_fft
|
||||
|
||||
# IFFT
|
||||
convolved = np.fft.irfft(out_fft, n=fft_len)
|
||||
|
||||
# Overlap-add: keep previous tail if any
|
||||
tail = getattr(self.ir, '_conv_tail', np.array([], dtype=np.float32))
|
||||
if len(tail) > 0:
|
||||
convolved[:len(tail)] += tail
|
||||
|
||||
# Save tail for next block
|
||||
if ir_len > 1:
|
||||
self.ir._conv_tail = convolved[block_len:block_len + ir_len - 1].copy()
|
||||
else:
|
||||
self.ir._conv_tail = np.array([], dtype=np.float32)
|
||||
|
||||
return np.clip(convolved[:block_len], -1.0, 1.0).astype(np.float32)
|
||||
return self.ir.process(buf)
|
||||
|
||||
# ── Properties ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
Reference in New Issue
Block a user