jyoti / example.py
voidash's picture
jyoti 2.3 kbps speech codec: development candidate C (p=0.50) + variant B (p=0.25), inference bundle
b1882a9 verified
Raw History Blame Contribute Delete
2.24 kB
#!/usr/bin/env python3
"""Minimal jyoti usage: encode a 16 kHz mono file to a 2.3 kbps stream and decode it back, in process.
python example.py input_16k_mono.wav output.wav [--checkpoint codec.pt]
The stream is bound to the checkpoint's SHA-256; decoding with a different checkpoint raises an error.
"""
import argparse
import hashlib
from pathlib import Path
import sys
import soundfile as sf
import torch
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
from tinyasr.codec_factory_v3 import load_codec_checkpoint # noqa: E402
from tinyasr.codec_live_v1 import LiveEncoder # noqa: E402
from tinyasr.codec_phase_streaming_v1 import PhaseLiveDecoder # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument('input', type=Path)
ap.add_argument('output', type=Path)
ap.add_argument('--checkpoint', type=Path, default=ROOT / 'codec.pt')
args = ap.parse_args()
if args.output.exists():
raise FileExistsError(args.output)
audio, rate = sf.read(args.input, dtype='float32', always_2d=False)
if rate != 16000 or audio.ndim != 1:
raise ValueError(f'Expected 16 kHz mono input, got {rate} Hz with shape {audio.shape}')
torch.set_num_threads(1)
model_sha256 = hashlib.sha256(args.checkpoint.read_bytes()).hexdigest()
codec = load_codec_checkpoint(args.checkpoint)
# Encode in 20 ms (320-sample) blocks, as a live sender would.
encoder = LiveEncoder(codec, model_sha256)
stream = bytearray(encoder.start())
x = torch.from_numpy(audio)
for start in range(0, x.numel(), 320):
stream += encoder.push(x[start:start + 320])
stream += encoder.finish()
# Decode the byte stream; finish() checks the footer and CRC.
decoder = PhaseLiveDecoder(codec, model_sha256)
parts = [decoder.push(bytes(stream))]
tail, end = decoder.finish()
parts.append(tail)
decoded = torch.cat(parts).numpy()
sf.write(args.output, decoded, 16000, subtype='FLOAT')
seconds = audio.size / 16000
print(f'{seconds:.2f} s audio -> {len(stream)} bytes ({8 * len(stream) / seconds:.0f} bit/s incl. framing); '
f'decoded {decoded.size} samples -> {args.output}')
if __name__ == '__main__':
main()