File size: 6,052 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 | """Versioned standalone container for the independent codec's scalar indices.
Frame-major, stage-major, coordinate-major fixed-width symbols occupy exactly
46 bits/frame, packed MSB first. A fixed header stores model SHA-256, waveform
length and frame count. CRC32 covers the header fields and packed payload.
"""
from __future__ import annotations
from dataclasses import dataclass
import struct
import zlib
import numpy as np
MAGIC = b"TCSPCH1\0"
VERSION = 1
SAMPLE_RATE = 16000
HOP = 320
BITS_PER_FRAME = 46
LEVELS = np.array([[8, 8, 8, 4, 4, 4]] * 2 + [[4, 4, 2, 2, 2, 2]] * 2)
WIDTHS = np.log2(LEVELS).astype(np.int64).reshape(-1)
HEADER_WITHOUT_CRC = struct.Struct("<8sBBHIQQ32s")
CRC = struct.Struct("<I")
HEADER_BYTES = HEADER_WITHOUT_CRC.size + CRC.size
DEFAULT_MAX_FRAMES = 1000000
@dataclass(frozen=True)
class PackedSpeech:
codes: np.ndarray
samples: int
model_sha256: str
file_bytes: int
sample_rate: int = SAMPLE_RATE
@property
def payload_bits(self):
return self.codes.shape[-1] * BITS_PER_FRAME
def accounting(self):
seconds = self.samples / self.sample_rate
payload_bytes = (self.payload_bits + 7) // 8
return {"samples": self.samples, "frames": self.codes.shape[-1],
"seconds": seconds, "sample_rate": self.sample_rate,
"model_sha256": self.model_sha256, "payload_bits": self.payload_bits,
"payload_bytes": payload_bytes,
"padding_bits": payload_bytes * 8 - self.payload_bits,
"header_bytes": HEADER_BYTES, "file_bytes": self.file_bytes,
"payload_bits_per_second": self.payload_bits / seconds,
"file_bits_per_second": self.file_bytes * 8 / seconds,
"format": "tinycall-independent-packed-v1"}
def model_bytes(fingerprint):
if not isinstance(fingerprint, str) or len(fingerprint) != 64:
raise ValueError("Model identity must be a SHA-256 hex string")
try:
value = bytes.fromhex(fingerprint)
except ValueError as error:
raise ValueError("Model identity must be a SHA-256 hex string") from error
if len(value) != 32:
raise ValueError("Model SHA-256 must contain exactly 32 bytes")
return value
def pack_codes(codes, samples, model_sha256):
values = np.asarray(codes)
if values.ndim != 3 or values.shape[:2] != (4, 6) or values.shape[-1] < 1:
raise ValueError("Expected nonempty [4,6,frames] scalar indices")
if values.dtype.kind not in "iu":
raise ValueError("Scalar indices must have integer dtype")
if np.any(values < 0) or np.any(values >= LEVELS[:, :, None]):
raise ValueError("Scalar index outside its grid")
if isinstance(samples, bool) or not isinstance(samples, (int, np.integer)) or samples <= 0:
raise ValueError("Sample count must be a positive integer")
frames = values.shape[-1]
if frames > DEFAULT_MAX_FRAMES or frames != (int(samples) + HOP - 1) // HOP:
raise ValueError("Frame count does not match waveform duration or exceeds format limit")
words = np.zeros(frames, dtype=np.uint64)
shift = BITS_PER_FRAME
for coordinate, width in zip(values.reshape(24, frames), WIDTHS):
shift -= int(width)
words |= coordinate.astype(np.uint64) << np.uint64(shift)
shifts = np.arange(BITS_PER_FRAME - 1, -1, -1, dtype=np.uint64)
bits = ((words[:, None] >> shifts) & 1).astype(np.uint8).reshape(-1)
payload = np.packbits(bits, bitorder="big").tobytes()
header = HEADER_WITHOUT_CRC.pack(MAGIC, VERSION, 0, BITS_PER_FRAME,
SAMPLE_RATE, int(samples), frames,
model_bytes(model_sha256))
checksum = zlib.crc32(payload, zlib.crc32(header))
return header + CRC.pack(checksum) + payload
def unpack_codes(blob, expected_model_sha256=None, max_frames=DEFAULT_MAX_FRAMES):
if len(blob) < HEADER_BYTES:
raise ValueError("Truncated codec header")
header = blob[:HEADER_WITHOUT_CRC.size]
magic, version, flags, width, rate, samples, frames, fingerprint = HEADER_WITHOUT_CRC.unpack(header)
if magic != MAGIC or version != VERSION or flags != 0:
raise ValueError("Unsupported codec format, version or flags")
if width != BITS_PER_FRAME or rate != SAMPLE_RATE:
raise ValueError("Unsupported scalar layout or sample rate")
if samples == 0 or frames == 0 or frames != (samples + HOP - 1) // HOP:
raise ValueError("Invalid waveform length or frame count")
if frames > max_frames or frames > DEFAULT_MAX_FRAMES:
raise ValueError("Declared stream exceeds the frame limit")
size = (frames * BITS_PER_FRAME + 7) // 8
if len(blob) != HEADER_BYTES + size:
raise ValueError("Truncated payload or unexpected trailing bytes")
payload = blob[HEADER_BYTES:]
expected_crc, = CRC.unpack(blob[HEADER_WITHOUT_CRC.size:HEADER_BYTES])
if zlib.crc32(payload, zlib.crc32(header)) != expected_crc:
raise ValueError("Codec stream checksum mismatch")
if expected_model_sha256 is not None and fingerprint != model_bytes(expected_model_sha256):
raise ValueError("Codec stream requires a different model checkpoint")
padding = size * 8 - frames * BITS_PER_FRAME
if padding and payload[-1] & ((1 << padding) - 1):
raise ValueError("Nonzero unused padding bits")
bits = np.unpackbits(np.frombuffer(payload, dtype=np.uint8), bitorder="big")
bits = bits[:frames * BITS_PER_FRAME].reshape(frames, BITS_PER_FRAME)
shifts = np.arange(BITS_PER_FRAME - 1, -1, -1, dtype=np.uint64)
words = np.sum(bits.astype(np.uint64) << shifts, axis=1, dtype=np.uint64)
values = np.empty((24, frames), dtype=np.int64)
shift = BITS_PER_FRAME
for index, width in enumerate(WIDTHS):
shift -= int(width)
values[index] = (words >> np.uint64(shift)) & np.uint64((1 << int(width)) - 1)
return PackedSpeech(values.reshape(4, 6, frames), int(samples), fingerprint.hex(), len(blob))
|