from __future__ import annotations import os import logging import asyncio import json import time from dataclasses import dataclass, field import numpy as np from aiortc import RTCPeerConnection, RTCSessionDescription from dotenv import load_dotenv from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from pydantic import BaseModel from config import settings from core.pipeline import ChatPipeline from core.state_machine import Arbitrator from services.asr import ASRService from services.avatar import AvatarService from services.llm import LLMService from services.tts import TTSService from services.vad import VADService from webrtc.tracks import AudioBus, AvatarAudioTrack, AvatarVideoTrack load_dotenv() if settings.hf_token: os.environ["HF_TOKEN"] = settings.hf_token logging.basicConfig(level=logging.INFO) app = FastAPI(title="Visual Voice Chat") app.mount("/web", StaticFiles(directory="web"), name="web") app.add_middleware( CORSMiddleware, allow_origins=[o.strip() for o in settings.cors_origins.split(",")] if settings.cors_origins else ["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) arbitrator = Arbitrator() asr_service = ASRService() llm_service = LLMService() tts_service = TTSService() vad_service = VADService() avatar_service = AvatarService() pipeline = ChatPipeline(arbitrator, asr_service, llm_service, tts_service) audio_bus = AudioBus(sample_rate=48000) pcs: set[RTCPeerConnection] = set() subtitle_clients: set[WebSocket] = set() @dataclass class RuntimeState: last_user_text: str = "" last_reply_text: str = "" pipeline_runs: int = 0 last_latency_ms: int = 0 last_asr_latency_ms: int = 0 last_llm_latency_ms: int = 0 last_llm_source: str = "unknown" last_tts_latency_ms: int = 0 last_tts_first_chunk_ms: int = 0 last_barge_in_ms: int = 0 barge_in_count: int = 0 last_input_mode: str = "none" vad_start_count: int = 0 vad_end_count: int = 0 busy: bool = False buffer_16k: list[np.ndarray] = field(default_factory=list) runtime = RuntimeState() class TextChatRequest(BaseModel): text: str async def _broadcast_subtitle( role: str, text: str, source: str, *, partial: bool = False, final: bool = True, ) -> None: if not text: return payload = { "role": role, "text": text, "source": source, "partial": partial, "final": final, "ts_ms": int(time.time() * 1000), } stale: list[WebSocket] = [] for ws in list(subtitle_clients): try: await ws.send_text(json.dumps(payload, ensure_ascii=False)) except Exception: stale.append(ws) for ws in stale: subtitle_clients.discard(ws) def _resample_mono(audio: np.ndarray, from_sr: int, to_sr: int) -> np.ndarray: if from_sr == to_sr: return audio.astype(np.float32, copy=False) in_len = audio.shape[0] out_len = max(1, int(in_len * to_sr / from_sr)) x_old = np.linspace(0.0, 1.0, in_len, endpoint=False) x_new = np.linspace(0.0, 1.0, out_len, endpoint=False) return np.interp(x_new, x_old, audio).astype(np.float32) def _resample_to_16k_mono(pcm: np.ndarray, sample_rate: int, channels: int) -> np.ndarray: # aiortc audio ndarray is typically (channels, samples) if pcm.ndim == 2: mono = pcm.mean(axis=0) else: mono = pcm mono = mono.astype(np.float32) / 32768.0 if sample_rate == 16000: return mono in_len = mono.shape[0] out_len = max(1, int(in_len * 16000 / sample_rate)) x_old = np.linspace(0.0, 1.0, in_len, endpoint=False) x_new = np.linspace(0.0, 1.0, out_len, endpoint=False) return np.interp(x_new, x_old, mono).astype(np.float32) async def _run_pipeline_from_buffer() -> None: if runtime.busy: return if not runtime.buffer_16k: return runtime.busy = True start_t = time.perf_counter() try: runtime.last_input_mode = "voice" audio_16k = np.concatenate(runtime.buffer_16k, axis=0) runtime.buffer_16k.clear() async def on_audio_chunk(seg_audio: np.ndarray, seg_sr: int) -> None: audio_48k_seg = _resample_mono(seg_audio.astype(np.float32, copy=False), seg_sr, 48000) pcm_seg = np.clip(audio_48k_seg * 32767.0, -32768, 32767).astype(np.int16) await audio_bus.enqueue(pcm_seg) async def on_user_text_now(user_text: str) -> None: runtime.last_user_text = user_text await _broadcast_subtitle("user", user_text, "voice") async def on_reply_text_now(reply_text: str) -> None: runtime.last_reply_text = reply_text async def on_reply_segment_now(seg_text: str, idx: int, total: int) -> None: await _broadcast_subtitle("ai", seg_text, "voice", partial=True, final=(idx >= total - 1)) result, audio, sr = await pipeline.process_turn( audio_16k, on_audio_chunk=on_audio_chunk, on_user_text=on_user_text_now, on_reply_text=on_reply_text_now, on_reply_segment=on_reply_segment_now, ) asyncio.create_task(avatar_service.enqueue_musetalk_from_audio(audio, sr)) runtime.last_user_text = result["user_text"] runtime.last_reply_text = result["reply_text"] runtime.pipeline_runs += 1 runtime.last_latency_ms = int((time.perf_counter() - start_t) * 1000) runtime.last_asr_latency_ms = int(result.get("asr_latency_ms", 0)) runtime.last_llm_latency_ms = int(result.get("llm_latency_ms", 0)) runtime.last_llm_source = str(result.get("llm_source", "unknown")) runtime.last_tts_latency_ms = int(result.get("tts_latency_ms", 0)) runtime.last_tts_first_chunk_ms = int(result.get("tts_first_chunk_ms", 0)) # Fallback: if chunk callback path produced nothing, enqueue full audio. if result.get("audio_samples", 0) <= 1: audio_48k = _resample_mono(audio.astype(np.float32, copy=False), sr, 48000) pcm_int16 = np.clip(audio_48k * 32767.0, -32768, 32767).astype(np.int16) await audio_bus.enqueue(pcm_int16) finally: runtime.busy = False async def _consume_user_audio(track) -> None: vad_buffer = np.zeros(0, dtype=np.float32) while True: frame = await track.recv() pcm = frame.to_ndarray() sample_rate = getattr(frame, "sample_rate", 48000) or 48000 layout = getattr(frame, "layout", None) channels = len(layout.channels) if layout and layout.channels else 1 mono_16k = _resample_to_16k_mono(pcm, sample_rate, channels) runtime.buffer_16k.append(mono_16k) vad_buffer = np.concatenate([vad_buffer, mono_16k], axis=0) # Silero VAD requires exact 512 samples at 16k. while vad_buffer.shape[0] >= 512: chunk = vad_buffer[:512] vad_buffer = vad_buffer[512:] try: event = vad_service.push_chunk(chunk) except Exception as exc: logging.exception("VAD chunk processing failed: %s", exc) event = None if event and "start" in event: barge_t0 = time.perf_counter() was_avatar_speaking = str(arbitrator.state) == "avatar_speaking" await arbitrator.on_speech_start() await audio_bus.clear() if was_avatar_speaking: runtime.last_barge_in_ms = int((time.perf_counter() - barge_t0) * 1000) runtime.barge_in_count += 1 runtime.vad_start_count += 1 if event and "end" in event: await arbitrator.on_speech_end() runtime.vad_end_count += 1 asyncio.create_task(_run_pipeline_from_buffer()) @app.get("/health") async def health(): return { "state": arbitrator.state, "peers": len(pcs), "pipeline_runs": runtime.pipeline_runs, "last_latency_ms": runtime.last_latency_ms, "last_asr_latency_ms": runtime.last_asr_latency_ms, "last_llm_latency_ms": runtime.last_llm_latency_ms, "last_llm_source": runtime.last_llm_source, "last_tts_latency_ms": runtime.last_tts_latency_ms, "last_tts_first_chunk_ms": runtime.last_tts_first_chunk_ms, "last_barge_in_ms": runtime.last_barge_in_ms, "barge_in_count": runtime.barge_in_count, "last_input_mode": runtime.last_input_mode, "vad_start_count": runtime.vad_start_count, "vad_end_count": runtime.vad_end_count, "last_user_text": runtime.last_user_text, "last_reply_text": runtime.last_reply_text, "pipeline_busy": runtime.busy, "llm": llm_service.health, "asr": asr_service.health, "tts": tts_service.health, "vad": vad_service.health, "avatar": avatar_service.health, } @app.get("/meta") async def meta(): return { "stun_url": settings.stun_url, "https_enabled": bool(settings.ssl_certfile and settings.ssl_keyfile), "host": settings.webrtc_host, "port": settings.webrtc_port, } @app.get("/events") async def events(): async def gen(): while True: payload = { "state": str(arbitrator.state), "peers": len(pcs), "pipeline_runs": runtime.pipeline_runs, "last_latency_ms": runtime.last_latency_ms, "last_asr_latency_ms": runtime.last_asr_latency_ms, "last_llm_latency_ms": runtime.last_llm_latency_ms, "last_llm_source": runtime.last_llm_source, "last_tts_latency_ms": runtime.last_tts_latency_ms, "last_tts_first_chunk_ms": runtime.last_tts_first_chunk_ms, "last_barge_in_ms": runtime.last_barge_in_ms, "barge_in_count": runtime.barge_in_count, "last_input_mode": runtime.last_input_mode, "vad_start_count": runtime.vad_start_count, "vad_end_count": runtime.vad_end_count, "last_user_text": runtime.last_user_text, "last_reply_text": runtime.last_reply_text, "pipeline_busy": runtime.busy, "avatar": avatar_service.health, } yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" await asyncio.sleep(1.0) return StreamingResponse(gen(), media_type="text/event-stream") @app.websocket("/ws/subtitles") async def ws_subtitles(websocket: WebSocket): await websocket.accept() subtitle_clients.add(websocket) try: while True: await websocket.send_text(json.dumps({"type": "ping"}, ensure_ascii=False)) await asyncio.sleep(15) except WebSocketDisconnect: subtitle_clients.discard(websocket) except Exception: subtitle_clients.discard(websocket) @app.post("/demo/run-once") async def run_once(): # 3 seconds of silence placeholder; replace with real microphone stream path. fake_audio = np.zeros(16000 * 3, dtype=np.float32) await arbitrator.on_speech_start() result = await pipeline.run_once(fake_audio) return result @app.post("/chat/text") async def chat_text(req: TextChatRequest): text = req.text.strip() if not text: return JSONResponse({"ok": False, "message": "text is empty"}, status_code=400) if runtime.busy: return JSONResponse({"ok": False, "message": "pipeline busy"}, status_code=429) runtime.busy = True start_t = time.perf_counter() try: async def on_audio_chunk(seg_audio: np.ndarray, seg_sr: int) -> None: audio_48k_seg = _resample_mono(seg_audio.astype(np.float32, copy=False), seg_sr, 48000) pcm_seg = np.clip(audio_48k_seg * 32767.0, -32768, 32767).astype(np.int16) await audio_bus.enqueue(pcm_seg) async def on_reply_text_now(reply_text: str) -> None: runtime.last_reply_text = reply_text async def on_reply_segment_now(seg_text: str, idx: int, total: int) -> None: await _broadcast_subtitle("ai", seg_text, "text", partial=True, final=(idx >= total - 1)) await _broadcast_subtitle("user", text, "text") result, audio, sr = await pipeline.process_text_turn( text, on_audio_chunk=on_audio_chunk, on_reply_text=on_reply_text_now, on_reply_segment=on_reply_segment_now, ) asyncio.create_task(avatar_service.enqueue_musetalk_from_audio(audio, sr)) runtime.last_input_mode = "text" runtime.last_user_text = result["user_text"] runtime.last_reply_text = result["reply_text"] runtime.pipeline_runs += 1 runtime.last_latency_ms = int((time.perf_counter() - start_t) * 1000) runtime.last_asr_latency_ms = 0 runtime.last_llm_latency_ms = int(result.get("llm_latency_ms", 0)) runtime.last_llm_source = str(result.get("llm_source", "unknown")) runtime.last_tts_latency_ms = int(result.get("tts_latency_ms", 0)) runtime.last_tts_first_chunk_ms = int(result.get("tts_first_chunk_ms", 0)) # Fallback: if no chunk was enqueued for some reason, enqueue full audio. if result.get("audio_samples", 0) <= 1: audio_48k = _resample_mono(audio.astype(np.float32, copy=False), sr, 48000) pcm_int16 = np.clip(audio_48k * 32767.0, -32768, 32767).astype(np.int16) await audio_bus.enqueue(pcm_int16) return {"ok": True, **result} finally: runtime.busy = False @app.post("/chat/reset") async def chat_reset(): llm_service.clear_history() runtime.last_user_text = "" runtime.last_reply_text = "" runtime.pipeline_runs = 0 runtime.last_latency_ms = 0 runtime.last_asr_latency_ms = 0 runtime.last_llm_latency_ms = 0 runtime.last_llm_source = "unknown" runtime.last_tts_latency_ms = 0 runtime.last_tts_first_chunk_ms = 0 runtime.last_barge_in_ms = 0 runtime.barge_in_count = 0 runtime.last_input_mode = "none" runtime.vad_start_count = 0 runtime.vad_end_count = 0 return {"ok": True} @app.get("/demo/frame-idle-shape") async def frame_idle_shape(): frame = await avatar_service.idle_frame() return {"shape": list(frame.shape)} @app.get("/") async def root(): return FileResponse("web/index.html") @app.post("/webrtc/offer") async def webrtc_offer(request: Request): params = await request.json() offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"]) # Single-connection mode: kick existing peers before accepting new one. replaced_previous = len(pcs) > 0 if replaced_previous: old_peers = list(pcs) await asyncio.gather(*(peer.close() for peer in old_peers), return_exceptions=True) for peer in old_peers: pcs.discard(peer) pc = RTCPeerConnection() pcs.add(pc) @pc.on("connectionstatechange") async def on_connectionstatechange(): if pc.connectionState in {"failed", "closed", "disconnected"}: await pc.close() pcs.discard(pc) @pc.on("track") def on_track(track): if track.kind == "audio": asyncio.create_task(_consume_user_audio(track)) pc.addTrack(AvatarVideoTrack(avatar_service, arbitrator, audio_bus, settings.avatar_fps)) pc.addTrack(AvatarAudioTrack(audio_bus=audio_bus)) await pc.setRemoteDescription(offer) answer = await pc.createAnswer() await pc.setLocalDescription(answer) return JSONResponse( { "sdp": pc.localDescription.sdp, "type": pc.localDescription.type, "replaced_previous": replaced_previous, } ) @app.on_event("shutdown") async def on_shutdown(): await asyncio.gather(*(pc.close() for pc in list(pcs)), return_exceptions=True) pcs.clear() if __name__ == "__main__": import uvicorn uvicorn_kwargs = { "app": app, "host": settings.webrtc_host, "port": settings.webrtc_port, } if settings.ssl_certfile and settings.ssl_keyfile: uvicorn_kwargs["ssl_certfile"] = settings.ssl_certfile uvicorn_kwargs["ssl_keyfile"] = settings.ssl_keyfile uvicorn.run(**uvicorn_kwargs)