| """Bounded inference for the frozen TinyCall-inspired codec. |
| |
| Finite context is recomputed on each 20 ms hop. This favors a small, auditable |
| state over speed; model parameters and caller-owned input/output are excluded |
| from the persistent state accounting. Only batch-one CPU inference is supported. |
| """ |
| import math |
|
|
| import torch |
|
|
|
|
| class StreamingEncoder: |
| hop = 320 |
| |
| |
| history_samples = 1280 |
|
|
| def __init__(self, codec): |
| if next(codec.parameters()).device.type != 'cpu': |
| raise ValueError('Streaming v1 supports CPU inference only') |
| self.codec = codec |
| self.history = torch.empty(0) |
| self.pending = torch.empty(0) |
| self.samples = 0 |
| self.frames = 0 |
| self.finished = False |
|
|
| @property |
| def state_bytes(self): |
| return sum(t.untyped_storage().nbytes() for t in [self.history, self.pending]) |
|
|
| def _frame(self, hop): |
| context = torch.cat((self.history, hop)) |
| features = self.codec.encoder(context[None])[..., -1:] |
| _, codes = self.codec.quantizer(features) |
| self.history = context[-self.history_samples:].clone() |
| self.frames += 1 |
| return codes |
|
|
| @torch.inference_mode() |
| def push(self, audio): |
| if self.finished: |
| raise ValueError('Encoder stream already finished') |
| if audio.ndim != 1 or audio.device.type != 'cpu' or audio.dtype != torch.float32: |
| raise ValueError('Expected one-dimensional CPU float32 waveform') |
| if not torch.isfinite(audio).all(): |
| raise ValueError('Nonfinite waveform') |
| self.samples += audio.numel() |
| |
| |
| data = torch.cat((self.pending, audio)) |
| complete = data.numel() // self.hop * self.hop |
| frames = [self._frame(data[i:i + self.hop]) for i in range(0, complete, self.hop)] |
| self.pending = data[complete:].clone() |
| return torch.cat(frames, -1) if frames else torch.empty((1, 4, 6, 0), dtype=torch.long) |
|
|
| @torch.inference_mode() |
| def finish(self): |
| if self.finished or not self.samples: |
| raise ValueError('Empty or already finished encoder stream') |
| self.finished = True |
| if self.pending.numel(): |
| |
| |
| |
| result = self._frame(self.pending) |
| else: |
| result = torch.empty((1, 4, 6, 0), dtype=torch.long) |
| self.pending = torch.empty(0) |
| return result |
|
|
|
|
| class StreamingDecoder: |
| hop = 320 |
| history_frames = 30 |
|
|
| def __init__(self, codec): |
| if next(codec.parameters()).device.type != 'cpu': |
| raise ValueError('Streaming v1 supports CPU inference only') |
| self.codec = codec |
| self.history = torch.empty((1, 64, 0)) |
| self.previous_right = torch.empty(0) |
| window = codec.decoder.window |
| self.denominator = window[:self.hop].square() + window[self.hop:].square() |
| self.frames = 0 |
| self.emitted = 0 |
| self.finished = False |
|
|
| @property |
| def state_bytes(self): |
| return sum(t.untyped_storage().nbytes() for t in [self.history, self.previous_right, self.denominator]) |
|
|
| def _frame(self, features): |
| context = torch.cat((self.history, features), -1) |
| decoder = self.codec.decoder |
| values = decoder.spectrum(decoder.blocks(decoder.input(context)))[..., -1] |
| log_magnitude, phase = values.chunk(2, dim=1) |
| magnitude = log_magnitude.clamp(-10, 5).exp() |
| phase = math.pi * phase.tanh() |
| spectrum = torch.complex(magnitude * phase.cos(), magnitude * phase.sin()) |
| frame = torch.fft.irfft(spectrum, n=decoder.n_fft)[0] * decoder.window |
| self.history = context[..., -self.history_frames:].clone() |
| output = ((self.previous_right + frame[:self.hop]) / self.denominator |
| if self.previous_right.numel() else torch.empty(0)) |
| self.previous_right = frame[self.hop:].clone() |
| self.emitted += output.numel() |
| return output |
|
|
| @torch.inference_mode() |
| def push(self, codes): |
| if self.finished: |
| raise ValueError('Decoder stream already finished') |
| if codes.ndim != 4 or codes.shape[:3] != (1, 4, 6) or codes.dtype != torch.long or codes.device.type != 'cpu': |
| raise ValueError('Expected CPU int64 codes [1,4,6,frames]') |
| frames = [] |
| for i in range(codes.shape[-1]): |
| features = self.codec.quantizer.decode(codes[..., i:i + 1]) |
| frames.append(self._frame(features)) |
| self.frames += 1 |
| return torch.cat(frames) if frames else torch.empty(0) |
|
|
| @torch.inference_mode() |
| def finish(self, samples): |
| if self.finished or not self.frames: |
| raise ValueError('Empty or already finished decoder stream') |
| if not isinstance(samples, int) or not (self.frames - 1) * self.hop < samples <= self.frames * self.hop: |
| raise ValueError('Sample count inconsistent with received code frames') |
| self.finished = True |
| |
| |
| remaining = samples - self.emitted |
| output = self._frame(torch.zeros((1, 64, 1)))[:remaining] |
| self.emitted = samples |
| return output |
|
|
