mirror of
https://github.com/facefusion/facefusion.git
synced 2026-06-02 10:51:39 +02:00
Refactor RTC structure (#1113)
* refactor rtc part1 * skip for macos * merge create spd and create sdp offer * fix lint * add test for create_sdp_offer * better naming for negotiate method as we get an answer * extend tests based on mutations * remove dead code * rename rtc store and related methods * clean store, move sender logic to stream helper under apis * generate tests for rtc store
This commit is contained in:
Henry Ruhs
@@ -3,9 +3,9 @@ from starlette.responses import Response
|
||||
from starlette.status import HTTP_201_CREATED, HTTP_404_NOT_FOUND
|
||||
from starlette.websockets import WebSocket
|
||||
|
||||
from facefusion import rtc_store, session_context, session_manager
|
||||
from facefusion import session_context, session_manager
|
||||
from facefusion.apis.session_helper import extract_access_token
|
||||
from facefusion.apis.stream_helper import handle_image_stream, handle_video_stream
|
||||
from facefusion.apis.stream_helper import add_rtc_viewer, handle_image_stream, handle_video_stream
|
||||
|
||||
|
||||
# TODO: reject websocket with invalid or missing stream mode
|
||||
@@ -27,7 +27,7 @@ async def post_stream(request : Request) -> Response:
|
||||
|
||||
if session_id:
|
||||
sdp_offer = await request.body()
|
||||
sdp_answer = rtc_store.add_rtc_viewer(session_id, sdp_offer.decode())
|
||||
sdp_answer = add_rtc_viewer(session_id, sdp_offer.decode())
|
||||
|
||||
if sdp_answer:
|
||||
return Response(sdp_answer, status_code = HTTP_201_CREATED, media_type = 'application/sdp')
|
||||
|
||||
@@ -1,20 +1,20 @@
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Tuple, cast, get_args
|
||||
from typing import Optional, Tuple, cast, get_args
|
||||
|
||||
import cv2
|
||||
import numpy
|
||||
from starlette.websockets import WebSocket, WebSocketState
|
||||
|
||||
from facefusion import rtc_store, session_context, session_manager, state_manager
|
||||
from facefusion import rtc, 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.codecs.aom import create_aom_encoder, destroy_aom_encoder, encode_aom_buffer
|
||||
from facefusion.codecs.opus import create_opus_encoder, destroy_opus_encoder, encode_opus_buffer
|
||||
from facefusion.codecs.vpx import create_vpx_encoder, destroy_vpx_encoder, encode_vpx_buffer
|
||||
from facefusion.streamer import process_vision_frame
|
||||
from facefusion.types import Resolution, SessionId, VideoCodec, VisionFrame
|
||||
from facefusion.types import AudioCodec, PeerConnection, Resolution, RtcAudioTrack, RtcPeer, RtcVideoTrack, SdpAnswer, SdpOffer, SessionId, VideoCodec, VisionFrame
|
||||
|
||||
|
||||
async def receive_stream_frames(websocket : WebSocket) -> AsyncIterator[Tuple[int, bytes]]:
|
||||
@@ -42,6 +42,40 @@ async def receive_vision_frames(websocket : WebSocket) -> AsyncIterator[VisionFr
|
||||
websocket_event = await websocket.receive()
|
||||
|
||||
|
||||
# TODO: clean up peer connection on failed sdp negotiation, wrap in run_in_executor to avoid blocking async event loop
|
||||
def add_rtc_viewer(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpAnswer]:
|
||||
rtc_peers = rtc_store.get_rtc_peers(session_id)
|
||||
|
||||
if rtc_peers is not None:
|
||||
payload_types = rtc.parse_sdp_payload_types(sdp_offer)
|
||||
peer_connection : PeerConnection = rtc.create_peer_connection()
|
||||
audio_codec : AudioCodec = 'opus'
|
||||
audio_track : RtcAudioTrack = rtc.add_audio_track(peer_connection, 'sendonly', audio_codec, payload_types.get(audio_codec, 111))
|
||||
|
||||
#TODO: Fix me via resolve method
|
||||
video_codec : VideoCodec = 'av1'
|
||||
if payload_types.get('av1'):
|
||||
video_codec = 'av1'
|
||||
if payload_types.get('vp8'):
|
||||
video_codec = 'vp8'
|
||||
|
||||
video_track : RtcVideoTrack = rtc.add_video_track(peer_connection, 'sendonly', video_codec, payload_types.get(video_codec, 96))
|
||||
local_sdp = rtc.negotiate_sdp_answer(peer_connection, sdp_offer)
|
||||
|
||||
if local_sdp:
|
||||
rtc_peer : RtcPeer =\
|
||||
{
|
||||
'peer_connection': peer_connection,
|
||||
'video_track': video_track,
|
||||
'audio_track': audio_track
|
||||
}
|
||||
rtc_peers.append(rtc_peer)
|
||||
|
||||
return local_sdp
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def run_aom_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : SessionId, initial_resolution : Resolution, keyframe_interval : int) -> None:
|
||||
aom_encoder = create_aom_encoder(initial_resolution, 4500, 8, 10)
|
||||
current_resolution = initial_resolution
|
||||
@@ -65,7 +99,10 @@ def run_aom_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : Se
|
||||
frame_buffer = encode_aom_buffer(aom_encoder, yuv_frame.tobytes(), frame_resolution, pts)
|
||||
|
||||
if frame_buffer:
|
||||
rtc_store.send_rtc_video(session_id, frame_buffer)
|
||||
rtc_peers = rtc_store.get_rtc_peers(session_id)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.send_video_to_peers(rtc_peers, frame_buffer)
|
||||
|
||||
pts += 1
|
||||
|
||||
@@ -96,7 +133,10 @@ def run_vp8_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : Se
|
||||
frame_buffer = encode_vpx_buffer(vpx_encoder, yuv_frame.tobytes(), frame_resolution, pts)
|
||||
|
||||
if frame_buffer:
|
||||
rtc_store.send_rtc_video(session_id, frame_buffer)
|
||||
rtc_peers = rtc_store.get_rtc_peers(session_id)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.send_video_to_peers(rtc_peers, frame_buffer)
|
||||
|
||||
pts += 1
|
||||
|
||||
@@ -159,7 +199,7 @@ async def handle_video_stream(websocket : WebSocket) -> None:
|
||||
audio_timestamp = 0
|
||||
|
||||
vision_frame_deque.append(first_vision_frame)
|
||||
rtc_store.create_rtc_stream(session_id)
|
||||
rtc_store.create_rtc_peers(session_id)
|
||||
|
||||
event_loop = asyncio.get_running_loop()
|
||||
encode_loop = run_aom_encode_loop
|
||||
@@ -186,7 +226,10 @@ async def handle_video_stream(websocket : WebSocket) -> None:
|
||||
audio_buffer = encode_opus_buffer(opus_encoder, audio_chunk.tobytes(), 960)
|
||||
|
||||
if audio_buffer:
|
||||
rtc_store.send_rtc_audio(session_id, audio_buffer, audio_timestamp)
|
||||
rtc_peers = rtc_store.get_rtc_peers(session_id)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.send_audio_to_peers(rtc_peers, audio_buffer, audio_timestamp)
|
||||
|
||||
audio_timestamp += 960
|
||||
|
||||
@@ -196,7 +239,7 @@ async def handle_video_stream(websocket : WebSocket) -> None:
|
||||
if opus_encoder:
|
||||
destroy_opus_encoder(opus_encoder)
|
||||
|
||||
rtc_store.destroy_rtc_stream(session_id)
|
||||
rtc_store.destroy_rtc_peers(session_id)
|
||||
|
||||
if websocket.client_state == WebSocketState.CONNECTED:
|
||||
await websocket.close()
|
||||
|
||||
+109
-126
@@ -43,115 +43,9 @@ def create_peer_connection(
|
||||
return datachannel_library.rtcCreatePeerConnection(ctypes.byref(rtc_configuration))
|
||||
|
||||
|
||||
#TODO: I think we dont need this method at all
|
||||
def build_media_description(media_type : str, payload_type : int, rtp_codec : str, media_direction : MediaDirection, media_id : int) -> bytes:
|
||||
lines =\
|
||||
[
|
||||
'm=' + media_type + ' 9 UDP/TLS/RTP/SAVPF ' + str(payload_type),
|
||||
'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec,
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack',
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack pli',
|
||||
'a=' + media_direction,
|
||||
'a=mid:' + str(media_id),
|
||||
'a=rtcp-mux',
|
||||
''
|
||||
]
|
||||
return '\r\n'.join(lines).encode()
|
||||
|
||||
|
||||
def parse_sdp_payload_types(sdp_offer : SdpOffer) -> Dict[str, int]:
|
||||
payload_types : Dict[str, int] = {}
|
||||
|
||||
# TODO: consider having a codec helper to resolve these
|
||||
for line in sdp_offer.splitlines():
|
||||
if line.startswith('a=rtpmap:') and 'AV1/90000' in line and not payload_types.get('av1'):
|
||||
payload_types['av1'] = int(line.split(':')[1].split(' ')[0])
|
||||
if line.startswith('a=rtpmap:') and 'VP8/90000' in line and not payload_types.get('vp8'):
|
||||
payload_types['vp8'] = int(line.split(':')[1].split(' ')[0])
|
||||
if line.startswith('a=rtpmap:') and 'opus/48000/2' in line and not payload_types.get('opus'):
|
||||
payload_types['opus'] = int(line.split(':')[1].split(' ')[0])
|
||||
|
||||
return payload_types
|
||||
|
||||
|
||||
def add_audio_track(peer_connection : PeerConnection, media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> RtcAudioTrack:
|
||||
# TODO: check if sleep is needed
|
||||
def create_sdp_offer(peer_connection : PeerConnection) -> Optional[SdpOffer]:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
|
||||
# TODO: Fix me via resolve method
|
||||
rtp_codec = 'opus/48000/2'
|
||||
if audio_codec == 'opus':
|
||||
rtp_codec = 'opus/48000/2'
|
||||
|
||||
media_description = build_media_description('audio', payload_type, rtp_codec, media_direction, 1)
|
||||
|
||||
audio_track = datachannel_library.rtcAddTrack(peer_connection, media_description)
|
||||
|
||||
audio_packetizer = datachannel_module.define_rtc_packetizer_init()
|
||||
audio_packetizer.ssrc = 43
|
||||
audio_packetizer.cname = b'audio'
|
||||
audio_packetizer.payloadType = payload_type
|
||||
audio_packetizer.clockRate = 48000
|
||||
|
||||
if audio_codec == 'opus':
|
||||
datachannel_library.rtcSetOpusPacketizer(audio_track, ctypes.byref(audio_packetizer))
|
||||
|
||||
datachannel_library.rtcChainRtcpSrReporter(audio_track)
|
||||
|
||||
return audio_track
|
||||
|
||||
|
||||
def add_video_track(peer_connection : PeerConnection, media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> RtcVideoTrack:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
|
||||
#TODO: Fix me via resolve method
|
||||
rtp_codec = 'AV1/90000'
|
||||
if video_codec == 'av1':
|
||||
rtp_codec = 'AV1/90000'
|
||||
if video_codec == 'vp8':
|
||||
rtp_codec = 'VP8/90000'
|
||||
|
||||
media_description = build_media_description('video', payload_type, rtp_codec, media_direction, 0)
|
||||
|
||||
video_track = datachannel_library.rtcAddTrack(peer_connection, media_description)
|
||||
|
||||
video_packetizer = datachannel_module.define_rtc_packetizer_init()
|
||||
video_packetizer.ssrc = 42
|
||||
video_packetizer.cname = b'video'
|
||||
video_packetizer.payloadType = payload_type
|
||||
video_packetizer.clockRate = 90000
|
||||
video_packetizer.maxFragmentSize = 1200
|
||||
|
||||
if video_codec == 'av1':
|
||||
video_packetizer.obuPacketization = 1
|
||||
datachannel_library.rtcSetAV1Packetizer(video_track, ctypes.byref(video_packetizer))
|
||||
if video_codec == 'vp8':
|
||||
datachannel_library.rtcSetVP8Packetizer(video_track, ctypes.byref(video_packetizer))
|
||||
|
||||
datachannel_library.rtcChainRtcpSrReporter(video_track)
|
||||
datachannel_library.rtcChainRtcpNackResponder(video_track, 512)
|
||||
|
||||
return video_track
|
||||
|
||||
|
||||
def create_sdp(peer_connection : PeerConnection) -> Optional[SdpOffer]:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
datachannel_library.rtcSetLocalDescription(peer_connection, b'offer')
|
||||
buffer_size = 8192
|
||||
buffer_string = ctypes.create_string_buffer(buffer_size)
|
||||
|
||||
if datachannel_library.rtcGetLocalDescription(peer_connection, buffer_string, buffer_size) > 0:
|
||||
return buffer_string.value.decode()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# TODO: move from testing suite helper to rtc.py - belongs here to complete the rtc flow
|
||||
def create_sdp_offer() -> Optional[SdpOffer]:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
peer_connection = create_peer_connection(disable_auto_negotiation = True)
|
||||
|
||||
datachannel_library.rtcAddTrack(peer_connection, build_media_description('video', 96, 'VP8/90000', 'recvonly', 0))
|
||||
datachannel_library.rtcAddTrack(peer_connection, build_media_description('audio', 111, 'opus/48000/2', 'recvonly', 1))
|
||||
datachannel_library.rtcSetLocalDescription(peer_connection, b'offer')
|
||||
|
||||
buffer_size = 16384
|
||||
@@ -160,31 +54,15 @@ def create_sdp_offer() -> Optional[SdpOffer]:
|
||||
|
||||
while time.monotonic() < wait_limit:
|
||||
if datachannel_library.rtcGetLocalDescription(peer_connection, buffer_string, buffer_size) > 0:
|
||||
sdp = buffer_string.value.decode()
|
||||
datachannel_library.rtcDeletePeerConnection(peer_connection)
|
||||
return sdp
|
||||
return buffer_string.value.decode()
|
||||
|
||||
time.sleep(0.05)
|
||||
|
||||
datachannel_library.rtcDeletePeerConnection(peer_connection)
|
||||
return None
|
||||
|
||||
|
||||
@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_void_p)
|
||||
def on_sdp_ready(peer_connection : int, sdp : Optional[bytes], sdp_type : int, user_pointer : Optional[int]) -> None:
|
||||
ctypes.cast(user_pointer, ctypes.py_object).value.set()
|
||||
|
||||
|
||||
# TODO: unused callback, remove or wire up
|
||||
@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int, ctypes.c_void_p)
|
||||
def on_ice_complete(peer_connection : int, state : int, user_pointer : Optional[int]) -> None:
|
||||
if state == 2:
|
||||
context = ctypes.cast(user_pointer, ctypes.py_object).value
|
||||
context['event'].set()
|
||||
|
||||
|
||||
# TODO: sanitize sdp_offer, wrap in run_in_executor, track peer connection state
|
||||
def negotiate_sdp(peer_connection : PeerConnection, sdp_offer : SdpOffer) -> Optional[SdpAnswer]:
|
||||
def negotiate_sdp_answer(peer_connection : PeerConnection, sdp_offer : SdpOffer) -> Optional[SdpAnswer]:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
sdp_event = threading.Event()
|
||||
sdp_event_pointer = ctypes.cast(id(sdp_event), ctypes.c_void_p)
|
||||
@@ -249,3 +127,108 @@ def delete_peers(rtc_peers : List[RtcPeer]) -> None:
|
||||
datachannel_library.rtcDeletePeerConnection(peer_connection_id)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def add_audio_track(peer_connection : PeerConnection, media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> RtcAudioTrack:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
media_description = create_audio_description(media_direction, audio_codec, payload_type)
|
||||
|
||||
audio_track = datachannel_library.rtcAddTrack(peer_connection, media_description)
|
||||
|
||||
audio_packetizer = datachannel_module.define_rtc_packetizer_init()
|
||||
audio_packetizer.ssrc = 43
|
||||
audio_packetizer.cname = b'audio'
|
||||
audio_packetizer.payloadType = payload_type
|
||||
audio_packetizer.clockRate = 48000
|
||||
|
||||
if audio_codec == 'opus':
|
||||
datachannel_library.rtcSetOpusPacketizer(audio_track, ctypes.byref(audio_packetizer))
|
||||
|
||||
datachannel_library.rtcChainRtcpSrReporter(audio_track)
|
||||
|
||||
return audio_track
|
||||
|
||||
|
||||
def add_video_track(peer_connection : PeerConnection, media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> RtcVideoTrack:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
media_description = create_video_description(media_direction, video_codec, payload_type)
|
||||
|
||||
video_track = datachannel_library.rtcAddTrack(peer_connection, media_description)
|
||||
|
||||
video_packetizer = datachannel_module.define_rtc_packetizer_init()
|
||||
video_packetizer.ssrc = 42
|
||||
video_packetizer.cname = b'video'
|
||||
video_packetizer.payloadType = payload_type
|
||||
video_packetizer.clockRate = 90000
|
||||
video_packetizer.maxFragmentSize = 1200
|
||||
|
||||
if video_codec == 'av1':
|
||||
video_packetizer.obuPacketization = 1
|
||||
datachannel_library.rtcSetAV1Packetizer(video_track, ctypes.byref(video_packetizer))
|
||||
if video_codec == 'vp8':
|
||||
datachannel_library.rtcSetVP8Packetizer(video_track, ctypes.byref(video_packetizer))
|
||||
|
||||
datachannel_library.rtcChainRtcpSrReporter(video_track)
|
||||
datachannel_library.rtcChainRtcpNackResponder(video_track, 512)
|
||||
|
||||
return video_track
|
||||
|
||||
|
||||
def create_audio_description(media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> bytes:
|
||||
rtp_codec = 'opus/48000/2'
|
||||
if audio_codec == 'opus':
|
||||
rtp_codec = 'opus/48000/2'
|
||||
|
||||
lines =\
|
||||
[
|
||||
'm=audio 9 UDP/TLS/RTP/SAVPF ' + str(payload_type),
|
||||
'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec,
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack',
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack pli',
|
||||
'a=' + media_direction,
|
||||
'a=mid:1',
|
||||
'a=rtcp-mux',
|
||||
''
|
||||
]
|
||||
return '\r\n'.join(lines).encode()
|
||||
|
||||
|
||||
def create_video_description(media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> bytes:
|
||||
rtp_codec = 'AV1/90000'
|
||||
if video_codec == 'av1':
|
||||
rtp_codec = 'AV1/90000'
|
||||
if video_codec == 'vp8':
|
||||
rtp_codec = 'VP8/90000'
|
||||
|
||||
lines =\
|
||||
[
|
||||
'm=video 9 UDP/TLS/RTP/SAVPF ' + str(payload_type),
|
||||
'a=rtpmap:' + str(payload_type) + ' ' + rtp_codec,
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack',
|
||||
'a=rtcp-fb:' + str(payload_type) + ' nack pli',
|
||||
'a=' + media_direction,
|
||||
'a=mid:0',
|
||||
'a=rtcp-mux',
|
||||
''
|
||||
]
|
||||
return '\r\n'.join(lines).encode()
|
||||
|
||||
|
||||
def parse_sdp_payload_types(sdp_offer : SdpOffer) -> Dict[str, int]:
|
||||
payload_types : Dict[str, int] = {}
|
||||
|
||||
# TODO: consider having a codec helper to resolve these
|
||||
for line in sdp_offer.splitlines():
|
||||
if line.startswith('a=rtpmap:') and 'AV1/90000' in line and not payload_types.get('av1'):
|
||||
payload_types['av1'] = int(line.split(':')[1].split(' ')[0])
|
||||
if line.startswith('a=rtpmap:') and 'VP8/90000' in line and not payload_types.get('vp8'):
|
||||
payload_types['vp8'] = int(line.split(':')[1].split(' ')[0])
|
||||
if line.startswith('a=rtpmap:') and 'opus/48000/2' in line and not payload_types.get('opus'):
|
||||
payload_types['opus'] = int(line.split(':')[1].split(' ')[0])
|
||||
|
||||
return payload_types
|
||||
|
||||
|
||||
@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_void_p)
|
||||
def on_sdp_ready(peer_connection : int, sdp : Optional[bytes], sdp_type : int, user_pointer : Optional[int]) -> None:
|
||||
ctypes.cast(user_pointer, ctypes.py_object).value.set()
|
||||
|
||||
+9
-57
@@ -1,70 +1,22 @@
|
||||
from typing import List, Optional
|
||||
from typing import List
|
||||
|
||||
from facefusion import rtc
|
||||
from facefusion.types import AudioCodec, PeerConnection, RtcAudioTrack, RtcPeer, RtcStreamStore, RtcVideoTrack, SdpAnswer, SdpOffer, SessionId, VideoCodec
|
||||
from facefusion.types import RtcPeer, RtcStore, SessionId
|
||||
|
||||
|
||||
# TODO: aint this a peer store?
|
||||
RTC_STREAM_STORE : RtcStreamStore = {}
|
||||
RTC_STORE : RtcStore = {}
|
||||
|
||||
|
||||
def get_rtc_stream(session_id : SessionId) -> Optional[List[RtcPeer]]:
|
||||
return RTC_STREAM_STORE.get(session_id)
|
||||
def get_rtc_peers(session_id : SessionId) -> List[RtcPeer]:
|
||||
return RTC_STORE.get(session_id)
|
||||
|
||||
|
||||
def create_rtc_stream(session_id : SessionId) -> None:
|
||||
RTC_STREAM_STORE[session_id] = []
|
||||
def create_rtc_peers(session_id : SessionId) -> None:
|
||||
RTC_STORE[session_id] = []
|
||||
|
||||
|
||||
def destroy_rtc_stream(session_id : SessionId) -> None:
|
||||
rtc_peers = RTC_STREAM_STORE.pop(session_id, None)
|
||||
def destroy_rtc_peers(session_id : SessionId) -> None:
|
||||
rtc_peers = RTC_STORE.pop(session_id, None)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.delete_peers(rtc_peers)
|
||||
|
||||
|
||||
# TODO: clean up peer connection on failed sdp negotiation, wrap in run_in_executor to avoid blocking async event loop
|
||||
def add_rtc_viewer(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpAnswer]:
|
||||
if session_id in RTC_STREAM_STORE:
|
||||
payload_types = rtc.parse_sdp_payload_types(sdp_offer)
|
||||
peer_connection : PeerConnection = rtc.create_peer_connection()
|
||||
audio_codec : AudioCodec = 'opus'
|
||||
audio_track : RtcAudioTrack = rtc.add_audio_track(peer_connection, 'sendonly', audio_codec, payload_types.get(audio_codec, 111))
|
||||
|
||||
#TODO: Fix me via resolve method
|
||||
video_codec : VideoCodec = 'av1'
|
||||
if payload_types.get('av1'):
|
||||
video_codec = 'av1'
|
||||
if payload_types.get('vp8'):
|
||||
video_codec = 'vp8'
|
||||
|
||||
video_track : RtcVideoTrack = rtc.add_video_track(peer_connection, 'sendonly', video_codec, payload_types.get(video_codec, 96))
|
||||
local_sdp = rtc.negotiate_sdp(peer_connection, sdp_offer)
|
||||
|
||||
if local_sdp:
|
||||
rtc_peer : RtcPeer =\
|
||||
{
|
||||
'peer_connection': peer_connection,
|
||||
'video_track': video_track,
|
||||
'audio_track': audio_track
|
||||
}
|
||||
RTC_STREAM_STORE[session_id].append(rtc_peer)
|
||||
|
||||
return local_sdp
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# TODO: detect and remove dead peers
|
||||
def send_rtc_video(session_id : SessionId, frame_buffer : bytes) -> None:
|
||||
rtc_peers = get_rtc_stream(session_id)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.send_video_to_peers(rtc_peers, frame_buffer)
|
||||
|
||||
|
||||
def send_rtc_audio(session_id : SessionId, audio_buffer : bytes, audio_pts : int) -> None:
|
||||
rtc_peers = get_rtc_stream(session_id)
|
||||
|
||||
if rtc_peers:
|
||||
rtc.send_audio_to_peers(rtc_peers, audio_buffer, audio_pts)
|
||||
|
||||
+1
-1
@@ -286,7 +286,7 @@ RtcPeer = TypedDict('RtcPeer',
|
||||
'audio_track': RtcAudioTrack,
|
||||
})
|
||||
|
||||
RtcStreamStore : TypeAlias = Dict[SessionId, List[RtcPeer]]
|
||||
RtcStore : TypeAlias = Dict[SessionId, List[RtcPeer]]
|
||||
|
||||
ModelOptions : TypeAlias = Dict[str, Any]
|
||||
ModelSet : TypeAlias = Dict[str, ModelOptions]
|
||||
|
||||
@@ -13,6 +13,7 @@ from facefusion.apis.core import create_api
|
||||
from facefusion.core import common_pre_check
|
||||
from facefusion.download import conditional_download
|
||||
from facefusion.hash_helper import create_hash
|
||||
from facefusion.libraries import datachannel as datachannel_module
|
||||
from .assert_helper import get_test_example_file, get_test_examples_directory
|
||||
|
||||
|
||||
@@ -123,7 +124,7 @@ def test_stream_video(test_client : TestClient, create_event : threading.Event,
|
||||
'Authorization': 'Bearer ' + access_token
|
||||
})
|
||||
|
||||
with patch('facefusion.rtc_store.send_rtc_video', side_effect = partial(set_event, event = create_event)):
|
||||
with patch('facefusion.rtc.send_video_to_peers', side_effect = partial(set_event, event = create_event)):
|
||||
with test_client.websocket_connect('/stream?mode=video&codec=' + video_codec, subprotocols =
|
||||
[
|
||||
'access_token.' + access_token
|
||||
@@ -131,7 +132,11 @@ def test_stream_video(test_client : TestClient, create_event : threading.Event,
|
||||
websocket.send_bytes(chr(1).encode() + source_content)
|
||||
websocket.receive_text()
|
||||
|
||||
sdp_offer = rtc.create_sdp_offer()
|
||||
peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True)
|
||||
rtc.add_video_track(peer_connection, 'recvonly', 'vp8', 96)
|
||||
rtc.add_audio_track(peer_connection, 'recvonly', 'opus', 111)
|
||||
sdp_offer = rtc.create_sdp_offer(peer_connection)
|
||||
datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection)
|
||||
stream_response = test_client.post('/stream', content = sdp_offer, headers =
|
||||
{
|
||||
'Authorization': 'Bearer ' + access_token,
|
||||
|
||||
+51
-20
@@ -16,12 +16,6 @@ def before_all() -> None:
|
||||
vpx_module.pre_check()
|
||||
|
||||
|
||||
# TODO: add test_parse_sdp_payload_types
|
||||
def test_build_media_description() -> None:
|
||||
assert rtc.build_media_description('audio', 111, 'opus/48000/2', 'sendonly', 1) == b'm=audio 9 UDP/TLS/RTP/SAVPF 111\r\na=rtpmap:111 opus/48000/2\r\na=rtcp-fb:111 nack\r\na=rtcp-fb:111 nack pli\r\na=sendonly\r\na=mid:1\r\na=rtcp-mux\r\n'
|
||||
assert rtc.build_media_description('video', 96, 'VP8/90000', 'recvonly', 0) == b'm=video 9 UDP/TLS/RTP/SAVPF 96\r\na=rtpmap:96 VP8/90000\r\na=rtcp-fb:96 nack\r\na=rtcp-fb:96 nack pli\r\na=recvonly\r\na=mid:0\r\na=rtcp-mux\r\n'
|
||||
|
||||
|
||||
def test_create_peer_connection() -> None:
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
@@ -30,40 +24,41 @@ def test_create_peer_connection() -> None:
|
||||
assert datachannel_library.rtcDeletePeerConnection(peer_connection) == 0
|
||||
|
||||
|
||||
def test_add_audio_track() -> None:
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
def test_create_sdp_offer() -> None:
|
||||
peer_connection = rtc.create_peer_connection(disable_auto_negotiation = True)
|
||||
rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96)
|
||||
rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111)
|
||||
sdp_offer = rtc.create_sdp_offer(peer_connection)
|
||||
|
||||
assert rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111) > 0
|
||||
assert sdp_offer
|
||||
assert 'm=video' in sdp_offer
|
||||
assert 'VP8/90000' in sdp_offer
|
||||
assert 'm=audio' in sdp_offer
|
||||
assert 'opus/48000/2' in sdp_offer
|
||||
|
||||
datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection)
|
||||
|
||||
|
||||
def test_add_video_track() -> None:
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
|
||||
assert rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96) > 0
|
||||
|
||||
datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection)
|
||||
|
||||
|
||||
def test_negotiate_sdp() -> None:
|
||||
def test_negotiate_sdp_answer() -> None:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
|
||||
sender_connection = rtc.create_peer_connection()
|
||||
rtc.add_video_track(sender_connection, 'sendonly', 'vp8', 96)
|
||||
rtc.add_audio_track(sender_connection, 'sendonly', 'opus', 111)
|
||||
sdp_offer = rtc.create_sdp(sender_connection)
|
||||
sdp_offer = rtc.create_sdp_offer(sender_connection)
|
||||
|
||||
receiver_connection = rtc.create_peer_connection()
|
||||
rtc.add_video_track(receiver_connection, 'recvonly', 'vp8', 96)
|
||||
rtc.add_audio_track(receiver_connection, 'recvonly', 'opus', 111)
|
||||
sdp_answer = rtc.negotiate_sdp(receiver_connection, sdp_offer)
|
||||
sdp_answer = rtc.negotiate_sdp_answer(receiver_connection, sdp_offer)
|
||||
|
||||
assert sdp_answer
|
||||
assert 'm=video' in sdp_answer
|
||||
assert 'VP8/90000' in sdp_answer
|
||||
assert 'm=audio' in sdp_answer
|
||||
assert 'opus/48000/2' in sdp_answer
|
||||
# TODO: review
|
||||
assert 'a=recvonly' in sdp_answer
|
||||
|
||||
assert datachannel_library.rtcDeletePeerConnection(sender_connection) == 0
|
||||
assert datachannel_library.rtcDeletePeerConnection(receiver_connection) == 0
|
||||
@@ -84,3 +79,39 @@ def test_delete_peers() -> None:
|
||||
rtc.delete_peers(rtc_peers)
|
||||
|
||||
assert datachannel_library.rtcDeletePeerConnection(peer_connection) < 0
|
||||
|
||||
|
||||
def test_add_audio_track() -> None:
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
audio_track = rtc.add_audio_track(peer_connection, 'sendonly', 'opus', 111)
|
||||
|
||||
assert audio_track > 0
|
||||
|
||||
# TODO: review
|
||||
sdp_offer = rtc.create_sdp_offer(peer_connection)
|
||||
|
||||
assert 'opus/48000/2' in sdp_offer
|
||||
|
||||
datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection)
|
||||
|
||||
|
||||
def test_add_video_track() -> None:
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
video_track = rtc.add_video_track(peer_connection, 'sendonly', 'vp8', 96)
|
||||
|
||||
assert video_track > 0
|
||||
|
||||
# TODO: review
|
||||
sdp_offer = rtc.create_sdp_offer(peer_connection)
|
||||
|
||||
assert 'VP8/90000' in sdp_offer
|
||||
|
||||
datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection)
|
||||
|
||||
|
||||
def test_create_audio_description() -> None:
|
||||
assert rtc.create_audio_description('sendonly', 'opus', 111) == b'm=audio 9 UDP/TLS/RTP/SAVPF 111\r\na=rtpmap:111 opus/48000/2\r\na=rtcp-fb:111 nack\r\na=rtcp-fb:111 nack pli\r\na=sendonly\r\na=mid:1\r\na=rtcp-mux\r\n'
|
||||
|
||||
|
||||
def test_create_video_description() -> None:
|
||||
assert rtc.create_video_description('recvonly', 'vp8', 96) == b'm=video 9 UDP/TLS/RTP/SAVPF 96\r\na=rtpmap:96 VP8/90000\r\na=rtcp-fb:96 nack\r\na=rtcp-fb:96 nack pli\r\na=recvonly\r\na=mid:0\r\na=rtcp-mux\r\n'
|
||||
|
||||
+50
-7
@@ -1,7 +1,10 @@
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
|
||||
from facefusion import state_manager
|
||||
from facefusion import rtc, rtc_store, state_manager
|
||||
from facefusion.libraries import datachannel as datachannel_module, opus as opus_module, vpx as vpx_module
|
||||
from facefusion.types import RtcPeer
|
||||
|
||||
|
||||
@pytest.fixture(scope = 'module', autouse = True)
|
||||
@@ -13,11 +16,51 @@ def before_all() -> None:
|
||||
vpx_module.pre_check()
|
||||
|
||||
|
||||
# TODO: test create_rtc_stream, get_rtc_stream, destroy_rtc_stream lifecycle
|
||||
def test_rtc_stream_lifecycle() -> None:
|
||||
pass
|
||||
@pytest.fixture(autouse = True)
|
||||
def before_each() -> None:
|
||||
rtc_store.RTC_STORE.clear()
|
||||
|
||||
|
||||
# TODO: test add_rtc_viewer with valid session and sdp offer
|
||||
def test_add_rtc_viewer() -> None:
|
||||
pass
|
||||
# TODO: needs review
|
||||
def test_create_rtc_peers() -> None:
|
||||
rtc_store.create_rtc_peers('test-session')
|
||||
|
||||
assert rtc_store.RTC_STORE.get('test-session') == []
|
||||
|
||||
|
||||
# TODO: needs review
|
||||
def test_get_rtc_peers() -> None:
|
||||
assert rtc_store.get_rtc_peers('test-session') is None
|
||||
|
||||
rtc_store.create_rtc_peers('test-session')
|
||||
|
||||
assert rtc_store.get_rtc_peers('test-session') == []
|
||||
|
||||
|
||||
# TODO: needs review
|
||||
def test_destroy_rtc_peers() -> None:
|
||||
rtc_store.create_rtc_peers('test-session')
|
||||
rtc_store.destroy_rtc_peers('test-session')
|
||||
|
||||
assert rtc_store.get_rtc_peers('test-session') is None
|
||||
|
||||
|
||||
# TODO: needs review
|
||||
def test_destroy_rtc_peers_with_connections() -> None:
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
peer_connection = rtc.create_peer_connection()
|
||||
rtc_store.create_rtc_peers('test-session')
|
||||
rtc_peers : List[RtcPeer] =\
|
||||
[
|
||||
{
|
||||
'peer_connection': peer_connection,
|
||||
'video_track': 0,
|
||||
'audio_track': 0
|
||||
}
|
||||
]
|
||||
rtc_store.RTC_STORE['test-session'] = rtc_peers
|
||||
|
||||
rtc_store.destroy_rtc_peers('test-session')
|
||||
|
||||
assert rtc_store.get_rtc_peers('test-session') is None
|
||||
assert datachannel_library.rtcDeletePeerConnection(peer_connection) < 0
|
||||
|
||||
Reference in New Issue
Block a user