save code
This commit is contained in:
@@ -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}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user