Upload 2 files

asr.py

@@ -0,0 +1,233 @@
+from typing import *
+import logging
+import time
+import logging
+import sherpa_onnx
+import os
+import asyncio
+import numpy as np
+
+logger = logging.getLogger(__file__)
+_asr_engines = {}
+
+
+class ASRResult:
+    def __init__(self, text: str, finished: bool, idx: int):
+        self.text = text
+        self.finished = finished
+        self.idx = idx
+
+    def to_dict(self):
+        return {"text": self.text, "finished": self.finished, "idx": self.idx}
+
+
+class ASRStream:
+    def __init__(self, recognizer: Union[sherpa_onnx.OnlineRecognizer | sherpa_onnx.OfflineRecognizer], sample_rate: int) -> None:
+        self.recognizer = recognizer
+        self.inbuf = asyncio.Queue()
+        self.outbuf = asyncio.Queue()
+        self.sample_rate = sample_rate
+        self.is_closed = False
+        self.online = isinstance(recognizer, sherpa_onnx.OnlineRecognizer)
+
+    async def start(self):
+        if self.online:
+            asyncio.create_task(self.run_online())
+        else:
+            asyncio.create_task(self.run_offline())
+
+    async def run_online(self):
+        stream = self.recognizer.create_stream()
+        last_result = ""
+        segment_id = 0
+        logger.info('asr: start real-time recognizer')
+        while not self.is_closed:
+            samples = await self.inbuf.get()
+            stream.accept_waveform(self.sample_rate, samples)
+            while self.recognizer.is_ready(stream):
+                self.recognizer.decode_stream(stream)
+
+            is_endpoint = self.recognizer.is_endpoint(stream)
+            result = self.recognizer.get_result(stream)
+
+            if result and (last_result != result):
+                last_result = result
+                logger.info(f' > {segment_id}:{result}')
+                self.outbuf.put_nowait(
+                    ASRResult(result, False, segment_id))
+
+            if is_endpoint:
+                if result:
+                    logger.info(f'{segment_id}: {result}')
+                    self.outbuf.put_nowait(
+                        ASRResult(result, True, segment_id))
+                    segment_id += 1
+                self.recognizer.reset(stream)
+
+    async def run_offline(self):
+        vad = _asr_engines['vad']
+        segment_id = 0
+        st = None
+        while not self.is_closed:
+            samples = await self.inbuf.get()
+            vad.accept_waveform(samples)
+            while not vad.empty():
+                if not st:
+                    st = time.time()
+                stream = self.recognizer.create_stream()
+                stream.accept_waveform(self.sample_rate, vad.front.samples)
+
+                vad.pop()
+                self.recognizer.decode_stream(stream)
+
+                result = stream.result.text.strip()
+                if result:
+                    duration = time.time() - st
+                    logger.info(f'{segment_id}:{result} ({duration:.2f}s)')
+                    self.outbuf.put_nowait(ASRResult(result, True, segment_id))
+                    segment_id += 1
+            st = None
+
+    async def close(self):
+        self.is_closed = True
+        self.outbuf.put_nowait(None)
+
+    async def write(self, pcm_bytes: bytes):
+        pcm_data = np.frombuffer(pcm_bytes, dtype=np.int16)
+        samples = pcm_data.astype(np.float32) / 32768.0
+        self.inbuf.put_nowait(samples)
+
+    async def read(self) -> ASRResult:
+        return await self.outbuf.get()
+
+
+def create_zipformer(samplerate: int, args) -> sherpa_onnx.OnlineRecognizer:
+    d = os.path.join(
+        args.models_root, 'sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20')
+    if not os.path.exists(d):
+        raise ValueError(f"asr: model not found {d}")
+
+    encoder = os.path.join(d, "encoder-epoch-99-avg-1.onnx")
+    decoder = os.path.join(d, "decoder-epoch-99-avg-1.onnx")
+    joiner = os.path.join(d, "joiner-epoch-99-avg-1.onnx")
+    tokens = os.path.join(d, "tokens.txt")
+
+    recognizer = sherpa_onnx.OnlineRecognizer.from_transducer(
+        tokens=tokens,
+        encoder=encoder,
+        decoder=decoder,
+        joiner=joiner,
+        provider=args.asr_provider,
+        num_threads=args.threads,
+        sample_rate=samplerate,
+        feature_dim=80,
+        enable_endpoint_detection=True,
+        rule1_min_trailing_silence=2.4,
+        rule2_min_trailing_silence=1.2,
+        rule3_min_utterance_length=20,  # it essentially disables this rule
+    )
+    return recognizer
+
+
+def create_sensevoice(samplerate: int, args) -> sherpa_onnx.OfflineRecognizer:
+    d = os.path.join(args.models_root,
+                     'sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17')
+
+    if not os.path.exists(d):
+        raise ValueError(f"asr: model not found {d}")
+
+    recognizer = sherpa_onnx.OfflineRecognizer.from_sense_voice(
+        model=os.path.join(d, 'model.onnx'),
+        tokens=os.path.join(d, 'tokens.txt'),
+        num_threads=args.threads,
+        sample_rate=samplerate,
+        use_itn=True,
+        debug=0,
+        language=args.asr_lang,
+    )
+    return recognizer
+
+
+def create_paraformer_trilingual(samplerate: int, args) -> sherpa_onnx.OnlineRecognizer:
+    d = os.path.join(
+        args.models_root, 'sherpa-onnx-paraformer-trilingual-zh-cantonese-en')
+    if not os.path.exists(d):
+        raise ValueError(f"asr: model not found {d}")
+
+    recognizer = sherpa_onnx.OfflineRecognizer.from_paraformer(
+        paraformer=os.path.join(d, 'model.onnx'),
+        tokens=os.path.join(d, 'tokens.txt'),
+        num_threads=args.threads,
+        sample_rate=samplerate,
+        debug=0,
+        provider=args.asr_provider,
+    )
+    return recognizer
+
+
+def create_paraformer_en(samplerate: int, args) -> sherpa_onnx.OnlineRecognizer:
+    d = os.path.join(
+        args.models_root, 'sherpa-onnx-paraformer-en')
+    if not os.path.exists(d):
+        raise ValueError(f"asr: model not found {d}")
+
+    recognizer = sherpa_onnx.OfflineRecognizer.from_paraformer(
+        paraformer=os.path.join(d, 'model.onnx'),
+        tokens=os.path.join(d, 'tokens.txt'),
+        num_threads=args.threads,
+        sample_rate=samplerate,
+        use_itn=True,
+        debug=0,
+        provider=args.asr_provider,
+    )
+    return recognizer
+
+
+def load_asr_engine(samplerate: int, args) -> sherpa_onnx.OnlineRecognizer:
+    cache_engine = _asr_engines.get(args.asr_model)
+    if cache_engine:
+        return cache_engine
+    st = time.time()
+    if args.asr_model == 'zipformer-bilingual':
+        cache_engine = create_zipformer(samplerate, args)
+    elif args.asr_model == 'sensevoice':
+        cache_engine = create_sensevoice(samplerate, args)
+        _asr_engines['vad'] = load_vad_engine(samplerate, args)
+    elif args.asr_model == 'paraformer-trilingual':
+        cache_engine = create_paraformer_trilingual(samplerate, args)
+        _asr_engines['vad'] = load_vad_engine(samplerate, args)
+    elif args.asr_model == 'paraformer-en':
+        cache_engine = create_paraformer_en(samplerate, args)
+        _asr_engines['vad'] = load_vad_engine(samplerate, args)
+    else:
+        raise ValueError(f"asr: unknown model {args.asr_model}")
+    _asr_engines[args.asr_model] = cache_engine
+    logger.info(f"asr: engine loaded in {time.time() - st:.2f}s")
+    return cache_engine
+
+
+def load_vad_engine(samplerate: int, args, min_silence_duration: float = 0.25, buffer_size_in_seconds: int = 100) -> sherpa_onnx.VoiceActivityDetector:
+    config = sherpa_onnx.VadModelConfig()
+    d = os.path.join(args.models_root, 'silero_vad')
+    if not os.path.exists(d):
+        raise ValueError(f"vad: model not found {d}")
+
+    config.silero_vad.model = os.path.join(d, 'silero_vad.onnx')
+    config.silero_vad.min_silence_duration = min_silence_duration
+    config.sample_rate = samplerate
+    config.provider = args.asr_provider
+    config.num_threads = args.threads
+
+    vad = sherpa_onnx.VoiceActivityDetector(
+        config,
+        buffer_size_in_seconds=buffer_size_in_seconds)
+    return vad
+
+
+async def start_asr_stream(samplerate: int, args) -> ASRStream:
+    """
+    Start a ASR stream
+    """
+    stream = ASRStream(load_asr_engine(samplerate, args), samplerate)
+    await stream.start()
+    return stream

tts.py

@@ -0,0 +1,216 @@
+from typing import *
+import os
+import time
+import sherpa_onnx
+import logging
+import numpy as np
+import asyncio
+import time
+import soundfile
+from scipy.signal import resample
+import io
+import re
+
+logger = logging.getLogger(__file__)
+
+splitter = re.compile(r'[,，。.!?！？;；、\n]')
+_tts_engines = {}
+
+tts_configs = {
+    'vits-zh-hf-theresa': {
+        'model': 'theresa.onnx',
+        'lexicon': 'lexicon.txt',
+        'dict_dir': 'dict',
+        'tokens': 'tokens.txt',
+        'sample_rate': 22050,
+        # 'rule_fsts': ['phone.fst', 'date.fst', 'number.fst'],
+    },
+    'vits-melo-tts-zh_en': {
+        'model': 'model.onnx',
+        'lexicon': 'lexicon.txt',
+        'dict_dir': 'dict',
+        'tokens': 'tokens.txt',
+        'sample_rate': 44100,
+        'rule_fsts': ['phone.fst', 'date.fst', 'number.fst'],
+    },
+}
+
+
+def load_tts_model(name: str, model_root: str, provider: str, num_threads: int = 1, max_num_sentences: int = 20) -> sherpa_onnx.OfflineTtsConfig:
+    cfg = tts_configs[name]
+    fsts = []
+    model_dir = os.path.join(model_root, name)
+    for f in cfg.get('rule_fsts', ''):
+        fsts.append(os.path.join(model_dir, f))
+    tts_rule_fsts = ','.join(fsts) if fsts else ''
+
+    model_config = sherpa_onnx.OfflineTtsModelConfig(
+        vits=sherpa_onnx.OfflineTtsVitsModelConfig(
+            model=os.path.join(model_dir, cfg['model']),
+            lexicon=os.path.join(model_dir, cfg['lexicon']),
+            dict_dir=os.path.join(model_dir, cfg['dict_dir']),
+            tokens=os.path.join(model_dir, cfg['tokens']),
+        ),
+        provider=provider,
+        debug=0,
+        num_threads=num_threads,
+    )
+    tts_config = sherpa_onnx.OfflineTtsConfig(
+        model=model_config,
+        rule_fsts=tts_rule_fsts,
+        max_num_sentences=max_num_sentences)
+
+    if not tts_config.validate():
+        raise ValueError("tts: invalid config")
+
+    return tts_config
+
+
+def get_tts_engine(args) -> Tuple[sherpa_onnx.OfflineTts, int]:
+    sample_rate = tts_configs[args.tts_model]['sample_rate']
+    cache_engine = _tts_engines.get(args.tts_model)
+    if cache_engine:
+        return cache_engine, sample_rate
+    st = time.time()
+    tts_config = load_tts_model(
+        args.tts_model, args.models_root, args.tts_provider)
+
+    cache_engine = sherpa_onnx.OfflineTts(tts_config)
+    elapsed = time.time() - st
+    logger.info(f"tts: loaded {args.tts_model} in {elapsed:.2f}s")
+    _tts_engines[args.tts_model] = cache_engine
+
+    return cache_engine, sample_rate
+
+
+class TTSResult:
+    def __init__(self, pcm_bytes: bytes, finished: bool):
+        self.pcm_bytes = pcm_bytes
+        self.finished = finished
+        self.progress: float = 0.0
+        self.elapsed: float = 0.0
+        self.audio_duration: float = 0.0
+        self.audio_size: int = 0
+
+    def to_dict(self):
+        return {
+            "progress": self.progress,
+            "elapsed": f'{int(self.elapsed * 1000)}ms',
+            "duration": f'{self.audio_duration:.2f}s',
+            "size": self.audio_size
+        }
+
+
+class TTSStream:
+    def __init__(self, engine, sid: int, speed: float = 1.0, sample_rate: int = 16000, original_sample_rate: int = 16000):
+        self.engine = engine
+        self.sid = sid
+        self.speed = speed
+        self.outbuf: asyncio.Queue[TTSResult | None] = asyncio.Queue()
+        self.is_closed = False
+        self.target_sample_rate = sample_rate
+        self.original_sample_rate = original_sample_rate
+
+    def on_process(self, chunk: np.ndarray, progress: float):
+        if self.is_closed:
+            return 0
+
+        # resample to target sample rate
+        if self.target_sample_rate != self.original_sample_rate:
+            num_samples = int(
+                len(chunk) * self.target_sample_rate / self.original_sample_rate)
+            resampled_chunk = resample(chunk, num_samples)
+            chunk = resampled_chunk.astype(np.float32)
+
+        scaled_chunk = chunk * 32768.0
+        clipped_chunk = np.clip(scaled_chunk, -32768, 32767)
+        int16_chunk = clipped_chunk.astype(np.int16)
+        samples = int16_chunk.tobytes()
+        self.outbuf.put_nowait(TTSResult(samples, False))
+        return self.is_closed and 0 or 1
+
+    async def write(self, text: str, split: bool, pause: float = 0.2):
+        start = time.time()
+        if split:
+            texts = re.split(splitter, text)
+        else:
+            texts = [text]
+
+        audio_duration = 0.0
+        audio_size = 0
+
+        for idx, text in enumerate(texts):
+            text = text.strip()
+            if not text:
+                continue
+            sub_start = time.time()
+
+            audio = await asyncio.to_thread(self.engine.generate,
+                                            text, self.sid, self.speed,
+                                            self.on_process)
+
+            if not audio or not audio.sample_rate or not audio.samples:
+                logger.error(f"tts: failed to generate audio for "
+                             f"'{text}' (audio={audio})")
+                continue
+
+            if split and idx < len(texts) - 1:  # add a pause between sentences
+                noise = np.zeros(int(audio.sample_rate * pause))
+                self.on_process(noise, 1.0)
+                audio.samples = np.concatenate([audio.samples, noise])
+
+            audio_duration += len(audio.samples) / audio.sample_rate
+            audio_size += len(audio.samples)
+            elapsed_seconds = time.time() - sub_start
+            logger.info(f"tts: generated audio for '{text}', "
+                        f"audio duration: {audio_duration:.2f}s, "
+                        f"elapsed: {elapsed_seconds:.2f}s")
+
+        elapsed_seconds = time.time() - start
+        logger.info(f"tts: generated audio in {elapsed_seconds:.2f}s, "
+                    f"audio duration: {audio_duration:.2f}s")
+
+        r = TTSResult(None, True)
+        r.elapsed = elapsed_seconds
+        r.audio_duration = audio_duration
+        r.progress = 1.0
+        r.finished = True
+        await self.outbuf.put(r)
+
+    async def close(self):
+        self.is_closed = True
+        self.outbuf.put_nowait(None)
+        logger.info("tts: stream closed")
+
+    async def read(self) -> TTSResult:
+        return await self.outbuf.get()
+
+    async def generate(self,  text: str) -> io.BytesIO:
+        start = time.time()
+        audio = await asyncio.to_thread(self.engine.generate,
+                                        text, self.sid, self.speed)
+        elapsed_seconds = time.time() - start
+        audio_duration = len(audio.samples) / audio.sample_rate
+
+        logger.info(f"tts: generated audio in {elapsed_seconds:.2f}s, "
+                    f"audio duration: {audio_duration:.2f}s, "
+                    f"sample rate: {audio.sample_rate}")
+
+        if self.target_sample_rate != audio.sample_rate:
+            audio.samples = resample(audio.samples,
+                                     int(len(audio.samples) * self.target_sample_rate / audio.sample_rate))
+            audio.sample_rate = self.target_sample_rate
+
+        output = io.BytesIO()
+        soundfile.write(output,
+                        audio.samples,
+                        samplerate=audio.sample_rate,
+                        subtype="PCM_16",
+                        format="WAV")
+        output.seek(0)
+        return output
+
+
+async def start_tts_stream(sid: int, sample_rate: int, speed: float, args) -> TTSStream:
+    engine, original_sample_rate = get_tts_engine(args)
+    return TTSStream(engine, sid, speed, sample_rate, original_sample_rate)