Fix/video codecs patching (#245)

* Fix publish_codec_preferences() not called with real transceivers

In PublisherPeerConnection, publish_codec_preferences() was called on init, when the list of transceivers was still empty, so it was not applied, and vp8 codec was still used.
Moved the func call to `addTrack()` to fix that.

* Set "ultrafast" preset for h264 codec by default

- Updated StreamH264Encoder to pass "preset" to codec options
- Added STREAM_PATCH_AIORTC_H264_PRESET env variable with "ultrafast" by default. To reset to default "medium", pass "medium" or empty value

* Update StreamH264Encoder._encode_frame to patch the codec preset more robustly in case of resolution change

* fix ty check

Author: dangusev

getstream/video/rtc/encoders_patches.py

@@ -1,8 +1,11 @@
+import fractions
 import logging
 import os
 from typing import Optional
 
+import av
 from aiortc import RTCRtpCodecParameters
+from aiortc.codecs.h264 import MAX_FRAME_RATE as H264_MAX_FRAME_RATE
 from aiortc.codecs.h264 import H264Encoder
 from aiortc.codecs.vpx import Vp8Encoder
 from aiortc.rtcrtpsender import RTCEncodedFrame, RTCRtpSender
@@ -26,6 +29,11 @@
     "off",
 )
 
+# libx264 preset for the StreamH264Encoder. Defaults to "ultrafast" for
+# real-time CPU; can be raised (e.g. "veryfast", "medium") for better
+# quality-per-bit at the cost of more CPU per frame.
+STREAM_H264_PRESET = os.getenv("STREAM_PATCH_AIORTC_H264_PRESET", "ultrafast").strip()
+
 
 try:
     # Verify the name-mangled attributes we depend on still exist.
@@ -51,11 +59,13 @@ def target_bitrate(self, bitrate: int) -> None:
             self._Vp8Encoder__target_bitrate = bitrate
 
     class StreamH264Encoder(H264Encoder):
-        """H264Encoder subclass with higher bitrate bounds for Stream calls."""
+        """H264Encoder subclass with higher bitrate bounds and a real-time
+        libx264 preset for Stream calls."""
 
         def __init__(self) -> None:
             super().__init__()
             self._H264Encoder__target_bitrate = STREAM_VIDEO_DEFAULT_BITRATE
+            self.preset = STREAM_H264_PRESET
 
         @property
         def target_bitrate(self) -> int:
@@ -68,10 +78,45 @@ def target_bitrate(self, bitrate: int) -> None:
             )
             self._H264Encoder__target_bitrate = bitrate
 
+        def _encode_frame(self, frame, force_keyframe):
+            # Mirror parent's invalidation policy so we own codec creation in
+            # both first-init AND recreation (resolution/bitrate change) paths;
+            # otherwise parent's `if self.codec is None` branch silently
+            # reverts our preset to libx264's default `medium`.
+            codec = self.codec
+            if (
+                codec is not None
+                and codec.bit_rate
+                and (
+                    frame.width != codec.width
+                    or frame.height != codec.height
+                    or abs(self.target_bitrate - codec.bit_rate) / codec.bit_rate > 0.1
+                )
+            ):
+                self.buffer_data = b""
+                self.buffer_pts = None
+                self.codec = None
+
+            if self.codec is None:
+                self.codec = av.CodecContext.create("libx264", "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(H264_MAX_FRAME_RATE, 1)
+                self.codec.time_base = fractions.Fraction(1, H264_MAX_FRAME_RATE)
+                self.codec.options = {
+                    "level": "31",
+                    "tune": "zerolatency",
+                    "preset": self.preset,
+                }
+                self.codec.profile = "Baseline"
+            yield from super()._encode_frame(frame, force_keyframe)
+
 except Exception:
     logger.warning(
-        "Failed to patch aiortc video encoder subclasses with Stream bitrate values (aiortc internals may have changed), "
-        "falling back to default aiortc bitrates. \n"
+        "Failed to patch aiortc encoder subclasses with Stream values (aiortc internals may have changed), "
+        "falling back to default aiortc encoders. \n"
         "Set STREAM_PATCH_AIORTC_BITRATES=0 to disable patching.",
         exc_info=True,
     )
@@ -80,10 +125,11 @@ def target_bitrate(self, bitrate: int) -> None:
 
 
 def patch_sender_encoder(sender: RTCRtpSender) -> None:
-    """Patch a video sender to use Stream's higher-bitrate encoders.
+    """Patch a sender to use Stream's tuned encoders for the negotiated codec.
 
-    If anything goes wrong (e.g. aiortc internals changed), the sender
-    is left untouched and will use the stock encoder via get_encoder().
+    Works for video (VP8/H264) senders. If anything
+    goes wrong (e.g. aiortc internals changed), the sender is left untouched
+    and will use the stock encoder via get_encoder().
     """
     if StreamVp8Encoder is None or StreamH264Encoder is None:
         return
@@ -105,8 +151,8 @@ async def _next_with_stream_encoder(
         sender._next_encoded_frame = _next_with_stream_encoder  # type: ignore[method-assign]
     except Exception:
         logger.warning(
-            "Failed to patch aiortc video encoder subclasses with Stream bitrate values (aiortc internals may have changed), "
-            "falling back to default aiortc bitrates. \n"
+            "Failed to patch aiortc encoder subclasses with Stream values (aiortc internals may have changed), "
+            "falling back to default aiortc encoders. \n"
             "Set STREAM_PATCH_AIORTC_BITRATES=0 to disable patching.",
             exc_info=True,
         )

getstream/video/rtc/pc.py

@@ -53,10 +53,6 @@ def __init__(
         self._closed = False
         self._connected_event = asyncio.Event()
 
-        for transceiver in self.getTransceivers():
-            if transceiver.kind == "video":
-                transceiver.setCodecPreferences(publish_codec_preferences())
-
         @self.on("icegatheringstatechange")
         def on_icegatheringstatechange():
             logger.info(
@@ -79,8 +75,13 @@ def on_connectionstatechange():
 
     def addTrack(self, track: MediaStreamTrack) -> RTCRtpSender:
         sender = super().addTrack(track)
-        if track.kind == "video" and not BITRATE_PATCH_DISABLED:
-            patch_sender_encoder(sender)
+        if track.kind == "video":
+            for transceiver in self.getTransceivers():
+                if transceiver.sender is sender:
+                    transceiver.setCodecPreferences(publish_codec_preferences())
+                    break
+            if not BITRATE_PATCH_DISABLED:
+                patch_sender_encoder(sender)
         return sender
 
     async def handle_answer(self, response):

tests/rtc/test_encoders_patches.py

@@ -93,6 +93,39 @@ def test_disabled_via_env(self, monkeypatch, env_val):
         assert mod.BITRATE_PATCH_DISABLED is True
 
 
+class TestStreamH264Preset:
+    def test_default_is_ultrafast(self, monkeypatch):
+        """STREAM_H264_PRESET defaults to 'ultrafast' when env var is unset."""
+        monkeypatch.delenv("STREAM_PATCH_AIORTC_H264_PRESET", raising=False)
+
+        import getstream.video.rtc.encoders_patches as mod
+
+        importlib.reload(mod)
+        assert mod.STREAM_H264_PRESET == "ultrafast"
+        assert mod.StreamH264Encoder().preset == "ultrafast"
+
+    @pytest.mark.parametrize("env_val", ["medium", "veryfast", "superfast"])
+    def test_overridden_via_env(self, monkeypatch, env_val):
+        """STREAM_H264_PRESET reflects the env var and is stored on the instance."""
+        monkeypatch.setenv("STREAM_PATCH_AIORTC_H264_PRESET", env_val)
+
+        import getstream.video.rtc.encoders_patches as mod
+
+        importlib.reload(mod)
+        assert mod.STREAM_H264_PRESET == env_val
+        assert mod.StreamH264Encoder().preset == env_val
+
+    def test_whitespace_is_stripped(self, monkeypatch):
+        """Leading/trailing whitespace in the env value is stripped."""
+        monkeypatch.setenv("STREAM_PATCH_AIORTC_H264_PRESET", "  medium  ")
+
+        import getstream.video.rtc.encoders_patches as mod
+
+        importlib.reload(mod)
+        assert mod.STREAM_H264_PRESET == "medium"
+        assert mod.StreamH264Encoder().preset == "medium"
+
+
 class TestPatchSenderEncoder:
     @pytest.mark.asyncio
     async def test_installs_vp8_encoder(self):