"""Real-time audio conversation with WebSockets. This module provides WebSocket endpoints for real-time audio conversation using the CSM-1B model and WhisperX for transcription. """ import os import io import base64 import json import time import asyncio import logging import tempfile from enum import Enum from typing import Dict, List, Optional, Any, Union import numpy as np import torch import torchaudio from pydub import AudioSegment import whisperx from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException, Request from fastapi.responses import JSONResponse # Set up logging logger = logging.getLogger(__name__) router = APIRouter(prefix="/realtime", tags=["Real-time Conversation"]) # Audio processing constants SAMPLE_RATE = 16000 # Sample rate for audio processing CHUNK_SIZE = 4096 # Chunk size for audio processing MAX_AUDIO_DURATION = 10 # Maximum audio duration in seconds SILENCE_THRESHOLD = 400 # Threshold for detecting silence (RMS) MIN_SILENCE_DURATION = 0.5 # Minimum silence duration to consider a pause # WebSocket message types class MessageType(str, Enum): AUDIO_CHUNK = "audio_chunk" TRANSCRIPT = "transcript" RESPONSE = "response" START_SPEAKING = "start_speaking" STOP_SPEAKING = "stop_speaking" ERROR = "error" STATUS = "status" # WhisperX model cache for performance _whisperx_model = None _whisperx_model_lock = asyncio.Lock() # Connection manager for websockets class ConnectionManager: def __init__(self): self.active_connections: Dict[str, WebSocket] = {} self.conversation_contexts: Dict[str, List] = {} self.voice_preferences: Dict[str, int] = {} # Store voice preferences by client_id async def connect(self, websocket: WebSocket, client_id: str): """Connect a client to the WebSocket""" await websocket.accept() self.active_connections[client_id] = websocket self.conversation_contexts[client_id] = [] self.voice_preferences[client_id] = 1 # Default to echo voice logger.info(f"Client {client_id} connected, active connections: {len(self.active_connections)}") def disconnect(self, client_id: str): """Disconnect a client from the WebSocket""" if client_id in self.active_connections: del self.active_connections[client_id] if client_id in self.conversation_contexts: del self.conversation_contexts[client_id] if client_id in self.voice_preferences: del self.voice_preferences[client_id] logger.info(f"Client {client_id} disconnected, active connections: {len(self.active_connections)}") def set_voice_preference(self, client_id: str, speaker_id: int): """Set voice preference for a client""" self.voice_preferences[client_id] = speaker_id def get_voice_preference(self, client_id: str) -> int: """Get voice preference for a client""" return self.voice_preferences.get(client_id, 1) # Default to echo (speaker_id=1) async def send_message(self, client_id: str, message_type: MessageType, data: Any): """Send a message to a client""" if client_id in self.active_connections: message = { "type": message_type, "data": data, "timestamp": time.time() } await self.active_connections[client_id].send_json(message) def add_to_context(self, client_id: str, speaker: int, text: str, audio: Union[torch.Tensor, bytes]): """Add a message to the conversation context""" if client_id in self.conversation_contexts: # Convert audio tensor to base64 if needed if isinstance(audio, torch.Tensor): audio_bytes = convert_tensor_to_wav_bytes(audio) audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') elif isinstance(audio, bytes): audio_base64 = base64.b64encode(audio).decode('utf-8') else: raise ValueError(f"Unsupported audio type: {type(audio)}") # Add to context, limiting size to last 5 exchanges self.conversation_contexts[client_id].append({ "speaker": speaker, "text": text, "audio": audio_base64 }) # Limit context size (keep last 5 exchanges to prevent context growing too large) if len(self.conversation_contexts[client_id]) > 5: self.conversation_contexts[client_id] = self.conversation_contexts[client_id][-5:] def get_context(self, client_id: str) -> List[Dict]: """Get the conversation context for a client""" return self.conversation_contexts.get(client_id, []) # Initialize connection manager manager = ConnectionManager() async def load_whisperx_model(compute_type="float16"): """Load WhisperX model if not already loaded""" global _whisperx_model # Use lock to ensure model loading is thread-safe async with _whisperx_model_lock: # Load WhisperX model if not already loaded if _whisperx_model is None: logger.info("Loading WhisperX model for real-time transcription") device = "cuda" if torch.cuda.is_available() else "cpu" # Use small model for lower latency _whisperx_model = whisperx.load_model( "small", # Small model for faster processing in real-time device, compute_type=compute_type, asr_options={"beam_size": 5, "vad_onset": 0.5, "vad_offset": 0.5} ) logger.info(f"WhisperX model loaded on {device} with compute_type={compute_type}") return _whisperx_model def convert_tensor_to_wav_bytes(audio_tensor: torch.Tensor) -> bytes: """Convert audio tensor to WAV bytes""" buf = io.BytesIO() if len(audio_tensor.shape) == 1: audio_tensor = audio_tensor.unsqueeze(0) torchaudio.save(buf, audio_tensor.cpu(), SAMPLE_RATE, format="wav") buf.seek(0) return buf.read() def convert_audio_data(audio_data: bytes) -> torch.Tensor: """Convert audio data to tensor""" with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp: temp.write(audio_data) temp.flush() # Load audio try: # First try with torchaudio waveform, sample_rate = torchaudio.load(temp.name) # Convert to mono if needed if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample if needed if sample_rate != SAMPLE_RATE: waveform = torchaudio.functional.resample( waveform, orig_freq=sample_rate, new_freq=SAMPLE_RATE ) return waveform.squeeze(0) except: # Fallback to pydub if torchaudio fails audio = AudioSegment.from_file(temp.name) # Convert to mono if needed if audio.channels > 1: audio = audio.set_channels(1) # Resample if needed if audio.frame_rate != SAMPLE_RATE: audio = audio.set_frame_rate(SAMPLE_RATE) # Convert to numpy array samples = np.array(audio.get_array_of_samples(), dtype=np.float32) / 32768.0 # Convert to tensor waveform = torch.tensor(samples, dtype=torch.float32) return waveform async def transcribe_audio(audio_data: bytes, language: Optional[str] = None) -> Dict: """Transcribe audio using WhisperX""" # Load WhisperX model model = await load_whisperx_model() # Save audio to temporary file with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp: temp.write(audio_data) temp.flush() # Transcribe with WhisperX device = "cuda" if torch.cuda.is_available() else "cpu" result = model.transcribe( temp.name, language=language, batch_size=16 if device == "cuda" else 1 ) return result async def generate_response(app, text: str, speaker_id: int, context: List[Dict]) -> torch.Tensor: """Generate response using CSM-1B model""" generator = app.state.generator # Validate model availability if generator is None: raise RuntimeError("TTS model not loaded") # Setup context segments segments = [] for ctx in context: if 'speaker' not in ctx or 'text' not in ctx or 'audio' not in ctx: continue # Decode base64 audio audio_data = base64.b64decode(ctx['audio']) # Convert to tensor audio_tensor = convert_audio_data(audio_data) # Create segment segments.append({ "speaker": ctx['speaker'], "text": ctx['text'], "audio": audio_tensor }) # Format text for better voice consistency from app.prompt_engineering import format_text_for_voice # Determine voice name from speaker_id voice_names = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"] voice_name = voice_names[speaker_id] if 0 <= speaker_id < len(voice_names) else "alloy" formatted_text = format_text_for_voice(text, voice_name) # Generate audio with context audio = generator.generate( text=formatted_text, speaker=speaker_id, context=segments, max_audio_length_ms=10000, # 10 seconds max for low latency temperature=0.65, # Lower temperature for more stable output topk=40, ) # Process audio for better quality from app.voice_enhancement import process_generated_audio processed_audio = process_generated_audio( audio, voice_name, generator.sample_rate, text ) return processed_audio def is_silence(audio_data: bytes, threshold=SILENCE_THRESHOLD) -> bool: """Check if audio is silence""" with io.BytesIO(audio_data) as buf: try: audio = AudioSegment.from_file(buf) # Get RMS (root mean square) amplitude rms = audio.rms return rms < threshold except: # If can't process, assume not silent return False @router.websocket("/conversation/{client_id}") async def websocket_conversation(websocket: WebSocket, client_id: str): """WebSocket endpoint for real-time audio conversation""" await manager.connect(websocket, client_id) # Get access to app state through the websocket app = websocket.app # Validate model availability if not hasattr(app.state, "generator") or app.state.generator is None: await manager.send_message(client_id, MessageType.ERROR, {"message": "TTS model not available"}) manager.disconnect(client_id) return # Initialize audio buffer and state audio_buffer = io.BytesIO() is_speaking = False silence_start = None try: # Tell client we're ready await manager.send_message(client_id, MessageType.STATUS, {"status": "ready", "message": "Connection established"}) # Process messages async for message in websocket.iter_json(): message_type = message.get("type") if message_type == "audio_chunk": # Get audio data audio_data = base64.b64decode(message["data"]) # Check if silence or speech current_is_silence = is_silence(audio_data) # Handle silence detection for end of speech if current_is_silence: if not silence_start: silence_start = time.time() elif time.time() - silence_start > MIN_SILENCE_DURATION and is_speaking: # End of speech detected is_speaking = False # Get audio from buffer audio_buffer.seek(0) full_audio = audio_buffer.read() # Reset buffer audio_buffer = io.BytesIO() # Process the complete audio asynchronously asyncio.create_task(process_complete_audio( app, client_id, full_audio )) # Notify client of end of speech await manager.send_message(client_id, MessageType.STOP_SPEAKING, {}) else: # Reset silence detection on new speech silence_start = None # Start of speech if not already speaking if not is_speaking: is_speaking = True await manager.send_message(client_id, MessageType.START_SPEAKING, {}) # Add chunk to buffer if speaking if is_speaking: audio_buffer.write(audio_data) elif message_type == "end_audio": # Explicit end of audio from client if audio_buffer.tell() > 0: # Get audio from buffer audio_buffer.seek(0) full_audio = audio_buffer.read() # Reset buffer audio_buffer = io.BytesIO() is_speaking = False # Process the complete audio asynchronously asyncio.create_task(process_complete_audio( app, client_id, full_audio )) elif message_type == "set_voice": # Set the voice for the response voice = message.get("voice", "alloy") # Map voice string to speaker ID voice_to_speaker = {"alloy": 0, "echo": 1, "fable": 2, "onyx": 3, "nova": 4, "shimmer": 5} speaker_id = voice_to_speaker.get(voice, 0) # Store in client state manager.set_voice_preference(client_id, speaker_id) # Send confirmation to client await manager.send_message(client_id, MessageType.STATUS, {"status": "voice_set", "voice": voice, "speaker_id": speaker_id}) elif message_type == "clear_context": # Clear the conversation context if client_id in manager.conversation_contexts: manager.conversation_contexts[client_id] = [] await manager.send_message(client_id, MessageType.STATUS, {"status": "context_cleared"}) except WebSocketDisconnect: logger.info(f"Client {client_id} disconnected") except Exception as e: logger.error(f"Error in websocket conversation: {e}", exc_info=True) try: await manager.send_message(client_id, MessageType.ERROR, {"message": str(e)}) except: pass finally: manager.disconnect(client_id) async def process_complete_audio(app, client_id: str, audio_data: bytes): """Process complete audio chunk from WebSocket""" try: # Transcribe audio transcription = await transcribe_audio(audio_data) # Get the text text = transcription.get("text", "").strip() # Send transcription to client await manager.send_message(client_id, MessageType.TRANSCRIPT, {"text": text, "segments": transcription.get("segments", [])}) # Skip if empty text if not text: return # Add user message to context (user is always speaker 0) manager.add_to_context(client_id, 0, text, audio_data) # Get current context context = manager.get_context(client_id) # Generate response voice_id = manager.get_voice_preference(client_id) response_audio = await generate_response(app, text, voice_id, context) # Convert to bytes response_bytes = convert_tensor_to_wav_bytes(response_audio) response_base64 = base64.b64encode(response_bytes).decode('utf-8') # Send response to client await manager.send_message(client_id, MessageType.RESPONSE, { "audio": response_base64, "speaker_id": voice_id }) # Add assistant response to context manager.add_to_context(client_id, voice_id, text, response_audio) except Exception as e: logger.error(f"Error processing audio: {e}", exc_info=True) await manager.send_message(client_id, MessageType.ERROR, {"message": f"Error processing audio: {str(e)}"})