File size: 5,581 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 | #!/usr/bin/env python3
"""Unknown-duration waveform/byte streaming through files or raw stdin/stdout."""
import argparse
from contextlib import ExitStack
import hashlib
import json
import os
from pathlib import Path
import sys
import time
import numpy as np
import soundfile as sf
import torch
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT))
from tinyasr.codec_factory_v3 import load_codec_checkpoint
from tinyasr.codec_live_v1 import LiveDecoder,LiveEncoder
from tinyasr.codec_phase_streaming_v1 import PhaseLiveDecoder
from tinyasr.codec_phase_residual_v1 import PhaseResidualDecoder
def chunks(source,size):
read=getattr(source,'read1',source.read)
while True:
data=read(size)
if not data:break
yield data
def raw_audio(source,format,block_samples):
width=2 if format=='s16le' else 4
pending=b''
for data in chunks(source,block_samples*width):
data=pending+data;complete=len(data)//width*width;pending=data[complete:]
if complete:
x=np.frombuffer(data[:complete],dtype='<i2' if width==2 else '<f4').astype('float32')
if width==2:x/=32768
yield torch.from_numpy(x)
if pending:raise ValueError('Raw input ended inside a PCM sample')
def main():
ap=argparse.ArgumentParser()
ap.add_argument('operation',choices=['encode','decode'])
ap.add_argument('--checkpoint',type=Path,required=True)
ap.add_argument('--input',required=True)
ap.add_argument('--output',required=True)
ap.add_argument('--input-format',choices=['audio','s16le','f32le'],default='audio')
ap.add_argument('--output-format',choices=['wav','f32le'],default='wav')
ap.add_argument('--block-samples',type=int,default=320)
ap.add_argument('--threads',type=int,default=1)
args=ap.parse_args()
if not 1 <= args.threads <= 4:raise ValueError('Threads must be between one and four')
if args.block_samples<1:raise ValueError('Block size must be positive')
if args.input=='-' and args.operation=='encode' and args.input_format=='audio':
raise ValueError('Raw live stdin requires --input-format s16le or f32le')
if args.output=='-' and args.operation=='decode' and args.output_format!='f32le':
raise ValueError('Live stdout requires --output-format f32le')
target=Path(args.output) if args.output!='-' else None
partial=target.with_name(target.name+'.partial') if target else None
if target and (target.exists() or partial.exists()):raise FileExistsError('Output or its partial file exists')
torch.set_num_threads(args.threads)
fingerprint=hashlib.sha256(args.checkpoint.read_bytes()).hexdigest()
codec=load_codec_checkpoint(args.checkpoint)
started=time.monotonic()
with ExitStack() as stack:
if args.operation=='encode':
sender=LiveEncoder(codec,fingerprint)
output=stack.enter_context(partial.open('xb')) if target else sys.stdout.buffer
output.write(sender.start());output.flush()
if args.input_format=='audio':
source=stack.enter_context(sf.SoundFile(args.input))
if source.samplerate!=16000 or source.channels!=1:raise ValueError('Expected 16 kHz mono input')
audio=(torch.from_numpy(x) for x in source.blocks(blocksize=args.block_samples,dtype='float32'))
else:
source=stack.enter_context(Path(args.input).open('rb')) if args.input!='-' else sys.stdin.buffer
audio=raw_audio(source,args.input_format,args.block_samples)
for x in audio:
output.write(sender.push(x));output.flush()
output.write(sender.finish());output.flush()
samples=sender.audio.samples;state_bytes=sender.float_state_bytes
result=dict(samples=samples,frames=sender.audio.frames,payload_bits=sender.audio.frames*46,
framing_bits=sender.audio.frames*2,wire_bytes=48+6*sender.audio.frames+8)
else:
receiver=(PhaseLiveDecoder if isinstance(codec.decoder,PhaseResidualDecoder) else LiveDecoder)(codec,fingerprint)
source=stack.enter_context(Path(args.input).open('rb')) if args.input!='-' else sys.stdin.buffer
if args.output_format=='wav':
output=stack.enter_context(sf.SoundFile(partial,mode='w',samplerate=16000,channels=1,format='WAV',subtype='FLOAT'))
write=lambda x:output.write(x.numpy())
else:
output=stack.enter_context(partial.open('xb')) if target else sys.stdout.buffer
def write(x):
output.write(x.numpy().astype('<f4',copy=False).tobytes());output.flush()
for data in chunks(source,4096):
write(receiver.push(data))
if args.input=='-' and receiver.transport.end is not None:break
tail,end=receiver.finish();write(tail)
samples=end.samples;state_bytes=receiver.float_state_bytes
result=end.accounting()
# File outputs become visible at the requested path only after complete
# finalization/CRC checks. Raw stdout has already delivered provisional data.
if target:
os.link(partial,target)
partial.unlink()
elapsed=time.monotonic()-started
result.update(operation=args.operation,elapsed_seconds=elapsed,rtf=elapsed/(samples/16000),
persistent_float_state_bytes=state_bytes,model_sha256=fingerprint,quality_qualified=False)
print(json.dumps(result),file=sys.stderr)
if __name__=='__main__':main()
|