""" WebRTC streaming for MuseTalk Uses aiortc to stream generated frames as H.264 video """ import asyncio import numpy as np import cv2 import time import fractions from av import VideoFrame from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCConfiguration, RTCIceServer from aiortc.contrib.media import MediaRelay # ICE Server configuration (STUN + TURN for NAT traversal) ICE_SERVERS = RTCConfiguration(iceServers=[ # STUN servers RTCIceServer(urls=["stun:stun.l.google.com:19302"]), RTCIceServer(urls=["stun:stun1.l.google.com:19302"]), # OpenRelay TURN servers (free) RTCIceServer( urls=["turn:openrelay.metered.ca:80"], username="openrelayproject", credential="openrelayproject" ), RTCIceServer( urls=["turn:openrelay.metered.ca:443?transport=tcp"], username="openrelayproject", credential="openrelayproject" ), RTCIceServer( urls=["turns:openrelay.metered.ca:443?transport=tcp"], username="openrelayproject", credential="openrelayproject" ), ]) # Video track that streams frames from a queue class MuseTalkVideoTrack(VideoStreamTrack): """ Video track that reads frames from an async queue. Frames are numpy arrays (BGR). """ def __init__(self, fps: int = 30): super().__init__() self.fps = fps self.frame_queue = asyncio.Queue(maxsize=60) # Buffer up to 60 frames self.frame_count = 0 self.start_time = None self._running = True self.last_frame = None # Keep last frame for idle async def recv(self): """Called by aiortc to get the next frame.""" if self.start_time is None: self.start_time = time.time() # Calculate expected timestamp pts = self.frame_count time_base = fractions.Fraction(1, self.fps) try: # Try to get frame from queue with timeout frame_bgr = await asyncio.wait_for( self.frame_queue.get(), timeout=1.0 / self.fps # Wait up to one frame period ) self.last_frame = frame_bgr except asyncio.TimeoutError: # No frame available, use last frame or black if self.last_frame is not None: frame_bgr = self.last_frame else: # Create black frame frame_bgr = np.zeros((480, 640, 3), dtype=np.uint8) # Convert BGR to RGB for av frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) # Create VideoFrame video_frame = VideoFrame.from_ndarray(frame_rgb, format="rgb24") video_frame.pts = pts video_frame.time_base = time_base self.frame_count += 1 return video_frame async def put_frame(self, frame: np.ndarray): """Add a frame to the queue.""" if self._running: try: # Non-blocking put, drop frame if queue is full self.frame_queue.put_nowait(frame) if self.frame_count < 5: print(f"[WebRTC Track] Frame queued, queue size: {self.frame_queue.qsize()}") except asyncio.QueueFull: # Drop oldest frame and add new one try: self.frame_queue.get_nowait() self.frame_queue.put_nowait(frame) print(f"[WebRTC Track] Frame dropped and replaced") except: pass def stop(self): """Stop the track.""" self._running = False class WebRTCManager: """ Manages WebRTC connections for streaming video. """ def __init__(self): self.connections = {} # session_id -> RTCPeerConnection self.video_tracks = {} # session_id -> MuseTalkVideoTrack self.connection_ready = {} # session_id -> asyncio.Event self.relay = MediaRelay() async def create_connection(self, session_id: str, fps: int = 30) -> RTCPeerConnection: """Create a new WebRTC connection with TURN servers for NAT traversal.""" pc = RTCPeerConnection(configuration=ICE_SERVERS) # Create video track video_track = MuseTalkVideoTrack(fps=fps) pc.addTrack(video_track) # Store references self.connections[session_id] = pc self.video_tracks[session_id] = video_track self.connection_ready[session_id] = asyncio.Event() # Handle connection state changes @pc.on("connectionstatechange") async def on_connectionstatechange(): print(f"[WebRTC] Connection state: {pc.connectionState}") if pc.connectionState == "connected": # Signal that connection is ready for streaming if session_id in self.connection_ready: self.connection_ready[session_id].set() print(f"[WebRTC] Connection ready for streaming: {session_id}") elif pc.connectionState == "failed" or pc.connectionState == "closed": # Only close if session still exists if session_id in self.connections: await self.close_connection(session_id) @pc.on("iceconnectionstatechange") async def on_iceconnectionstatechange(): print(f"[WebRTC] ICE connection state: {pc.iceConnectionState}") @pc.on("icegatheringstatechange") async def on_icegatheringstatechange(): print(f"[WebRTC] ICE gathering state: {pc.iceGatheringState}") @pc.on("signalingstatechange") async def on_signalingstatechange(): print(f"[WebRTC] Signaling state: {pc.signalingState}") return pc async def handle_offer(self, session_id: str, sdp: str, type: str, fps: int = 30) -> dict: """ Handle incoming SDP offer from client. Returns answer SDP. """ # Create connection if not exists if session_id not in self.connections: pc = await self.create_connection(session_id, fps) else: pc = self.connections[session_id] # Set remote description (offer from client) offer = RTCSessionDescription(sdp=sdp, type=type) await pc.setRemoteDescription(offer) # Create answer answer = await pc.createAnswer() await pc.setLocalDescription(answer) return { "sdp": pc.localDescription.sdp, "type": pc.localDescription.type } async def wait_for_connection(self, session_id: str, timeout: float = 10.0) -> bool: """Wait for WebRTC connection to be established.""" if session_id not in self.connection_ready: print(f"[WebRTC] No connection event for {session_id}") return False try: await asyncio.wait_for( self.connection_ready[session_id].wait(), timeout=timeout ) print(f"[WebRTC] Connection established for {session_id}") return True except asyncio.TimeoutError: print(f"[WebRTC] Connection timeout for {session_id}") return False def is_connected(self, session_id: str) -> bool: """Check if WebRTC connection is established.""" if session_id in self.connections: return self.connections[session_id].connectionState == "connected" return False async def send_frame(self, session_id: str, frame: np.ndarray): """Send a frame to the specified session.""" if session_id in self.video_tracks: await self.video_tracks[session_id].put_frame(frame) async def close_connection(self, session_id: str): """Close a WebRTC connection.""" try: if session_id in self.video_tracks: self.video_tracks[session_id].stop() del self.video_tracks[session_id] if session_id in self.connections: pc = self.connections.pop(session_id) await pc.close() if session_id in self.connection_ready: del self.connection_ready[session_id] print(f"[WebRTC] Closed connection: {session_id}") except Exception as e: print(f"[WebRTC] Error closing connection {session_id}: {e}") def get_video_track(self, session_id: str) -> MuseTalkVideoTrack: """Get the video track for a session.""" return self.video_tracks.get(session_id) # Global manager instance webrtc_manager = WebRTCManager()