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