|
|
@@ -0,0 +1,214 @@
|
|
|
+"""
|
|
|
+Full Pipeline Worker: VAD → ASR → LLM → TTS
|
|
|
+ VAD: energy-based voice activity detection (local)
|
|
|
+ VP: AEC/ANS/AGC handled by LiveKit WebRTC
|
|
|
+ ASR: Qwen3-ASR (local model)
|
|
|
+ LLM: vLLM (local, or Mimo fallback)
|
|
|
+ TTS: Mimo API (remote)
|
|
|
+"""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import asyncio
|
|
|
+import base64
|
|
|
+import json
|
|
|
+import logging
|
|
|
+import os
|
|
|
+import sys
|
|
|
+import time
|
|
|
+
|
|
|
+import aiohttp
|
|
|
+import jwt
|
|
|
+import numpy as np
|
|
|
+from livekit import rtc
|
|
|
+
|
|
|
+_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "whisper-qt-client", "whisper_asr")
|
|
|
+if not os.path.isdir(_root):
|
|
|
+ _root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "whisper_asr")
|
|
|
+sys.path.insert(0, _root)
|
|
|
+from audio_processor import AudioBuffer, VADProcessor
|
|
|
+from qwen_engine import QwenASREngine
|
|
|
+
|
|
|
+logger = logging.getLogger("worker")
|
|
|
+
|
|
|
+LIVEKIT_URL = os.environ.get("LIVEKIT_URL", "ws://localhost:7888")
|
|
|
+ASR_MODEL = os.environ.get("ASR_MODEL_PATH", "Qwen/Qwen3-ASR-0.6B")
|
|
|
+VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8000/v1")
|
|
|
+MIMO_KEY = os.environ.get("MIMO_KEY", "")
|
|
|
+
|
|
|
+# VAD params
|
|
|
+VAD_THRESHOLD = 0.015 # energy threshold
|
|
|
+MIN_SPEECH_S = 0.3 # minimum speech before considering utterance
|
|
|
+MIN_SILENCE_S = 1.2 # silence to end utterance
|
|
|
+
|
|
|
+
|
|
|
+def _token(room, ident):
|
|
|
+ n = int(time.time())
|
|
|
+ return jwt.encode({"iss": "devkey", "sub": ident, "name": ident, "nbf": n - 60, "exp": n + 6 * 3600,
|
|
|
+ "video": {"roomJoin": True, "room": room, "canPublish": True, "canSubscribe": True, "canPublishData": True}},
|
|
|
+ "secretsecretsecretsecretsecret12", algorithm="HS256")
|
|
|
+
|
|
|
+
|
|
|
+async def _llm(prompt: str, hist: list[dict]) -> str:
|
|
|
+ msgs = [{"role": "system", "content": "你是友好的中文语音助手,回答简洁自然,2-3句即可。"}]
|
|
|
+ msgs.extend(hist); msgs.append({"role": "user", "content": prompt})
|
|
|
+ async with aiohttp.ClientSession() as s:
|
|
|
+ async with s.post(f"{VLLM_URL}/chat/completions",
|
|
|
+ json={"model": "default", "messages": msgs, "max_tokens": 128, "temperature": 0.7},
|
|
|
+ timeout=aiohttp.ClientTimeout(total=15)) as r:
|
|
|
+ return (await r.json())["choices"][0]["message"]["content"]
|
|
|
+
|
|
|
+
|
|
|
+async def _tts(text: str) -> np.ndarray | None:
|
|
|
+ key = os.environ.get("MIMO_KEY", "")
|
|
|
+ if not key: return None
|
|
|
+ H = {"api-key": key, "Content-Type": "application/json"}
|
|
|
+ async with aiohttp.ClientSession() as s:
|
|
|
+ async with s.post("https://token-plan-cn.xiaomimimo.com/v1/chat/completions",
|
|
|
+ json={"model": "mimo-v2.5-tts", "messages": [
|
|
|
+ {"role": "user", "content": "请用自然语速朗读。"},
|
|
|
+ {"role": "assistant", "content": text}], "audio": {"format": "wav", "voice": "Chloe"}},
|
|
|
+ headers=H, timeout=aiohttp.ClientTimeout(total=20)) as r:
|
|
|
+ try:
|
|
|
+ b64 = (await r.json())["choices"][0]["message"]["audio"]["data"]
|
|
|
+ return np.frombuffer(base64.b64decode(b64), dtype=np.int16)
|
|
|
+ except: return None
|
|
|
+
|
|
|
+
|
|
|
+class Worker:
|
|
|
+ def __init__(self, room, identity="asr-bot"):
|
|
|
+ self.rn, self.id = room, identity
|
|
|
+ self.room = rtc.Room()
|
|
|
+ self.tasks: dict = {}
|
|
|
+ self.hist: list[dict] = []
|
|
|
+
|
|
|
+ async def run(self):
|
|
|
+ logger.info("Loading ASR model...")
|
|
|
+ self.asr = QwenASREngine(model_id=ASR_MODEL, language=None)
|
|
|
+ logger.info("ASR ready")
|
|
|
+
|
|
|
+ self.room.on("track_subscribed", self._on_track)
|
|
|
+ self.room.on("participant_disconnected", self._on_off)
|
|
|
+ self.room.on("participant_connected", lambda p: None)
|
|
|
+ self.room.on("track_published", lambda pub, p: None)
|
|
|
+ await self.room.connect(LIVEKIT_URL, _token(self.rn, self.id))
|
|
|
+ logger.info("Worker ready (VAD+ASR+LLM+TTS)")
|
|
|
+
|
|
|
+ for p in self.room.remote_participants.values():
|
|
|
+ for pub in p.track_publications.values():
|
|
|
+ if pub.track and pub.kind == rtc.TrackKind.KIND_AUDIO:
|
|
|
+ self._on_track(pub.track, pub, p)
|
|
|
+ try: await asyncio.Future()
|
|
|
+ finally: await self.room.disconnect()
|
|
|
+
|
|
|
+ def _on_track(self, t, _, p):
|
|
|
+ s = p.sid
|
|
|
+ if s in self.tasks: return
|
|
|
+ self.tasks[s] = asyncio.create_task(self._run(s, t))
|
|
|
+
|
|
|
+ def _on_off(self, p):
|
|
|
+ x = self.tasks.pop(p.sid, None)
|
|
|
+ if x: x.cancel()
|
|
|
+
|
|
|
+ async def _run(self, sid, track):
|
|
|
+ buf = AudioBuffer(max_duration=10, sample_rate=16000)
|
|
|
+ vad = VADProcessor(
|
|
|
+ threshold=VAD_THRESHOLD,
|
|
|
+ min_speech_duration=MIN_SPEECH_S,
|
|
|
+ min_silence_duration=MIN_SILENCE_S,
|
|
|
+ sample_rate=16000,
|
|
|
+ )
|
|
|
+ busy = False
|
|
|
+ cur_tts = None
|
|
|
+
|
|
|
+ async def _acc():
|
|
|
+ async for ev in rtc.AudioStream(track, sample_rate=16000, num_channels=1):
|
|
|
+ if not busy:
|
|
|
+ data = ev.frame.data
|
|
|
+ arr = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
|
|
|
+ buf.append(arr)
|
|
|
+
|
|
|
+ # Feed VAD with audio frames
|
|
|
+ in_speech = vad.in_speech
|
|
|
+ segments = vad.process(arr)
|
|
|
+ # Speech just ended
|
|
|
+ if not vad.in_speech and in_speech and len(buf.buffer) > 1600 * MIN_SPEECH_S:
|
|
|
+ nonlocal cur_tts
|
|
|
+ if cur_tts: cur_tts = None
|
|
|
+ busy = True
|
|
|
+ await self._transcribe_and_respond(buf)
|
|
|
+ busy = False
|
|
|
+
|
|
|
+ async def _loop():
|
|
|
+ # Fallback: periodic transcription when VAD detects speech
|
|
|
+ last_len = 0
|
|
|
+ while True:
|
|
|
+ await asyncio.sleep(1.5)
|
|
|
+ if busy: continue
|
|
|
+ if len(buf.buffer) <= 1600: continue
|
|
|
+ if not vad.in_speech: continue
|
|
|
+ # Only transcribe if buffer grew (new speech happened)
|
|
|
+ if len(buf.buffer) == last_len: continue
|
|
|
+ last_len = len(buf.buffer)
|
|
|
+
|
|
|
+ a = buf.get_last(4.0)
|
|
|
+ if len(a) <= 1600: continue
|
|
|
+ loop = asyncio.get_running_loop()
|
|
|
+ result = await loop.run_in_executor(None, self.asr.transcribe_array, a)
|
|
|
+ txt = result.get("text", "").strip()
|
|
|
+ if txt:
|
|
|
+ await self._send({"type": "partial", "text": txt, "seq": 0})
|
|
|
+
|
|
|
+ t = asyncio.create_task(_loop())
|
|
|
+ try: await _acc()
|
|
|
+ finally: t.cancel()
|
|
|
+
|
|
|
+ async def _transcribe_and_respond(self, buf: AudioBuffer):
|
|
|
+ all_audio = buf.get_all()
|
|
|
+ buf.clear()
|
|
|
+ if len(all_audio) <= 1600: return
|
|
|
+
|
|
|
+ loop = asyncio.get_running_loop()
|
|
|
+ result = await loop.run_in_executor(None, self.asr.transcribe_array, all_audio)
|
|
|
+ txt = result.get("text", "").strip()
|
|
|
+ if not txt: return
|
|
|
+
|
|
|
+ await self._send({"type": "utterance", "text": txt, "seq": 0})
|
|
|
+ await self._reply(txt)
|
|
|
+
|
|
|
+ async def _reply(self, text: str):
|
|
|
+ try: rep = await _llm(text, self.hist)
|
|
|
+ except: rep = "抱歉。"
|
|
|
+ await self._send({"type": "reply", "text": rep})
|
|
|
+ try:
|
|
|
+ a = await _tts(rep)
|
|
|
+ if a is not None and len(a) > 0:
|
|
|
+ await self._play(a, 24000)
|
|
|
+ except (asyncio.CancelledError, Exception):
|
|
|
+ pass
|
|
|
+ self.hist.extend([{"role": "user", "content": text}, {"role": "assistant", "content": rep}])
|
|
|
+ self.hist[:] = self.hist[-20:]
|
|
|
+
|
|
|
+ async def _play(self, audio: np.ndarray, sr: int):
|
|
|
+ src = rtc.AudioSource(sr, 1)
|
|
|
+ tk = rtc.LocalAudioTrack.create_audio_track("tts", src)
|
|
|
+ await self.room.local_participant.publish_track(tk)
|
|
|
+ for i in range(0, len(audio), sr // 50):
|
|
|
+ c = audio[i : i + sr // 50]
|
|
|
+ await src.capture_frame(rtc.AudioFrame(data=c.tobytes(), sample_rate=sr, num_channels=1, samples_per_channel=len(c)))
|
|
|
+ await asyncio.sleep(0.018)
|
|
|
+ return tk
|
|
|
+
|
|
|
+ async def _send(self, msg: dict):
|
|
|
+ try: await self.room.local_participant.publish_data(json.dumps(msg, ensure_ascii=False).encode(), reliable=True, topic="transcription")
|
|
|
+ except: pass
|
|
|
+
|
|
|
+
|
|
|
+async def main():
|
|
|
+ import argparse; p = argparse.ArgumentParser()
|
|
|
+ p.add_argument("--room", required=True); a = p.parse_args()
|
|
|
+ logging.basicConfig(level=logging.INFO)
|
|
|
+ await Worker(a.room).run()
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ asyncio.run(main())
|