| """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] |
|
|
