funasr_wss_server.py

import asyncio
import json
import websockets
import time
import numpy as np
import argparse
import ssl
import os
import wave
import functools
from concurrent.futures import ThreadPoolExecutor
from scipy.spatial.distance import cosine

import torch  # 保留不影响


def to_python(obj):
    """递归地把 numpy / torch 等类型转成纯 Python，可 JSON 序列化。"""
    try:
        import numpy as np  # noqa
        import torch  # noqa
    except Exception:
        np = None
        torch = None

    if np is not None and isinstance(obj, np.generic):
        return obj.item()
    if np is not None and isinstance(obj, np.ndarray):
        return obj.tolist()
    if torch is not None and isinstance(obj, torch.Tensor):
        return obj.cpu().tolist()

    if isinstance(obj, dict):
        return {k: to_python(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [to_python(v) for v in obj]

    return obj


parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="0.0.0.0", required=False, help="host ip")
parser.add_argument("--port", type=int, default=10095, required=False, help="grpc server port")

parser.add_argument(
    "--asr_model",
    type=str,
    default="iic/speech_paraformer-large-contextual_asr_nat-zh-cn-16k-common-vocab8404",
    help="model from modelscope",
)
parser.add_argument("--asr_model_revision", type=str, default="v2.0.4", help="")

parser.add_argument(
    "--asr_model_online",
    type=str,
    default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
    help="model from modelscope",
)
parser.add_argument("--asr_model_online_revision", type=str, default="v2.0.4", help="")

parser.add_argument(
    "--vad_model",
    type=str,
    default="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
    help="model from modelscope",
)
parser.add_argument("--vad_model_revision", type=str, default="v2.0.4", help="")

parser.add_argument(
    "--punc_model",
    type=str,
    default="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727",
    help="model from modelscope",
)
parser.add_argument("--punc_model_revision", type=str, default="v2.0.4", help="")

parser.add_argument("--ngpu", type=int, default=1, help="0 for cpu, 1 for gpu")
parser.add_argument("--device", type=str, default="cuda", help="cuda, cpu")
parser.add_argument("--ncpu", type=int, default=4, help="cpu cores")

parser.add_argument(
    "--certfile",
    type=str,
    default="../../ssl_key/server.crt",
    required=False,
    help="certfile for ssl",
)
parser.add_argument(
    "--keyfile",
    type=str,
    default="../../ssl_key/server.key",
    required=False,
    help="keyfile for ssl",
)

# ====== 保存 2pass 离线阶段送入 ASR 的音频片段（排查 VAD 切分）======
parser.add_argument(
    "--save_offline_segments",
    action="store_true",
    help="Save each offline (2pass) audio segment sent to offline ASR as wav for debugging VAD split.",
)
parser.add_argument(
    "--save_offline_segments_dir",
    type=str,
    default="./offline_segments",
    help="Directory to save offline wav segments when --save_offline_segments is enabled.",
)

# ====== 并发控制：核心新增 ======
parser.add_argument(
    "--worker_threads",
    type=int,
    default=max(4, (os.cpu_count() or 4)),
    help="ThreadPoolExecutor max_workers. Used to offload blocking inference so event loop won't be blocked.",
)
parser.add_argument("--concurrent_vad", type=int, default=4, help="Max concurrent VAD generate() calls.")
parser.add_argument("--concurrent_asr_online", type=int, default=4, help="Max concurrent streaming ASR generate() calls.")
parser.add_argument("--concurrent_asr_offline", type=int, default=2, help="Max concurrent offline ASR generate() calls.")
parser.add_argument("--concurrent_punc", type=int, default=1, help="Max concurrent punctuation generate() calls.")
parser.add_argument("--concurrent_sv", type=int, default=1, help="Max concurrent speaker verification generate() calls.")
parser.add_argument(
    "--speaker_db_reload_sec",
    type=int,
    default=5,
    help="Reload speaker_db.json at most once every N seconds (avoid frequent disk IO).",
)

args = parser.parse_args()

websocket_users = set()
SPEAKER_DB_PATH = os.path.join(os.path.dirname(__file__), "speaker_db.json")


def _ensure_dir(p: str):
    try:
        os.makedirs(p, exist_ok=True)
    except Exception:
        pass


def _pcm_duration_ms(pcm_bytes: bytes, fs: int, ch: int = 1, sampwidth: int = 2) -> int:
    """根据 fs/ch/sampwidth 计算 PCM 时长，避免写死 16k -> 32 bytes/ms。"""
    if not pcm_bytes:
        return 0
    bytes_per_ms = (fs * ch * sampwidth) / 1000.0
    if bytes_per_ms <= 0:
        return 0
    return int(len(pcm_bytes) / bytes_per_ms)


def _safe_int(v, default):
    try:
        return int(v)
    except Exception:
        return default


# ========= speaker db：加缓存，避免每段都读盘 =========
_SPEAKER_DB_CACHE = {}
_SPEAKER_DB_CACHE_TS = 0.0


def _load_speaker_db_sync():
    if not os.path.exists(SPEAKER_DB_PATH):
        return {}
    try:
        with open(SPEAKER_DB_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
            return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def get_speaker_db_cached(now_ts: float, reload_sec: int):
    global _SPEAKER_DB_CACHE, _SPEAKER_DB_CACHE_TS
    if (now_ts - _SPEAKER_DB_CACHE_TS) >= max(1, int(reload_sec)):
        _SPEAKER_DB_CACHE = _load_speaker_db_sync()
        _SPEAKER_DB_CACHE_TS = now_ts
    return _SPEAKER_DB_CACHE or {}


def _save_wav_sync(out_path: str, audio_bytes: bytes, fs: int, ch: int, sampwidth: int):
    with wave.open(out_path, "wb") as wf:
        wf.setnchannels(ch)
        wf.setsampwidth(sampwidth)
        wf.setframerate(fs)
        wf.writeframes(audio_bytes)


def save_offline_wav_segment_sync(websocket, audio_bytes: bytes, reason: str = "offline"):
    """
    保存离线阶段送入 ASR 的音频片段，方便人工试听排查 VAD 切分是否正确。
    约定：audio_bytes 为 单声道 PCM16 little-endian（默认 16k）。
    （注意：这是同步函数，外层会放线程池执行）
    """
    if not getattr(websocket, "save_offline_segments", False):
        return
    if "2pass" not in (getattr(websocket, "mode", "") or ""):
        return
    if not audio_bytes:
        return

    fs = int(getattr(websocket, "audio_fs", 16000) or 16000)
    ch = 1
    sampwidth = 2  # int16

    # int16 对齐
    if len(audio_bytes) % 2 == 1:
        audio_bytes = audio_bytes[:-1]
        if not audio_bytes:
            return

    seg_idx = int(getattr(websocket, "offline_seg_idx", 0))
    websocket.offline_seg_idx = seg_idx + 1

    duration_ms = _pcm_duration_ms(audio_bytes, fs=fs, ch=ch, sampwidth=sampwidth)

    base_dir = getattr(websocket, "offline_save_dir", args.save_offline_segments_dir)
    _ensure_dir(base_dir)

    wav_name = (getattr(websocket, "wav_name", "microphone") or "microphone").replace("/", "_")
    ts = int(time.time() * 1000)
    fname = f"{wav_name}_{ts}_seg{seg_idx:04d}_{reason}_{duration_ms}ms.wav"
    out_path = os.path.join(base_dir, fname)

    try:
        _save_wav_sync(out_path, audio_bytes, fs=fs, ch=ch, sampwidth=sampwidth)
        print(f"[SAVE_OFFLINE_SEG] {out_path} ({duration_ms} ms, {len(audio_bytes)} bytes)")
    except Exception as e:
        print(f"[SAVE_OFFLINE_SEG] failed: {e}")


print("model loading")
from funasr import AutoModel  # noqa

# ====== 离线 ASR ======
model_asr = AutoModel(
    model="paraformer-zh",
    model_revision="v2.0.4",
    ngpu=args.ngpu,
    ncpu=args.ncpu,
    device=args.device,
    disable_pbar=True,
    disable_log=True,
)

# streaming asr
model_asr_streaming = AutoModel(
    model=args.asr_model_online,
    model_revision=args.asr_model_online_revision,
    ngpu=args.ngpu,
    ncpu=args.ncpu,
    device=args.device,
    disable_pbar=True,
    disable_log=True,
)

# vad
model_vad = AutoModel(
    model=args.vad_model,
    model_revision=args.vad_model_revision,
    ngpu=args.ngpu,
    ncpu=args.ncpu,
    device=args.device,
    disable_pbar=True,
    disable_log=True,
)

# punc
if args.punc_model != "":
    model_punc = AutoModel(
        model=args.punc_model,
        model_revision=args.punc_model_revision,
        ngpu=args.ngpu,
        ncpu=args.ncpu,
        device=args.device,
        disable_pbar=True,
        disable_log=True,
    )
else:
    model_punc = None

# sv
model_sv = AutoModel(
    model="iic/speech_campplus_sv_zh-cn_16k-common",
    ngpu=args.ngpu,
    device=args.device,
    disable_pbar=True,
    disable_log=True,
)

print("model loaded! (now supports multi-client with non-blocking inference)")


# ====== 线程池 + 并发阈值（核心）======
EXECUTOR = ThreadPoolExecutor(max_workers=int(args.worker_threads))

SEM_VAD = asyncio.Semaphore(max(1, int(args.concurrent_vad)))
SEM_ASR_ONLINE = asyncio.Semaphore(max(1, int(args.concurrent_asr_online)))
SEM_ASR_OFFLINE = asyncio.Semaphore(max(1, int(args.concurrent_asr_offline)))
SEM_PUNC = asyncio.Semaphore(max(1, int(args.concurrent_punc)))
SEM_SV = asyncio.Semaphore(max(1, int(args.concurrent_sv)))
SEM_WAV = asyncio.Semaphore(max(1, 4))  # 保存 wav 一般不需要太大


async def run_blocking(fn, *a, sem: asyncio.Semaphore | None = None, **kw):
    """
    把阻塞函数丢线程池执行，避免卡 event loop。
    sem 用于限流（避免 GPU / 模型被打爆）。
    """
    loop = asyncio.get_running_loop()
    call = functools.partial(fn, *a, **kw)
    if sem is None:
        return await loop.run_in_executor(EXECUTOR, call)
    async with sem:
        return await loop.run_in_executor(EXECUTOR, call)


def _generate_sync(model, audio_or_text, status_dict):
    # 注意：status_dict 里包含 cache，会被 generate 更新
    return model.generate(input=audio_or_text, **status_dict)


async def ws_reset(websocket):
    print("ws reset now, total num is ", len(websocket_users))

    websocket.status_dict_asr_online["cache"] = {}
    websocket.status_dict_asr_online["is_final"] = True
    websocket.status_dict_vad["cache"] = {}
    websocket.status_dict_vad["is_final"] = True
    websocket.status_dict_punc["cache"] = {}

    await websocket.close()


async def clear_websocket():
    for websocket in list(websocket_users):
        await ws_reset(websocket)
    websocket_users.clear()


async def ws_serve(websocket, path=None):
    # websockets 新版本不会传 path，这里做兼容
    if path is None:
        path = getattr(websocket, "path", None)
    frames = []
    frames_asr = []
    frames_asr_online = []
    pending_offline_audio = []
    global websocket_users
    websocket_users.add(websocket)

    websocket.status_dict_asr = {}  # hotword 等
    websocket.status_dict_asr_online = {"cache": {}, "is_final": False}
    websocket.status_dict_vad = {"cache": {}, "is_final": False}
    websocket.status_dict_punc = {"cache": {}}

    websocket.chunk_interval = 10
    websocket.vad_pre_idx = 0
    speech_start = False
    speech_end_i = -1
    online_needs_finalization = False
    session_errors = []

    websocket.wav_name = "microphone"
    websocket.mode = "2pass"
    websocket.is_speaking = True  # ✅ 默认初始化，避免 AttributeError

    # 保存离线片段
    websocket.audio_fs = 16000
    websocket.offline_seg_idx = 0
    websocket.save_offline_segments = bool(args.save_offline_segments)
    websocket.offline_save_dir = args.save_offline_segments_dir
    if websocket.save_offline_segments:
        _ensure_dir(websocket.offline_save_dir)
        print(f"[SAVE_OFFLINE_SEG] enabled, dir={websocket.offline_save_dir}")

    print("new user connected", flush=True)

    def record_error(message):
        if message not in session_errors:
            session_errors.append(message)

    async def finalize_online_segment():
        nonlocal frames_asr_online, online_needs_finalization

        if websocket.mode not in ("2pass", "online") or not online_needs_finalization:
            return

        websocket.status_dict_asr_online["is_final"] = True
        try:
            await async_asr_online(websocket, b"".join(frames_asr_online))
        except Exception as e:
            print("error in final asr streaming:", e)
            record_error(f"online inference failed: {e}")

        frames_asr_online = []
        websocket.status_dict_asr_online["cache"] = {}
        websocket.status_dict_asr_online["is_final"] = False
        online_needs_finalization = False

    async def finish_input(send_end_ack):
        nonlocal frames, frames_asr, frames_asr_online, pending_offline_audio
        nonlocal speech_start, speech_end_i, online_needs_finalization

        await finalize_online_segment()

        if websocket.mode in ("2pass", "offline"):
            audio_in = b"".join(frames_asr)
            if not audio_in:
                audio_in = b"".join(pending_offline_audio)

            if audio_in:
                if websocket.save_offline_segments and audio_in:
                    try:
                        await run_blocking(
                            save_offline_wav_segment_sync,
                            websocket,
                            audio_in,
                            "not_speaking",
                            sem=SEM_WAV,
                        )
                    except Exception as e:
                        print("[SAVE_OFFLINE_SEG] async failed:", e)

                try:
                    await async_asr(websocket, audio_in)
                    pending_offline_audio = []
                except Exception as e:
                    print("error in final asr offline:", e)
                    record_error(f"offline inference failed: {e}")

        errors = list(session_errors)

        frames = []
        frames_asr = []
        frames_asr_online = []
        pending_offline_audio = []
        speech_start = False
        speech_end_i = -1
        online_needs_finalization = False
        websocket.vad_pre_idx = 0
        websocket.status_dict_asr_online["cache"] = {}
        websocket.status_dict_vad["cache"] = {}

        if send_end_ack:
            acknowledgement = {
                "mode": websocket.mode,
                "wav_name": websocket.wav_name,
                "is_final": not errors,
                "is_end": True,
            }
            if errors:
                acknowledgement["error"] = "; ".join(errors)
            await websocket.send(
                json.dumps(acknowledgement, ensure_ascii=False)
            )
            session_errors.clear()
        elif errors:
            raise RuntimeError("; ".join(errors))

    try:
        async for message in websocket:
            # ========== 1) 先处理“文本配置消息” ==========
            if isinstance(message, str):
                try:
                    messagejson = json.loads(message)
                except Exception as e:
                    print("bad json message:", e, message[:200])
                    continue

                print("=============messagejson============", messagejson)

                end_of_input = False
                if "is_speaking" in messagejson:
                    websocket.is_speaking = bool(messagejson["is_speaking"])
                    websocket.status_dict_asr_online["is_final"] = (not websocket.is_speaking)
                    end_of_input = not websocket.is_speaking

                if "chunk_interval" in messagejson:
                    websocket.chunk_interval = _safe_int(
                        messagejson["chunk_interval"], websocket.chunk_interval
                    )

                if "wav_name" in messagejson:
                    websocket.wav_name = messagejson.get("wav_name") or websocket.wav_name

                if "chunk_size" in messagejson:
                    chunk_size = messagejson["chunk_size"]
                    if isinstance(chunk_size, str):
                        chunk_size = [x.strip() for x in chunk_size.split(",") if x.strip()]
                    websocket.status_dict_asr_online["chunk_size"] = [int(x) for x in chunk_size]

                if "encoder_chunk_look_back" in messagejson:
                    websocket.status_dict_asr_online["encoder_chunk_look_back"] = messagejson[
                        "encoder_chunk_look_back"
                    ]

                if "decoder_chunk_look_back" in messagejson:
                    websocket.status_dict_asr_online["decoder_chunk_look_back"] = messagejson[
                        "decoder_chunk_look_back"
                    ]

                if "hotwords" in messagejson:
                    hotword_data = messagejson["hotwords"]
                    websocket.status_dict_asr["hotword"] = hotword_data
                    websocket.status_dict_asr_online["hotword"] = hotword_data
                    print(f"热词已更新: {hotword_data}")

                if "mode" in messagejson:
                    requested_mode = messagejson["mode"]
                    if requested_mode and requested_mode not in ("online", "offline", "2pass"):
                        websocket.mode = requested_mode
                        record_error(f"unsupported mode: {requested_mode!r}")
                    else:
                        websocket.mode = requested_mode or websocket.mode

                if "audio_fs" in messagejson:
                    websocket.audio_fs = _safe_int(messagejson["audio_fs"], 16000)

                if end_of_input:
                    await finish_input(send_end_ack=bool(messagejson.get("is_end")))

                continue

            # ========== 2) 处理“二进制音频消息” ==========
            if websocket.mode not in ("online", "offline", "2pass"):
                continue

            if "chunk_size" not in websocket.status_dict_asr_online:
                print("[WARN] chunk_size not set yet, skip audio frame (send config first).")
                record_error("audio frame discarded: chunk_size is not configured")
                continue

            try:
                websocket.status_dict_vad["chunk_size"] = int(
                    websocket.status_dict_asr_online["chunk_size"][1] * 60 / websocket.chunk_interval
                )
            except Exception as e:
                print("[WARN] set vad chunk_size failed:", e)
                record_error(f"audio frame discarded: invalid VAD chunk_size: {e}")
                continue

            pcm = message
            frames.append(pcm)
            if websocket.mode in ("2pass", "offline"):
                pending_offline_audio.append(pcm)

            duration_ms = _pcm_duration_ms(pcm, fs=websocket.audio_fs, ch=1, sampwidth=2)
            websocket.vad_pre_idx += duration_ms

            # online asr
            frames_asr_online.append(pcm)
            if websocket.mode in ("2pass", "online"):
                online_needs_finalization = True
            websocket.status_dict_asr_online["is_final"] = (speech_end_i != -1)

            if (len(frames_asr_online) % websocket.chunk_interval == 0) or websocket.status_dict_asr_online["is_final"]:
                if websocket.mode in ("2pass", "online"):
                    audio_in = b"".join(frames_asr_online)
                    try:
                        await async_asr_online(websocket, audio_in)
                    except Exception as e:
                        print(f"error in asr streaming, {websocket.status_dict_asr_online}")
                        record_error(f"online inference failed: {e}")
                frames_asr_online = []

            if speech_start:
                frames_asr.append(pcm)

            # vad online
            try:
                speech_start_i, speech_end_i = await async_vad(websocket, pcm)
            except Exception as e:
                print("error in vad:", e)
                record_error(f"vad inference failed: {e}")
                speech_start_i, speech_end_i = -1, -1

            if speech_start_i != -1:
                speech_start = True
                if duration_ms > 0:
                    beg_bias = (websocket.vad_pre_idx - speech_start_i) // duration_ms
                else:
                    beg_bias = 0
                frames_pre = frames[-beg_bias:] if beg_bias > 0 else []
                frames_asr = []
                frames_asr.extend(frames_pre)

            # ========== 3) 2pass：离线阶段触发点 ==========
            if (speech_end_i != -1) or (not websocket.is_speaking):
                await finalize_online_segment()

                if websocket.mode in ("2pass", "offline"):
                    audio_in = b"".join(frames_asr)
                    if not audio_in and speech_end_i != -1:
                        audio_in = b"".join(pending_offline_audio)
                    reason = "vad_end" if speech_end_i != -1 else "not_speaking"

                    # 保存 wav：放线程池，避免磁盘 IO 卡 loop
                    if websocket.save_offline_segments and audio_in:
                        try:
                            await run_blocking(
                                save_offline_wav_segment_sync,
                                websocket,
                                audio_in,
                                reason,
                                sem=SEM_WAV,
                            )
                        except Exception as e:
                            print("[SAVE_OFFLINE_SEG] async failed:", e)

                    if audio_in:
                        try:
                            await async_asr(websocket, audio_in)
                            pending_offline_audio = []
                        except Exception as e:
                            print("error in asr offline:", e)
                            record_error(f"offline inference failed: {e}")

                frames_asr = []
                speech_start = False
                frames_asr_online = []
                websocket.status_dict_asr_online["cache"] = {}
                websocket.status_dict_asr_online["is_final"] = False
                online_needs_finalization = False
                speech_end_i = -1

                if not websocket.is_speaking:
                    websocket.vad_pre_idx = 0
                    frames = []
                    websocket.status_dict_vad["cache"] = {}
                else:
                    frames = frames[-20:]

    except websockets.ConnectionClosed:
        print("ConnectionClosed...", websocket_users, flush=True)
        await ws_reset(websocket)
        if websocket in websocket_users:
            websocket_users.remove(websocket)
    except websockets.InvalidState:
        print("InvalidState...")
        try:
            await ws_reset(websocket)
        except Exception:
            pass
        websocket_users.discard(websocket)
    except Exception as e:
        print("Exception:", e)
        try:
            await ws_reset(websocket)
        except Exception:
            pass
        if websocket in websocket_users:
            websocket_users.remove(websocket)


# ===================== 推理：全部改为“线程池 + 限流” =====================

async def async_vad(websocket, audio_in: bytes):
    # model_vad.generate 是阻塞的，必须 offload
    out = await run_blocking(_generate_sync, model_vad, audio_in, websocket.status_dict_vad, sem=SEM_VAD)
    segments_result = out[0].get("value", [])

    speech_start = -1
    speech_end = -1

    if len(segments_result) == 0 or len(segments_result) > 1:
        return speech_start, speech_end
    if segments_result[0][0] != -1:
        speech_start = segments_result[0][0]
    if segments_result[0][1] != -1:
        speech_end = segments_result[0][1]
    return speech_start, speech_end


def _sv_and_match_sync(audio_in: bytes, reload_sec: int):
    """
    同步执行：SV embedding + speaker_db 匹配
    返回 (spk_name, best_score)
    """
    spk_name = "unknown"
    best_score = 0.0

    sv_out = model_sv.generate(input=audio_in, embedding=True)[0]
    embedding = sv_out["spk_embedding"][0].cpu().numpy()

    now_ts = time.time()
    local_speaker_db = get_speaker_db_cached(now_ts, reload_sec=reload_sec)
    if local_speaker_db:
        for name, ref_embedding in local_speaker_db.items():
            if ref_embedding is None:
                continue
            arr = np.array(ref_embedding, dtype=np.float32)
            similarity = 1.0 - cosine(embedding, arr)
            print("sv similarity with {}: {}".format(name, similarity))
            if similarity > best_score and similarity > 0.2:
                best_score = similarity
                spk_name = name

    return spk_name, float(best_score)


async def async_asr(websocket, audio_in: bytes):
    mode = "2pass-offline" if "2pass" in (websocket.mode or "") else websocket.mode

    if len(audio_in) <= 0:
        message = {
            "mode": mode,
            "text": "",
            "wav_name": websocket.wav_name,
            "is_final": True,
        }
        await websocket.send(json.dumps(message, ensure_ascii=False))
        return

    # 1) ASR（阻塞，线程池执行）
    rec_result_list = await run_blocking(
        _generate_sync,
        model_asr,
        audio_in,
        websocket.status_dict_asr,
        sem=SEM_ASR_OFFLINE,
    )
    rec_result = rec_result_list[0]

    print("offline_asr, raw:", rec_result)
    print("offline_asr, keys:", rec_result.keys())

    text = rec_result.get("text", "")
    timestamp = rec_result.get("timestamp", None)
    sentence_info = rec_result.get("sentence_info", None)

    # 2) 声纹识别（阻塞，线程池执行）
    spk_name = "unknown"
    best_score = 0.0
    try:
        spk_name, best_score = await run_blocking(
            _sv_and_match_sync,
            audio_in,
            int(args.speaker_db_reload_sec),
            sem=SEM_SV,
        )
    except Exception as e:
        print(f"声纹识别失败: {e}")

    # 3) 标点（阻塞，线程池执行）
    punc_array = None
    if model_punc is not None and len(text) > 0:
        try:
            # punc 只对文本处理
            punc_out = await run_blocking(
                _generate_sync,
                model_punc,
                text,
                websocket.status_dict_punc,
                sem=SEM_PUNC,
            )
            punc_result = punc_out[0]
            print("offline, after punc", punc_result)

            if "text" in punc_result and punc_result["text"]:
                text = punc_result["text"]
            if "punc_array" in punc_result:
                punc_array = punc_result["punc_array"]
        except Exception as e:
            print("punc failed:", e)

    # 4) 构造最终 message
    if len(text) > 0:
        print("======offline final text:", text)
        message = {
            "mode": mode,
            "spk_name": spk_name,
            "spk_score": float(best_score),
            "text": text,
            "wav_name": websocket.wav_name,
            "is_final": True,
        }
        if timestamp is not None:
            message["timestamp"] = to_python(timestamp)
        if sentence_info is not None:
            message["sentence_info"] = to_python(sentence_info)
        if punc_array is not None:
            message["punc_array"] = to_python(punc_array)

        await websocket.send(json.dumps(message, ensure_ascii=False))
    else:
        message = {
            "mode": mode,
            "spk_name": spk_name,
            "spk_score": float(best_score),
            "text": "",
            "wav_name": websocket.wav_name,
            "is_final": True,
        }
        await websocket.send(json.dumps(message, ensure_ascii=False))


async def async_asr_online(websocket, audio_in: bytes):
    if len(audio_in) <= 0 and not websocket.status_dict_asr_online.get("is_final", False):
        return

    # streaming generate 也是阻塞：线程池执行
    rec_out = await run_blocking(
        _generate_sync,
        model_asr_streaming,
        audio_in,
        websocket.status_dict_asr_online,
        sem=SEM_ASR_ONLINE,
    )
    rec_result = rec_out[0]
    print("online, ", rec_result)

    # 2pass：online 只要 partial，不发 final（final 交给 offline）
    if websocket.mode == "2pass" and websocket.status_dict_asr_online.get("is_final", False):
        return

    if rec_result.get("text"):
        mode = "2pass-online" if "2pass" in (websocket.mode or "") else websocket.mode
        message = {
            "mode": mode,
            "text": rec_result["text"],
            "wav_name": websocket.wav_name,
            "is_final": bool(
                websocket.status_dict_asr_online.get("is_final", False) or (not websocket.is_speaking)
            ),
        }
        await websocket.send(json.dumps(message, ensure_ascii=False))


# ===================== 启动服务 =====================

async def main():
    if len(args.certfile) > 0:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(args.certfile, keyfile=args.keyfile)
        server = await websockets.serve(
            ws_serve,
            args.host,
            args.port,
            subprotocols=["binary"],
            ping_interval=None,
            ssl=ssl_context,
        )
    else:
        server = await websockets.serve(
            ws_serve,
            args.host,
            args.port,
            subprotocols=["binary"],
            ping_interval=None,
        )

    print(f"WS server started at ws(s)://{args.host}:{args.port}")
    await server.wait_closed()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    finally:
        try:
            EXECUTOR.shutdown(wait=False, cancel_futures=True)
        except Exception:
            pass