"""Pure-numpy host DSP for LavaSR (STFT / mel / iSTFT / resample / LR merge). Everything here mirrors the original PyTorch pipeline: * denoiser (UL-UNAS): 16 kHz, n_fft=512, hop=256, center=True STFT * enhancer (Vocos): 48 kHz samples, 44.1 kHz mel config, n_fft=2048, hop=512, padding="same", 80 slaney mels * FastLRMerge: keep the low band of the original 48 kHz signal and the high band generated by the model Only numpy is required. The same math is used by the converter, the SDK and the calibration generator, so board accuracy can be compared apples-to-apples. """ import math import os import numpy as np _ASSET_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") def _load_asset(name): return np.load(os.path.join(_ASSET_DIR, name)).astype(np.float32) class Assets: def __init__(self, asset_dir=None): d = asset_dir or _ASSET_DIR self.mel_fb = np.load(os.path.join(d, "mel_fb.npy")).astype(np.float32) self.mel_window = np.load(os.path.join(d, "mel_window.npy")).astype(np.float32) self.den_window = np.load(os.path.join(d, "den_window.npy")).astype(np.float32) self.erb_fc = np.load(os.path.join(d, "erb_fc.npy")).astype(np.float32) self.ierb_fc = np.load(os.path.join(d, "ierb_fc.npy")).astype(np.float32) def resample(x, orig_freq, new_freq, lowpass_filter_width=6, rolloff=0.99): """Bit-faithful numpy port of torchaudio.functional.resample (sinc_interp_hann).""" x = np.asarray(x, dtype=np.float32) if orig_freq == new_freq: return x gcd = math.gcd(int(orig_freq), int(new_freq)) o = int(orig_freq) // gcd n = int(new_freq) // gcd base = min(o, n) * rolloff width = int(math.ceil(lowpass_filter_width * o / base)) idx = np.arange(-width, width + o, dtype=np.float64) / o t = np.arange(0, -n, -1, dtype=np.float64)[:, None] / n + idx[None, :] t = np.clip(t * base, -lowpass_filter_width, lowpass_filter_width) window = np.cos(t * math.pi / lowpass_filter_width / 2) ** 2 kernel = np.sinc(t) * window * (base / o) kernel = kernel.astype(np.float32) k = kernel.shape[1] padded = np.pad(x, (width, width + o)) frames = np.lib.stride_tricks.sliding_window_view(padded, k)[::o] y = frames @ kernel.T target = int(math.ceil(n * len(x) / o)) return np.ascontiguousarray(y.reshape(-1)[:target], dtype=np.float32) def _frames(x, n_fft, hop): return np.lib.stride_tricks.sliding_window_view(x, n_fft)[::hop] def den_stft(x, window): n_fft = window.shape[0] pad = n_fft // 2 xp = np.pad(x, (pad, pad), mode="reflect") return np.fft.rfft(_frames(xp, n_fft, pad) * window, n=n_fft, axis=-1) def den_istft(spec, window, out_len): n_fft = window.shape[0] hop = n_fft // 2 frames = np.fft.irfft(spec, n=n_fft, axis=-1) * window total = (len(frames) - 1) * hop + n_fft y = np.zeros(total, dtype=np.float32) env = np.zeros(total, dtype=np.float32) for i, frame in enumerate(frames): s = i * hop y[s:s + n_fft] += frame env[s:s + n_fft] += window * window pad = hop y = y[pad:total - pad] / env[pad:total - pad] if len(y) < out_len: y = np.pad(y, (0, out_len - len(y))) return np.ascontiguousarray(y[:out_len], dtype=np.float32) def den_features(spec, assets): """magnitude -> log10 -> ERB band merge; returns (1,1,T,129) float32.""" feat = np.log10(np.clip(np.abs(spec), 1e-12, None)) low_bins = 257 - assets.erb_fc.shape[1] low = feat[:, :low_bins] high = feat[:, low_bins:] @ assets.erb_fc.T feat = np.concatenate([low, high], axis=1) return np.ascontiguousarray(feat[None, None], dtype=np.float32) def den_mask_expand(mask, assets): """ERB mask (T,129) -> full 257-bin mask (T,257).""" low_bins = 257 - assets.ierb_fc.shape[0] low = mask[:, :low_bins] high = mask[:, low_bins:] @ assets.ierb_fc.T return np.concatenate([low, high], axis=1) def mel_features(x, assets): """Original Vocos MelSpectrogramFeatures(padding='same'); returns (1,80,T).""" n_fft = assets.mel_window.shape[0] hop = n_fft // 4 pad = (n_fft - hop) // 2 xp = np.pad(x, (pad, pad), mode="reflect") spec = np.fft.rfft(_frames(xp, n_fft, hop) * assets.mel_window, n=n_fft, axis=-1) mel = np.abs(spec) @ assets.mel_fb feat = np.log(np.clip(mel, 1e-5, None)) return np.ascontiguousarray(feat.T[None], dtype=np.float32) def bwe_istft(spec, window): n_fft = window.shape[0] hop = n_fft // 4 pad = (n_fft - hop) // 2 frames = np.fft.irfft(spec, n=n_fft, axis=-1) * window total = (len(frames) - 1) * hop + n_fft y = np.zeros(total, dtype=np.float32) env = np.zeros(total, dtype=np.float32) for i, frame in enumerate(frames): s = i * hop y[s:s + n_fft] += frame env[s:s + n_fft] += window * window y = y[pad:total - pad] / env[pad:total - pad] return np.ascontiguousarray(y, dtype=np.float32) def lr_mask(n_bins, cutoff, transition_bins=1024, sample_rate=48000): cutoff_bin = int((cutoff / (sample_rate / 2)) * n_bins) mask = np.ones(n_bins, dtype=np.float32) half = transition_bins // 2 start = max(0, cutoff_bin - half) end = min(n_bins, cutoff_bin + half) x = np.linspace(-1.0, 1.0, transition_bins, dtype=np.float64) t = (x + 1.0) / 2.0 fade = (3.0 * t * t - 2.0 * t * t * t).astype(np.float32) mask[:start] = 0.0 mask[start:end] = fade[: end - start] mask[end:] = 1.0 return mask def lr_merge(generated, original, cutoff, transition_bins=1024, sample_rate=48000, _cache={}): """Linkwitz-Riley style low/high band merge. The exact full-length FFT is very slow on the board when the sample count has a large prime factor (e.g. 851456 = 2**9 * 1663). Zero-pad to the next power of two, scale the transition width so the crossover stays in Hz, and trim back. This is ~60 dB SNR vs. the exact merge and ~100x faster. """ n = len(generated) m = 1 << (n + 16384 - 1).bit_length() tb = max(2, int(round(transition_bins * m / n)) // 2 * 2) key = (m, cutoff, tb, sample_rate) mask = _cache.get(key) if mask is None: mask = lr_mask(m // 2 + 1, cutoff, tb, sample_rate) _cache[key] = mask spec1 = np.fft.rfft(generated, n=m) spec2 = np.fft.rfft(original, n=m) merged = spec2 + (spec1 - spec2) * mask return np.ascontiguousarray(np.fft.irfft(merged, n=m)[:n], dtype=np.float32) def crossfade(chunks, hop, chunk_len, out_len): """Overlap-add chunks with a Hann crossfade; edges are not attenuated.""" total = chunk_len + hop * (len(chunks) - 1) out = np.zeros(total, dtype=np.float32) wsum = np.zeros(total, dtype=np.float32) win = np.hanning(chunk_len).astype(np.float32) for i, c in enumerate(chunks): s = i * hop w = win.copy() if i == 0: w[:hop] = 1.0 if i == len(chunks) - 1: w[-hop:] = 1.0 out[s:s + chunk_len] += c * w wsum[s:s + chunk_len] += w return out[:out_len] / wsum[:out_len]