Implement basic webrtc stream (#1054)

* implement basic webrtc_stream

* add aiortc to requirements.txt

* update aiortc version

* rename variables with rtc_ prefix

* changes

* changes

* change helper to assert_helper and stream_helper

* rename variables with rtc_ prefix

* add error handling

* return whole connection

* remove monkey patch and some cleaning

* cleanup

* tiny adjustments

* tiny adjustments

* proper typing and naming for rtc offer set

* - remove async from on_video_track method
- rename source -> target
- add audio

* audio always before video

---------

Co-authored-by: henryruhs <info@henryruhs.com>

facefusion/apis/endpoints/stream.py

@@ -1,11 +1,17 @@
+from functools import partial
+
 import cv2
 import numpy
+from aiortc import RTCPeerConnection, RTCSessionDescription
 from starlette.requests import Request
-from starlette.websockets import WebSocket, WebSocketDisconnect
+from starlette.responses import JSONResponse, Response
+from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
+from starlette.websockets import WebSocket
 
 from facefusion import session_context, session_manager, state_manager
 from facefusion.apis.api_helper import get_sec_websocket_protocol
-from facefusion.apis.endpoints.session import extract_access_token
+from facefusion.apis.session_helper import extract_access_token
+from facefusion.apis.stream_helper import on_video_track
 from facefusion.streamer import process_stream_frame
 
 
@@ -21,8 +27,8 @@ async def websocket_stream(websocket : WebSocket) -> None:
 
 	if source_paths:
 		try:
-			image_bytes = await websocket.receive_bytes()
-			target_vision_frame = cv2.imdecode(numpy.frombuffer(image_bytes, numpy.uint8), cv2.IMREAD_COLOR)
+			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_stream_frame(target_vision_frame)
@@ -31,12 +37,32 @@ async def websocket_stream(websocket : WebSocket) -> None:
 				if is_success:
 					await websocket.send_bytes(output_vision_frame.tobytes())
 
-		except (WebSocketDisconnect, OSError):
+		except Exception:
 			pass
 		return
 
 	await websocket.close()
 
 
-async def webrtc_stream(request : Request) -> None: # TODO: implement webrtc streaming
-	pass
+async def webrtc_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()
+		rtc_offer = RTCSessionDescription(sdp = body.get('sdp'), type = body.get('type'))
+		rtc_connection = RTCPeerConnection()
+
+		rtc_connection.on('track', partial(on_video_track, rtc_connection))
+
+		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

@@ -0,0 +1,31 @@
+import asyncio
+from typing import cast
+
+from aiortc import MediaStreamTrack, RTCPeerConnection, VideoStreamTrack
+from av import VideoFrame
+
+from facefusion.streamer import process_stream_frame
+
+
+def create_output_track(target_track : MediaStreamTrack) -> VideoStreamTrack:
+	output_track = VideoStreamTrack()
+
+	async def read_stream_frame() -> VideoFrame:
+		target_stream_frame = cast(VideoFrame, await target_track.recv())
+		output_vision_frame = await asyncio.get_running_loop().run_in_executor(None, process_stream_frame, target_stream_frame.to_ndarray(format = 'bgr24'))
+		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
+
+	output_track.recv = read_stream_frame
+	return output_track
+
+
+def on_video_track(rtc_connection : RTCPeerConnection, target_track : MediaStreamTrack) -> None:
+	if target_track.kind == 'audio':
+		rtc_connection.addTrack(target_track)
+
+	if target_track.kind == 'video':
+		output_track = create_output_track(target_track)
+		rtc_connection.addTrack(output_track)

facefusion/types.py

@@ -257,6 +257,12 @@ BenchmarkCycleSet = TypedDict('BenchmarkCycleSet',
 WebcamMode = Literal['inline', 'udp', 'v4l2']
 StreamMode = Literal['udp', 'v4l2']
 
+RtcOfferSet = TypedDict('RtcOfferSet',
+{
+	'sdp': str,
+	'type': str
+})
+
 ModelOptions : TypeAlias = Dict[str, Any]
 ModelSet : TypeAlias = Dict[str, ModelOptions]
 ModelInitializer : TypeAlias = NDArray[Any]