From 2854bf1db0df8e2b012dade90fc84fe76dfbce25 Mon Sep 17 00:00:00 2001 From: Henry Ruhs Date: Mon, 15 Jun 2026 11:37:25 +0200 Subject: [PATCH] Resolve final stream TODOs (#1155) * refactor a lot * refactor a lot * fix test * fix test --- facefusion/apis/stream_audio.py | 13 +++++---- facefusion/apis/stream_event.py | 8 +++++- facefusion/apis/stream_video.py | 48 +++++---------------------------- facefusion/rtc.py | 14 +++++++++- facefusion/types.py | 4 +-- tests/test_api_stream_audio.py | 28 +++++++++---------- tests/test_api_stream_video.py | 30 ++++++++++----------- 7 files changed, 64 insertions(+), 81 deletions(-) diff --git a/facefusion/apis/stream_audio.py b/facefusion/apis/stream_audio.py index 4fc27694..c4e027c4 100644 --- a/facefusion/apis/stream_audio.py +++ b/facefusion/apis/stream_audio.py @@ -1,5 +1,3 @@ -import ctypes -import time from functools import partial from queue import Queue from typing import Optional, Tuple @@ -13,6 +11,7 @@ from facefusion.types import AudioCodec, AudioFrame, Buffer, OpusDecoder, RtcPee def run_audio_encode_loop(rtc_peer : RtcPeer, audio_queue : Queue[Tuple[Time, AudioFrame]]) -> None: + audio_codec = rtc_peer.get('audio').get('codec') temp_audio_time, temp_audio_frame = audio_queue.get() audio_encoder = opus_encoder.create(48000, 2) @@ -20,7 +19,7 @@ def run_audio_encode_loop(rtc_peer : RtcPeer, audio_queue : Queue[Tuple[Time, Au audio_buffer = opus_encoder.encode(audio_encoder, temp_audio_frame.tobytes(), 960) if audio_buffer: - audio_timestamp = int(temp_audio_time * 48000) + audio_timestamp = rtc.convert_time_to_timestamp(audio_codec, temp_audio_time) rtc.send_audio(rtc_peer, audio_buffer, audio_timestamp) temp_audio_time, temp_audio_frame = audio_queue.get() @@ -59,10 +58,10 @@ def destroy_audio_decoder(audio_codec : AudioCodec, audio_decoder : OpusDecoder) opus_decoder.destroy(audio_decoder) -def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[Time, AudioFrame]], track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None: - audio_buffer = ctypes.string_at(data, size) +def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[Time, AudioFrame]], audio_buffer : Buffer, audio_timestamp : int) -> None: audio_frame = decode_audio_frame(audio_codec, audio_decoder, audio_buffer) if audio_frame: - temp_audio_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32) - audio_queue.put((time.monotonic(), temp_audio_frame)) + audio_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32) + audio_time = rtc.convert_timestamp_to_time(audio_codec, audio_timestamp) + audio_queue.put((audio_time, audio_frame)) diff --git a/facefusion/apis/stream_event.py b/facefusion/apis/stream_event.py index aafb1ac6..5da36040 100644 --- a/facefusion/apis/stream_event.py +++ b/facefusion/apis/stream_event.py @@ -10,7 +10,7 @@ def create_receive_event(track : int, frame_handler : FrameHandler) -> threading datachannel_library = datachannel_module.create_static_library() receive_event = threading.Event() - frame_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p)(frame_handler) + frame_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p)(partial(dispatch_frame, frame_handler)) close_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p)(partial(dispatch_event, receive_event)) datachannel_library.rtcSetFrameCallback(track, frame_callback) datachannel_library.rtcSetClosedCallback(track, close_callback) @@ -20,5 +20,11 @@ def create_receive_event(track : int, frame_handler : FrameHandler) -> threading return receive_event +def dispatch_frame(frame_handler : FrameHandler, track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None: + frame_buffer = ctypes.string_at(data, size) + frame_timestamp = ctypes.cast(info, ctypes.POINTER(ctypes.c_uint32)).contents.value + frame_handler(frame_buffer, frame_timestamp) + + def dispatch_event(event : threading.Event, track : int, pointer : ctypes.c_void_p) -> None: event.set() diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py index 451ecf67..947a3cf5 100644 --- a/facefusion/apis/stream_video.py +++ b/facefusion/apis/stream_video.py @@ -1,5 +1,3 @@ -import ctypes -import time from concurrent.futures import Future, ThreadPoolExecutor from functools import partial from queue import Queue @@ -26,32 +24,26 @@ def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[Time, Fu temp_resolution : Resolution = video_resolution temp_bitrate : BitRate = 8000 video_encoder = create_video_encoder(video_codec, temp_resolution, temp_bitrate) - temp_video_time = video_time frame_index = 0 while video_buffer: - encode_start = time.monotonic() - sender_bitrate = calculate_sender_bitrate(rtc_peer, temp_bitrate) + sender_bitrate = rtc_peer.get('sender_bitrate').value if video_resolution[0] - temp_resolution[0] or video_resolution[1] - temp_resolution[1]: temp_resolution = video_resolution update_video_encoder_resolution(video_codec, video_encoder, temp_resolution) - if sender_bitrate - temp_bitrate: + if sender_bitrate > 0 and sender_bitrate - temp_bitrate: temp_bitrate = sender_bitrate update_video_encoder_bitrate(video_codec, video_encoder, temp_bitrate) __video_buffer__ = encode_video_frame(video_codec, video_encoder, video_buffer, temp_resolution, frame_index) if __video_buffer__: - video_timestamp = int(video_time * 90000) + video_timestamp = rtc.convert_time_to_timestamp(video_codec, video_time) rtc.send_video(rtc_peer, __video_buffer__, video_timestamp) - encode_time = time.monotonic() - encode_start - frame_interval = video_time - temp_video_time - temp_video_time = video_time - - receiver_bitrate = calculate_receiver_bitrate(rtc_peer, encode_time, frame_interval) + receiver_bitrate = rtc_peer.get('receiver_bitrate').value rtc.adapt_receiver_bitrate(rtc_peer, receiver_bitrate) frame_index += 1 @@ -86,31 +78,6 @@ def process_video_frame(input_vision_frame : VisionFrame) -> BufferPack: return BufferPack(buffer = output_buffer, resolution = output_resolution) -def calculate_receiver_bitrate(rtc_peer : RtcPeer, encode_time : float, frame_interval : float) -> BitRate: - min_bitrate : BitRate = 500 - max_bitrate : BitRate = 8000 - bitrate : BitRate = rtc_peer.get('receiver_bitrate').value - - if frame_interval > 0: - scale = frame_interval / encode_time - bitrate = int(bitrate * scale) - bitrate = max(min_bitrate, min(max_bitrate, bitrate)) - - return bitrate - - -#todo: does not feel final as this is an clamp and not calculate -def calculate_sender_bitrate(rtc_peer : RtcPeer, bitrate : BitRate) -> BitRate: - min_bitrate : BitRate = 500 - max_bitrate : BitRate = 8000 - peer_bitrate : BitRate = rtc_peer.get('sender_bitrate').value - - if peer_bitrate > 0: - bitrate = max(min_bitrate, min(max_bitrate, peer_bitrate)) - - return bitrate - - def decode_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, input_buffer : Buffer) -> Optional[VisionFrame]: if video_codec == 'av1': aom_pointer = aom_decoder.decode(video_decoder, input_buffer) @@ -199,11 +166,10 @@ def update_video_encoder_bitrate(video_codec : VideoCodec, video_encoder : VpxEn return False -#todo: we can remove the dead args or pass audio buffer -def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None: - video_buffer = ctypes.string_at(data, size) +def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, video_buffer : Buffer, video_timestamp : int) -> None: vision_frame = decode_video_frame(video_codec, video_decoder, video_buffer) if numpy.any(vision_frame) and video_queue.qsize() < video_queue.maxsize: video_future = video_executor.submit(process_video_frame, vision_frame) - video_queue.put((time.monotonic(), video_future)) + video_time = rtc.convert_timestamp_to_time(video_codec, video_timestamp) + video_queue.put((video_time, video_future)) diff --git a/facefusion/rtc.py b/facefusion/rtc.py index 743d1138..8af13e90 100644 --- a/facefusion/rtc.py +++ b/facefusion/rtc.py @@ -2,7 +2,7 @@ import ctypes from typing import List, Optional from facefusion.libraries import datachannel as datachannel_module -from facefusion.types import AudioCodec, BitRate, Buffer, MediaDirection, PeerConnection, RtcAudioTrack, RtcPeer, RtcTrackInit, RtcVideoTrack, SdpAnswer, SdpOffer, VideoCodec +from facefusion.types import AudioCodec, BitRate, Buffer, MediaDirection, PeerConnection, RtcAudioTrack, RtcPeer, RtcTrackInit, RtcVideoTrack, SdpAnswer, SdpOffer, Time, VideoCodec def create_peer_connection() -> PeerConnection: @@ -239,6 +239,18 @@ def get_payload_type(sdp_offer : SdpOffer, codec : AudioCodec | VideoCodec) -> i return 0 +def convert_time_to_timestamp(codec : AudioCodec | VideoCodec, time : Time) -> int: + if codec == 'opus': + return int(time * 48000) + return int(time * 90000) + + +def convert_timestamp_to_time(codec : AudioCodec | VideoCodec, timestamp : int) -> Time: + if codec == 'opus': + return timestamp / 48000 + return timestamp / 90000 + + @ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_uint, ctypes.c_void_p) def handle_sender_bitrate(_ : int, bitrate : BitRate, pointer : int) -> None: ctypes.cast(pointer, ctypes.POINTER(ctypes.c_uint)).contents.value = bitrate // 1000 diff --git a/facefusion/types.py b/facefusion/types.py index 6b7fc0c5..60bf7477 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -102,8 +102,6 @@ VideoCodec : TypeAlias = Literal['av1', 'vp8', 'vp9'] AomVideoCodec : TypeAlias = Literal['av1'] VxpVideoCodec : TypeAlias = Literal['vp8', 'vp9'] -FrameHandler : TypeAlias = Callable[..., None] - AomEncoder : TypeAlias = ctypes.Array[ctypes.c_char] AomDecoder : TypeAlias = ctypes.Array[ctypes.c_char] OpusEncoder : TypeAlias = ctypes.c_void_p @@ -118,6 +116,8 @@ BufferPack = TypedDict('BufferPack', 'resolution' : Resolution }) +FrameHandler : TypeAlias = Callable[[Buffer, int], None] + Args : TypeAlias = Dict[str, Any] Choice : TypeAlias = Union[int | str] diff --git a/tests/test_api_stream_audio.py b/tests/test_api_stream_audio.py index dc20293f..6e04d0da 100644 --- a/tests/test_api_stream_audio.py +++ b/tests/test_api_stream_audio.py @@ -14,7 +14,7 @@ from facefusion.download import conditional_download from facefusion.ffmpeg import read_audio_buffer from facefusion.hash_helper import create_hash from facefusion.libraries import datachannel as datachannel_module, opus as opus_module -from facefusion.types import AudioCodec, AudioFrame, FrameHandler, RtcPeer, RtcPeerAudio +from facefusion.types import AudioCodec, AudioFrame, Buffer, FrameHandler, RtcPeer, RtcPeerAudio from .assert_helper import get_test_example_file, get_test_examples_directory @@ -37,8 +37,11 @@ def before_each() -> None: rtc_store.clear() -def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None: +def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event: + frame_handler(buffer, 0) + ready_event = threading.Event() ready_event.set() + return ready_event def test_run_audio_encode_loop() -> None: @@ -48,6 +51,12 @@ def test_run_audio_encode_loop() -> None: rtc_peer : RtcPeer =\ { 'peer_connection': peer_connection, + 'audio': + { + 'sender_track': 0, + 'receiver_track': 0, + 'codec': 'opus' + }, 'video': { 'sender_track': 0, @@ -82,11 +91,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None: audio_frame = numpy.frombuffer(audio_buffer, dtype = numpy.int16).astype(numpy.float32) / 32768.0 audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300) - datachannel_mock = MagicMock() - ready_event = threading.Event() - datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event) - - with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock): + with patch('facefusion.apis.stream_audio.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))): with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()): rtc_peer_audio : RtcPeerAudio =\ { @@ -94,12 +99,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None: 'receiver_track': 0, 'codec': audio_codec } - audio_receiver_thread = threading.Thread(target = receive_audio_frames, args = (rtc_peer_audio, audio_queue), daemon = True) - audio_receiver_thread.start() - ready_event.wait(timeout = 5.0) - datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None) - datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None) - audio_receiver_thread.join(timeout = 5.0) + receive_audio_frames(rtc_peer_audio, audio_queue) _, temp_audio_frame = audio_queue.get_nowait() @@ -113,7 +113,7 @@ def test_handle_audio_frame() -> None: audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300) with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()): - handle_audio_frame('opus', audio_decoder_mock, audio_queue, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p()) + handle_audio_frame('opus', audio_decoder_mock, audio_queue, bytes(), 0) _, temp_audio_frame = audio_queue.get_nowait() diff --git a/tests/test_api_stream_video.py b/tests/test_api_stream_video.py index d741a324..a1e22244 100644 --- a/tests/test_api_stream_video.py +++ b/tests/test_api_stream_video.py @@ -5,7 +5,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from functools import partial from queue import Queue from typing import Tuple -from unittest.mock import MagicMock, patch +from unittest.mock import patch import cv2 import numpy @@ -18,7 +18,7 @@ from facefusion.common_helper import is_linux, is_macos, is_windows from facefusion.download import conditional_download from facefusion.hash_helper import create_hash from facefusion.libraries import aom as aom_module, datachannel as datachannel_module, vpx as vpx_module -from facefusion.types import BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec +from facefusion.types import Buffer, BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec from facefusion.vision import read_video_frame from .assert_helper import get_test_example_file, get_test_examples_directory @@ -44,8 +44,11 @@ def before_each() -> None: rtc_store.clear() -def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None: +def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event: + frame_handler(buffer, 0) + ready_event = threading.Event() ready_event.set() + return ready_event @pytest.mark.parametrize('video_codec, payload_type', [ ('av1', 35), ('vp8', 96), ('vp9', 98) ]) @@ -57,6 +60,12 @@ def test_run_video_encode_loop(video_codec : VideoCodec, payload_type : int) -> rtc_peer : RtcPeer =\ { 'peer_connection': peer_connection, + 'audio': + { + 'sender_track': 0, + 'receiver_track': 0, + 'codec': 'opus' + }, 'video': { 'sender_track': video_sender_track, @@ -101,12 +110,8 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None: video_frame = read_video_frame(get_test_example_file('target-240p.mp4')) video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = 30) - datachannel_mock = MagicMock() - ready_event = threading.Event() - datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event) - with ThreadPoolExecutor(max_workers = 1) as executor: - with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock): + with patch('facefusion.apis.stream_video.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))): with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame): with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))): rtc_peer_video : RtcPeerVideo =\ @@ -115,12 +120,7 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None: 'receiver_track': 0, 'codec': video_codec } - video_receiver_thread = threading.Thread(target = receive_video_frames, args = (rtc_peer_video, video_queue, executor), daemon = True) - video_receiver_thread.start() - ready_event.wait(timeout = 5.0) - datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None) - datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None) - video_receiver_thread.join(timeout = 5.0) + receive_video_frames(rtc_peer_video, video_queue, executor) _, video_future = video_queue.get_nowait() video_buffer = video_future.result().get('buffer') @@ -267,7 +267,7 @@ def test_handle_video_frame(video_codec : VideoCodec) -> None: with ThreadPoolExecutor(max_workers = 1) as executor: with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame): with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))): - handle_video_frame(video_codec, video_decoder, video_queue, executor, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p()) + handle_video_frame(video_codec, video_decoder, video_queue, executor, bytes(), 0) _, video_future = video_queue.get_nowait() video_buffer = video_future.result().get('buffer')