大功告成

main.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-import os
-import logging
 import asyncio
 import json
+import logging
+import os
 import time
 from dataclasses import dataclass, field
 
@@ -11,8 +11,8 @@ 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.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 
@@ -24,7 +24,7 @@ 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
+from webrtc.tracks import AudioBus, AvatarAudioTrack
 
 load_dotenv()
 if settings.hf_token:
@@ -51,6 +51,7 @@ 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
@@ -69,6 +70,8 @@ class RuntimeState:
     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)
 
@@ -80,6 +83,18 @@ 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,
@@ -90,22 +105,54 @@ async def _broadcast_subtitle(
 ) -> None:
     if not text:
         return
-    payload = {
-        "role": role,
-        "text": text,
-        "source": source,
-        "partial": partial,
-        "final": final,
-        "ts_ms": int(time.time() * 1000),
+    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,
     }
-    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:
@@ -119,7 +166,6 @@ def _resample_mono(audio: np.ndarray, from_sr: int, to_sr: int) -> np.ndarray:
 
 
 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:
@@ -135,10 +181,28 @@ def _resample_to_16k_mono(pcm: np.ndarray, sample_rate: int, channels: int) -> n
     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:
-        return
-    if not runtime.buffer_16k:
+    if runtime.busy or not runtime.buffer_16k:
         return
 
     runtime.busy = True
@@ -147,10 +211,6 @@ async def _run_pipeline_from_buffer() -> None:
         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
@@ -164,12 +224,11 @@ async def _run_pipeline_from_buffer() -> None:
 
         result, audio, sr = await pipeline.process_turn(
             audio_16k,
-            on_audio_chunk=on_audio_chunk,
+            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,
         )
-        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
@@ -179,11 +238,11 @@ async def _run_pipeline_from_buffer() -> None:
         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)
+            await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
     finally:
         runtime.busy = False
 
@@ -200,7 +259,6 @@ async def _consume_user_audio(track) -> None:
         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:]
@@ -214,6 +272,7 @@ async def _consume_user_audio(track) -> None:
                 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
@@ -226,30 +285,7 @@ async def _consume_user_audio(track) -> None:
 
 @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,
-    }
+    return _runtime_payload()
 
 
 @app.get("/meta")
@@ -259,34 +295,20 @@ async def meta():
         "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:
-            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"
+            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")
@@ -306,9 +328,24 @@ async def ws_subtitles(websocket: WebSocket):
         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():
-    # 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)
@@ -326,11 +363,6 @@ async def chat_text(req: TextChatRequest):
     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
 
@@ -338,13 +370,13 @@ async def chat_text(req: TextChatRequest):
             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=on_audio_chunk,
+            on_audio_chunk=_enqueue_audio_and_animation,
             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"]
@@ -355,11 +387,11 @@ async def chat_text(req: TextChatRequest):
         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)
+            await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
         return {"ok": True, **result}
     finally:
         runtime.busy = False
@@ -382,13 +414,14 @@ async def chat_reset():
     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():
-    frame = await avatar_service.idle_frame()
-    return {"shape": list(frame.shape)}
+    return await avatar_service.build_idle_payload(reason="demo")
 
 
 @app.get("/")
@@ -401,7 +434,6 @@ 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)
@@ -423,7 +455,6 @@ async def webrtc_offer(request: Request):
         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)