server/webrtc_stream.py

"""
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()