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:
+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
|
||||
Reference in New Issue
Block a user