diff --git a/getstream/video/rtc/encoders_patches.py b/getstream/video/rtc/encoders_patches.py index 63ed03b0..7c93445f 100644 --- a/getstream/video/rtc/encoders_patches.py +++ b/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, ) diff --git a/getstream/video/rtc/pc.py b/getstream/video/rtc/pc.py index 0df0d777..5b221b59 100644 --- a/getstream/video/rtc/pc.py +++ b/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): diff --git a/tests/rtc/test_encoders_patches.py b/tests/rtc/test_encoders_patches.py index 7d0a45d3..80807f3b 100644 --- a/tests/rtc/test_encoders_patches.py +++ b/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):