"""Streaming host half of the iSTFT codec decoder: overlap-add iSTFT + FIR mixer, then the 44160 -> 44100 Hz resampler. Line-for-line NumPy ports of the web server's `IstftHead` (servers/web/src/istft.ts) and `Resampler` (servers/web/src/resample.ts), which port zac.cpp's head_push / ola_finalize / mix_emit / head_flush / resample. Both return only final samples: after F head frames pushed, the iSTFT stream has emitted everything the whole `torch.istft(center=True)` + mixer decode would produce up to the mixer's right-hand taps; `flush()` emits the rest. Arithmetic is float64; outputs are float32. """ import numpy as np from numpy.lib.stride_tricks import sliding_window_view _EMPTY = np.zeros(0, np.float32) class IstftHeadStream: """Two-branch centered iSTFT (shared hop) followed by the linear FIR mixer, fed head frames incrementally.""" def __init__(self, n_ffts, hop, windows, mixer, max_log_mag): self.n_ffts = [int(n) for n in n_ffts] self.hop = int(hop) self.windows = [np.asarray(w, np.float64) for w in windows] self.mixer = np.asarray(mixer, np.float64) if self.mixer.ndim != 2 or self.mixer.shape[0] != 2 or self.mixer.shape[1] % 2 != 1: raise ValueError("istft: mixer kernel must be odd") for n, w in zip(self.n_ffts, self.windows): if len(w) != n or n % self.hop: raise ValueError(f"istft: window of {len(w)} taps for n_fft {n}, hop {self.hop}") self.max_log_mag = float(max_log_mag) self.K = self.mixer.shape[1] self.h = (self.K - 1) // 2 self.ola = [np.zeros(0), np.zeros(0)] self.env = [np.zeros(0), np.zeros(0)] self.ola_base = [0, 0] self.mix_in = [np.zeros(self.h), np.zeros(self.h)] # index 0 sits at absolute mix_base self.mix_base = self.mix_next = 0 self.frames = 0 self.done = False def push(self, y0, y1): """y_i [cols, n_i + 2]: log-magnitude for n_i//2+1 bins, then phase. Returns the 44160 Hz samples now final.""" if self.done: raise RuntimeError("istft: push after flush") cols = None for i, y in enumerate((y0, y1)): n, hop = self.n_ffts[i], self.hop bins = n // 2 + 1 y = np.asarray(y, np.float64) if y.ndim != 2 or y.shape[1] != n + 2 or (cols is not None and y.shape[0] != cols): raise ValueError(f"istft: branch {i} has shape {y.shape}, expected [cols, {n + 2}]") cols = y.shape[0] if cols == 0: return _EMPTY mag = np.exp(np.minimum(y[:, :bins], self.max_log_mag)) ph = y[:, bins:] imag = mag * np.sin(ph) imag[:, 0] = 0.0 imag[:, -1] = 0.0 win = self.windows[i] frames = np.fft.irfft(mag * np.cos(ph) + 1j * imag, n, axis=-1) * win r = n // hop span = (cols + r - 1) * hop add = np.zeros(span) env = np.zeros(span) parts = frames.reshape(cols, r, hop) w2 = np.broadcast_to((win * win).reshape(r, hop), (cols, r, hop)) for j in range(r): add[j * hop : (j + cols) * hop] += parts[:, j].ravel() env[j * hop : (j + cols) * hop] += w2[:, j].ravel() o = self.frames * hop - self.ola_base[i] need = o + span if len(self.ola[i]) < need: grow = need - len(self.ola[i]) self.ola[i] = np.concatenate([self.ola[i], np.zeros(grow)]) self.env[i] = np.concatenate([self.env[i], np.zeros(grow)]) self.ola[i][o:need] += add self.env[i][o:need] += env self.frames += cols for i in range(2): self._finalize(i, self.frames * self.hop) return self._mix_emit() def flush(self): """End of stream: every sample m < frames*hop is final; the mixer sees h zeros after the last sample.""" if self.done or self.frames == 0: return _EMPTY self.done = True L = self.frames * self.hop for i in range(2): self._finalize(i, self.ola_base[i] + len(self.ola[i])) keep = self.h + L - self.mix_base m = self.mix_in[i] m = m[:keep] if len(m) >= keep else np.concatenate([m, np.zeros(keep - len(m))]) self.mix_in[i] = np.concatenate([m, np.zeros(self.h)]) return self._mix_emit() def _finalize(self, i, pf): """Moves overlap-add positions [ola_base, pf) of branch i into the mixer input, dropping the centre padding.""" n = pf - self.ola_base[i] if n <= 0: return skip = min(n, max(0, self.n_ffts[i] // 2 - self.ola_base[i])) self.mix_in[i] = np.concatenate([self.mix_in[i], self.ola[i][skip:n] / self.env[i][skip:n]]) self.ola[i] = self.ola[i][n:] self.env[i] = self.env[i][n:] self.ola_base[i] = pf def _mix_emit(self): """y[m] = sum_c sum_k mixer[c][k] * mix_in_c[m + k] for every m whose taps are available.""" K = self.K m1 = min(len(m) for m in self.mix_in) + self.mix_base - K + 1 if m1 <= self.mix_next: return _EMPTY take = m1 - self.mix_base + K - 1 y = sum(np.correlate(self.mix_in[c][:take], self.mixer[c], "valid") for c in range(2)) drop = m1 - self.mix_base self.mix_in = [m[drop:] for m in self.mix_in] self.mix_base = self.mix_next = m1 return y.astype(np.float32) class ResamplerStream: """Streaming torchaudio sinc resampler: out[n*new + i] = sum_k kernel[i][k] * x[n*orig + k - width].""" def __init__(self, kernel, orig, new, width): self.orig, self.new, self.width = int(orig), int(new), int(width) self.kl = 2 * self.width + self.orig self.kernel = np.asarray(kernel, np.float64) if self.kernel.shape != (self.new, self.kl): raise ValueError(f"resampler: kernel shape {self.kernel.shape}, expected {(self.new, self.kl)}") self.buf = np.zeros(0) self.base = 0 # absolute index of buf[0] self.n_in = 0 # samples pushed self.next_block = 0 def push(self, x): x = np.asarray(x, np.float64).ravel() self.buf = np.concatenate([self.buf, x]) self.n_in += len(x) return self._run(False) def flush(self): return self._run(True) def _run(self, flush): O, N, W, KL = self.orig, self.new, self.width, self.kl if flush: total = -(-N * self.n_in // O) end = -(-total // N) else: reach = self.n_in + W - KL end = reach // O + 1 if reach >= 0 else 0 out = _EMPTY if end > self.next_block: lo = self.next_block * O - W hi = (end - 1) * O - W + KL seg = np.zeros(hi - lo) a, b = max(lo, 0), min(hi, self.n_in) if b > a: seg[a - lo : b - lo] = self.buf[a - self.base : b - self.base] win = sliding_window_view(seg, KL)[::O][: end - self.next_block] y = (win @ self.kernel.T).ravel() if flush: y = y[: total - self.next_block * N] out = y.astype(np.float32) self.next_block = end keep_from = min(self.n_in, max(0, self.next_block * O - W)) if keep_from > self.base: self.buf = self.buf[keep_from - self.base :] self.base = keep_from return out