shared/live_inference.py

"""Shared live inference utilities for Gradio demos.

Provides Twilio RTC configuration, a reusable WebRTC tab builder,
and a FastRTC frame-queue patch to prevent unbounded memory growth.
"""

import asyncio
import os
import time
from typing import Any, Callable

import gradio as gr
import numpy as np
from log_utils import setup_logger

InferenceFn = Callable[[np.ndarray], np.ndarray]

logger = setup_logger("LiveInference")

TAB_SWITCH_AUTO_STOP_JS = """\
<script>
(function () {
  function setup() {
    var col = document.getElementById("webrtc-stream-col");
    if (!col) return setTimeout(setup, 500);
    var panel = col.closest(".tabitem") || col.closest('[role="tabpanel"]');
    if (!panel) return;
    new MutationObserver(function () {
      if (panel.style.display === "none" || panel.hidden) {
        col.querySelectorAll("button").forEach(function (b) {
          if (b.textContent.indexOf("Stop") >= 0) b.click();
        });
      }
    }).observe(panel, { attributes: true, attributeFilter: ["style", "hidden", "class"] });
  }
  setTimeout(setup, 1000);
})();
</script>"""

_TOKEN_REFRESH_SECONDS = 12 * 3600  # 12 hours


def get_rtc_configuration() -> dict[str, Any]:
    """Build a WebRTC ICE configuration using Twilio TURN/STUN servers.

    Reads ``TWILIO_ACCOUNT_SID`` and ``TWILIO_AUTH_TOKEN`` from the
    environment. Returns an empty config when credentials are absent,
    which causes WebRTC to fall back to a direct peer-to-peer connection
    (works on LAN, may fail behind symmetric NATs).

    Returns:
        ICE server configuration dict (empty when no Twilio credentials).
    """
    account_sid = os.environ.get("TWILIO_ACCOUNT_SID")
    auth_token = os.environ.get("TWILIO_AUTH_TOKEN")
    if not account_sid or not auth_token:
        logger.info(
            "Twilio credentials not set -- WebRTC will use direct connection "
            "(set TWILIO_ACCOUNT_SID and TWILIO_AUTH_TOKEN for TURN relay)"
        )
        return {}

    from twilio.rest import Client

    client = Client(account_sid, auth_token)
    token = client.tokens.create()
    return {
        "iceServers": token.ice_servers,
        "iceTransportPolicy": "relay",
    }


class RtcConfigProvider:
    """Lazily creates and caches a Twilio ICE configuration, refreshing on expiry.

    Twilio tokens typically last 24 h. This class re-creates the token after
    ``_TOKEN_REFRESH_SECONDS`` (12 h by default) so that new WebRTC connections
    always get a valid config.
    """

    def __init__(self) -> None:
        self._config: dict[str, Any] | None = None
        self._created_at: float = 0.0

    def get(self) -> dict[str, Any]:
        """Return the ICE configuration, creating or refreshing as needed."""
        now = time.monotonic()
        if self._config is None or (now - self._created_at) > _TOKEN_REFRESH_SECONDS:
            self._config = get_rtc_configuration()
            self._created_at = now
        return self._config


RtcConfigurationInput = Callable[[], dict[str, Any]] | dict[str, Any] | None
"""Accepted types for WebRTC ICE configuration.

A callable is invoked per-connection by FastRTC (allowing credential refresh)
and **must** return a dict (not ``None``) — FastRTC's client-side JS crashes on
``null``. A dict is used as-is for every connection; ``None`` falls back to
direct P2P (handled safely by FastRTC's ``or {}`` fallback).
"""


def build_live_inference_tab(
    rtc_configuration: RtcConfigurationInput,
    description_html: str = "",
    tab_label: str = "Live Inference",
    width: int = 640,
    height: int = 360,
) -> tuple[gr.TabItem, Any]:
    """Build a Live Inference tab containing a WebRTC video stream.

    Must be called inside a ``gr.Blocks`` / ``gr.Tabs`` context. The caller
    is responsible for wiring the stream event.

    Args:
        rtc_configuration: ICE configuration dict, callable that returns one,
            or ``None`` for direct connection. A callable is called per-connection
            by FastRTC, enabling credential refresh.
        description_html: Optional HTML shown above the stream.
        tab_label: Label for the tab.
        width: Camera frame width in pixels.
        height: Camera frame height in pixels.

    Returns:
        Tuple of ``(tab, webrtc_stream)``.
    """
    with gr.TabItem(tab_label) as tab:
        if description_html:
            gr.HTML(description_html)
        webrtc_stream = build_webrtc_stream(rtc_configuration, width=width, height=height)

    return tab, webrtc_stream


def build_webrtc_stream(
    rtc_configuration: RtcConfigurationInput,
    max_fps: int = 15,
    width: int = 640,
    height: int = 360,
) -> Any:
    """Create a WebRTC video-stream component.

    Must be called inside a Gradio layout context (e.g. ``gr.Row``,
    ``gr.Column``, ``gr.TabItem``).

    Args:
        rtc_configuration: ICE configuration dict, callable that returns one,
            or ``None`` for direct connection. A callable is called per-connection
            by FastRTC, enabling credential refresh.
        max_fps: Maximum frame rate requested from the browser camera.
        width: Camera frame width in pixels.
        height: Camera frame height in pixels.

    Returns:
        A ``fastrtc.WebRTC`` component instance.
    """
    from fastrtc import WebRTC

    with gr.Row():
        with gr.Column(
            elem_classes="webrtc-stream-col",
            elem_id="webrtc-stream-col",
        ):
            webrtc_stream = WebRTC(
                label="Live Detection",
                mode="send-receive",
                modality="video",
                rtc_configuration=rtc_configuration,
                container=True,
                show_label=True,
                mirror_webcam=False,
                track_constraints={
                    "width": {"exact": width},
                    "height": {"exact": height},
                    "frameRate": {"max": max_fps},
                },
                #rtp_params={"degradationPreference": "maintain-resolution"},
                full_screen=False,
                button_labels={
                    "start": "Start Inference",
                    "stop": "Stop Inference",
                },
            )
    return webrtc_stream


def patch_fastrtc_frame_queue() -> None:
    """Replace FastRTC's unbounded frame queue with a latest-frame-only queue.

    FastRTC's ``VideoCallback`` uses an unbounded ``asyncio.Queue`` for
    incoming WebRTC frames.  When the handler can't keep up (e.g. pipe
    round-trip ~35-50 ms at 30 fps), frames accumulate faster than
    they're consumed.  Each queued 1280×720 ``VideoFrame`` is ~1.38 MB
    of C memory (invisible to ``tracemalloc``), causing >1 MB/s growth
    on Windows.

    We can never "catch up" on skipped frames, so always discard stale
    ones and keep only the most recent frame.

    Must be called **before** any ``WebRTC`` component is created
    (typically at the top of ``__main__``).
    """
    from typing import cast as _cast

    import fastrtc.tracks as frt
    from aiortc.mediastreams import MediaStreamError

    _orig_init = frt.VideoCallback.__init__

    def _patched_init(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        _orig_init(self, *args, **kwargs)
        self.frame_queue = asyncio.Queue(maxsize=1)

    async def _latest_frame_accept_input(self):  # type: ignore[no-untyped-def]
        self.has_started = True
        while not self.thread_quit.is_set():
            try:
                frame = _cast(frt.VideoFrame, await self.track.recv())
                self.latest_frame = frame
                # Flush any stale frame — we only ever want the latest
                while not self.frame_queue.empty():
                    try:
                        self.frame_queue.get_nowait()
                    except asyncio.QueueEmpty:
                        break
                self.frame_queue.put_nowait(frame)
            except MediaStreamError:
                self.stop()
                return

    frt.VideoCallback.__init__ = _patched_init  # type: ignore[assignment]
    frt.VideoCallback.accept_input = _latest_frame_accept_input  # type: ignore[assignment]

    # Guard against None payload during client switches.
    # FastRTC's WebRTC.preprocess assumes payload is never None, but Gradio
    # can deliver a None when the outgoing connection drops mid-handoff.
    from fastrtc import WebRTC as _WebRTC

    _orig_preprocess = _WebRTC.preprocess

    def _safe_preprocess(self, payload):  # type: ignore[no-untyped-def]
        if payload is None:
            return None
        return _orig_preprocess(self, payload)

    _WebRTC.preprocess = _safe_preprocess  # type: ignore[assignment]


def patch_aiortc_h264_nvenc() -> None:
    """Patch aiortc's H264Encoder to use GPU encoding (h264_nvenc) when available.

    The default aiortc H264Encoder hardcodes ``libx264`` (CPU). On machines
    with an NVIDIA GPU and NVENC-enabled FFmpeg, this patch swaps in
    ``h264_nvenc`` to offload encoding from the CPU.

    Falls back silently if NVENC is not available (no patch applied).

    Must be called **before** any WebRTC component is created.
    """
    import fractions
    from collections.abc import Iterator

    import av
    from aiortc.codecs.h264 import MAX_FRAME_RATE, H264Encoder

    try:
        ctx = av.CodecContext.create("h264_nvenc", "w")
        ctx.width = 64
        ctx.height = 64
        ctx.pix_fmt = "yuv420p"
        ctx.open()
        ctx.close()
    except Exception:
        logger.info("WebRTC: h264_nvenc unavailable, keeping libx264")
        return

    logger.info("WebRTC: patching H264Encoder to use h264_nvenc")

    def _encode_frame_nvenc(
        self: H264Encoder, frame: av.VideoFrame, force_keyframe: bool
    ) -> Iterator[bytes]:
        if self.codec and (
            frame.width != self.codec.width
            or frame.height != self.codec.height
            or abs(self.target_bitrate - self.codec.bit_rate) / self.codec.bit_rate > 0.1
        ):
            self.buffer_data = b""
            self.buffer_pts = None
            self.codec = None

        if force_keyframe:
            frame.pict_type = av.video.frame.PictureType.I
        else:
            frame.pict_type = av.video.frame.PictureType.NONE

        if self.codec is None:
            self.codec = av.CodecContext.create("h264_nvenc", "w")
            self.codec.width = frame.width
            self.codec.height = frame.height
            self.codec.bit_rate = self.target_bitrate
            self.codec.pix_fmt = "yuv420p"
            self.codec.framerate = fractions.Fraction(MAX_FRAME_RATE, 1)
            self.codec.time_base = fractions.Fraction(1, MAX_FRAME_RATE)
            self.codec.options = {"preset": "p1", "delay": "0"}
            self.codec.profile = "Baseline"

        data_to_send = b""
        for package in self.codec.encode(frame):
            data_to_send += bytes(package)

        if data_to_send:
            yield from self._split_bitstream(data_to_send)

    H264Encoder._encode_frame = _encode_frame_nvenc  # type: ignore[assignment]


def patch_aioice_stun_transaction() -> None:
    """Guard aioice's STUN Transaction.__retry against already-resolved futures.

    aioice schedules ``Transaction.__retry`` via ``loop.call_later``.  When the
    retry fires after ``response_received`` has already resolved the future,
    ``set_exception(TransactionTimeout())`` raises ``InvalidStateError``.  This
    is a known upstream race condition (aiortc/aioice#78).

    The patch adds a ``future.done()`` guard identical to the one already
    present in ``response_received``, and logs a debug message when the race is
    hit.

    Must be called **before** any WebRTC component is created.
    """
    from aioice.stun import Transaction

    _orig_retry = Transaction._Transaction__retry  # type: ignore[attr-defined]

    def _safe_retry(self) -> None:  # type: ignore[no-untyped-def]
        if self._Transaction__future.done():
            logger.debug(
                "aioice STUN transaction already resolved — "
                "suppressing stale retry (aiortc/aioice#78)"
            )
            return
        _orig_retry(self)

    Transaction._Transaction__retry = _safe_retry  # type: ignore[attr-defined]


def patch_fastrtc_yuv420p_output() -> None:
    """Pre-convert outbound frames to yuv420p to avoid BufferError on Python 3.13.

    PyAV's ``VideoFrame.from_ndarray()`` wraps numpy memory via internal
    ``BytesIO`` objects.  When aiortc's VP8 encoder later calls
    ``frame.reformat(format="yuv420p")``, the intermediate ``BytesIO``
    cleanup raises ``BufferError`` because the numpy array still holds an
    active buffer export (Python 3.13's stricter buffer protocol).

    By converting BGR→YUV420p in numpy-space (via OpenCV) and constructing
    the ``VideoFrame`` directly in ``yuv420p``, the VP8 encoder sees the
    target format and skips ``reformat()`` entirely — avoiding the conflict.

    Must be called **before** any ``WebRTC`` component is created.
    """
    import cv2
    from av import VideoFrame

    import fastrtc.tracks as frt

    def _yuv420p_array_to_frame(self, array: np.ndarray) -> VideoFrame:  # type: ignore[no-untyped-def]
        yuv = cv2.cvtColor(array, cv2.COLOR_BGR2YUV_I420)
        return VideoFrame.from_ndarray(yuv, format="yuv420p")

    frt.VideoCallback.array_to_frame = _yuv420p_array_to_frame  # type: ignore[assignment]
    frt.ServerToClientVideo.array_to_frame = _yuv420p_array_to_frame  # type: ignore[assignment]