491 lines
17 KiB
Python
491 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
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.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
|
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
|
|
|
|
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()
|
|
animation_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
|
|
last_animation_mode: str = settings.avatar_driver_mode
|
|
last_animation_frame_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_json(clients: set[WebSocket], payload: dict) -> None:
|
|
stale: list[WebSocket] = []
|
|
message = json.dumps(payload, ensure_ascii=False)
|
|
for ws in list(clients):
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
stale.append(ws)
|
|
for ws in stale:
|
|
clients.discard(ws)
|
|
|
|
|
|
async def _broadcast_subtitle(
|
|
role: str,
|
|
text: str,
|
|
source: str,
|
|
*,
|
|
partial: bool = False,
|
|
final: bool = True,
|
|
) -> None:
|
|
if not text:
|
|
return
|
|
await _broadcast_json(
|
|
subtitle_clients,
|
|
{
|
|
"role": role,
|
|
"text": text,
|
|
"source": source,
|
|
"partial": partial,
|
|
"final": final,
|
|
"ts_ms": int(time.time() * 1000),
|
|
},
|
|
)
|
|
|
|
|
|
async def _broadcast_animation(payload: dict) -> None:
|
|
runtime.last_animation_mode = str(payload.get("driver", settings.avatar_driver_mode))
|
|
runtime.last_animation_frame_count = int(payload.get("frame_count", 0))
|
|
await _broadcast_json(animation_clients, payload)
|
|
|
|
|
|
def _runtime_payload() -> dict:
|
|
return {
|
|
"state": str(arbitrator.state),
|
|
"peers": len(pcs),
|
|
"subtitle_clients": len(subtitle_clients),
|
|
"animation_clients": len(animation_clients),
|
|
"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,
|
|
"last_animation_mode": runtime.last_animation_mode,
|
|
"last_animation_frame_count": runtime.last_animation_frame_count,
|
|
"pipeline_busy": runtime.busy,
|
|
"llm": llm_service.health,
|
|
"asr": asr_service.health,
|
|
"tts": tts_service.health,
|
|
"vad": vad_service.health,
|
|
"avatar": avatar_service.health,
|
|
}
|
|
|
|
|
|
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:
|
|
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 _enqueue_audio_and_animation(
|
|
seg_audio: np.ndarray,
|
|
seg_sr: int,
|
|
seg_text: str,
|
|
idx: int,
|
|
total: 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)
|
|
animation_payload = await avatar_service.build_controls_from_audio(
|
|
seg_audio,
|
|
seg_sr,
|
|
text=seg_text,
|
|
chunk_index=idx,
|
|
total_chunks=total,
|
|
)
|
|
await _broadcast_animation(animation_payload)
|
|
|
|
|
|
async def _run_pipeline_from_buffer() -> None:
|
|
if runtime.busy or 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_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=_enqueue_audio_and_animation,
|
|
on_user_text=on_user_text_now,
|
|
on_reply_text=on_reply_text_now,
|
|
on_reply_segment=on_reply_segment_now,
|
|
)
|
|
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))
|
|
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)
|
|
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
|
|
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)
|
|
|
|
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()
|
|
await _broadcast_animation(await avatar_service.build_reset_payload(reason="speech-start"))
|
|
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 _runtime_payload()
|
|
|
|
|
|
@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,
|
|
"avatar_protocol": settings.avatar_control_protocol,
|
|
}
|
|
|
|
|
|
@app.get("/avatar/schema")
|
|
async def avatar_schema():
|
|
return avatar_service.schema
|
|
|
|
|
|
@app.get("/events")
|
|
async def events():
|
|
async def gen():
|
|
while True:
|
|
yield f"data: {json.dumps(_runtime_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.websocket("/ws/animation")
|
|
async def ws_animation(websocket: WebSocket):
|
|
await websocket.accept()
|
|
animation_clients.add(websocket)
|
|
try:
|
|
await websocket.send_text(json.dumps({"type": "animation_ready", **avatar_service.schema}, ensure_ascii=False))
|
|
await websocket.send_text(json.dumps(await avatar_service.build_idle_payload(reason="client-connected"), ensure_ascii=False))
|
|
while True:
|
|
await websocket.send_text(json.dumps({"type": "ping"}, ensure_ascii=False))
|
|
await asyncio.sleep(15)
|
|
except WebSocketDisconnect:
|
|
animation_clients.discard(websocket)
|
|
except Exception:
|
|
animation_clients.discard(websocket)
|
|
|
|
|
|
@app.post("/demo/run-once")
|
|
async def run_once():
|
|
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_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")
|
|
await _broadcast_animation(await avatar_service.build_reset_payload(reason="text-turn-start"))
|
|
result, audio, sr = await pipeline.process_text_turn(
|
|
text,
|
|
on_audio_chunk=_enqueue_audio_and_animation,
|
|
on_reply_text=on_reply_text_now,
|
|
on_reply_segment=on_reply_segment_now,
|
|
)
|
|
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))
|
|
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)
|
|
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
|
|
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
|
|
runtime.last_animation_frame_count = 0
|
|
await _broadcast_animation(await avatar_service.build_reset_payload(reason="chat-reset"))
|
|
return {"ok": True}
|
|
|
|
|
|
@app.get("/demo/frame-idle-shape")
|
|
async def frame_idle_shape():
|
|
return await avatar_service.build_idle_payload(reason="demo")
|
|
|
|
|
|
@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"])
|
|
|
|
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(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)
|