Replace aiortc with libdatachannel direct pipeline (#1083)

* fix stdin close error

* Refactor stream endpoint, fix encoder thread safety and improve tests

* fix and improve test

* remove not None

* use Enum

* use Enum and add todo

* remove poll

tests/stream_helper.py

@@ -1,21 +1,58 @@
-from aiortc import RTCPeerConnection, VideoStreamTrack
+import ctypes
+import os
+import threading
+import time
+from typing import Optional
 
-from facefusion.types import RtcOfferSet
+from starlette.testclient import TestClient
+
+from facefusion import rtc
+from facefusion.types import RtcSdpOffer
 
 
-async def create_rtc_offer() -> RtcOfferSet:
-	rtc_connection = RTCPeerConnection()
-	rtc_connection.addTrack(VideoStreamTrack())
-	rtc_offer = await rtc_connection.createOffer()
+def create_sdp_offer() -> Optional[RtcSdpOffer]:
+	rtc_library = rtc.create_static_rtc_library()
+	peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True)
 
-	await rtc_connection.setLocalDescription(rtc_offer)
+	media_video = os.linesep.join(
+	[
+		'm=video 9 UDP/TLS/RTP/SAVPF 96',
+		'a=rtpmap:96 VP8/90000',
+		'a=recvonly',
+		'a=mid:0',
+		''
+	]).encode()
+	media_audio = os.linesep.join(
+	[
+		'm=audio 9 UDP/TLS/RTP/SAVPF 111',
+		'a=rtpmap:111 opus/48000/2',
+		'a=recvonly',
+		'a=mid:1',
+		''
+	]).encode()
 
-	rtc_offer_set : RtcOfferSet =\
-	{
-		'sdp': rtc_connection.localDescription.sdp,
-		'type': rtc_connection.localDescription.type
-	}
+	rtc_library.rtcAddTrack(peer_connection, media_video)
+	rtc_library.rtcAddTrack(peer_connection, media_audio)
+	rtc_library.rtcSetLocalDescription(peer_connection, b'offer')
 
-	await rtc_connection.close()
+	buffer_size = 16384
+	buffer_string = ctypes.create_string_buffer(buffer_size)
+	wait_limit = time.monotonic() + 5
 
-	return rtc_offer_set
+	while time.monotonic() < wait_limit:
+		if rtc_library.rtcGetLocalDescription(peer_connection, buffer_string, buffer_size) > 0:
+			sdp = buffer_string.value.decode()
+			rtc_library.rtcDeletePeerConnection(peer_connection)
+			return sdp
+		time.sleep(0.05)
+
+	rtc_library.rtcDeletePeerConnection(peer_connection)
+	return None
+
+
+def open_websocket_stream(test_client : TestClient, subprotocols : list[str], source_content : bytes, ready_event : threading.Event, stop_event : threading.Event) -> None:
+	with test_client.websocket_connect('/stream', subprotocols = subprotocols) as websocket:
+		websocket.send_bytes(source_content)
+		websocket.receive_text()
+		ready_event.set()
+		stop_event.wait()

tests/test_api_stream.py

@@ -1,5 +1,5 @@
-import asyncio
 import tempfile
+import threading
 from typing import Iterator
 
 import cv2
@@ -13,7 +13,7 @@ from facefusion.apis.core import create_api
 from facefusion.core import common_pre_check, processors_pre_check
 from facefusion.download import conditional_download
 from .assert_helper import get_test_example_file, get_test_examples_directory
-from .stream_helper import create_rtc_offer
+from .stream_helper import create_sdp_offer, open_websocket_stream
 
 
 @pytest.fixture(scope = 'module', autouse = True)
@@ -92,7 +92,8 @@ def test_stream_image(test_client : TestClient) -> None:
 
 	with test_client.websocket_connect('/stream', subprotocols =
 	[
-		'access_token.' + access_token
+		'access_token.' + access_token,
+		'image'
 	]) as websocket:
 		websocket.send_bytes(source_content)
 		output_bytes = websocket.receive_bytes()
@@ -129,12 +130,21 @@ def test_stream_video(test_client : TestClient) -> None:
 		'Authorization': 'Bearer ' + access_token
 	})
 
-	rtc_offer = asyncio.run(create_rtc_offer())
-	stream_response = test_client.post('/stream', json = rtc_offer, headers =
+	ready_event = threading.Event()
+	stop_event = threading.Event()
+	stream_thread = threading.Thread(target = open_websocket_stream, args = (test_client, [ 'access_token.' + access_token, 'video' ], source_content, ready_event, stop_event))
+	stream_thread.start()
+	ready_event.wait()
+
+	sdp_offer = create_sdp_offer()
+	stream_response = test_client.post('/stream', content = sdp_offer, headers =
 	{
-		'Authorization': 'Bearer ' + access_token
+		'Authorization': 'Bearer ' + access_token,
+		'Content-Type': 'application/sdp'
 	})
 
-	assert stream_response.status_code == 200
-	assert stream_response.json().get('type') == 'answer'
-	assert stream_response.json().get('sdp')
+	assert stream_response.status_code == 201
+	assert stream_response.text
+
+	stop_event.set()
+	stream_thread.join()

tests/test_stream_helper.py

@@ -1,9 +1,6 @@
 import os
-import subprocess
 
-from facefusion import ffmpeg_builder
-from facefusion.apis.stream_helper import calculate_bitrate, calculate_buffer_size, forward_stream_frame, get_websocket_stream_mode, read_pipe_buffer
-from facefusion.vision import pack_resolution
+from facefusion.apis.stream_helper import calculate_bitrate, calculate_buffer_size, get_websocket_stream_mode, read_pipe_buffer
 
 
 def make_scope(protocol : str) -> dict[str, object]:
@@ -30,7 +27,7 @@ def test_calculate_buffer_size() -> None:
 	assert calculate_buffer_size((3840, 2160)) == 14000
 
 
-def test_get_websocket_stream_mode() -> None:
+def test_get_stream_mode() -> None:
 	assert get_websocket_stream_mode(make_scope('image')) == 'image'
 	assert get_websocket_stream_mode(make_scope('video')) == 'video'
 
@@ -47,23 +44,4 @@ def test_read_pipe_buffer() -> None:
 	os.close(read_fd)
 
 
-def test_forward_frames() -> None:
-	resolution = (320, 240)
-	frame_size = resolution[0] * resolution[1] * 3
-	commands = ffmpeg_builder.run(ffmpeg_builder.chain(
-		ffmpeg_builder.capture_video(),
-		ffmpeg_builder.set_media_resolution(pack_resolution(resolution)),
-		ffmpeg_builder.set_input_fps(30),
-		ffmpeg_builder.set_input('-'),
-		ffmpeg_builder.set_video_encoder('libvpx'),
-		ffmpeg_builder.set_encoder_deadline('realtime'),
-		ffmpeg_builder.set_stream_quality(400),
-		ffmpeg_builder.set_muxer('ivf'),
-		ffmpeg_builder.set_output('-')
-	))
-	encoder = subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE, stderr = subprocess.PIPE)
-	encoder.stdin.write(bytes(frame_size))
-	encoder.stdin.close()
-
-	for stream_buffer in forward_stream_frame(encoder):
-		assert 0 < len(stream_buffer) < frame_size
+# TODO: add remaining tests