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))