Spaces:
Running on Zero
Running on Zero
| """Streaming iSTFT head + resampler on the codec's device: the torch twins of `zonos3.codec.stream_host`. | |
| Same algorithm and float64 arithmetic as the NumPy ports (overlap-add iSTFT of the two head branches, the FIR mixer, | |
| the 44160 -> 44100 Hz sinc resampler), on the tensors the trunk leaves on its device, so a chunk needs no host round | |
| trip before its final samples. Positions and lengths are host integers (they follow from the frame counts), so the | |
| only device sync is the caller's copy of the output. Outputs are float32 like the NumPy ports' (the head's output is | |
| rounded to float32 before the resampler, as there); the overlap-add sums in another order, so the samples match the | |
| NumPy path to float32 rounding, not bit for bit. | |
| """ | |
| import torch | |
| from torch.nn import functional as F | |
| 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, device): | |
| self.n_ffts = [int(n) for n in n_ffts] | |
| self.hop = int(hop) | |
| self.device = torch.device(device) | |
| f64 = dict(dtype=torch.float64, device=self.device) | |
| self.windows = [torch.as_tensor(w).to(**f64) for w in windows] | |
| self.mixer = torch.as_tensor(mixer).to(**f64) | |
| 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._f64 = f64 | |
| self.ola = [torch.zeros(0, **f64), torch.zeros(0, **f64)] | |
| self.env = [torch.zeros(0, **f64), torch.zeros(0, **f64)] | |
| self.ola_base = [0, 0] | |
| self.mix_in = [torch.zeros(self.h, **f64), torch.zeros(self.h, **f64)] # index 0 sits at absolute mix_base | |
| self.mix_base = self.mix_next = 0 | |
| self.frames = 0 | |
| self.done = False | |
| self._env: dict[tuple[int, int], torch.Tensor] = {} # (branch, cols) -> overlap-added squared window | |
| def copy(self) -> "IstftHeadStream": | |
| """An independent twin (the overlap-add buffers are updated in place; everything else is replaced).""" | |
| twin = object.__new__(IstftHeadStream) | |
| twin.__dict__.update(self.__dict__) | |
| twin.ola, twin.env = [t.clone() for t in self.ola], [t.clone() for t in self.env] | |
| twin.ola_base, twin.mix_in = list(self.ola_base), list(self.mix_in) | |
| return twin | |
| def _overlap_add(self, frames: torch.Tensor, n: int) -> torch.Tensor: | |
| """Frames [cols, n] placed every hop samples and summed -> [(cols - 1) * hop + n].""" | |
| cols = frames.shape[0] | |
| span = (cols - 1) * self.hop + n | |
| return F.fold( | |
| frames.T.unsqueeze(0), output_size=(1, span), kernel_size=(1, n), stride=(1, self.hop) | |
| ).reshape(span) | |
| def push(self, y0: torch.Tensor, y1: torch.Tensor) -> torch.Tensor: | |
| """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 = self.n_ffts[i] | |
| bins = n // 2 + 1 | |
| y = y.to(**self._f64) | |
| 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 {tuple(y.shape)}, expected [cols, {n + 2}]") | |
| cols = y.shape[0] | |
| if cols == 0: | |
| return torch.zeros(0, dtype=torch.float32, device=self.device) | |
| mag = torch.exp(torch.clamp(y[:, :bins], max=self.max_log_mag)) | |
| ph = y[:, bins:] | |
| imag = mag * torch.sin(ph) | |
| imag[:, 0] = 0.0 | |
| imag[:, -1] = 0.0 | |
| win = self.windows[i] | |
| frames = torch.fft.irfft(torch.complex(mag * torch.cos(ph), imag), n, dim=-1) * win | |
| add = self._overlap_add(frames, n) | |
| env = self._env.get((i, cols)) | |
| if env is None: | |
| env = self._env[(i, cols)] = self._overlap_add((win * win).expand(cols, n), n) | |
| span = add.shape[0] | |
| o = self.frames * self.hop - self.ola_base[i] | |
| need = o + span | |
| if self.ola[i].shape[0] < need: | |
| grow = torch.zeros(need - self.ola[i].shape[0], **self._f64) | |
| self.ola[i] = torch.cat([self.ola[i], grow]) | |
| self.env[i] = torch.cat([self.env[i], 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) -> torch.Tensor: | |
| """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 torch.zeros(0, dtype=torch.float32, device=self.device) | |
| self.done = True | |
| L = self.frames * self.hop | |
| for i in range(2): | |
| self._finalize(i, self.ola_base[i] + self.ola[i].shape[0]) | |
| keep = self.h + L - self.mix_base | |
| m = self.mix_in[i] | |
| m = m[:keep] if m.shape[0] >= keep else torch.cat([m, torch.zeros(keep - m.shape[0], **self._f64)]) | |
| self.mix_in[i] = torch.cat([m, torch.zeros(self.h, **self._f64)]) | |
| 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] = torch.cat([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) -> torch.Tensor: | |
| """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(m.shape[0] for m in self.mix_in) + self.mix_base - K + 1 | |
| if m1 <= self.mix_next: | |
| return torch.zeros(0, dtype=torch.float32, device=self.device) | |
| take = m1 - self.mix_base + K - 1 | |
| x = torch.stack([m[:take] for m in self.mix_in]).unsqueeze(0) # [1, 2, take] | |
| y = F.conv1d(x, self.mixer.unsqueeze(0)).reshape(-1) # cross-correlation summed over both channels | |
| 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.to(torch.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, device): | |
| self.orig, self.new, self.width = int(orig), int(new), int(width) | |
| self.kl = 2 * self.width + self.orig | |
| self.device = torch.device(device) | |
| self._f64 = dict(dtype=torch.float64, device=self.device) | |
| self.kernel = torch.as_tensor(kernel).to(**self._f64) | |
| if tuple(self.kernel.shape) != (self.new, self.kl): | |
| raise ValueError(f"resampler: kernel shape {tuple(self.kernel.shape)}, expected {(self.new, self.kl)}") | |
| self._weight = self.kernel.unsqueeze(1) # [new, 1, kl] | |
| self.buf = torch.zeros(0, **self._f64) | |
| self.base = 0 # absolute index of buf[0] | |
| self.n_in = 0 # samples pushed | |
| self.next_block = 0 | |
| def copy(self) -> "ResamplerStream": | |
| """An independent twin (the buffer is only ever replaced, never written in place).""" | |
| twin = object.__new__(ResamplerStream) | |
| twin.__dict__.update(self.__dict__) | |
| return twin | |
| def push(self, x: torch.Tensor) -> torch.Tensor: | |
| x = x.to(**self._f64).reshape(-1) | |
| self.buf = torch.cat([self.buf, x]) | |
| self.n_in += x.shape[0] | |
| return self._run(False) | |
| def flush(self) -> torch.Tensor: | |
| 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 = torch.zeros(0, dtype=torch.float32, device=self.device) | |
| if end > self.next_block: | |
| lo = self.next_block * O - W | |
| hi = (end - 1) * O - W + KL | |
| a, b = max(lo, 0), min(hi, self.n_in) | |
| if b > a: | |
| seg = F.pad(self.buf[a - self.base : b - self.base], (a - lo, hi - b)) | |
| else: | |
| seg = torch.zeros(hi - lo, **self._f64) | |
| # windows of KL samples every O samples, each dotted with the N kernel rows: [N, blocks] -> block-major | |
| y = F.conv1d(seg.view(1, 1, -1), self._weight, stride=O)[0].T.reshape(-1) | |
| if flush: | |
| y = y[: total - self.next_block * N] | |
| out = y.to(torch.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 | |