File size: 7,766 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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """Self-delimiting live transport: 46 payload bits in each 48-bit record.
Each six-byte frame has a leading one bit and a zero padding bit. A footer
starts with a zero high bit and holds a 31-bit sample count plus CRC32. The
48-byte model header and eight-byte footer add 56 fixed bytes. Payload is
2300 bit/s; framing adds 100 bit/s. CRC is verified at EOS, so earlier indices
are provisional. No seek, full-file pass or duration-sized buffer is needed.
"""
from dataclasses import dataclass
import struct
import zlib
import numpy as np
from tinyasr.codec_bitstream import BITS_PER_FRAME, DEFAULT_MAX_FRAMES, HOP, LEVELS, SAMPLE_RATE, WIDTHS, model_bytes
MAGIC=b'TCLIVE1\0'
HEADER=struct.Struct('<8sBBHI32s')
FOOTER_PREFIX=struct.Struct('>I')
CRC=struct.Struct('<I')
FOOTER_BYTES=FOOTER_PREFIX.size+CRC.size
RECORD_BYTES=6
def frame_limit(value):
if isinstance(value,bool) or not isinstance(value,int) or not 1<=value<=DEFAULT_MAX_FRAMES:
raise ValueError('Invalid live transport frame limit')
return value
def empty_codes():return np.empty((4,6,0),dtype=np.int64)
@dataclass(frozen=True)
class LiveEnd:
samples:int
frames:int
file_bytes:int
model_sha256:str
def accounting(self):
seconds=self.samples/SAMPLE_RATE
return dict(samples=self.samples,frames=self.frames,seconds=seconds,
payload_bits=self.frames*BITS_PER_FRAME,framing_bits=self.frames*2,
file_bytes=self.file_bytes,fixed_container_bytes=HEADER.size+FOOTER_BYTES,
file_bits_per_second=self.file_bytes*8/seconds,
model_sha256=self.model_sha256,format='tinycall-live-v1')
class LiveBitWriter:
def __init__(self,model_sha256,max_frames=DEFAULT_MAX_FRAMES):
self.max_frames=frame_limit(max_frames)
self.header=HEADER.pack(MAGIC,1,1,BITS_PER_FRAME,SAMPLE_RATE,model_bytes(model_sha256))
self.checksum=zlib.crc32(self.header);self.frames=0;self.finished=False
@property
def retained_bit_reservoir_bytes(self):return 0
def push(self,codes):
if self.finished:raise ValueError('Live writer already finished')
codes=np.asarray(codes)
if codes.ndim!=3 or codes.shape[:2]!=(4,6) or codes.dtype.kind not in 'iu':
raise ValueError('Expected integer scalar codes [4,6,frames]')
if codes.shape[-1]+self.frames>self.max_frames:raise ValueError('Live frame limit exceeded')
if np.any(codes<0) or np.any(codes>=LEVELS[:,:,None]):raise ValueError('Scalar index outside its grid')
output=bytearray()
for index in range(codes.shape[-1]):
word=0
for value,width in zip(codes[:,:,index].reshape(-1),WIDTHS):word=(word<<int(width))|int(value)
output.extend(((1<<47)|(word<<1)).to_bytes(RECORD_BYTES,'big'));self.frames+=1
self.checksum=zlib.crc32(output,self.checksum)
return bytes(output)
def finish(self,samples):
if self.finished:raise ValueError('Live writer already finished')
if isinstance(samples,bool) or not isinstance(samples,(int,np.integer)) or not 0<samples<2**31:
raise ValueError('Expected positive 31-bit sample count')
if not self.frames or self.frames!=(int(samples)+HOP-1)//HOP:
raise ValueError('Sample count does not match received frames')
prefix=FOOTER_PREFIX.pack(int(samples));self.finished=True
return prefix+CRC.pack(zlib.crc32(prefix,self.checksum))
class LiveBitReader:
def __init__(self,expected_model_sha256,max_frames=DEFAULT_MAX_FRAMES):
self.expected_model=model_bytes(expected_model_sha256);self.max_frames=frame_limit(max_frames)
self.header_buffer=bytearray();self.record_buffer=bytearray();self.record_size=0
self.header_complete=False;self.finished=False;self.failed=False;self.end=None
self.frames=0;self.received_bytes=0;self.checksum=0
@property
def retained_transport_buffer_bytes(self):return len(self.header_buffer)+len(self.record_buffer)
def _record(self,record):
if self.record_size==RECORD_BYTES:
value=int.from_bytes(record,'big')
if value&1:raise ValueError('Nonzero live frame padding bit')
if self.frames>=self.max_frames:raise ValueError('Live frame limit exceeded')
word=(value>>1)&((1<<BITS_PER_FRAME)-1);shift=BITS_PER_FRAME;values=[]
for width in WIDTHS:
shift-=int(width);values.append((word>>shift)&((1<<int(width))-1))
self.frames+=1;self.checksum=zlib.crc32(record,self.checksum)
return np.array(values,dtype=np.int64).reshape(4,6)
prefix=record[:FOOTER_PREFIX.size];samples=FOOTER_PREFIX.unpack(prefix)[0]
if not 0<samples<2**31 or self.frames!=(samples+HOP-1)//HOP:
raise ValueError('Live frame/sample counts disagree')
if zlib.crc32(prefix,self.checksum)!=CRC.unpack(record[FOOTER_PREFIX.size:])[0]:
raise ValueError('Live checksum mismatch')
self.end=LiveEnd(samples,self.frames,self.received_bytes,self.expected_model.hex())
return None
def push(self,data):
if self.finished or self.failed:raise ValueError('Live reader is closed or failed')
if not isinstance(data,(bytes,bytearray,memoryview)):raise ValueError('Expected incoming bytes')
try:data=memoryview(data).cast('B')
except (TypeError,ValueError) as error:raise ValueError('Expected contiguous bytes') from error
limit=HEADER.size+FOOTER_BYTES+self.max_frames*RECORD_BYTES
if self.received_bytes+len(data)>limit:
self.failed=True;raise ValueError('Live byte limit exceeded')
self.received_bytes+=len(data)
try:
if self.end is not None:
if len(data):raise ValueError('Unexpected trailing live data')
return empty_codes()
if not self.header_complete:
need=HEADER.size-len(self.header_buffer)
used=min(need,len(data));self.header_buffer.extend(data[:used]);data=data[used:]
if len(self.header_buffer)<HEADER.size:return empty_codes()
magic,version,flags,bits,rate,fingerprint=HEADER.unpack(self.header_buffer)
if (magic,version,flags,bits,rate)!=(MAGIC,1,1,BITS_PER_FRAME,SAMPLE_RATE):raise ValueError('Unsupported live header')
if fingerprint!=self.expected_model:raise ValueError('Live stream requires a different model checkpoint')
self.checksum=zlib.crc32(self.header_buffer);self.header_buffer=bytearray();self.header_complete=True
frames=[];position=0
while position<len(data):
if self.end is not None:raise ValueError('Unexpected trailing live data')
if not self.record_buffer:self.record_size=RECORD_BYTES if data[position]&128 else FOOTER_BYTES
used=min(self.record_size-len(self.record_buffer),len(data)-position)
self.record_buffer.extend(data[position:position+used]);position+=used
if len(self.record_buffer)==self.record_size:
value=self._record(bytes(self.record_buffer));self.record_buffer=bytearray()
if value is not None:frames.append(value)
return np.stack(frames,-1) if frames else empty_codes()
except (ValueError,struct.error):
self.failed=True;raise
def finish(self):
if self.finished or self.failed:raise ValueError('Live reader is closed or failed')
if not self.header_complete or self.record_buffer or self.end is None:
self.failed=True;raise ValueError('Truncated live stream')
self.finished=True;return self.end
|