Resolve final stream TODOs (#1155)

* refactor a lot

* refactor a lot

* fix test

* fix test

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:
+def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[Time, AudioFrame]], audio_buffer : Buffer, audio_timestamp : int) -> None:
-	audio_buffer = ctypes.string_at(data, size)
 	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_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32)
-		audio_queue.put((time.monotonic(), temp_audio_frame))
+		audio_time = rtc.convert_timestamp_to_time(audio_codec, audio_timestamp)
+		audio_queue.put((audio_time, audio_frame))

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

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 = rtc_peer.get('sender_bitrate').value
-			sender_bitrate = calculate_sender_bitrate(rtc_peer, temp_bitrate)
 
 			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
+			receiver_bitrate = rtc_peer.get('receiver_bitrate').value
-			frame_interval = video_time - temp_video_time
-			temp_video_time = video_time
-
-			receiver_bitrate = calculate_receiver_bitrate(rtc_peer, encode_time, frame_interval)
 			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, video_buffer : Buffer, video_timestamp : int) -> None:
-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)
 	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))

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

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]

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()
+	with patch('facefusion.apis.stream_audio.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))):
-	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.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)
+			receive_audio_frames(rtc_peer_audio, audio_queue)
-			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)
 
 	_, 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()
 

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)
+					receive_video_frames(rtc_peer_video, video_queue, executor)
-					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)
 					_, 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')