mirror of
https://github.com/facefusion/facefusion.git
synced 2026-07-06 10:37:54 +02:00
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
This commit is contained in:
Harisreedhar
@@ -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)
|
||||
|
||||
@@ -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.status import HTTP_500_INTERNAL_SERVER_ERROR
|
||||
from starlette.responses import Response
|
||||
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.apis.api_helper import get_sec_websocket_protocol
|
||||
from facefusion import rtc_store, session_context, session_manager
|
||||
from facefusion.apis.session_helper import extract_access_token
|
||||
from facefusion.apis.stream_helper import create_output_track, on_video_track
|
||||
from facefusion.streamer import process_vision_frame
|
||||
from facefusion.apis.stream_helper import get_websocket_stream_mode, handle_image_stream, handle_video_stream
|
||||
|
||||
|
||||
async def websocket_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)
|
||||
stream_mode = get_websocket_stream_mode(websocket.scope)
|
||||
|
||||
session_context.set_session_id(session_id)
|
||||
source_paths = state_manager.get_item('source_paths')
|
||||
if stream_mode == 'image':
|
||||
await handle_image_stream(websocket)
|
||||
|
||||
await websocket.accept(subprotocol = subprotocol)
|
||||
|
||||
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()
|
||||
if stream_mode == 'video':
|
||||
await handle_video_stream(websocket)
|
||||
|
||||
|
||||
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()
|
||||
buffer_size = int(body.get('buffer_size', 30))
|
||||
bitrate_init = int(body.get('bitrate_init', 100000))
|
||||
bitrate_min = int(body.get('bitrate_min', 100000))
|
||||
bitrate_max = int(body.get('bitrate_max', 4000000))
|
||||
sdp_offer = (await request.body()).decode()
|
||||
sdp_answer = rtc_store.add_rtc_viewer(session_id, sdp_offer)
|
||||
|
||||
rtc_offer = RTCSessionDescription(sdp = body.get('sdp'), type = body.get('type'))
|
||||
rtc_connection = RTCPeerConnection()
|
||||
if sdp_answer:
|
||||
return Response(sdp_answer, status_code = HTTP_201_CREATED, media_type = 'application/sdp')
|
||||
|
||||
output_track, sender = create_output_track(rtc_connection, buffer_size)
|
||||
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)
|
||||
return Response(status_code = HTTP_404_NOT_FOUND)
|
||||
|
||||
@@ -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
|
||||
from aiortc.mediastreams import MediaStreamError
|
||||
from av import VideoFrame
|
||||
import cv2
|
||||
import numpy
|
||||
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
|
||||
|
||||
|
||||
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))
|
||||
from facefusion.types import Resolution, SessionId, VisionFrame, WebSocketStreamMode
|
||||
|
||||
|
||||
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]:
|
||||
pipe_handle = process.stdout.fileno()
|
||||
async def receive_vision_frames(websocket : WebSocket) -> AsyncIterator[VisionFrame]:
|
||||
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()
|
||||
|
||||
+13
-15
@@ -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]]:
|
||||
return RTC_STREAMS.get(stream_path)
|
||||
def get_rtc_stream(session_id : SessionId) -> Optional[List[RtcPeer]]:
|
||||
return RTC_STREAMS.get(session_id)
|
||||
|
||||
|
||||
def create_rtc_stream(stream_path : str) -> None:
|
||||
RTC_STREAMS[stream_path] = []
|
||||
def create_rtc_stream(session_id : SessionId) -> None:
|
||||
RTC_STREAMS[session_id] = []
|
||||
|
||||
|
||||
def destroy_rtc_stream(stream_path : str) -> None:
|
||||
peers = RTC_STREAMS.pop(stream_path, None)
|
||||
def destroy_rtc_stream(session_id : SessionId) -> 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]:
|
||||
peers = get_rtc_stream(stream_path)
|
||||
|
||||
if peers:
|
||||
return rtc.handle_whep_offer(peers, sdp_offer)
|
||||
def add_rtc_viewer(session_id : SessionId, sdp_offer : RtcSdpOffer) -> Optional[RtcSdpAnswer]:
|
||||
if session_id in RTC_STREAMS:
|
||||
return rtc.handle_whep_offer(RTC_STREAMS.get(session_id), sdp_offer)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def send_rtc_frame(stream_path : str, frame_data : bytes) -> None:
|
||||
peers = get_rtc_stream(stream_path)
|
||||
def send_rtc_frame(session_id : SessionId, frame_data : bytes) -> None:
|
||||
peers = get_rtc_stream(session_id)
|
||||
|
||||
if peers:
|
||||
rtc.send_to_peers(peers, frame_data)
|
||||
|
||||
@@ -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',
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user