ai-companion/p0/stt_tts_bench/bench_concurrent.py

#!/usr/bin/env python3
"""P0 基准 · STT 与 TTS 并发。

为什么必须测：

  真实系统里助手说话时麦克风并没有关——要检测打断，STT 必须同时在跑。
  这台机器只有 4 个核，且 bench_threads.py 已经证明 STT 超过 2 线程就会急剧退化。
  那么"STT(2线程) + TTS(2线程)"极可能互相踩踏。

  本脚本把两种负载跑在一起，和各自单独跑的结果对比，量化退化幅度。
  只看单独指标会得出"两个都能实时"的错误结论。
"""

import os
import statistics
import threading
import time
import wave

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))
STT = os.path.join(HERE, "models", "stt")
WAVS = os.path.join(HERE, "wavs")
SR = 16000
CHUNK = int(SR * 0.1)

TTS_SENTENCES = [
    "我查了一下，明天上午十点有个会议。",
    "地点在三号会议室，需要提前十分钟到。",
    "另外提醒你，后天是妈妈的生日。",
    "要不要我提前帮你订一束花？",
    "还有你上周说要买的那本书现在有货了。",
    "价格是四十九块八，要不要下单？",
]


def read_wav(p):
    with wave.open(p, "rb") as w:
        return np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).astype(np.float32) / 32768.0


def build_stt(nt):
    return sherpa_onnx.OnlineRecognizer.from_transducer(
        tokens=os.path.join(STT, "tokens.txt"),
        encoder=os.path.join(STT, "encoder-epoch-99-avg-1.int8.onnx"),
        decoder=os.path.join(STT, "decoder-epoch-99-avg-1.int8.onnx"),
        joiner=os.path.join(STT, "joiner-epoch-99-avg-1.int8.onnx"),
        num_threads=nt, sample_rate=SR, feature_dim=80,
        decoding_method="greedy_search", enable_endpoint_detection=False)


def build_tts(which, nt):
    if which == "matcha":
        d = os.path.join(HERE, "models", "matcha")
        matcha = sherpa_onnx.OfflineTtsMatchaModelConfig(
            acoustic_model=os.path.join(d, "model-steps-3.onnx"),
            vocoder=os.path.join(HERE, "models", "matcha_vocoder", "hifigan_v2.onnx"),
            lexicon=os.path.join(d, "lexicon.txt"),
            tokens=os.path.join(d, "tokens.txt"),
            dict_dir=os.path.join(d, "dict"))
        return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
            model=sherpa_onnx.OfflineTtsModelConfig(matcha=matcha, num_threads=nt, provider="cpu"),
            rule_fsts=",".join(os.path.join(d, f) for f in ("date.fst", "number.fst", "phone.fst")),
            max_num_sentences=1))
    if which == "kokoro_int8":
        d = os.path.join(HERE, "models", "kokoro_int8")
        kokoro = sherpa_onnx.OfflineTtsKokoroModelConfig(
            model=os.path.join(d, "model.int8.onnx"),
            voices=os.path.join(d, "voices.bin"),
            tokens=os.path.join(d, "tokens.txt"),
            data_dir=os.path.join(d, "espeak-ng-data"),
            dict_dir=os.path.join(d, "dict"),
            lexicon=",".join(os.path.join(d, f) for f in
                             ("lexicon-us-en.txt", "lexicon-gb-en.txt", "lexicon-zh.txt")))
        return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
            model=sherpa_onnx.OfflineTtsModelConfig(kokoro=kokoro, num_threads=nt, provider="cpu"),
            max_num_sentences=1))
    if which == "piper":
        d = os.path.join(HERE, "models", "tts_piper_xlow")
        vits = sherpa_onnx.OfflineTtsVitsModelConfig(
            model=os.path.join(d, "zh_CN-huayan-x_low.onnx"),
            lexicon=os.path.join(d, "lexicon.txt"),
            tokens=os.path.join(d, "tokens.txt"))
        return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
            model=sherpa_onnx.OfflineTtsModelConfig(vits=vits, num_threads=nt, provider="cpu"),
            max_num_sentences=1))
    d = os.path.join(HERE, "models", "tts")
    vits = sherpa_onnx.OfflineTtsVitsModelConfig(
        model=os.path.join(d, "model.onnx"), lexicon=os.path.join(d, "lexicon.txt"),
        tokens=os.path.join(d, "tokens.txt"), dict_dir=os.path.join(d, "dict"))
    return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(vits=vits, num_threads=nt, provider="cpu"),
        rule_fsts=",".join(os.path.join(d, f) for f in
                           ("date.fst", "number.fst", "phone.fst", "new_heteronym.fst")),
        max_num_sentences=1))


def stt_pass(rec, samples):
    """喂一遍音频，返回 (rtf, 单块解码耗时列表)。"""
    s = rec.create_stream()
    durs = []
    pos = 0
    t_all = time.perf_counter()
    while pos < len(samples):
        block = samples[pos:pos + CHUNK]
        pos += CHUNK
        t0 = time.perf_counter()
        s.accept_waveform(SR, block)
        while rec.is_ready(s):
            rec.decode_stream(s)
        durs.append(time.perf_counter() - t0)
    total = time.perf_counter() - t_all
    return total / (len(samples) / SR), durs


def tts_pass(tts, sents):
    """连续合成若干句，返回 (rtf, 每句耗时ms)。"""
    per = []
    audio_s = 0.0
    t_all = time.perf_counter()
    for s in sents:
        t0 = time.perf_counter()
        a = tts.generate(s, sid=0, speed=1.0)
        per.append((time.perf_counter() - t0) * 1000)
        audio_s += len(a.samples) / a.sample_rate
    total = time.perf_counter() - t_all
    return total / audio_s, per


def main():
    stt_nt = int(os.environ.get("STT_THREADS", "2"))
    tts_nt = int(os.environ.get("TTS_THREADS", "2"))
    tts_kind = os.environ.get("TTS_MODEL", "piper")
    wav = read_wav(os.path.join(WAVS, "stt_0.wav"))
    dur = len(wav) / SR

    print(f"音频 {dur:.2f}s   STT线程={stt_nt}   TTS={tts_kind}线程={tts_nt}\n")

    stt = build_stt(stt_nt)
    tts = build_tts(tts_kind, tts_nt)
    stt_pass(stt, wav[: SR * 2])          # 预热
    tts.generate("预热", sid=0, speed=1.0)

    # 单独
    solo_stt, _ = stt_pass(stt, wav)
    solo_tts, _ = tts_pass(tts, TTS_SENTENCES)
    print(f"单独:  STT RTF={solo_stt:.2f}   TTS RTF={solo_tts:.2f}")

    # 并发：TTS 在后台线程持续合成，主线程跑 STT
    stop = threading.Event()
    tts_stats = {"runs": 0, "rtf": [], "done": False}

    def tts_worker():
        try:
            while not stop.is_set():
                rtf, _ = tts_pass(tts, TTS_SENTENCES)
                tts_stats["rtf"].append(rtf)
                tts_stats["runs"] += 1
                if tts_stats["runs"] >= 3:
                    break
        finally:
            tts_stats["done"] = True

    th = threading.Thread(target=tts_worker)
    th.start()
    time.sleep(0.2)  # 让 TTS 先跑起来，制造真实争抢
    conc_rtfs, conc_durs = [], []
    for _ in range(3):
        r, d = stt_pass(stt, wav)
        conc_rtfs.append(r)
        conc_durs.extend(d)
    stop.set()
    th.join()

    conc_tts = statistics.median(tts_stats["rtf"]) if tts_stats["rtf"] else float("nan")
    stt_med = statistics.median(conc_rtfs)
    ratios = sorted(x / (CHUNK / SR) for x in conc_durs)

    print(f"并发:  STT RTF={stt_med:.2f} (×{stt_med/solo_stt:.2f})   "
          f"TTS RTF={conc_tts:.2f} (×{conc_tts/solo_tts:.2f})")
    print(f"       并发下单块解码 p50={statistics.median(ratios):.2f} "
          f"p95={ratios[int(len(ratios)*0.95)]:.2f} max={ratios[-1]:.2f}（>1 表示积压）")

    # 判定口径：看 RTF（聚合力，<1 表示跟得上）。单块 p95 只作观察项——
    # 流式模型并非每块都解码（实测约每 3 块解一次），所以单块 >1 不等于积压。
    # 留 0.7 的余量，因为还要给视觉检测等其它负载留 CPU。
    if stt_med < 0.7 and conc_tts < 0.7:
        verdict = f"✅ 并发可实时（STT 余量 {1/stt_med:.1f}x，TTS 余量 {1/conc_tts:.1f}x）"
    elif stt_med < 1 and conc_tts < 1:
        verdict = "△ 并发勉强够用，视觉检测再加负载有风险"
    else:
        verdict = "⚠ 并发跟不上实时"
    print(f"\n判定: {verdict}")

    import json
    out = os.path.join(HERE, f"results_concurrent_{tts_kind}.json")
    with open(out, "w") as f:
        json.dump({"stt_threads": stt_nt, "tts_threads": tts_nt, "tts_model": tts_kind,
                   "solo_stt_rtf": solo_stt, "solo_tts_rtf": solo_tts,
                   "conc_stt_rtf": stt_med, "conc_tts_rtf": conc_tts,
                   "stt_degrade": stt_med / solo_stt, "tts_degrade": conc_tts / solo_tts,
                   "conc_chunk_p95": ratios[int(len(ratios) * 0.95)],
                   "conc_chunk_max": ratios[-1]}, f, ensure_ascii=False, indent=2)
    print(f"结果已写入 {out}")


if __name__ == "__main__":
    main()