828 lines
31 KiB
Python
828 lines
31 KiB
Python
import asyncio
|
|
import numpy as np
|
|
import time
|
|
import os
|
|
from typing import Dict, Optional, Callable, Awaitable, Any, Union
|
|
import numpy.typing as npt
|
|
|
|
|
|
# Core dependencies
|
|
import librosa
|
|
from shared.logger import logger
|
|
from aiortc import MediaStreamTrack
|
|
from aiortc.mediastreams import MediaStreamError
|
|
from av import AudioFrame, VideoFrame
|
|
import cv2
|
|
import fractions
|
|
from time import perf_counter
|
|
import wave
|
|
|
|
# Import shared models for chat functionality
|
|
import sys
|
|
|
|
sys.path.append(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
)
|
|
from shared.models import ChatMessageModel
|
|
from voicebot.models import Peer
|
|
|
|
from .vibevoicetts import text_to_speech
|
|
|
|
# Global configuration and constants
|
|
AGENT_NAME = "Text To Speech Bot"
|
|
AGENT_DESCRIPTION = (
|
|
"Real-time speech generation- converts text to speech on Intel Arc B580"
|
|
)
|
|
SAMPLE_RATE = 16000 # Whisper expects 16kHz
|
|
|
|
|
|
_audio_data = None
|
|
|
|
class TTSModel:
|
|
"""TTS via MSFT VibeVoice"""
|
|
|
|
def __init__(
|
|
self,
|
|
peer_name: str,
|
|
send_chat_func: Callable[[ChatMessageModel], Awaitable[None]],
|
|
create_chat_message_func: Callable[[str, Optional[str]], ChatMessageModel],
|
|
):
|
|
self.peer_name = peer_name
|
|
self.send_chat_func = send_chat_func
|
|
self.create_chat_message_func = create_chat_message_func
|
|
|
|
|
|
_send_chat_func: Optional[Callable[[ChatMessageModel], Awaitable[None]]] = None
|
|
_create_chat_message_func: Optional[
|
|
Callable[[str, Optional[str]], ChatMessageModel]
|
|
] = None
|
|
|
|
class MediaClock:
|
|
"""Simple monotonic clock for media tracks."""
|
|
|
|
def __init__(self) -> None:
|
|
self.t0 = perf_counter()
|
|
|
|
def now(self) -> float:
|
|
return perf_counter() - self.t0
|
|
|
|
|
|
class WaveformVideoTrack(MediaStreamTrack):
|
|
"""Video track that renders a live waveform of the incoming audio.
|
|
|
|
The track reads the most-active `OptimizedAudioProcessor` in
|
|
`_audio_processors` and renders the last ~2s of its `current_phrase_audio`.
|
|
If no audio is available, the track will display a "No audio" message.
|
|
"""
|
|
|
|
kind = "video"
|
|
|
|
# Shared buffer for audio data
|
|
buffer: Dict[str, npt.NDArray[np.float32]] = {}
|
|
speech_status: Dict[str, Dict[str, Any]] = {}
|
|
|
|
def __init__(
|
|
self, session_name: str, width: int = 640, height: int = 480, fps: int = 15
|
|
) -> None:
|
|
super().__init__()
|
|
self.session_name = session_name
|
|
self.width = int(width)
|
|
self.height = int(height)
|
|
self.fps = int(fps)
|
|
self.clock = MediaClock()
|
|
self._next_frame_index = 0
|
|
|
|
async def next_timestamp(self) -> tuple[int, float]:
|
|
pts = int(self._next_frame_index * (1 / self.fps) * 90000)
|
|
time_base = 1 / 90000
|
|
return pts, time_base
|
|
|
|
async def recv(self) -> VideoFrame:
|
|
pts, _ = await self.next_timestamp()
|
|
|
|
# schedule frame according to clock
|
|
target_t = self._next_frame_index / self.fps
|
|
now = self.clock.now()
|
|
if target_t > now:
|
|
await asyncio.sleep(target_t - now)
|
|
|
|
self._next_frame_index += 1
|
|
|
|
frame_array: npt.NDArray[np.uint8] = np.zeros(
|
|
(self.height, self.width, 3), dtype=np.uint8
|
|
)
|
|
|
|
# Display model loading status prominently
|
|
status_text = "Initializing..."
|
|
progress = 0.1
|
|
|
|
# Draw status background (increased height for larger text)
|
|
cv2.rectangle(frame_array, (0, 0), (self.width, 80), (0, 0, 0), -1)
|
|
|
|
# Draw progress bar if loading
|
|
if progress < 1.0 and "Ready" not in status_text:
|
|
bar_width = int(progress * (self.width - 40))
|
|
cv2.rectangle(frame_array, (20, 55), (20 + bar_width, 70), (0, 255, 0), -1)
|
|
cv2.rectangle(
|
|
frame_array, (20, 55), (self.width - 20, 70), (255, 255, 255), 2
|
|
)
|
|
|
|
# Draw status text (larger font)
|
|
cv2.putText(
|
|
frame_array,
|
|
f"Status: {status_text}",
|
|
(10, 35),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.2,
|
|
(255, 255, 255),
|
|
3,
|
|
)
|
|
|
|
# Draw clock in lower right corner, right justified
|
|
current_time = time.strftime("%H:%M:%S")
|
|
(text_width, _), _ = cv2.getTextSize(
|
|
current_time, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2
|
|
)
|
|
clock_x = self.width - text_width - 10 # 10px margin from right edge
|
|
clock_y = self.height - 30 # Move to 450 for height=480
|
|
cv2.putText(
|
|
frame_array,
|
|
current_time,
|
|
(clock_x, clock_y),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.0,
|
|
(255, 255, 255),
|
|
2,
|
|
)
|
|
|
|
# Select the most active audio buffer and get its speech status
|
|
best_proc = None
|
|
best_rms = 0.0
|
|
speech_info = None
|
|
|
|
try:
|
|
for pname, arr in self.__class__.buffer.items():
|
|
try:
|
|
if len(arr) == 0:
|
|
rms = 0.0
|
|
else:
|
|
rms = float(np.sqrt(np.mean(arr**2)))
|
|
if rms > best_rms:
|
|
best_rms = rms
|
|
best_proc = (pname, arr.copy())
|
|
speech_info = self.__class__.speech_status.get(pname, {})
|
|
except Exception:
|
|
continue
|
|
except Exception:
|
|
best_proc = None
|
|
|
|
if best_proc is not None:
|
|
pname, arr = best_proc
|
|
|
|
# Use the last 2 second of audio data, padded with zeros if less
|
|
samples_needed = SAMPLE_RATE * 2 # 2 second(s)
|
|
if len(arr) <= 0:
|
|
arr_segment = np.zeros(samples_needed, dtype=np.float32)
|
|
elif len(arr) >= samples_needed:
|
|
arr_segment = arr[-samples_needed:].copy()
|
|
else:
|
|
# Pad with zeros at the beginning
|
|
arr_segment = np.concatenate(
|
|
[np.zeros(samples_needed - len(arr), dtype=np.float32), arr]
|
|
)
|
|
|
|
# Single normalization code path: normalize based on the historical
|
|
# peak observed for this stream (proc.max_observed_amplitude). This
|
|
# ensures the waveform display is consistent over time and avoids
|
|
# using the instantaneous buffer peak.
|
|
proc = None
|
|
norm = arr_segment.astype(np.float32)
|
|
|
|
# Map audio samples to pixels across the width
|
|
if norm.size < self.width:
|
|
padded = np.zeros(self.width, dtype=np.float32)
|
|
if norm.size > 0:
|
|
padded[-norm.size :] = norm
|
|
norm = padded
|
|
else:
|
|
block = int(np.ceil(norm.size / self.width))
|
|
norm = np.array(
|
|
[
|
|
np.mean(norm[i * block : min((i + 1) * block, norm.size)])
|
|
for i in range(self.width)
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
# For display we use the same `norm` computed above (single code
|
|
# path). Use `display_norm` alias to avoid confusion later in the
|
|
# code but don't recompute normalization.
|
|
display_norm = norm
|
|
|
|
# Draw waveform with color coding for speech detection
|
|
points: list[tuple[int, int]] = []
|
|
colors: list[tuple[int, int, int]] = [] # Color for each point
|
|
|
|
for x in range(self.width):
|
|
v = (
|
|
float(display_norm[x])
|
|
if x < display_norm.size and not np.isnan(display_norm[x])
|
|
else 0.0
|
|
)
|
|
y = int((1.0 - ((v + 1.0) / 2.0)) * (self.height - 120)) + 100
|
|
points.append((x, y))
|
|
|
|
# Color based on speech detection status
|
|
is_speech = (
|
|
speech_info.get("is_speech", False) if speech_info else False
|
|
)
|
|
energy_check = (
|
|
speech_info.get("energy_check", False) if speech_info else False
|
|
)
|
|
|
|
if is_speech:
|
|
colors.append((0, 255, 0)) # Green for detected speech
|
|
elif energy_check:
|
|
colors.append((255, 255, 0)) # Yellow for energy but not speech
|
|
else:
|
|
colors.append((128, 128, 128)) # Gray for background noise
|
|
|
|
# Draw colored waveform
|
|
if len(points) > 1:
|
|
for i in range(len(points) - 1):
|
|
cv2.line(frame_array, points[i], points[i + 1], colors[i], 1)
|
|
|
|
# Draw historical peak indicator (horizontal lines at +/-(target_peak))
|
|
try:
|
|
if proc is not None and getattr(proc, "normalization_enabled", False):
|
|
target_peak = float(getattr(proc, "normalization_target_peak", 0.0))
|
|
# Ensure target_peak is within [0, 1]
|
|
target_peak = max(0.0, min(1.0, target_peak))
|
|
|
|
def _amp_to_y(a: float) -> int:
|
|
return (
|
|
int((1.0 - ((a + 1.0) / 2.0)) * (self.height - 120)) + 100
|
|
)
|
|
|
|
top_y = _amp_to_y(target_peak)
|
|
bot_y = _amp_to_y(-target_peak)
|
|
|
|
# Draw thin magenta lines across the waveform area
|
|
cv2.line(
|
|
frame_array,
|
|
(0, top_y),
|
|
(self.width - 1, top_y),
|
|
(255, 0, 255),
|
|
1,
|
|
)
|
|
cv2.line(
|
|
frame_array,
|
|
(0, bot_y),
|
|
(self.width - 1, bot_y),
|
|
(255, 0, 255),
|
|
1,
|
|
)
|
|
|
|
# Label the peak with small text near the right edge
|
|
label = f"Peak:{target_peak:.2f}"
|
|
(tw, _), _ = cv2.getTextSize(
|
|
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
|
|
)
|
|
lx = max(10, self.width - tw - 12)
|
|
ly = max(12, top_y - 6)
|
|
cv2.putText(
|
|
frame_array,
|
|
label,
|
|
(lx, ly),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.5,
|
|
(255, 0, 255),
|
|
1,
|
|
)
|
|
except Exception:
|
|
# Non-critical: ignore any drawing errors
|
|
pass
|
|
|
|
# Add speech detection status overlay
|
|
if speech_info:
|
|
self._draw_speech_status(frame_array, speech_info, pname)
|
|
|
|
cv2.putText(
|
|
frame_array,
|
|
f"Waveform: {pname}",
|
|
(10, self.height - 15),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.0,
|
|
(255, 255, 255),
|
|
2,
|
|
)
|
|
|
|
frame = VideoFrame.from_ndarray(frame_array, format="bgr24")
|
|
frame.pts = pts
|
|
frame.time_base = fractions.Fraction(1 / 90000).limit_denominator(1000000)
|
|
return frame
|
|
|
|
def _draw_speech_status(
|
|
self,
|
|
frame_array: npt.NDArray[np.uint8],
|
|
speech_info: Dict[str, Any],
|
|
pname: str,
|
|
):
|
|
"""Draw speech detection status information."""
|
|
|
|
y_offset = 100
|
|
|
|
# Main status
|
|
is_speech = speech_info.get("is_speech", False)
|
|
status_text = "SPEECH" if is_speech else "NOISE"
|
|
status_color = (0, 255, 0) if is_speech else (128, 128, 128)
|
|
|
|
adaptive_thresh = speech_info.get("adaptive_threshold", 0)
|
|
cv2.putText(
|
|
frame_array,
|
|
f"{pname}: {status_text} (thresh: {adaptive_thresh:.4f})",
|
|
(10, y_offset),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
status_color,
|
|
2,
|
|
)
|
|
|
|
# Detailed metrics (smaller text)
|
|
metrics = [
|
|
f"Energy: {speech_info.get('energy', 0):.3f} ({'Y' if speech_info.get('energy_check', False) else 'N'})",
|
|
f"ZCR: {speech_info.get('zcr', 0):.3f} ({'Y' if speech_info.get('zcr_check', False) else 'N'})",
|
|
f"Spectral: {'Y' if 300 < speech_info.get('centroid', 0) < 3400 else 'N'}/{'Y' if speech_info.get('rolloff', 0) < 2000 else 'N'}/{'Y' if speech_info.get('flux', 0) > 0.01 else 'N'} ({'Y' if speech_info.get('spectral_check', False) else 'N'})",
|
|
f"Harmonic: {speech_info.get('hamonicity', 0):.3f} ({'Y' if speech_info.get('harmonic_check', False) else 'N'})",
|
|
f"Temporal: ({'Y' if speech_info.get('temporal_consistency', False) else 'N'})",
|
|
]
|
|
|
|
for _, metric in enumerate(metrics):
|
|
cv2.putText(
|
|
frame_array,
|
|
metric,
|
|
(320, y_offset),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.4,
|
|
(255, 255, 255),
|
|
1,
|
|
)
|
|
y_offset += 15
|
|
|
|
logic_result = "E:" + ("Y" if speech_info.get("energy_check", False) else "N")
|
|
logic_result += " Z:" + ("Y" if speech_info.get("zcr_check", False) else "N")
|
|
logic_result += " S:" + (
|
|
"Y" if speech_info.get("spectral_check", False) else "N"
|
|
)
|
|
logic_result += " H:" + (
|
|
"Y" if speech_info.get("harmonic_check", False) else "N"
|
|
)
|
|
logic_result += " T:" + (
|
|
"Y" if speech_info.get("temporal_consistency", False) else "N"
|
|
)
|
|
|
|
cv2.putText(
|
|
frame_array,
|
|
logic_result,
|
|
(320, y_offset + 5),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(255, 255, 255),
|
|
1,
|
|
)
|
|
|
|
# Noise floor indicator
|
|
noise_floor = speech_info.get("noise_floor_energy", 0)
|
|
cv2.putText(
|
|
frame_array,
|
|
f"Noise Floor: {noise_floor:.4f}",
|
|
(10, y_offset + 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(200, 200, 200),
|
|
1,
|
|
)
|
|
|
|
|
|
async def handle_track_received(peer: Peer, track: MediaStreamTrack) -> None:
|
|
"""Handle incoming audio tracks from WebRTC peers."""
|
|
logger.info(
|
|
f"handle_track_received called for {peer.peer_name} with track kind: {track.kind}"
|
|
)
|
|
global _audio_processors, _send_chat_func
|
|
|
|
if track.kind != "audio":
|
|
logger.info(f"Ignoring non-audio track from {peer.peer_name}: {track.kind}")
|
|
return
|
|
|
|
# Initialize raw audio buffer for immediate graphing
|
|
if peer.peer_name not in WaveformVideoTrack.buffer:
|
|
WaveformVideoTrack.buffer[peer.peer_name] = np.array([], dtype=np.float32)
|
|
|
|
# Start processing frames immediately (before processor is ready)
|
|
logger.info(f"Starting frame processing loop for {peer.peer_name}")
|
|
frame_count = 0
|
|
while True:
|
|
try:
|
|
frame = await track.recv()
|
|
frame_count += 1
|
|
|
|
if frame_count % 100 == 0:
|
|
logger.debug(f"Received {frame_count} frames from {peer.peer_name}")
|
|
|
|
except MediaStreamError as e:
|
|
logger.info(f"Audio stream ended for {peer.peer_name}: {e}")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error receiving frame from {peer.peer_name}: {e}")
|
|
break
|
|
|
|
if isinstance(frame, AudioFrame):
|
|
try:
|
|
# Convert frame to numpy array
|
|
audio_data = frame.to_ndarray()
|
|
|
|
# Handle audio format conversion
|
|
audio_data = _process_audio_frame(audio_data, frame)
|
|
|
|
# Resample if needed
|
|
if frame.sample_rate != SAMPLE_RATE:
|
|
audio_data = _resample_audio(
|
|
audio_data, frame.sample_rate, SAMPLE_RATE
|
|
)
|
|
|
|
# Convert to float32
|
|
audio_data_float32 = audio_data.astype(np.float32)
|
|
|
|
# Update visualization buffer immediately
|
|
WaveformVideoTrack.buffer[peer.peer_name] = np.concatenate(
|
|
[WaveformVideoTrack.buffer[peer.peer_name], audio_data_float32]
|
|
)
|
|
# Limit buffer size to last 10 seconds
|
|
max_samples = SAMPLE_RATE * 10
|
|
if len(WaveformVideoTrack.buffer[peer.peer_name]) > max_samples:
|
|
WaveformVideoTrack.buffer[peer.peer_name] = (
|
|
WaveformVideoTrack.buffer[peer.peer_name][-max_samples:]
|
|
)
|
|
|
|
# # Process with optimized processor if available
|
|
# if peer.peer_name in _audio_processors:
|
|
# audio_processor = _audio_processors[peer.peer_name]
|
|
# audio_processor.add_audio_data(audio_data_float32)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing audio frame for {peer.peer_name}: {e}")
|
|
continue
|
|
|
|
# If processor already exists, just continue processing
|
|
# audio_processor = _audio_processors[peer.peer_name]
|
|
logger.info(f"Continuing OpenVINO audio processing for {peer.peer_name}")
|
|
|
|
try:
|
|
frame_count = 0
|
|
logger.info(f"Entering frame processing loop for {peer.peer_name}")
|
|
while True:
|
|
try:
|
|
logger.debug(f"Waiting for frame from {peer.peer_name}")
|
|
frame = await track.recv()
|
|
frame_count += 1
|
|
|
|
if frame_count == 1:
|
|
logger.info(f"Received first frame from {peer.peer_name}")
|
|
elif frame_count % 50 == 0:
|
|
logger.info(f"Received {frame_count} frames from {peer.peer_name}")
|
|
|
|
except MediaStreamError as e:
|
|
logger.info(f"Audio stream ended for {peer.peer_name}: {e}")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error receiving frame from {peer.peer_name}: {e}")
|
|
break
|
|
|
|
if isinstance(frame, AudioFrame):
|
|
try:
|
|
# Convert frame to numpy array
|
|
audio_data = frame.to_ndarray()
|
|
|
|
# Handle audio format conversion
|
|
audio_data = _process_audio_frame(audio_data, frame)
|
|
|
|
# Resample if needed
|
|
if frame.sample_rate != SAMPLE_RATE:
|
|
audio_data = _resample_audio(
|
|
audio_data, frame.sample_rate, SAMPLE_RATE
|
|
)
|
|
|
|
# Convert to float32
|
|
audio_data_float32 = audio_data.astype(np.float32)
|
|
|
|
logger.debug(
|
|
f"Processed audio frame {frame_count} from {peer.peer_name}: {len(audio_data_float32)} samples"
|
|
)
|
|
|
|
# Process with optimized processor if available
|
|
# audio_processor.add_audio_data(audio_data_float32)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error processing audio frame for {peer.peer_name}: {e}"
|
|
)
|
|
continue
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Unexpected error in audio processing for {peer.peer_name}: {e}",
|
|
exc_info=True,
|
|
)
|
|
finally:
|
|
cleanup_peer_processor(peer.peer_name)
|
|
|
|
|
|
def _process_audio_frame(
|
|
audio_data: npt.NDArray[Any], frame: AudioFrame
|
|
) -> npt.NDArray[np.float32]:
|
|
"""Process audio frame format conversion."""
|
|
# Handle stereo to mono conversion
|
|
if audio_data.ndim == 2:
|
|
if audio_data.shape[0] == 1:
|
|
audio_data = audio_data.squeeze(0)
|
|
else:
|
|
audio_data = np.mean(
|
|
audio_data, axis=0 if audio_data.shape[0] > audio_data.shape[1] else 1
|
|
)
|
|
|
|
# Normalize based on data type
|
|
if audio_data.dtype == np.int16:
|
|
audio_data = audio_data.astype(np.float32) / 32768.0
|
|
elif audio_data.dtype == np.int32:
|
|
audio_data = audio_data.astype(np.float32) / 2147483648.0
|
|
|
|
return audio_data.astype(np.float32)
|
|
|
|
|
|
def _resample_audio(
|
|
audio_data: npt.NDArray[np.float32], orig_sr: int, target_sr: int
|
|
) -> npt.NDArray[np.float32]:
|
|
"""Resample audio efficiently."""
|
|
try:
|
|
# Handle stereo audio by converting to mono if necessary
|
|
if audio_data.ndim > 1:
|
|
audio_data = np.mean(audio_data, axis=1)
|
|
|
|
# Use high-quality resampling
|
|
resampled = librosa.resample( # type: ignore
|
|
audio_data.astype(np.float64),
|
|
orig_sr=orig_sr,
|
|
target_sr=target_sr,
|
|
res_type="kaiser_fast", # Good balance of quality and speed
|
|
)
|
|
return resampled.astype(np.float32) # type: ignore
|
|
except Exception as e:
|
|
logger.error(f"Resampling failed: {e}")
|
|
raise ValueError(
|
|
f"Failed to resample audio from {orig_sr} Hz to {target_sr} Hz: {e}"
|
|
)
|
|
|
|
|
|
# Public API functions
|
|
def agent_info() -> Dict[str, str]:
|
|
return {
|
|
"name": AGENT_NAME,
|
|
"description": AGENT_DESCRIPTION,
|
|
"has_media": "true",
|
|
"configurable": "true",
|
|
}
|
|
|
|
|
|
def get_config_schema() -> Dict[str, Any]:
|
|
"""Get the configuration schema for the Whisper bot"""
|
|
return {
|
|
"bot_name": AGENT_NAME,
|
|
"version": "1.0",
|
|
"parameters": [
|
|
],
|
|
"categories": [
|
|
],
|
|
}
|
|
|
|
|
|
def handle_config_update(lobby_id: str, config_values: Dict[str, Any]) -> bool:
|
|
"""Handle configuration update for a specific lobby"""
|
|
global _model_id, _device, _ov_config
|
|
|
|
try:
|
|
logger.info(f"Updating TTS config for lobby {lobby_id}: {config_values}")
|
|
|
|
config_applied = False
|
|
|
|
config_applied = True
|
|
|
|
return config_applied
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to apply Whisper config update: {e}")
|
|
return False
|
|
|
|
|
|
def create_agent_tracks(session_name: str) -> Dict[str, MediaStreamTrack]:
|
|
"""Create agent tracks. Provides a synthetic video waveform track and a silent audio track for compatibility."""
|
|
|
|
class SilentAudioTrack(MediaStreamTrack):
|
|
kind = "audio"
|
|
|
|
def __init__(
|
|
self, sample_rate: int = SAMPLE_RATE, channels: int = 1, fps: int = 50
|
|
):
|
|
super().__init__()
|
|
self.sample_rate = sample_rate
|
|
self.channels = channels
|
|
self.fps = fps
|
|
self.samples_per_frame = int(self.sample_rate / self.fps)
|
|
self._timestamp = 0
|
|
# Per-track playback buffer (float32, mono, -1..1)
|
|
self._play_buffer = np.array([], dtype=np.float32)
|
|
# Source sample rate of the buffered audio (if any)
|
|
self._play_src_sr = None
|
|
# Phase for synthetic sine tone (radians)
|
|
self._sine_phase = 0.0
|
|
|
|
async def recv(self) -> AudioFrame:
|
|
global _audio_data
|
|
# If new global TTS audio was produced, grab and prepare it for playback
|
|
logger.info("recv called with _audio_data: %s and size: %s", "set" if _audio_data is not None else "unset", getattr(_audio_data, 'size', 0) if _audio_data is not None else 'N/A')
|
|
try:
|
|
if _audio_data is not None and getattr(_audio_data, 'size', 0) > 0:
|
|
# Get source sample rate from TTS engine if available
|
|
try:
|
|
from .vibevoicetts import get_tts_engine
|
|
|
|
src_sr = get_tts_engine().get_sample_rate()
|
|
except Exception:
|
|
# Fallback: assume 24000 (VibeVoice default)
|
|
logger.info("Falling back to default source sample rate of 24000 Hz")
|
|
src_sr = 24000
|
|
|
|
# Ensure numpy float32
|
|
src_audio = np.asarray(_audio_data, dtype=np.float32)
|
|
|
|
# Clear global buffer (we consumed it into this track)
|
|
_audio_data = None
|
|
|
|
# If source sr differs from track sr, resample
|
|
if src_sr != self.sample_rate:
|
|
try:
|
|
resampled = librosa.resample(
|
|
src_audio.astype(np.float64),
|
|
orig_sr=src_sr,
|
|
target_sr=self.sample_rate,
|
|
res_type="kaiser_fast",
|
|
)
|
|
self._play_buffer = resampled.astype(np.float32)
|
|
except Exception:
|
|
# On failure, fallback to nearest simple scaling (not ideal)
|
|
logger.exception("Failed to resample TTS audio; using source as-is")
|
|
self._play_buffer = src_audio.astype(np.float32)
|
|
else:
|
|
self._play_buffer = src_audio.astype(np.float32)
|
|
|
|
self._play_src_sr = self.sample_rate
|
|
|
|
# Save a copy of the consumed audio to disk for inspection
|
|
try:
|
|
pcm_save = np.clip(self._play_buffer, -1.0, 1.0)
|
|
pcm_int16_save = (pcm_save * 32767.0).astype(np.int16)
|
|
sample_path = "./sample.wav"
|
|
with wave.open(sample_path, "wb") as wf:
|
|
wf.setnchannels(1)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(self.sample_rate)
|
|
wf.writeframes(pcm_int16_save.tobytes())
|
|
logger.info("Wrote TTS sample to %s (%d samples, %d Hz)", sample_path, len(pcm_int16_save), self.sample_rate)
|
|
except Exception:
|
|
logger.exception("Failed to write sample.wav")
|
|
|
|
# Prepare output samples for this frame
|
|
if self._play_buffer.size > 0:
|
|
take = min(self.samples_per_frame, int(self._play_buffer.size))
|
|
out = self._play_buffer[:take]
|
|
# Advance buffer
|
|
if take >= self._play_buffer.size:
|
|
self._play_buffer = np.array([], dtype=np.float32)
|
|
else:
|
|
self._play_buffer = self._play_buffer[take:]
|
|
else:
|
|
# No TTS audio buffered; output a 110 Hz sine tone at 0.1 volume
|
|
freq = 110.0
|
|
volume = 0.1
|
|
n = self.samples_per_frame
|
|
# incremental phase per sample
|
|
phase_inc = 2.0 * np.pi * freq / float(self.sample_rate)
|
|
samples_idx = np.arange(n, dtype=np.float32)
|
|
out = volume * np.sin(self._sine_phase + phase_inc * samples_idx)
|
|
# advance and wrap phase
|
|
self._sine_phase = (self._sine_phase + phase_inc * n) % (2.0 * np.pi)
|
|
logger.debug("No TTS audio: emitting 110Hz test tone (%d samples at %d Hz)", n, self.sample_rate)
|
|
|
|
# Convert float32 [-1.0,1.0] to int16 PCM
|
|
pcm = np.clip(out, -1.0, 1.0)
|
|
pcm_int16 = (pcm * 32767.0).astype(np.int16)
|
|
|
|
# aiortc AudioFrame.from_ndarray expects shape (channels, samples)
|
|
if self.channels == 1:
|
|
data = np.expand_dims(pcm_int16, axis=0)
|
|
layout = "mono"
|
|
else:
|
|
# Duplicate mono into stereo
|
|
data = np.vstack([pcm_int16, pcm_int16])
|
|
layout = "stereo"
|
|
|
|
frame = AudioFrame.from_ndarray(data, layout=layout)
|
|
frame.sample_rate = self.sample_rate
|
|
frame.pts = self._timestamp
|
|
frame.time_base = fractions.Fraction(1, self.sample_rate)
|
|
self._timestamp += self.samples_per_frame
|
|
|
|
# Pace the frame rate to avoid busy-looping
|
|
await asyncio.sleep(1 / self.fps)
|
|
|
|
return frame
|
|
except Exception as e:
|
|
logger.exception(f"Error in SilentAudioTrack.recv: {e}")
|
|
# On error, fall back to silent frame
|
|
data = np.zeros((self.channels, self.samples_per_frame), dtype=np.int16)
|
|
frame = AudioFrame.from_ndarray(data, layout="mono" if self.channels == 1 else "stereo")
|
|
frame.sample_rate = self.sample_rate
|
|
frame.pts = self._timestamp
|
|
frame.time_base = fractions.Fraction(1, self.sample_rate)
|
|
self._timestamp += self.samples_per_frame
|
|
await asyncio.sleep(1 / self.fps)
|
|
return frame
|
|
|
|
try:
|
|
video_track = WaveformVideoTrack(
|
|
session_name=session_name, width=640, height=480, fps=15
|
|
)
|
|
audio_track = SilentAudioTrack()
|
|
return {"video": video_track, "audio": audio_track}
|
|
except Exception as e:
|
|
logger.error(f"Failed to create agent tracks: {e}")
|
|
return {}
|
|
|
|
|
|
async def handle_chat_message(
|
|
chat_message: ChatMessageModel,
|
|
send_message_func: Callable[[Union[str, ChatMessageModel]], Awaitable[None]],
|
|
) -> Optional[str]:
|
|
"""Handle incoming chat messages."""
|
|
global _audio_data
|
|
logger.info(f"Received chat message: {chat_message.message}")
|
|
# This ends up blocking; spin it off into a thread
|
|
try:
|
|
from .vibevoicetts import get_tts_engine
|
|
|
|
engine = get_tts_engine()
|
|
_audio_data = engine.text_to_speech(text=f"Speaker 1: {chat_message.message}", speaker_names=None, verbose=True)
|
|
except Exception:
|
|
# Fallback to convenience function if direct engine call fails
|
|
_audio_data = text_to_speech(text=f"Speaker 1: {chat_message.message}", speaker_names=None)
|
|
return None
|
|
|
|
|
|
async def on_track_received(peer: Peer, track: MediaStreamTrack) -> None:
|
|
"""Callback when a new track is received from a peer."""
|
|
await handle_track_received(peer, track)
|
|
|
|
|
|
def get_track_handler() -> Callable[[Peer, MediaStreamTrack], Awaitable[None]]:
|
|
"""Return the track handler function."""
|
|
return on_track_received
|
|
|
|
|
|
def bind_send_chat_function(
|
|
send_chat_func: Callable[[ChatMessageModel], Awaitable[None]],
|
|
create_chat_message_func: Callable[[str, Optional[str]], ChatMessageModel],
|
|
) -> None:
|
|
"""Bind the send chat function."""
|
|
global _send_chat_func, _create_chat_message_func, _audio_processors
|
|
|
|
logger.info("Binding send chat function to OpenVINO whisper agent")
|
|
_send_chat_func = send_chat_func
|
|
_create_chat_message_func = create_chat_message_func
|
|
|
|
# Update existing processors
|
|
# for peer_name, processor in _audio_processors.items():
|
|
# processor.send_chat_func = send_chat_func
|
|
# processor.create_chat_message_func = create_chat_message_func
|
|
# logger.debug(f"Updated processor for {peer_name} with new send chat function")
|
|
|
|
|
|
def cleanup_peer_processor(peer_name: str) -> None:
|
|
"""Clean up processor for disconnected peer."""
|
|
global _audio_processors
|
|
|
|
if peer_name in WaveformVideoTrack.buffer:
|
|
del WaveformVideoTrack.buffer[peer_name]
|
|
|
|
|
|
|