File size: 5,800 Bytes
b1882a9 | 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 | """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
# Encoder receptive field is 1438 samples. Four complete past hops are
# sufficient, and preserving hop alignment preserves every strided phase.
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()
# The temporary concatenation is input-sized, but no view of it survives
# this call; persistent storage is independent of the input duration.
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():
# Let each strided layer pad its own final partial activation, as
# in whole-waveform inference. Padding the raw waveform instead
# would generate biased hidden activations beyond its true end.
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
# Match the offline decoder's single terminal zero latent. It is not a
# quantized code and does not carry an undeclared transmitted feature.
remaining = samples - self.emitted
output = self._frame(torch.zeros((1, 64, 1)))[:remaining]
self.emitted = samples
return output
|