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

facefusion/apis/core.py

@@ -9,7 +9,7 @@ from facefusion.apis.endpoints.metrics import get_metrics, websocket_metrics
 from facefusion.apis.endpoints.ping import websocket_ping
 from facefusion.apis.endpoints.session import create_session, destroy_session, get_session, refresh_session
 from facefusion.apis.endpoints.state import get_state, set_state
-from facefusion.apis.endpoints.stream import webrtc_stream, websocket_stream
+from facefusion.apis.endpoints.stream import post_stream, websocket_stream
 from facefusion.apis.middlewares.session import create_session_guard
 
 
@@ -29,10 +29,10 @@ def create_api() -> Starlette:
 			Route('/assets', delete_assets, methods = [ 'DELETE' ], middleware = [ session_guard ]),
 			Route('/capabilities', get_capabilities, methods = [ 'GET' ]),
 			Route('/metrics', get_metrics, methods = [ 'GET' ], middleware = [ session_guard ]),
-			Route('/stream', webrtc_stream, methods = ['POST'], middleware = [session_guard]),
+			Route('/stream', post_stream, methods = [ 'POST' ], middleware = [ session_guard ]),
 			WebSocketRoute('/metrics', websocket_metrics, middleware = [ session_guard ]),
 			WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ]),
-			WebSocketRoute('/stream', websocket_stream, middleware = [session_guard])
+			WebSocketRoute('/stream', websocket_stream, middleware = [ session_guard ])
 		]
 
 	api = Starlette(routes = routes)

facefusion/apis/endpoints/stream.py

@@ -1,76 +1,33 @@
-from functools import partial
-
-import cv2
-import numpy
-from aiortc import RTCPeerConnection, RTCSessionDescription
 from starlette.requests import Request
-from starlette.responses import JSONResponse, Response
+from starlette.responses import Response
-from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
+from starlette.status import HTTP_201_CREATED, HTTP_404_NOT_FOUND
 from starlette.websockets import WebSocket
 
-from facefusion import session_context, session_manager, state_manager
+from facefusion import rtc_store, session_context, session_manager
-from facefusion.apis.api_helper import get_sec_websocket_protocol
 from facefusion.apis.session_helper import extract_access_token
-from facefusion.apis.stream_helper import create_output_track, on_video_track
+from facefusion.apis.stream_helper import get_websocket_stream_mode, handle_image_stream, handle_video_stream
-from facefusion.streamer import process_vision_frame
 
 
 async def websocket_stream(websocket : WebSocket) -> None:
-	subprotocol = get_sec_websocket_protocol(websocket.scope)
+	stream_mode = get_websocket_stream_mode(websocket.scope)
-	access_token = extract_access_token(websocket.scope)
-	session_id = session_manager.find_session_id(access_token)
 
-	session_context.set_session_id(session_id)
+	if stream_mode == 'image':
-	source_paths = state_manager.get_item('source_paths')
+		await handle_image_stream(websocket)
 
-	await websocket.accept(subprotocol = subprotocol)
+	if stream_mode == 'video':
-
+		await handle_video_stream(websocket)
-	if source_paths:
-		try:
-			image_buffer = await websocket.receive_bytes()
-			target_vision_frame = cv2.imdecode(numpy.frombuffer(image_buffer, numpy.uint8), cv2.IMREAD_COLOR)
-
-			if numpy.any(target_vision_frame):
-				temp_vision_frame = process_vision_frame(target_vision_frame)
-				is_success, output_vision_frame = cv2.imencode('.jpg', temp_vision_frame)
-
-				if is_success:
-					await websocket.send_bytes(output_vision_frame.tobytes())
-
-		except Exception:
-			pass
-		return
-
-	await websocket.close()
 
 
-async def webrtc_stream(request : Request) -> Response:
+async def post_stream(request : Request) -> Response:
 	access_token = extract_access_token(request.scope)
 	session_id = session_manager.find_session_id(access_token)
 	session_context.set_session_id(session_id)
 
 	if session_id:
-		body = await request.json()
+		sdp_offer = (await request.body()).decode()
-		buffer_size = int(body.get('buffer_size', 30))
+		sdp_answer = rtc_store.add_rtc_viewer(session_id, sdp_offer)
-		bitrate_init = int(body.get('bitrate_init', 100000))
-		bitrate_min = int(body.get('bitrate_min', 100000))
-		bitrate_max = int(body.get('bitrate_max', 4000000))
 
-		rtc_offer = RTCSessionDescription(sdp = body.get('sdp'), type = body.get('type'))
+		if sdp_answer:
-		rtc_connection = RTCPeerConnection()
+			return Response(sdp_answer, status_code = HTTP_201_CREATED, media_type = 'application/sdp')
 
-		output_track, sender = create_output_track(rtc_connection, buffer_size)
+	return Response(status_code = HTTP_404_NOT_FOUND)
-		sender.configure_bitrate(bitrate_init, bitrate_min, bitrate_max)
-
-		rtc_connection.on('track', partial(on_video_track, rtc_connection, output_track))
-
-		await rtc_connection.setRemoteDescription(rtc_offer)
-		await rtc_connection.setLocalDescription(await rtc_connection.createAnswer())
-
-		return JSONResponse(
-		{
-			'sdp': rtc_connection.localDescription.sdp,
-			'type': rtc_connection.localDescription.type
-		})
-
-	return Response(status_code = HTTP_500_INTERNAL_SERVER_ERROR)

facefusion/apis/stream_helper.py

@@ -2,53 +2,23 @@ import asyncio
 import math
 import os
 import subprocess
-from typing import Iterator, Optional, Tuple, cast
+from collections import deque
+from collections.abc import AsyncIterator
+from typing import Optional, cast
 
-from aiortc import MediaStreamTrack, QueuedVideoStreamTrack, RTCPeerConnection, RTCRtpSender
+import cv2
-from aiortc.mediastreams import MediaStreamError
+import numpy
-from av import VideoFrame
 from starlette.datastructures import Headers
 from starlette.types import Scope
+from starlette.websockets import WebSocket, WebSocketState
 
+from facefusion import rtc_store, session_context, session_manager, state_manager
+from facefusion.apis.api_helper import get_sec_websocket_protocol
+from facefusion.apis.session_helper import extract_access_token
 from facefusion.common_helper import is_linux, is_macos
+from facefusion.ffmpeg import spawn_stream
 from facefusion.streamer import process_vision_frame
-from facefusion.types import Resolution, StreamBuffer, WebSocketStreamMode
+from facefusion.types import Resolution, SessionId, VisionFrame, WebSocketStreamMode
-
-
-def process_stream_frame(target_stream_frame : VideoFrame) -> VideoFrame:
-	target_vision_frame = target_stream_frame.to_ndarray(format = 'bgr24')
-	output_vision_frame = process_vision_frame(target_vision_frame)
-	output_stream_frame = VideoFrame.from_ndarray(output_vision_frame, format = 'bgr24')
-	output_stream_frame.pts = target_stream_frame.pts
-	output_stream_frame.time_base = target_stream_frame.time_base
-	return output_stream_frame
-
-
-def create_output_track(rtc_connection : RTCPeerConnection, buffer_size : int) -> Tuple[QueuedVideoStreamTrack, RTCRtpSender]:
-	output_track = QueuedVideoStreamTrack(buffer_size = buffer_size)
-	sender = rtc_connection.addTrack(output_track)
-	return output_track, sender
-
-
-async def process_and_enqueue(target_track : MediaStreamTrack, output_track : QueuedVideoStreamTrack) -> None:
-	loop = asyncio.get_running_loop()
-
-	while True:
-		try:
-			target_stream_frame = await target_track.recv()
-		except MediaStreamError:
-			pass
-
-		output_stream_frame = await loop.run_in_executor(None, process_stream_frame, target_stream_frame) #type:ignore[arg-type]
-		await output_track.put(output_stream_frame)
-
-
-def on_video_track(rtc_connection : RTCPeerConnection, output_track : QueuedVideoStreamTrack, target_track : MediaStreamTrack) -> None:
-	if target_track.kind == 'audio':
-		rtc_connection.addTrack(target_track)
-
-	if target_track.kind == 'video':
-		asyncio.create_task(process_and_enqueue(target_track, output_track))
 
 
 def calculate_bitrate(resolution : Resolution) -> int: # TODO : improve the bitrate calculation
@@ -89,8 +59,21 @@ def read_pipe_buffer(pipe_handle : int, size : int) -> Optional[bytes]:
 	return None
 
 
-def forward_stream_frame(process : subprocess.Popen[bytes]) -> Iterator[StreamBuffer]:
+async def receive_vision_frames(websocket : WebSocket) -> AsyncIterator[VisionFrame]:
-	pipe_handle = process.stdout.fileno()
+	websocket_event = await websocket.receive()
+
+	while websocket_event.get('type') == 'websocket.receive':
+		frame_buffer = websocket_event.get('bytes') or b''
+		vision_frame = cv2.imdecode(numpy.frombuffer(frame_buffer, numpy.uint8), cv2.IMREAD_COLOR)
+
+		if numpy.any(vision_frame):
+			yield vision_frame
+
+		websocket_event = await websocket.receive()
+
+
+def forward_rtc_frames(encoder : subprocess.Popen[bytes], session_id : SessionId) -> None:
+	pipe_handle = encoder.stdout.fileno()
 
 	if is_linux() or is_macos():
 		os.set_blocking(pipe_handle, True)
@@ -105,6 +88,83 @@ def forward_stream_frame(process : subprocess.Popen[bytes]) -> Iterator[StreamBu
 			frame_data = read_pipe_buffer(pipe_handle, frame_size)
 
 			if frame_data:
-				yield frame_data
+				rtc_store.send_rtc_frame(session_id, frame_data)
 
 			frame_header = read_pipe_buffer(pipe_handle, 12)
+
+
+def submit_encoder_frame(encoder : subprocess.Popen[bytes], vision_frame_deque : deque[VisionFrame]) -> None:
+	output_vision_frame = process_vision_frame(vision_frame_deque[-1])
+	encoder.stdin.write(cv2.cvtColor(output_vision_frame, cv2.COLOR_BGR2RGB).tobytes())
+	encoder.stdin.flush()
+
+
+def run_encode_loop(encoder : subprocess.Popen[bytes], vision_frame_deque : deque[VisionFrame]) -> None:
+	while vision_frame_deque:
+		submit_encoder_frame(encoder, vision_frame_deque)
+
+	encoder.stdin.close()
+	encoder.wait()
+
+
+async def handle_image_stream(websocket : WebSocket) -> None:
+	subprotocol = get_sec_websocket_protocol(websocket.scope)
+	access_token = extract_access_token(websocket.scope)
+	session_id = session_manager.find_session_id(access_token)
+	session_context.set_session_id(session_id)
+	source_paths = state_manager.get_item('source_paths')
+
+	await websocket.accept(subprotocol = subprotocol)
+
+	if source_paths:
+		capture_vision_frame = await anext(receive_vision_frames(websocket), None)
+
+		if numpy.any(capture_vision_frame):
+			output_vision_frame = process_vision_frame(capture_vision_frame)
+			is_success, output_frame_buffer = cv2.imencode('.jpg', output_vision_frame)
+
+			if is_success:
+				await websocket.send_bytes(output_frame_buffer.tobytes())
+
+	if websocket.client_state == WebSocketState.CONNECTED:
+		await websocket.close()
+
+
+async def handle_video_stream(websocket : WebSocket) -> None:
+	subprotocol = get_sec_websocket_protocol(websocket.scope)
+	access_token = extract_access_token(websocket.scope)
+	session_id = session_manager.find_session_id(access_token)
+	session_context.set_session_id(session_id)
+	source_paths = state_manager.get_item('source_paths')
+
+	await websocket.accept(subprotocol = subprotocol)
+
+	if session_id and source_paths:
+		output_video_fps = int(state_manager.get_item('output_video_fps') or 30) # TODO: resolve from target video fps
+		vision_frames = receive_vision_frames(websocket)
+		vision_frame = await anext(vision_frames, None)
+
+		if numpy.any(vision_frame):
+			resolution = (vision_frame.shape[1], vision_frame.shape[0])
+			encoder = spawn_stream(resolution, output_video_fps, calculate_bitrate(resolution), calculate_buffer_size(resolution))
+
+			vision_frame_deque : deque[VisionFrame] = deque(maxlen = 1)
+
+			vision_frame_deque.append(vision_frame)
+			rtc_store.create_rtc_stream(session_id)
+
+			event_loop = asyncio.get_running_loop()
+			await event_loop.run_in_executor(None, submit_encoder_frame, encoder, vision_frame_deque)
+			await websocket.send_text('ready')
+			encode_task = event_loop.run_in_executor(None, run_encode_loop, encoder, vision_frame_deque)
+			rtc_task = event_loop.run_in_executor(None, forward_rtc_frames, encoder, session_id)
+
+			async for vision_frame in vision_frames:
+				vision_frame_deque.append(vision_frame)
+
+			vision_frame_deque.clear()
+			await asyncio.gather(encode_task, rtc_task)
+			rtc_store.destroy_rtc_stream(session_id)
+
+	if websocket.client_state == WebSocketState.CONNECTED:
+		await websocket.close()

facefusion/rtc_store.py

@@ -1,37 +1,35 @@
 from typing import List, Optional
 
 from facefusion import rtc
-from facefusion.types import RtcPeer, RtcSdpAnswer, RtcSdpOffer, RtcStreamStore
+from facefusion.types import RtcPeer, RtcSdpAnswer, RtcSdpOffer, RtcStreamStore, SessionId
 
-RTC_STREAMS : RtcStreamStore = {} # TODO: tie lifetime to session_id so streams are cleaned up on session expiry
+RTC_STREAMS : RtcStreamStore = {}
 
 
-def get_rtc_stream(stream_path : str) -> Optional[List[RtcPeer]]:
+def get_rtc_stream(session_id : SessionId) -> Optional[List[RtcPeer]]:
-	return RTC_STREAMS.get(stream_path)
+	return RTC_STREAMS.get(session_id)
 
 
-def create_rtc_stream(stream_path : str) -> None:
+def create_rtc_stream(session_id : SessionId) -> None:
-	RTC_STREAMS[stream_path] = []
+	RTC_STREAMS[session_id] = []
 
 
-def destroy_rtc_stream(stream_path : str) -> None:
+def destroy_rtc_stream(session_id : SessionId) -> None:
-	peers = RTC_STREAMS.pop(stream_path, None)
+	peers = RTC_STREAMS.pop(session_id, None)
 
 	if peers:
 		rtc.delete_peers(peers)
 
 
-def add_rtc_viewer(stream_path : str, sdp_offer : RtcSdpOffer) -> Optional[RtcSdpAnswer]:
+def add_rtc_viewer(session_id : SessionId, sdp_offer : RtcSdpOffer) -> Optional[RtcSdpAnswer]:
-	peers = get_rtc_stream(stream_path)
+	if session_id in RTC_STREAMS:
-
+		return rtc.handle_whep_offer(RTC_STREAMS.get(session_id), sdp_offer)
-	if peers:
-		return rtc.handle_whep_offer(peers, sdp_offer)
 
 	return None
 
 
-def send_rtc_frame(stream_path : str, frame_data : bytes) -> None:
+def send_rtc_frame(session_id : SessionId, frame_data : bytes) -> None:
-	peers = get_rtc_stream(stream_path)
+	peers = get_rtc_stream(session_id)
 
 	if peers:
 		rtc.send_to_peers(peers, frame_data)

facefusion/types.py

@@ -265,7 +265,6 @@ BenchmarkCycleSet = TypedDict('BenchmarkCycleSet',
 WebcamMode = Literal['inline', 'udp', 'v4l2']
 StreamMode = Literal['udp', 'v4l2']
 WebSocketStreamMode = Literal['image', 'video']
-StreamBuffer : TypeAlias = bytes
 
 RtcOfferSet = TypedDict('RtcOfferSet',
 {

requirements.txt

@@ -7,7 +7,6 @@ nvidia-ml-py==13.590.48
 psutil==7.2.2
 tqdm==4.67.3
 scipy==1.16.3
-aiortc @ git+https://github.com/facefusion/aiortc.git@feat/dynamic-bitrate
 starlette==0.52.1
 uvicorn==0.41.0
 websockets==16.0

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:
+def create_sdp_offer() -> Optional[RtcSdpOffer]:
-	rtc_connection = RTCPeerConnection()
+	rtc_library = rtc.create_static_rtc_library()
-	rtc_connection.addTrack(VideoStreamTrack())
+	peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True)
-	rtc_offer = await rtc_connection.createOffer()
 
-	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 =\
+	rtc_library.rtcAddTrack(peer_connection, media_video)
-	{
+	rtc_library.rtcAddTrack(peer_connection, media_audio)
-		'sdp': rtc_connection.localDescription.sdp,
+	rtc_library.rtcSetLocalDescription(peer_connection, b'offer')
-		'type': rtc_connection.localDescription.type
-	}
 
-	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())
+	ready_event = threading.Event()
-	stream_response = test_client.post('/stream', json = rtc_offer, headers =
+	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.status_code == 201
-	assert stream_response.json().get('type') == 'answer'
+	assert stream_response.text
-	assert stream_response.json().get('sdp')
+
+	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, get_websocket_stream_mode, read_pipe_buffer
-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
 
 
 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:
+# TODO: add remaining tests
-	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