save code

This commit is contained in:
xsl
2026-03-29 23:01:39 +08:00
parent 36e838caec
commit 4dcae11c1b
30 changed files with 1094 additions and 203 deletions
+159 -36
View File
@@ -1,10 +1,12 @@
from __future__ import annotations
import asyncio
import io
import json
import logging
import os
import time
import wave
from dataclasses import dataclass, field
import numpy as np
@@ -52,7 +54,10 @@ audio_bus = AudioBus(sample_rate=48000)
pcs: set[RTCPeerConnection] = set()
subtitle_clients: set[WebSocket] = set()
animation_clients: set[WebSocket] = set()
audio_clients: set[WebSocket] = set()
MAX_BUFFERED_AUDIO_SAMPLES = 16000 * 45
MIN_VOICE_TURN_SAMPLES = 16000 // 2
MIN_VOICE_RMS = 0.012
@dataclass
@@ -131,6 +136,7 @@ def _runtime_payload() -> dict:
return {
"state": str(arbitrator.state),
"peers": len(pcs),
"audio_clients": len(audio_clients),
"subtitle_clients": len(subtitle_clients),
"animation_clients": len(animation_clients),
"pipeline_runs": runtime.pipeline_runs,
@@ -197,6 +203,42 @@ 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)
def _encode_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
pcm = np.clip(audio.astype(np.float32, copy=False) * 32767.0, -32768, 32767).astype(np.int16)
with io.BytesIO() as buffer:
with wave.open(buffer, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(sample_rate)
wav_file.writeframes(pcm.tobytes())
return buffer.getvalue()
async def _broadcast_audio_chunk(audio: np.ndarray, sample_rate: int) -> None:
if not audio_clients:
return
stale: list[WebSocket] = []
payload = _encode_wav_bytes(audio, sample_rate)
for ws in list(audio_clients):
try:
await ws.send_bytes(payload)
except Exception:
stale.append(ws)
for ws in stale:
audio_clients.discard(ws)
async def _broadcast_audio_reset(reason: str) -> None:
await _broadcast_json(
audio_clients,
{
"type": "audio_reset",
"reason": reason,
"ts_ms": int(time.time() * 1000),
},
)
async def _enqueue_audio_and_animation(
seg_audio: np.ndarray,
seg_sr: int,
@@ -204,9 +246,7 @@ async def _enqueue_audio_and_animation(
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)
await _broadcast_audio_chunk(seg_audio, seg_sr)
animation_payload = await avatar_service.build_controls_from_audio(
seg_audio,
seg_sr,
@@ -225,10 +265,18 @@ async def _run_pipeline_from_buffer() -> None:
runtime.pending_voice_turn = False
start_t = time.perf_counter()
try:
runtime.last_input_mode = "voice"
audio_16k = np.concatenate(runtime.buffer_16k, axis=0)
_clear_buffered_audio()
if audio_16k.shape[0] < MIN_VOICE_TURN_SAMPLES:
return
rms = float(np.sqrt(np.mean(audio_16k.astype(np.float32) ** 2))) if audio_16k.size else 0.0
if rms < MIN_VOICE_RMS:
return
runtime.last_input_mode = "voice"
async def on_user_text_now(user_text: str) -> None:
runtime.last_user_text = user_text
await _broadcast_subtitle("user", user_text, "voice")
@@ -246,6 +294,8 @@ async def _run_pipeline_from_buffer() -> None:
on_reply_text=on_reply_text_now,
on_reply_segment=on_reply_segment_now,
)
if not result["user_text"]:
return
runtime.last_user_text = result["user_text"]
runtime.last_reply_text = result["reply_text"]
runtime.pipeline_runs += 1
@@ -256,9 +306,7 @@ async def _run_pipeline_from_buffer() -> None:
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_audio_chunk(audio, sr)
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
finally:
runtime.busy = False
@@ -266,6 +314,38 @@ async def _run_pipeline_from_buffer() -> None:
asyncio.create_task(_run_pipeline_from_buffer())
async def _process_user_audio_chunk(mono_16k: np.ndarray, vad_buffer: np.ndarray) -> np.ndarray:
_append_buffered_audio(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 = arbitrator.state == SessionState.AVATAR_SPEAKING
await arbitrator.on_speech_start()
await audio_bus.clear()
await _broadcast_audio_reset("speech-start")
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
runtime.pending_voice_turn = True
asyncio.create_task(_run_pipeline_from_buffer())
return vad_buffer
async def _consume_user_audio(track) -> None:
vad_buffer = np.zeros(0, dtype=np.float32)
while True:
@@ -275,32 +355,7 @@ async def _consume_user_audio(track) -> None:
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)
_append_buffered_audio(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 = arbitrator.state == SessionState.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
runtime.pending_voice_turn = True
asyncio.create_task(_run_pipeline_from_buffer())
vad_buffer = await _process_user_audio_chunk(mono_16k, vad_buffer)
@app.get("/health")
@@ -364,6 +419,74 @@ async def ws_animation(websocket: WebSocket):
animation_clients.discard(websocket)
@app.websocket("/ws/audio")
async def ws_audio(websocket: WebSocket):
await websocket.accept()
replaced_previous = len(audio_clients) > 0
if replaced_previous:
old_clients = list(audio_clients)
await asyncio.gather(*(client.close(code=1012, reason="replaced") for client in old_clients), return_exceptions=True)
for client in old_clients:
audio_clients.discard(client)
audio_clients.add(websocket)
sample_rate = 48000
channels = 1
vad_buffer = np.zeros(0, dtype=np.float32)
await websocket.send_text(
json.dumps(
{
"type": "audio_ready",
"transport": "websocket_audio",
"replaced_previous": replaced_previous,
},
ensure_ascii=False,
)
)
try:
while True:
message = await websocket.receive()
message_type = message.get("type")
if message_type == "websocket.disconnect":
break
text = message.get("text")
if text is not None:
try:
payload = json.loads(text)
except json.JSONDecodeError:
payload = {}
event_type = str(payload.get("type", ""))
if event_type == "audio_config":
sample_rate = max(8000, int(payload.get("sample_rate", sample_rate) or sample_rate))
channels = max(1, int(payload.get("channels", channels) or channels))
elif event_type == "audio_reset":
_clear_buffered_audio()
runtime.pending_voice_turn = False
await audio_bus.clear()
await _broadcast_audio_reset("client-reset")
elif event_type == "ping":
await websocket.send_text(json.dumps({"type": "pong"}, ensure_ascii=False))
continue
data = message.get("bytes")
if not data:
continue
pcm = np.frombuffer(data, dtype=np.int16)
if pcm.size == 0:
continue
mono_16k = _resample_to_16k_mono(pcm, sample_rate, channels)
vad_buffer = await _process_user_audio_chunk(mono_16k, vad_buffer)
except WebSocketDisconnect:
pass
finally:
audio_clients.discard(websocket)
@app.post("/demo/run-once")
async def run_once():
fake_audio = np.zeros(16000 * 3, dtype=np.float32)
@@ -390,6 +513,7 @@ 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_audio_reset("text-turn-start")
await _broadcast_animation(await avatar_service.build_reset_payload(reason="text-turn-start"))
result, audio, sr = await pipeline.process_text_turn(
text,
@@ -408,9 +532,7 @@ async def chat_text(req: TextChatRequest):
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_audio_chunk(audio, sr)
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
return {"ok": True, **result}
finally:
@@ -438,6 +560,7 @@ async def chat_reset():
runtime.vad_end_count = 0
runtime.last_animation_frame_count = 0
await audio_bus.clear()
await _broadcast_audio_reset("chat-reset")
await _broadcast_animation(await avatar_service.build_reset_payload(reason="chat-reset"))
return {"ok": True}