From bd6e72142f1cca3c30b7651bf5fa567dfe969810 Mon Sep 17 00:00:00 2001 From: zhifu gao Date: Fri, 17 Jul 2026 23:14:42 +0800 Subject: [PATCH] Add client-driven realtime endpoints (#3245) Co-authored-by: LauraGPT --- .../fun_asr_nano/docs/realtime_demo.md | 19 ++ .../fun_asr_nano/serve_realtime_ws.py | 168 ++++++++++++++- tests/test_realtime_ws_service.py | 201 ++++++++++++++++++ 3 files changed, 378 insertions(+), 10 deletions(-) diff --git a/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md b/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md index 3f40f1741..6dba0226e 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md +++ b/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md @@ -15,6 +15,25 @@ pip install -r requirements.txt CUDA_VISIBLE_DEVICES=0 python serve_realtime_ws.py --port 10095 --language 中文 ``` +### 客户端 VAD / 低延迟对话模式 + +客户端已经持续执行 VAD 时，可以关闭服务端 VAD 模型，并由客户端明确提交每句话： + +```bash +CUDA_VISIBLE_DEVICES=0 python serve_realtime_ws.py \ + --port 10095 --language 中文 --endpoint-mode client +``` + +同一条 WebSocket 连接上的协议为： + +1. 发送文本 `START`； +2. 持续发送 16 kHz、16-bit PCM 二进制帧，服务端仍会返回 partial； +3. 客户端 VAD 检测到句末后发送文本 `COMMIT`； +4. 服务端立即返回一次 `is_final=true`，清理当前句状态但保持连接； +5. 重复步骤 2-4，整个会话结束时发送 `STOP`。 + +`COMMIT` 不受默认 960 ms final 门槛限制，因此短句也会提交。该模式信任客户端 endpoint，不运行 FSMN-VAD；需要服务端自动去噪和分句时继续使用默认的 `--endpoint-mode server`。 + ## 客户端 ```bash diff --git a/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py b/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py index 58355b3c0..db6888d09 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py @@ -3,6 +3,7 @@ Features: - Streaming VAD segmentation (fsmn-vad) +- Client-driven utterance endpoints (COMMIT, without server VAD) - Per-segment ASR decoding (Fun-ASR-Nano via vLLM) - Speaker diarization (eres2netv2 + ClusterBackend) - Hotword customization @@ -75,6 +76,20 @@ def _clean_asr_text(text): from funasr.models.fsmn_vad_streaming.dynamic_vad import DynamicStreamingVAD +class ClientEndpointVAD: + """Track an utterance start without loading or running a VAD model.""" + + def __init__(self): + self.current_speech_start = None + + def start_utterance(self, start_ms): + if self.current_speech_start is None: + self.current_speech_start = start_ms + + def reset(self): + self.current_speech_start = None + + class HybridSpeakerTracker: """Speaker diarization: streaming ClusterBackend + final re-clustering.""" @@ -312,10 +327,14 @@ class RealtimeASRSession: chunk_ms=960, partial_window_sec=15.0, audio_lookback_sec=5.0, + endpoint_mode="server", ): + if endpoint_mode not in {"server", "client"}: + raise ValueError(f"Unsupported endpoint mode: {endpoint_mode}") self.vllm_engine = vllm_engine self.asr_kwargs = asr_kwargs self.vad = vad + self.endpoint_mode = endpoint_mode self.sample_rate = sample_rate self.chunk_samples = int(sample_rate * chunk_ms / 1000) self.first_chunk_samples = int(sample_rate * 480 / 1000) @@ -347,11 +366,19 @@ class RealtimeASRSession: def add_audio(self, pcm_bytes): audio_int16 = np.frombuffer(pcm_bytes, dtype=np.int16) audio_float = audio_int16.astype(np.float32) / 32768.0 + chunk_start_sample = self.total_samples self.audio_buffer = np.concatenate([self.audio_buffer, audio_float]) self.total_samples += len(audio_float) if len(audio_float) > 0: - new_confirmed = self.vad.feed(torch.from_numpy(audio_float).float(), is_final=False) + if self.endpoint_mode == "client": + start_ms = int(chunk_start_sample * 1000 / self.sample_rate) + self.vad.start_utterance(start_ms) + new_confirmed = [] + else: + new_confirmed = self.vad.feed( + torch.from_numpy(audio_float).float(), is_final=False + ) for seg in new_confirmed: seg_text = self._decode_segment(seg) @@ -398,11 +425,23 @@ class RealtimeASRSession: @torch.no_grad() def decode(self, is_final=False): - if self.total_samples < self.chunk_samples: + if is_final and self.endpoint_mode == "client": + return self.commit() + + if self.endpoint_mode == "server" and self.total_samples < self.chunk_samples: if is_final: self._release_audio_buffer() return self._build_response(is_final) + if self.endpoint_mode == "client": + if self.vad.current_speech_start is None: + return self._build_response(is_final=False) + utterance_start_sample = int( + self.vad.current_speech_start * self.sample_rate / 1000 + ) + if self.total_samples - utterance_start_sample < self.first_chunk_samples: + return self._build_response(is_final=False) + if is_final: if self.vad.current_speech_start is not None: end_ms = int(self.total_samples * 1000 / self.sample_rate) @@ -460,6 +499,52 @@ class RealtimeASRSession: return self._build_response(is_final) + @torch.no_grad() + def commit(self): + """Finalize one client-delimited utterance while keeping the session open.""" + if self.endpoint_mode != "client": + raise RuntimeError("COMMIT requires endpoint_mode='client'") + + if self.vad.current_speech_start is not None: + end_ms = int(self.total_samples * 1000 / self.sample_rate) + seg = [self.vad.current_speech_start, end_ms] + seg_text = self._decode_segment(seg) + if seg_text.strip(): + sentence = { + "text": seg_text, + "start": int(seg[0]), + "end": int(seg[1]), + } + if self.spk_tracker: + s0 = int(seg[0] * self.sample_rate / 1000) + segment_audio = self._slice_audio(s0, self.total_samples).copy() + self.spk_tracker.assign_streaming( + segment_audio, + seg[0] / 1000, + seg[1] / 1000, + sentence, + ) + self.locked_sentences.append(sentence) + self.vad.current_speech_start = None + + if self.spk_tracker and self.locked_sentences: + self.locked_sentences = self.spk_tracker.finalize(self.locked_sentences) + + response = self._build_response(is_final=True) + self._reset_utterance_state() + return response + + def _reset_utterance_state(self): + self._release_audio_buffer() + self.first_decode_done = False + self.vad.reset() + self.prev_text = "" + self.last_partial_text = "" + self.last_partial_start_ms = 0 + self.last_decode_samples = self.total_samples + self.locked_sentences = [] + self.prev_seg_text = "" + def get_partial_decode_audio(self): """Return the bounded audio window used for unstable partial decoding.""" seg_start_sample = int(self.vad.current_speech_start * self.sample_rate / 1000) @@ -564,8 +649,14 @@ def load_models(args): _asr_kwargs["language"] = args.language logger.info(f"Language: {args.language}") - logger.info("Loading VAD: fsmn-vad (streaming)") - _vad_model = AutoModel(model="fsmn-vad", device=args.device, disable_update=True) + if getattr(args, "endpoint_mode", "server") == "server": + logger.info("Loading VAD: fsmn-vad (streaming)") + _vad_model = AutoModel( + model="fsmn-vad", device=args.device, disable_update=True + ) + else: + _vad_model = None + logger.info("Server VAD disabled; client COMMIT controls utterance endpoints") if getattr(args, "enable_spk", False): logger.info(f"Loading SPK: {args.spk_model}") @@ -597,7 +688,11 @@ async def run_session_work(args, operation, *operation_args, **operation_kwargs) async def handle_client(websocket, args): vllm_engine, asr_kwargs, vad_model, spk_model = load_models(args) - vad = DynamicStreamingVAD(vad_model) + endpoint_mode = getattr(args, "endpoint_mode", "server") + if endpoint_mode == "client": + vad = ClientEndpointVAD() + else: + vad = DynamicStreamingVAD(vad_model) spk_tracker = create_speaker_tracker(spk_model, args) session = RealtimeASRSession( vllm_engine, @@ -605,6 +700,7 @@ async def handle_client(websocket, args): vad, spk_tracker=spk_tracker, partial_window_sec=getattr(args, 'partial_window_sec', 15.0), + endpoint_mode=endpoint_mode, ) logger.info(f"Client connected: {websocket.remote_address}") @@ -633,12 +729,55 @@ async def handle_client(websocket, args): session.asr_kwargs["language"] = lang if lang else None await websocket.send(json.dumps({"event": "language_set", "language": lang})) logger.info(f"Language set: {lang}") - elif cmd.upper() == "STOP": - if session.is_active and session.total_samples > 0: - result = await run_session_work(args, session.decode, is_final=True) + elif cmd.upper() == "COMMIT": + if endpoint_mode != "client": + await websocket.send( + json.dumps( + { + "event": "error", + "error": "COMMIT requires --endpoint-mode client", + } + ) + ) + elif not session.is_active: + await websocket.send( + json.dumps( + { + "event": "error", + "error": "Session is not active; send START first", + } + ) + ) + else: + commit_started = time.perf_counter() + result = await run_session_work(args, session.commit) await websocket.send(json.dumps(result)) - logger.info(f"Final: {len(result['sentences'])} sentences") - session.is_active = False + last_decode_time = time.time() + elapsed_ms = (time.perf_counter() - commit_started) * 1000 + logger.info( + "Commit final: %d sentences in %.1fms", + len(result.get("sentences", [])), + elapsed_ms, + ) + elif cmd.upper() == "STOP": + if endpoint_mode == "client": + has_pending_audio = ( + session.total_samples > session.audio_buffer_start_sample + ) + else: + has_pending_audio = session.total_samples > 0 + if session.is_active and has_pending_audio: + if endpoint_mode == "client": + result = await run_session_work(args, session.commit) + else: + result = await run_session_work( + args, session.decode, is_final=True + ) + await websocket.send(json.dumps(result)) + logger.info( + "Final: %d sentences", len(result.get("sentences", [])) + ) + session.is_active = False await websocket.send(json.dumps({"event": "stopped"})) elif isinstance(message, bytes) and session.is_active: await run_session_work(args, session.add_audio, message) @@ -693,6 +832,15 @@ def build_arg_parser(): parser.add_argument("--use-context", action="store_true", default=True) parser.add_argument("--no-context", dest="use_context", action="store_false") parser.add_argument("--decode-interval", type=float, default=0.48) + parser.add_argument( + "--endpoint-mode", + choices=["server", "client"], + default="server", + help=( + "Use server VAD endpoints (default), or accept client COMMIT messages " + "without loading the VAD model." + ), + ) parser.add_argument("--partial-window-sec", type=float, default=15.0, help="Cap the interim partial re-decode window to the most recent N seconds. " "A long ongoing speech segment is otherwise re-encoded from its start on " diff --git a/tests/test_realtime_ws_service.py b/tests/test_realtime_ws_service.py index 00db6bafd..c7da0b994 100644 --- a/tests/test_realtime_ws_service.py +++ b/tests/test_realtime_ws_service.py @@ -1,5 +1,6 @@ import asyncio import importlib.util +import json import sys import threading import time @@ -38,6 +39,15 @@ def test_cli_defaults_disable_speaker_and_bound_partial_window(): assert args.enable_spk is False assert args.partial_window_sec == 15.0 + assert args.endpoint_mode == "server" + + +def test_cli_accepts_client_endpoint_mode(): + module = load_service_module() + + args = module.build_arg_parser().parse_args(["--endpoint-mode", "client"]) + + assert args.endpoint_mode == "client" def test_cli_accepts_float32_dtype_alias(): @@ -214,6 +224,197 @@ class DummyEngine: return [{"text": "hello"}] +def test_client_endpoint_mode_skips_fsmn_vad_model(monkeypatch): + module = load_service_module() + auto_model_calls = [] + + import funasr + + monkeypatch.setattr( + funasr, + "AutoModel", + lambda model, **kwargs: auto_model_calls.append(model) or object(), + ) + vllm_stub = types.ModuleType("funasr.auto.auto_model_vllm") + vllm_stub.AutoModelVLLM = lambda **kwargs: object() + monkeypatch.setitem(sys.modules, "funasr.auto.auto_model_vllm", vllm_stub) + + args = types.SimpleNamespace( + model="FunAudioLLM/Fun-ASR-Nano-2512", + hub="ms", + device="cpu", + dtype="fp32", + tensor_parallel_size=1, + gpu_memory_utilization=0.8, + max_model_len=2048, + hotword_file="", + language=None, + enable_spk=False, + endpoint_mode="client", + ) + + _, _, vad_model, _ = module.load_models(args) + + assert vad_model is None + assert auto_model_calls == [] + + +def make_client_endpoint_session(module): + assert hasattr(module, "ClientEndpointVAD") + return module.RealtimeASRSession( + vllm_engine=DummyEngine(), + asr_kwargs={}, + vad=module.ClientEndpointVAD(), + sample_rate=16000, + chunk_ms=960, + endpoint_mode="client", + ) + + +def test_client_endpoint_mode_emits_partial_at_first_decode_threshold(): + module = load_service_module() + session = make_client_endpoint_session(module) + audio = np.zeros(int(0.48 * session.sample_rate), dtype=np.int16).tobytes() + + session.add_audio(audio) + response = session.decode(is_final=False) + + assert session.should_decode() is False + assert response["partial"] == "hello" + assert response["partial_start_ms"] == 0 + assert response["is_final"] is False + assert session.vllm_engine.input_lengths == [int(0.48 * session.sample_rate)] + + +def test_client_commits_short_utterances_once_with_monotonic_timestamps(): + module = load_service_module() + session = make_client_endpoint_session(module) + + session.add_audio(np.zeros(int(0.4 * session.sample_rate), dtype=np.int16).tobytes()) + first = session.commit() + session.add_audio(np.zeros(int(0.3 * session.sample_rate), dtype=np.int16).tobytes()) + second = session.commit() + + assert first == { + "sentences": [{"text": "hello", "start": 0, "end": 400}], + "partial": "", + "partial_start_ms": 0, + "duration_ms": 400, + "is_final": True, + } + assert second == { + "sentences": [{"text": "hello", "start": 400, "end": 700}], + "partial": "", + "partial_start_ms": 0, + "duration_ms": 700, + "is_final": True, + } + assert session.vllm_engine.input_lengths == [6400, 4800] + assert len(session.audio_buffer) == 0 + assert session.audio_buffer_start_sample == session.total_samples + + +def test_handler_commit_returns_one_final_and_reuses_connection(monkeypatch): + module = load_service_module() + dynamic_vad_calls = [] + client_vad_calls = [] + received_audio = [] + + class ProtocolSession: + def __init__(self, *args, **kwargs): + self.is_active = False + self.total_samples = 0 + self.audio_buffer_start_sample = 0 + self.audio_buffer = np.array([], dtype=np.float32) + self.commit_count = 0 + + def reset(self): + self.total_samples = 0 + self.audio_buffer_start_sample = 0 + self.audio_buffer = np.array([], dtype=np.float32) + + def add_audio(self, message): + received_audio.append(message) + self.total_samples += 1 + self.audio_buffer = np.ones(1, dtype=np.float32) + + def should_decode(self): + return False + + def commit(self): + self.commit_count += 1 + self.audio_buffer = np.array([], dtype=np.float32) + self.audio_buffer_start_sample = self.total_samples + return {"is_final": True, "sequence": self.commit_count} + + def decode(self, is_final=False): + return self.commit() + + class FakeWebSocket: + remote_address = ("127.0.0.1", 12345) + + def __init__(self): + self.messages = iter( + [ + "START", + b"first", + "COMMIT", + b"second", + "COMMIT", + "STOP", + b"after-stop", + ] + ) + self.sent = [] + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.messages) + except StopIteration as error: + raise StopAsyncIteration from error + + async def send(self, message): + self.sent.append(json.loads(message)) + + monkeypatch.setattr(module, "load_models", lambda args: (object(), {}, None, None)) + monkeypatch.setattr( + module, + "DynamicStreamingVAD", + lambda model: dynamic_vad_calls.append(model) or object(), + ) + monkeypatch.setattr( + module, + "ClientEndpointVAD", + lambda: client_vad_calls.append(True) or object(), + raising=False, + ) + monkeypatch.setattr(module, "create_speaker_tracker", lambda model, args: None) + monkeypatch.setattr(module, "RealtimeASRSession", ProtocolSession) + websocket = FakeWebSocket() + args = types.SimpleNamespace( + device="cpu", + decode_interval=0.48, + partial_window_sec=15.0, + endpoint_mode="client", + ) + + asyncio.run(module.handle_client(websocket, args)) + + finals = [message for message in websocket.sent if message.get("is_final")] + events = [message.get("event") for message in websocket.sent if message.get("event")] + assert finals == [ + {"is_final": True, "sequence": 1}, + {"is_final": True, "sequence": 2}, + ] + assert events == ["started", "stopped"] + assert dynamic_vad_calls == [] + assert client_vad_calls == [True] + assert received_audio == [b"first", b"second"] + + def test_two_hour_session_keeps_audio_bounded_and_duration_absolute(): module = load_service_module() sample_rate = 10