File size: 7,167 Bytes
408e7ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """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]
|