From 2f388a4d1c887eca23962d05627d21a84e93b782 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Mon, 15 Jun 2026 15:18:21 +0200 Subject: [PATCH] vibe code to recent v4 changes --- facefusion/apis/core.py | 10 +- facefusion/apis/endpoints/stream.py | 52 ++ facefusion/apis/process.py | 830 ---------------------------- facefusion/apis/stream_audio.py | 71 +++ facefusion/apis/stream_event.py | 24 + facefusion/apis/stream_manager.py | 137 +++++ facefusion/apis/stream_video.py | 204 +++++++ facefusion/codecs/__init__.py | 0 facefusion/codecs/aom_decoder.py | 70 +++ facefusion/codecs/aom_encoder.py | 94 ++++ facefusion/codecs/opus_decoder.py | 36 ++ facefusion/codecs/opus_encoder.py | 36 ++ facefusion/codecs/vpx_decoder.py | 74 +++ facefusion/codecs/vpx_encoder.py | 104 ++++ facefusion/core.py | 7 +- facefusion/libraries/aom.py | 134 +++++ facefusion/libraries/datachannel.py | 294 ++++++++++ facefusion/libraries/opus.py | 121 ++++ facefusion/libraries/vpx.py | 134 +++++ facefusion/rtc.py | 263 +++++++++ facefusion/rtc_store.py | 33 ++ facefusion/types.py | 68 ++- 22 files changed, 1959 insertions(+), 837 deletions(-) create mode 100644 facefusion/apis/endpoints/stream.py delete mode 100644 facefusion/apis/process.py create mode 100644 facefusion/apis/stream_audio.py create mode 100644 facefusion/apis/stream_event.py create mode 100644 facefusion/apis/stream_manager.py create mode 100644 facefusion/apis/stream_video.py create mode 100644 facefusion/codecs/__init__.py create mode 100644 facefusion/codecs/aom_decoder.py create mode 100644 facefusion/codecs/aom_encoder.py create mode 100644 facefusion/codecs/opus_decoder.py create mode 100644 facefusion/codecs/opus_encoder.py create mode 100644 facefusion/codecs/vpx_decoder.py create mode 100644 facefusion/codecs/vpx_encoder.py create mode 100644 facefusion/libraries/aom.py create mode 100644 facefusion/libraries/datachannel.py create mode 100644 facefusion/libraries/opus.py create mode 100644 facefusion/libraries/vpx.py create mode 100644 facefusion/rtc.py create mode 100644 facefusion/rtc_store.py diff --git a/facefusion/apis/core.py b/facefusion/apis/core.py index b6cb07e1..a7b6f171 100644 --- a/facefusion/apis/core.py +++ b/facefusion/apis/core.py @@ -8,8 +8,8 @@ from facefusion.apis.endpoints.assets import delete_asset, delete_assets, get_as from facefusion.apis.endpoints.ping import websocket_ping from facefusion.apis.endpoints.session import create_session, create_session_guard, destroy_session, get_session, refresh_session from facefusion.apis.endpoints.state import get_state, set_state +from facefusion.apis.endpoints.stream import delete_stream, post_stream, websocket_stream from facefusion.apis.metrics import websocket_metrics -from facefusion.apis.process import webrtc_offer, webrtc_stream_offer, websocket_process from facefusion.apis.remote import remote from facefusion.apis.timeline import get_timeline from facefusion.apis.version import create_version_guard @@ -34,11 +34,11 @@ def create_api() -> Starlette: Route('/choices', get_choices, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), Route('/remote', remote, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), Route('/timeline/{count:int}', get_timeline, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), - Route('/webrtc/offer', webrtc_offer, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), - Route('/stream/webrtc/offer', webrtc_stream_offer, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + Route('/stream', post_stream, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + Route('/stream', delete_stream, methods = [ 'DELETE' ], middleware = [ version_guard, session_guard ]), + WebSocketRoute('/stream', websocket_stream, middleware = [ version_guard, session_guard ]), WebSocketRoute('/metrics', websocket_metrics, middleware = [ version_guard, session_guard ]), - WebSocketRoute('/ping', websocket_ping, middleware = [ version_guard, session_guard ]), - WebSocketRoute('/process', websocket_process, middleware = [ version_guard, session_guard ]) + WebSocketRoute('/ping', websocket_ping, middleware = [ version_guard, session_guard ]) ] api = Starlette(routes = routes) diff --git a/facefusion/apis/endpoints/stream.py b/facefusion/apis/endpoints/stream.py new file mode 100644 index 00000000..93b7398a --- /dev/null +++ b/facefusion/apis/endpoints/stream.py @@ -0,0 +1,52 @@ +from starlette.requests import Request +from starlette.responses import Response +from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_404_NOT_FOUND +from starlette.websockets import WebSocket, WebSocketState + +from facefusion import session_context, session_manager +from facefusion.apis.api_helper import get_sec_websocket_protocol +from facefusion.apis.endpoints.session import extract_access_token +from facefusion.apis.stream_manager import destroy_stream, process_image, process_video + + +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) + session_context.set_session_id(session_id) + + await websocket.accept(subprotocol = subprotocol) + await process_image(websocket) + + if websocket.client_state == WebSocketState.CONNECTED: + await websocket.close() + + +async def post_stream(request : Request) -> Response: + headers =\ + { + 'Location': request.url_for('delete_stream').path + } + content_type = request.headers.get('content-type') + 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 and content_type == 'application/sdp': + sdp_offer = await request.body() + sdp_answer = process_video(session_id, sdp_offer.decode()) + + if sdp_answer: + return Response(sdp_answer, status_code = HTTP_201_CREATED, media_type = 'application/sdp', headers = headers) + + return Response(status_code = HTTP_404_NOT_FOUND) + + +async def delete_stream(request : Request) -> Response: + access_token = extract_access_token(request.scope) + session_id = session_manager.find_session_id(access_token) + + if session_id and destroy_stream(session_id): + return Response(status_code = HTTP_200_OK) + + return Response(status_code = HTTP_404_NOT_FOUND) diff --git a/facefusion/apis/process.py b/facefusion/apis/process.py deleted file mode 100644 index 11a2645f..00000000 --- a/facefusion/apis/process.py +++ /dev/null @@ -1,830 +0,0 @@ -import asyncio -import fractions -import subprocess -from functools import partial -from typing import Any, List, Optional, Set, TypeAlias - -import cv2 -import numpy -from aiortc import AudioStreamTrack, RTCPeerConnection, RTCSessionDescription, VideoStreamTrack -from aiortc.codecs import h264 -from av import AudioFrame, VideoFrame -from starlette.datastructures import Headers -from starlette.requests import Request -from starlette.responses import JSONResponse -from starlette.websockets import WebSocket - -from facefusion import config, content_analyser, logger, state_manager -from facefusion.streamer import process_stream_frame -from facefusion.vision import obscure_frame - -PeerConnectionSet : TypeAlias = Set[RTCPeerConnection] -ResolutionTuple : TypeAlias = tuple[int, int] - -NSFW_LOCK = False - -pcs : PeerConnectionSet = set() - -RESOLUTION_MAP : dict[str, ResolutionTuple] =\ -{ - '480p': (640, 480), - '720p': (1280, 720), - '1080p': (1920, 1080) -} - - -def init_default_state() -> None: - if state_manager.get_item('execution_providers') is None: - state_manager.set_item('execution_providers', config.get_str_list('execution', 'execution_providers', 'cuda')) - if state_manager.get_item('execution_thread_count') is None: - state_manager.set_item('execution_thread_count', config.get_int_value('execution', 'execution_thread_count', '32')) - if state_manager.get_item('face_detector_model') is None: - state_manager.set_item('face_detector_model', config.get_str_value('face_detector', 'face_detector_model', 'yolo_face')) - if state_manager.get_item('face_detector_size') is None: - state_manager.set_item('face_detector_size', config.get_str_value('face_detector', 'face_detector_size', '640x640')) - if state_manager.get_item('face_detector_margin') is None: - state_manager.set_item('face_detector_margin', config.get_int_list('face_detector', 'face_detector_margin', '0 0 0 0')) - if state_manager.get_item('face_detector_angles') is None: - state_manager.set_item('face_detector_angles', config.get_int_list('face_detector', 'face_detector_angles', '0')) - if state_manager.get_item('face_detector_score') is None: - state_manager.set_item('face_detector_score', config.get_float_value('face_detector', 'face_detector_score', '0.5')) - if state_manager.get_item('face_landmarker_model') is None: - state_manager.set_item('face_landmarker_model', config.get_str_value('face_landmarker', 'face_landmarker_model', '2dfan4')) - if state_manager.get_item('face_landmarker_score') is None: - state_manager.set_item('face_landmarker_score', config.get_float_value('face_landmarker', 'face_landmarker_score', '0.5')) - if state_manager.get_item('face_selector_mode') is None: - state_manager.set_item('face_selector_mode', config.get_str_value('face_selector', 'face_selector_mode', 'many')) - if state_manager.get_item('face_selector_order') is None: - state_manager.set_item('face_selector_order', config.get_str_value('face_selector', 'face_selector_order', 'large-small')) - if state_manager.get_item('face_mask_types') is None: - state_manager.set_item('face_mask_types', config.get_str_list('face_masker', 'face_mask_types', 'occlusion')) - if state_manager.get_item('face_mask_blur') is None: - state_manager.set_item('face_mask_blur', config.get_float_value('face_masker', 'face_mask_blur', '0.3')) - if state_manager.get_item('face_mask_padding') is None: - state_manager.set_item('face_mask_padding', config.get_int_list('face_masker', 'face_mask_padding', '0 0 0 0')) - if state_manager.get_item('face_swapper_model') is None: - state_manager.set_item('face_swapper_model', config.get_str_value('processors', 'face_swapper_model', 'hyperswap_1a_256')) - if state_manager.get_item('face_swapper_pixel_boost') is None: - state_manager.set_item('face_swapper_pixel_boost', config.get_str_value('processors', 'face_swapper_pixel_boost', '256x256')) - if state_manager.get_item('face_swapper_weight') is None: - state_manager.set_item('face_swapper_weight', config.get_float_value('processors', 'face_swapper_weight', '0.5')) - if state_manager.get_item('face_enhancer_model') is None: - state_manager.set_item('face_enhancer_model', config.get_str_value('processors', 'face_enhancer_model', 'gfpgan_1.4')) - if state_manager.get_item('face_enhancer_blend') is None: - state_manager.set_item('face_enhancer_blend', config.get_int_value('processors', 'face_enhancer_blend', '80')) - if state_manager.get_item('frame_enhancer_model') is None: - state_manager.set_item('frame_enhancer_model', config.get_str_value('processors', 'frame_enhancer_model', 'real_esrgan_x2')) - if state_manager.get_item('frame_enhancer_blend') is None: - state_manager.set_item('frame_enhancer_blend', config.get_int_value('processors', 'frame_enhancer_blend', '80')) - if state_manager.get_item('face_debugger_items') is None: - state_manager.set_item('face_debugger_items', config.get_str_list('processors', 'face_debugger_items', 'kps')) - logger.debug(f'Initialized state - execution_providers: {state_manager.get_item("execution_providers")}', __name__) - - -def setup_bitrate_config(bitrate : int, encoder : str, mode_prefix : str = 'WebRTC') -> tuple[int, bool]: - if bitrate == 0: - bitrate_bps = 100000 - h264.DEFAULT_BITRATE = bitrate_bps - h264.MIN_BITRATE = 100000 - h264.MAX_BITRATE = 2000000 - logger.info( - f'{mode_prefix} setup: mode=auto, encoder={encoder}, ' - f'DEF={h264.DEFAULT_BITRATE / 1000} kbps, ' - f'MIN={h264.MIN_BITRATE / 1000} kbps, MAX={h264.MAX_BITRATE / 1000} kbps', - __name__ - ) - adaptive_bitrate = True - else: - bitrate_bps = bitrate * 1000 - h264.DEFAULT_BITRATE = bitrate_bps - h264.MIN_BITRATE = max(500000, bitrate_bps // 2) - h264.MAX_BITRATE = max(bitrate_bps * 2, 3000000) - logger.info( - f'{mode_prefix} setup: mode=manual, encoder={encoder}, ' - f'DEF={h264.DEFAULT_BITRATE / 1000} kbps, ' - f'MIN={h264.MIN_BITRATE / 1000} kbps, MAX={h264.MAX_BITRATE / 1000} kbps', - __name__ - ) - adaptive_bitrate = False - - return bitrate_bps, adaptive_bitrate - - -def create_video_stream_track(pc : RTCPeerConnection, bitrate_bps : int, adaptive_bitrate : bool, buffer_size : int = 30, mode_prefix : str = 'WebRTC') -> tuple[VideoStreamTrack, Any]: - logger.info(f'Creating {mode_prefix} output queue with buffer size: {buffer_size}', __name__) - - processed_track = VideoStreamTrack() - processed_track.frame_queue = asyncio.Queue(maxsize = buffer_size) - processed_track.recv = partial(recv_from_queue, processed_track.frame_queue) - processed_track.data_channel = None - processed_track.ready_sent = False - processed_track._target_bitrate = bitrate_bps - processed_track._adaptive = adaptive_bitrate - processed_track._current_bitrate = bitrate_bps - - sender = pc.addTrack(processed_track) - - return processed_track, sender - - -async def monitor_and_set_bitrate(sender : Any, bitrate_bps : int, adaptive_bitrate : bool, processed_track : VideoStreamTrack, buffer_size : int) -> None: - encoder_obj = None - for attempt in range(30): - if hasattr(sender, '_RTCRtpSender__encoder') and sender._RTCRtpSender__encoder: - encoder_obj = sender._RTCRtpSender__encoder - encoder_type = type(encoder_obj).__name__ - if hasattr(encoder_obj, 'target_bitrate'): - old_bitrate = encoder_obj.target_bitrate - encoder_obj.target_bitrate = bitrate_bps - logger.info( - f'Encoder type: {encoder_type}, updated bitrate from {old_bitrate / 1000} kbps to {bitrate_bps / 1000} kbps', - __name__ - ) - if hasattr(encoder_obj, 'codec') and encoder_obj.codec: - logger.info( - f'Encoder codec context: {encoder_obj.codec.name if hasattr(encoder_obj.codec, "name") else "unknown"}', - __name__ - ) - break - - if not encoder_obj: - logger.warn('Encoder not created after 9 seconds', __name__) - return - - if not adaptive_bitrate: - return - - stable_checks = 0 - INCREASE_STEP = 50000 - DECREASE_FACTOR = 0.9 - - while True: - await asyncio.sleep(0.5) - - if not hasattr(encoder_obj, 'target_bitrate'): - break - - queue_ratio = processed_track.frame_queue.qsize() / buffer_size - - if queue_ratio > 0.7: - new_bitrate = max(int(processed_track._current_bitrate * DECREASE_FACTOR), h264.MIN_BITRATE) - processed_track._current_bitrate = new_bitrate - encoder_obj.target_bitrate = new_bitrate - stable_checks = 0 - logger.info(f'Auto: decreased to {new_bitrate / 1000} kbps (congestion)', __name__) - elif queue_ratio < 0.3: - stable_checks += 1 - if stable_checks >= 4: - new_bitrate = min(processed_track._current_bitrate + INCREASE_STEP, h264.MAX_BITRATE) - if new_bitrate > processed_track._current_bitrate: - processed_track._current_bitrate = new_bitrate - encoder_obj.target_bitrate = new_bitrate - logger.info(f'Auto: increased to {new_bitrate / 1000} kbps (stable)', __name__) - stable_checks = 0 - else: - stable_checks = 0 - - -async def websocket_process(websocket : WebSocket) -> None: - subprotocol = get_requested_subprotocol(websocket) - await websocket.accept(subprotocol = subprotocol) - init_default_state() - - output_resolution = websocket.query_params.get('output_resolution', 'original') - - while True: - message = await websocket.receive() - - if message['type'] == 'websocket.disconnect': - logger.debug('Client disconnected', __name__) - break - - if message['type'] == 'websocket.receive': - if 'bytes' in message: - logger.debug(f'Received {len(message["bytes"])} bytes', __name__) - - target_vision_frame = cv2.imdecode(numpy.frombuffer(message['bytes'], numpy.uint8), cv2.IMREAD_COLOR) - if target_vision_frame is None: - logger.error('Failed to decode target image!', __name__) - continue - - logger.debug(f'Decoded target frame shape: {target_vision_frame.shape}', __name__) - - if output_resolution and output_resolution != 'original': - resolution_map =\ - { - '480p': (640, 480), - '720p': (1280, 720), - '1080p': (1920, 1080) - } - if output_resolution in resolution_map: - target_width, target_height = resolution_map[output_resolution] - current_height, current_width = target_vision_frame.shape[:2] - if current_width > target_width or current_height > target_height: - scale = min(target_width / current_width, target_height / current_height) - new_width = int(current_width * scale) - new_height = int(current_height * scale) - target_vision_frame = cv2.resize(target_vision_frame, (new_width, new_height), interpolation = cv2.INTER_AREA) - logger.debug(f'Downscaled target from {current_width}x{current_height} to {new_width}x{new_height}', __name__) - - temp_vision_frame = process_stream_frame(target_vision_frame) - if temp_vision_frame is None: - continue - - if content_analyser.analyse_frame(temp_vision_frame): - logger.warn('NSFW content detected in output, blurring frame', __name__) - temp_vision_frame = obscure_frame(temp_vision_frame) - - success, result_bytes = cv2.imencode('.jpg', temp_vision_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 50]) - if success: - await websocket.send_bytes(result_bytes.tobytes()) - - -async def process_incoming_video_track(track : Any, frame_queue : Any, output_track : Any = None, output_resolution : str = 'original') -> None: - from aiortc.mediastreams import MediaStreamError - - logger.debug(f'Track received: {track.kind}', __name__) - max_fps = 60 - min_frame_time = 1.0 / max_fps - last_process_time = 0.0 - frame_counter = 0 - frame_skip = 2 - last_processed_frame = None - - try: - global NSFW_LOCK - while True: - try: - frame = await track.recv() - except MediaStreamError: - logger.info('Media stream ended (connection closed)', __name__) - break - except asyncio.CancelledError: - logger.info('Video processing cancelled', __name__) - raise - except Exception as e: - logger.error(f'Error receiving frame: {e}', __name__) - break - - current_time = asyncio.get_event_loop().time() - time_since_last = current_time - last_process_time - - if time_since_last < min_frame_time: - continue - - img = frame.to_ndarray(format='bgr24') - logger.debug(f'Received frame shape: {img.shape}', __name__) - - if output_resolution and output_resolution != 'original': - if output_resolution in RESOLUTION_MAP: - target_width, target_height = RESOLUTION_MAP[output_resolution] - current_height, current_width = img.shape[:2] - if current_width > target_width or current_height > target_height: - scale = min(target_width / current_width, target_height / current_height) - new_width = int(current_width * scale) - new_height = int(current_height * scale) - img = cv2.resize(img, (new_width, new_height), interpolation = cv2.INTER_AREA) - logger.debug(f'Downscaled target from {current_width}x{current_height} to {new_width}x{new_height}', __name__) - - if content_analyser.analyse_stream(img, float(max_fps)): - NSFW_LOCK = True - - if NSFW_LOCK: - temp_vision_frame = obscure_frame(img) - else: - frame_counter += 1 - if frame_counter % frame_skip == 0: - temp_vision_frame = process_stream_frame(img) - last_processed_frame = temp_vision_frame - else: - if last_processed_frame is not None: - temp_vision_frame = last_processed_frame - else: - temp_vision_frame = process_stream_frame(img) - last_processed_frame = temp_vision_frame - - if temp_vision_frame is not None: - new_frame = VideoFrame.from_ndarray(temp_vision_frame, format='bgr24') - new_frame.pts = frame.pts - new_frame.time_base = frame.time_base - - if frame_queue.full(): - try: - frame_queue.get_nowait() - except asyncio.QueueEmpty: - pass - - try: - frame_queue.put_nowait(new_frame) - last_process_time = current_time - - if output_track and not output_track.ready_sent and frame_queue.qsize() >= int(frame_queue.maxsize * 0.5): - output_track.ready_sent = True - logger.info(f'Buffer ready ({frame_queue.qsize()}/{frame_queue.maxsize} frames), sending ready signal', __name__) - if output_track.data_channel and output_track.data_channel.readyState == 'open': - output_track.data_channel.send('ready') - logger.info('Ready signal sent to client', __name__) - - if output_track and output_track.data_channel and output_track.data_channel.readyState == 'open': - try: - output_track.data_channel.send(f'frame:{frame_counter}') - except Exception: - pass - except asyncio.QueueFull: - logger.debug('Frame queue full, frame dropped', __name__) - except Exception as e: - logger.error(f'Unexpected error in video processing: {e}', __name__) - finally: - logger.info('Video processing task completed', __name__) - - -async def recv_from_queue(frame_queue : Any) -> Any: - frame = await frame_queue.get() - return frame - - -def check_and_lock_nsfw(vision_frame : Any, fps : float) -> bool: - global NSFW_LOCK - - if NSFW_LOCK: - return True - - if content_analyser.analyse_stream(vision_frame, fps): - NSFW_LOCK = True - logger.warn('NSFW content detected, locking all future frames', __name__) - return True - - return False - - -def get_requested_subprotocol(websocket : WebSocket) -> Optional[str]: - headers = Headers(scope = websocket.scope) - protocol_header = headers.get('Sec-WebSocket-Protocol') - - if protocol_header: - protocol, _, _ = protocol_header.partition(',') - return protocol.strip() - - return None - - -async def webrtc_offer(request : Request) -> JSONResponse: - global NSFW_LOCK - init_default_state() - NSFW_LOCK = False - - params = await request.json() - offer = RTCSessionDescription(sdp=params['sdp'], type=params['type']) - bitrate = int(params.get('bitrate', 0)) - encoder = params.get('encoder', 'VP8') - buffer_size = int(params.get('stream_buffer_size', 30)) - output_resolution = params.get('output_resolution', 'original') - - bitrate_bps, adaptive_bitrate = setup_bitrate_config(bitrate, encoder, 'WebRTC') - - pc = RTCPeerConnection() - pcs.add(pc) - - processed_track, sender = create_video_stream_track(pc, bitrate_bps, adaptive_bitrate, buffer_size, 'WebRTC') - - asyncio.create_task(monitor_and_set_bitrate(sender, bitrate_bps, adaptive_bitrate, processed_track, buffer_size)) - - processing_tasks : List[Any] = [] - - @pc.on('connectionstatechange') - async def on_connectionstatechange() -> None: - logger.info(f'WebRTC connection state: {pc.connectionState}', __name__) - if pc.connectionState == 'failed' or pc.connectionState == 'closed': - logger.info('WebRTC connection closed, cleaning up', __name__) - pcs.discard(pc) - for task in processing_tasks: - task.cancel() - - @pc.on('datachannel') - def on_datachannel(channel : Any) -> None: - logger.info(f'Data channel received: {channel.label}', __name__) - processed_track.data_channel = channel - - await pc.setRemoteDescription(offer) - - for transceiver in pc.getTransceivers(): - if transceiver.receiver and transceiver.receiver.track: - track = transceiver.receiver.track - if track.kind == 'video': - logger.info('Found video track, starting processing', __name__) - video_task = asyncio.create_task(process_incoming_video_track(track, processed_track.frame_queue, processed_track, output_resolution)) - processing_tasks.append(video_task) - if track.kind == 'audio': - logger.info('Found audio track, forwarding as-is', __name__) - pc.addTrack(track) - - answer = await pc.createAnswer() - await pc.setLocalDescription(answer) - - return JSONResponse({'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}) - - -async def process_stream_from_url(stream_url : str, frame_queue : Any, output_track : Any = None, width : int = 1280, height : int = 720, target_fps : int = 30, duration : float = 0, output_resolution : str = 'original') -> None: - logger.info(f'Opening stream from URL: {stream_url[:100]}', __name__) - logger.info(f'Using metadata - Resolution: {width}x{height}, FPS: {target_fps}, Duration: {duration}s', __name__) - - frame_size = width * height * 3 - frame_interval = 1.0 / target_fps - - import threading - import time - - current_process = None - seek_position = [0.0] - stop_flag = [False] - lock = threading.Lock() - - def start_ffmpeg(start_time : float) -> Any: - ffmpeg_command =\ - [ - 'ffmpeg', - '-user_agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', - '-reconnect', '1', - '-reconnect_streamed', '1', - '-reconnect_delay_max', '5' - ] - - if start_time > 0: - ffmpeg_command.extend(['-ss', str(start_time)]) - - ffmpeg_command.extend([ - '-i', stream_url, - '-f', 'rawvideo', - '-pix_fmt', 'bgr24', - '-' - ]) - - try: - return subprocess.Popen(ffmpeg_command, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, bufsize = 10**8) - except Exception as e: - logger.error(f'Failed to start ffmpeg: {e}', __name__) - return None - - def read_stream() -> None: - nonlocal current_process - frame_count = int(seek_position[0] * target_fps) - - current_process = start_ffmpeg(seek_position[0]) - if not current_process: - return - - logger.info(f'FFmpeg stream started at position {seek_position[0]}s', __name__) - - last_frame_time = time.time() - local_process = current_process - - while not stop_flag[0]: - with lock: - if local_process != current_process: - logger.info('Process changed due to seek, switching to new process', __name__) - local_process = current_process - frame_count = int(seek_position[0] * target_fps) - if not local_process: - break - continue - - if local_process.poll() is not None: - logger.info('Stream process terminated', __name__) - break - - raw_frame = local_process.stdout.read(frame_size) - if not raw_frame or len(raw_frame) != frame_size: - with lock: - if local_process != current_process: - logger.info('Incomplete read due to seek, switching to new process', __name__) - local_process = current_process - frame_count = int(seek_position[0] * target_fps) - continue - logger.info('Stream ended or incomplete frame', __name__) - break - - frame = numpy.frombuffer(raw_frame, dtype = numpy.uint8).reshape((height, width, 3)) - - if output_resolution and output_resolution != 'original': - if output_resolution in RESOLUTION_MAP: - target_width, target_height = RESOLUTION_MAP[output_resolution] - current_height, current_width = frame.shape[:2] - if current_width > target_width or current_height > target_height: - scale = min(target_width / current_width, target_height / current_height) - new_width = int(current_width * scale) - new_height = int(current_height * scale) - frame = cv2.resize(frame, (new_width, new_height), interpolation = cv2.INTER_AREA) - - if check_and_lock_nsfw(frame, float(target_fps)): - processed_frame = obscure_frame(frame) - else: - processed_frame = process_stream_frame(frame) - - if processed_frame is not None: - current_time = time.time() - elapsed = current_time - last_frame_time - - if elapsed < frame_interval: - time.sleep(frame_interval - elapsed) - current_time = time.time() - - video_frame = VideoFrame.from_ndarray(processed_frame, format = 'bgr24') - video_frame.pts = frame_count - video_frame.time_base = fractions.Fraction(1, target_fps) - - if frame_queue.full(): - try: - frame_queue.get_nowait() - except asyncio.QueueEmpty: - pass - - try: - frame_queue.put_nowait(video_frame) - if output_track and not output_track.ready_sent and frame_queue.qsize() >= int(frame_queue.maxsize * 0.5): - output_track.ready_sent = True - logger.info(f'Buffer ready ({frame_queue.qsize()}/{frame_queue.maxsize} frames)', __name__) - if output_track.data_channel and output_track.data_channel.readyState == 'open': - output_track.data_channel.send('ready') - - if output_track and output_track.data_channel and output_track.data_channel.readyState == 'open': - try: - output_track.data_channel.send(f'frame:{frame_count}') - except Exception: - pass - except asyncio.QueueFull: - pass - - last_frame_time = current_time - frame_count += 1 - - if current_process: - current_process.terminate() - current_process.wait() - logger.info(f'Stream reading completed, {frame_count} frames processed', __name__) - - def handle_seek(new_position : float) -> None: - nonlocal current_process - with lock: - seek_position[0] = new_position - if current_process: - logger.info(f'Seeking to {new_position}s, restarting ffmpeg', __name__) - current_process.terminate() - current_process = start_ffmpeg(new_position) - - if output_track: - output_track.seek_handler = handle_seek - - stream_task = asyncio.get_event_loop().run_in_executor(None, read_stream) - - try: - await stream_task - except asyncio.CancelledError: - logger.info('Video stream task cancelled', __name__) - raise - finally: - stop_flag[0] = True - if current_process: - current_process.terminate() - logger.info('Video stream cleanup completed', __name__) - - -async def process_audio_from_url(stream_url : str, audio_queue : Any, video_output_track : Any = None) -> None: - logger.info('Opening audio stream from URL', __name__) - - import threading - - current_process = None - seek_position = [0.0] - stop_flag = [False] - lock = threading.Lock() - - sample_rate = 24000 - channels = 2 - frame_samples = 480 - bytes_per_sample = 2 - frame_size = frame_samples * channels * bytes_per_sample - - def start_ffmpeg_audio(start_time : float) -> Any: - ffmpeg_command =\ - [ - 'ffmpeg', - '-user_agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', - '-reconnect', '1', - '-reconnect_streamed', '1', - '-reconnect_delay_max', '5' - ] - - if start_time > 0: - ffmpeg_command.extend(['-ss', str(start_time)]) - - ffmpeg_command.extend([ - '-i', stream_url, - '-vn', - '-f', 's16le', - '-acodec', 'pcm_s16le', - '-ar', str(sample_rate), - '-ac', str(channels), - '-' - ]) - - try: - return subprocess.Popen(ffmpeg_command, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, bufsize = 10**8) - except Exception as e: - logger.error(f'Failed to start ffmpeg audio: {e}', __name__) - return None - - def read_audio_stream() -> None: - nonlocal current_process - import time - - current_process = start_ffmpeg_audio(seek_position[0]) - if not current_process: - return - - logger.info(f'FFmpeg audio stream started at position {seek_position[0]}s', __name__) - - local_process = current_process - frame_count = 0 - - frame_duration = frame_samples / sample_rate - start_time = time.time() - expected_time = start_time - - logger.info(f'Starting audio read loop, expecting {frame_size} bytes per frame', __name__) - - while not stop_flag[0]: - with lock: - if local_process != current_process: - logger.info('Audio process changed due to seek, switching to new process', __name__) - local_process = current_process - frame_count = 0 - start_time = time.time() - expected_time = start_time - if not local_process: - break - continue - - if local_process.poll() is not None: - logger.info('Audio stream process terminated', __name__) - break - - raw_audio = local_process.stdout.read(frame_size) - - if frame_count == 0: - logger.info('Successfully read first audio frame', __name__) - if not raw_audio or len(raw_audio) != frame_size: - with lock: - if local_process != current_process: - logger.info('Incomplete audio read due to seek, switching to new process', __name__) - local_process = current_process - continue - logger.info('Audio stream ended or incomplete frame', __name__) - break - - audio_array = numpy.frombuffer(raw_audio, dtype = numpy.int16) - - audio_frame = AudioFrame(format = 's16', layout = 'stereo', samples = frame_samples) - audio_frame.sample_rate = sample_rate - audio_frame.pts = frame_count * frame_samples - audio_frame.time_base = fractions.Fraction(1, sample_rate) - - for plane in audio_frame.planes: - plane.update(audio_array.tobytes()) - - expected_time += frame_duration - current_time = time.time() - sleep_time = expected_time - current_time - - if sleep_time > 0: - time.sleep(sleep_time) - - if audio_queue.full(): - try: - audio_queue.get_nowait() - except asyncio.QueueEmpty: - pass - - try: - audio_queue.put_nowait(audio_frame) - frame_count += 1 - if frame_count % 50 == 0: - logger.info(f'Audio frames queued: {frame_count}, queue depth: {audio_queue.qsize()}', __name__) - except asyncio.QueueFull: - pass - - if current_process: - current_process.terminate() - logger.info('Audio stream reading completed', __name__) - - def handle_audio_seek(new_position : float) -> None: - nonlocal current_process - with lock: - seek_position[0] = new_position - if current_process: - logger.info(f'Seeking audio to {new_position}s, restarting ffmpeg', __name__) - current_process.terminate() - current_process = start_ffmpeg_audio(new_position) - - if video_output_track: - video_output_track.audio_seek_handler = handle_audio_seek - - audio_task = asyncio.get_event_loop().run_in_executor(None, read_audio_stream) - - try: - await audio_task - except asyncio.CancelledError: - logger.info('Audio stream task cancelled', __name__) - raise - finally: - stop_flag[0] = True - if current_process: - current_process.terminate() - logger.info('Audio stream cleanup completed', __name__) - - -async def webrtc_stream_offer(request : Request) -> JSONResponse: - global NSFW_LOCK - init_default_state() - NSFW_LOCK = False - - params = await request.json() - offer = RTCSessionDescription(sdp = params['sdp'], type = params['type']) - bitrate = int(params.get('bitrate', 0)) - encoder = params.get('encoder', 'VP8') - is_remote_stream = params.get('is_remote_stream', False) - stream_url = params.get('stream_url') - target_width = int(params.get('target_width', 1280)) - target_height = int(params.get('target_height', 720)) - target_fps = int(params.get('target_fps', 30)) - target_duration = float(params.get('target_duration', 0)) - target_audio_path = params.get('target_audio_path') - buffer_size = int(params.get('stream_buffer_size', 30)) - output_resolution = params.get('output_resolution', 'original') - - logger.info(f'[WebRTC Stream] stream_url: {stream_url[:50] if stream_url else None}', __name__) - logger.info(f'[WebRTC Stream] is_remote_stream: {is_remote_stream}', __name__) - - if not stream_url: - logger.error('[WebRTC Stream] No stream URL provided', __name__) - return JSONResponse({'error': 'No stream URL provided in request'}, status_code = 400) - - if not is_remote_stream: - logger.error('[WebRTC Stream] is_remote_stream=False, use /webrtc/offer for local files', __name__) - return JSONResponse({'error': 'Use /webrtc/offer endpoint for local files, not /stream/webrtc/offer'}, status_code = 400) - - bitrate_bps, adaptive_bitrate = setup_bitrate_config(bitrate, encoder, 'WebRTC stream') - - pc = RTCPeerConnection() - pcs.add(pc) - - processed_track, sender = create_video_stream_track(pc, bitrate_bps, adaptive_bitrate, buffer_size, 'WebRTC stream') - - audio_track = AudioStreamTrack() - audio_track.audio_queue = asyncio.Queue(maxsize = buffer_size) - audio_track.recv = partial(recv_from_queue, audio_track.audio_queue) - pc.addTrack(audio_track) - - asyncio.create_task(monitor_and_set_bitrate(sender, bitrate_bps, adaptive_bitrate, processed_track, buffer_size)) - - stream_tasks : List[Any] = [] - - @pc.on('connectionstatechange') - async def on_connectionstatechange() -> None: - logger.info(f'WebRTC stream connection state: {pc.connectionState}', __name__) - if pc.connectionState == 'failed' or pc.connectionState == 'closed': - logger.info('WebRTC stream connection closed, stopping ffmpeg processes', __name__) - pcs.discard(pc) - for task in stream_tasks: - task.cancel() - - @pc.on('datachannel') - def on_datachannel(channel : Any) -> None: - logger.info(f'Data channel received: {channel.label}', __name__) - processed_track.data_channel = channel - - @channel.on('message') - def on_message(message : Any) -> None: - if isinstance(message, str) and message.startswith('seek:'): - try: - seek_time = float(message.split(':', 1)[1]) - logger.info(f'Received seek command: {seek_time}s', __name__) - if hasattr(processed_track, 'seek_handler') and processed_track.seek_handler: - processed_track.seek_handler(seek_time) - if hasattr(processed_track, 'audio_seek_handler') and processed_track.audio_seek_handler: - processed_track.audio_seek_handler(seek_time) - except Exception as e: - logger.error(f'Error handling seek command: {e}', __name__) - - await pc.setRemoteDescription(offer) - - audio_url = target_audio_path or stream_url - logger.info(f'Starting stream processing from URL (video: {stream_url[:50]}..., audio: {audio_url[:50]}...)', __name__) - - video_task = asyncio.create_task(process_stream_from_url(stream_url, processed_track.frame_queue, processed_track, target_width, target_height, target_fps, target_duration, output_resolution)) - audio_task = asyncio.create_task(process_audio_from_url(audio_url, audio_track.audio_queue, processed_track)) - stream_tasks.extend([video_task, audio_task]) - - answer = await pc.createAnswer() - await pc.setLocalDescription(answer) - - return JSONResponse({'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}) diff --git a/facefusion/apis/stream_audio.py b/facefusion/apis/stream_audio.py new file mode 100644 index 00000000..e677d753 --- /dev/null +++ b/facefusion/apis/stream_audio.py @@ -0,0 +1,71 @@ +import ctypes +import time +from functools import partial +from queue import Queue +from typing import Optional, Tuple + +import numpy + +from facefusion import rtc +from facefusion.apis.stream_event import create_receive_event +from facefusion.codecs import opus_decoder, opus_encoder +from facefusion.types import AudioCodec, AudioFrame, OpusDecoder, RtcPeer, RtcPeerAudio + + +def run_audio_encode_loop(rtc_peer : RtcPeer, audio_queue : Queue[Tuple[float, AudioFrame]]) -> None: + temp_audio_time, temp_audio_frame = audio_queue.get() + audio_encoder = opus_encoder.create(48000, 2) + audio_timestamp = 0 + + while numpy.any(temp_audio_frame): + audio_frame_size = len(temp_audio_frame) // 2 + audio_buffer = opus_encoder.encode(audio_encoder, temp_audio_frame.tobytes(), audio_frame_size) + + if audio_buffer: + rtc.send_audio(rtc_peer, audio_buffer, audio_timestamp) + audio_timestamp += audio_frame_size + + temp_audio_time, temp_audio_frame = audio_queue.get() + + opus_encoder.destroy(audio_encoder) + + +def receive_audio_frames(rtc_peer_audio : RtcPeerAudio, audio_queue : Queue[Tuple[float, AudioFrame]]) -> None: + audio_track = rtc_peer_audio.get('receiver_track') + audio_codec = rtc_peer_audio.get('codec') + audio_decoder = create_audio_decoder(audio_codec) + + audio_frame_handler = partial(handle_audio_frame, audio_codec, audio_decoder, audio_queue) + receive_event = create_receive_event(audio_track, audio_frame_handler) + receive_event.wait() + + empty_audio_frame = numpy.empty(0) + audio_queue.put((0.0, empty_audio_frame)) + destroy_audio_decoder(audio_codec, audio_decoder) + + +def decode_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, input_buffer : bytes) -> Optional[bytes]: + if audio_codec == 'opus': + return opus_decoder.decode(audio_decoder, input_buffer, 960, 2) + return None + + +def create_audio_decoder(audio_codec : AudioCodec) -> Optional[OpusDecoder]: + if audio_codec == 'opus': + return opus_decoder.create(48000, 2) + return None + + +def destroy_audio_decoder(audio_codec : AudioCodec, audio_decoder : OpusDecoder) -> None: + if audio_codec == 'opus': + opus_decoder.destroy(audio_decoder) + + +#todo: Alias Time for float +def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[float, AudioFrame]], track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None: + audio_buffer = ctypes.string_at(data, size) + audio_frame = decode_audio_frame(audio_codec, audio_decoder, audio_buffer) + + if audio_frame: + temp_audio_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32) + audio_queue.put((time.monotonic(), temp_audio_frame)) diff --git a/facefusion/apis/stream_event.py b/facefusion/apis/stream_event.py new file mode 100644 index 00000000..aafb1ac6 --- /dev/null +++ b/facefusion/apis/stream_event.py @@ -0,0 +1,24 @@ +import ctypes +import threading +from functools import partial + +from facefusion.libraries import datachannel as datachannel_module +from facefusion.types import FrameHandler + + +def create_receive_event(track : int, frame_handler : FrameHandler) -> threading.Event: + datachannel_library = datachannel_module.create_static_library() + receive_event = threading.Event() + + frame_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p)(frame_handler) + close_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p)(partial(dispatch_event, receive_event)) + datachannel_library.rtcSetFrameCallback(track, frame_callback) + datachannel_library.rtcSetClosedCallback(track, close_callback) + receive_event.frame_callback = frame_callback # type: ignore[attr-defined] + receive_event.close_callback = close_callback # type: ignore[attr-defined] + + return receive_event + + +def dispatch_event(event : threading.Event, track : int, pointer : ctypes.c_void_p) -> None: + event.set() diff --git a/facefusion/apis/stream_manager.py b/facefusion/apis/stream_manager.py new file mode 100644 index 00000000..ded9a5cc --- /dev/null +++ b/facefusion/apis/stream_manager.py @@ -0,0 +1,137 @@ +import ctypes +import threading +from collections.abc import AsyncIterator +from concurrent.futures import Future, ThreadPoolExecutor +from queue import Queue +from typing import Optional, Tuple + +import cv2 +import numpy +from starlette.websockets import WebSocket + +from facefusion import rtc, rtc_store, state_manager, streamer +from facefusion.apis.stream_audio import receive_audio_frames, run_audio_encode_loop +from facefusion.apis.stream_video import receive_video_frames, run_video_encode_loop +from facefusion.libraries import datachannel as datachannel_module +from facefusion.types import AudioCodec, AudioFrame, PeerConnection, Resolution, RtcPeer, RtcPeerAudio, SdpAnswer, SdpOffer, SessionId, VideoCodec, VisionFrame + + +async def process_image(websocket : WebSocket) -> None: + capture_vision_frame = await anext(receive_vision_frames(websocket), None) + + if numpy.any(capture_vision_frame): + output_vision_frame = streamer.process_stream_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()) + + +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 bytes() + 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 process_video(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpAnswer]: + video_codec : VideoCodec = 'vp8' + + if rtc.get_payload_type(sdp_offer, 'vp9'): + video_codec = 'vp9' + + if rtc.get_payload_type(sdp_offer, 'av1'): + video_codec = 'av1' + + video_payload_type = rtc.get_payload_type(sdp_offer, video_codec) + + if video_payload_type: + peer_connection : PeerConnection = rtc.create_peer_connection() + video_receiver_track = rtc.add_video_track(peer_connection, 'recvonly', video_codec, video_payload_type) + video_sender_track = rtc.add_video_track(peer_connection, 'sendonly', video_codec, video_payload_type) + + sender_bitrate = ctypes.c_uint(0) + receiver_bitrate = ctypes.c_uint(8000) + rtc.wire_sender_bitrate(video_sender_track, sender_bitrate) + + audio_codec : AudioCodec = 'opus' + audio_payload_type = rtc.get_payload_type(sdp_offer, audio_codec) + + if audio_payload_type: + audio_receiver_track = rtc.add_audio_track(peer_connection, 'recvonly', audio_codec, audio_payload_type) + audio_sender_track = rtc.add_audio_track(peer_connection, 'sendonly', audio_codec, audio_payload_type) + + rtc.set_remote_description(peer_connection, sdp_offer) + sdp_answer = rtc.create_sdp_answer(peer_connection) + + if sdp_answer: + rtc_peer : RtcPeer =\ + { + 'peer_connection': peer_connection, + 'video': + { + 'sender_track': video_sender_track, + 'receiver_track': video_receiver_track, + 'codec': video_codec + }, + 'sender_bitrate': sender_bitrate, + 'receiver_bitrate': receiver_bitrate + } + + if audio_payload_type: + rtc_peer['audio'] = RtcPeerAudio( + sender_track = audio_sender_track, + receiver_track = audio_receiver_track, + codec = audio_codec + ) + + rtc_store.init_peers(session_id) + rtc_store.get_peers(session_id).append(rtc_peer) + + threading.Thread(target = run_peer_loop, args = (session_id, rtc_peer), daemon = True).start() + + return sdp_answer + + datachannel_module.create_static_library().rtcDeletePeerConnection(peer_connection) + + return None + + +def run_peer_loop(session_id : SessionId, rtc_peer : RtcPeer) -> None: + execution_thread_count = state_manager.get_item('execution_thread_count') + #todo: is bytes, Resolution not a XXXPointer type + video_queue : Queue[Tuple[float, Future[Tuple[bytes, Resolution]]]] = Queue(maxsize = execution_thread_count) + audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = execution_thread_count * 10) + video_executor = ThreadPoolExecutor(max_workers = execution_thread_count) + + video_receiver_thread = threading.Thread(target = receive_video_frames, args = (rtc_peer.get('video'), video_queue, video_executor), daemon = True) + video_encoder_thread = threading.Thread(target = run_video_encode_loop, args = (rtc_peer, video_queue), daemon = True) + video_receiver_thread.start() + video_encoder_thread.start() + + if rtc_peer.get('audio'): + audio_receiver_thread = threading.Thread(target = receive_audio_frames, args = (rtc_peer.get('audio'), audio_queue), daemon = True) + audio_encoder_thread = threading.Thread(target = run_audio_encode_loop, args = (rtc_peer, audio_queue), daemon = True) + audio_receiver_thread.start() + audio_encoder_thread.start() + audio_receiver_thread.join() + audio_encoder_thread.join() + + video_receiver_thread.join() + video_encoder_thread.join() + video_executor.shutdown(wait = True) + rtc_store.delete_peers(session_id) + + +def destroy_stream(session_id : SessionId) -> bool: + if rtc_store.has_peers(session_id): + rtc_store.delete_peers(session_id) + return not rtc_store.has_peers(session_id) + + return False diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py new file mode 100644 index 00000000..b982cb58 --- /dev/null +++ b/facefusion/apis/stream_video.py @@ -0,0 +1,204 @@ +import ctypes +import time +from concurrent.futures import Future, ThreadPoolExecutor +from functools import partial +from queue import Queue +from typing import Optional, Tuple + +import cv2 +import numpy + +from facefusion import rtc, streamer +from facefusion.apis.stream_event import create_receive_event +from facefusion.codecs import aom_decoder, aom_encoder, vpx_decoder, vpx_encoder +from facefusion.types import AomDecoder, AomEncoder, AomPointer, BitRate, Resolution, RtcPeer, RtcPeerVideo, VideoCodec, VisionFrame, VpxDecoder, VpxEncoder, VpxPointer + + +def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[float, Future[Tuple[bytes, Resolution]]]]) -> None: + video_codec = rtc_peer.get('video').get('codec') + video_time, video_future = video_queue.get() + video_buffer, video_resolution = video_future.result() + + if video_buffer: + temp_resolution : Resolution = video_resolution + temp_bitrate : BitRate = 8000 + video_encoder = create_video_encoder(video_codec, temp_resolution, temp_bitrate) + temp_video_time = video_time + frame_index = 0 + + while video_buffer: + encode_start = time.monotonic() + sender_bitrate = calculate_sender_bitrate(rtc_peer, temp_bitrate) + + if video_resolution[0] - temp_resolution[0] or video_resolution[1] - temp_resolution[1]: + temp_resolution = video_resolution + update_video_encoder_resolution(video_codec, video_encoder, temp_resolution) + + if sender_bitrate - temp_bitrate: + temp_bitrate = sender_bitrate + update_video_encoder_bitrate(video_codec, video_encoder, temp_bitrate) + + __video_buffer__ = encode_video_frame(video_codec, video_encoder, video_buffer, temp_resolution, frame_index) + + if __video_buffer__: + video_timestamp = int(video_time * 90000) + rtc.send_video(rtc_peer, __video_buffer__, video_timestamp) + + encode_time = time.monotonic() - encode_start + frame_interval = video_time - temp_video_time + temp_video_time = video_time + + receiver_bitrate = calculate_receiver_bitrate(rtc_peer, encode_time, frame_interval) + rtc.adapt_receiver_bitrate(rtc_peer, receiver_bitrate) + frame_index += 1 + + video_time, video_future = video_queue.get() + video_buffer, video_resolution = video_future.result() + + destroy_video_encoder(video_codec, video_encoder) + rtc.clear_bitrate(rtc_peer) + + +def receive_video_frames(rtc_peer_video : RtcPeerVideo, video_queue : Queue[Tuple[float, Future[Tuple[bytes, Resolution]]]], video_executor : ThreadPoolExecutor) -> None: + video_track = rtc_peer_video.get('receiver_track') + video_codec = rtc_peer_video.get('codec') + video_decoder = create_video_decoder(video_codec) + + video_frame_handler = partial(handle_video_frame, video_codec, video_decoder, video_queue, video_executor) + receive_event = create_receive_event(video_track, video_frame_handler) + receive_event.wait() + + empty_future : Future[Tuple[bytes, Resolution]] = Future() + empty_future.set_result((bytes(), (0, 0))) + video_queue.put((0.0, empty_future)) + destroy_video_decoder(video_codec, video_decoder) + + +def process_video_frame(input_vision_frame : VisionFrame) -> Tuple[bytes, Resolution]: + output_vision_frame = streamer.process_stream_frame(input_vision_frame) + output_resolution : Resolution = (output_vision_frame.shape[1], output_vision_frame.shape[0]) + output_buffer = cv2.cvtColor(output_vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes() + return output_buffer, output_resolution + + +def calculate_receiver_bitrate(rtc_peer : RtcPeer, encode_time : float, frame_interval : float) -> BitRate: + min_bitrate : BitRate = 500 + max_bitrate : BitRate = 8000 + bitrate : BitRate = rtc_peer.get('receiver_bitrate').value + + if frame_interval > 0: + scale = frame_interval / encode_time + bitrate = int(bitrate * scale) + bitrate = max(min_bitrate, min(max_bitrate, bitrate)) + + return bitrate + + +#todo: does not feel final as this is an clamp and not calculate +def calculate_sender_bitrate(rtc_peer : RtcPeer, bitrate : BitRate) -> BitRate: + min_bitrate : BitRate = 500 + max_bitrate : BitRate = 8000 + peer_bitrate : BitRate = rtc_peer.get('sender_bitrate').value + + if peer_bitrate > 0: + bitrate = max(min_bitrate, min(max_bitrate, peer_bitrate)) + + return bitrate + + +def decode_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, input_buffer : bytes) -> Optional[VisionFrame]: + if video_codec == 'av1': + aom_pointer = aom_decoder.decode(video_decoder, input_buffer) + + if aom_pointer: + return normalize_vision_frame(aom_pointer) + + if video_codec in [ 'vp8', 'vp9' ]: + vpx_pointer = vpx_decoder.decode(video_decoder, input_buffer) + + if vpx_pointer: + return normalize_vision_frame(vpx_pointer) + + return None + + +def encode_video_frame(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder, input_buffer : bytes, frame_resolution : Resolution, frame_index : int) -> bytes: + if video_codec == 'av1': + return aom_encoder.encode(video_encoder, input_buffer, frame_resolution, frame_index) + + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_encoder.encode(video_encoder, input_buffer, frame_resolution, frame_index) + + return bytes() + + +def normalize_vision_frame(frame_pointer : AomPointer | VpxPointer) -> VisionFrame: + frame_width, frame_height = frame_pointer.get('resolution') + vision_frame = numpy.frombuffer(frame_pointer.get('buffer'), dtype = numpy.uint8).reshape((frame_height * 3 // 2, frame_width)) + return cv2.cvtColor(vision_frame, cv2.COLOR_YUV2BGR_I420) + + +def create_video_decoder(video_codec : VideoCodec) -> Optional[VpxDecoder | AomDecoder]: + if video_codec == 'av1': + return aom_decoder.create(8) + + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_decoder.create(video_codec, 8) + + return None + + +def create_video_encoder(video_codec : VideoCodec, frame_resolution : Resolution, bitrate : BitRate) -> Optional[VpxEncoder | AomEncoder]: + if video_codec == 'av1': + return aom_encoder.create(frame_resolution, bitrate, 8, 10) + + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_encoder.create(video_codec, frame_resolution, bitrate, 8, 10) + + return None + + +def destroy_video_decoder(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder) -> None: + if video_codec == 'av1': + aom_decoder.destroy(video_decoder) + + if video_codec in [ 'vp8', 'vp9' ]: + vpx_decoder.destroy(video_decoder) + + +def destroy_video_encoder(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder) -> None: + if video_codec == 'av1': + aom_encoder.destroy(video_encoder) + + if video_codec in [ 'vp8', 'vp9' ]: + vpx_encoder.destroy(video_encoder) + + +def update_video_encoder_resolution(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder, frame_resolution : Resolution) -> bool: + if video_codec == 'av1': + return aom_encoder.update_resolution(video_encoder, frame_resolution) + + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_encoder.update_resolution(video_encoder, frame_resolution) + + return False + + +def update_video_encoder_bitrate(video_codec : VideoCodec, video_encoder : VpxEncoder | AomEncoder, bitrate : BitRate) -> bool: + if video_codec == 'av1': + return aom_encoder.update_bitrate(video_encoder, bitrate) + + if video_codec in [ 'vp8', 'vp9' ]: + return vpx_encoder.update_bitrate(video_encoder, bitrate) + + return False + + +#todo: we can remove the dead args or pass audio buffer +def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[float, Future[Tuple[bytes, Resolution]]]], video_executor : ThreadPoolExecutor, track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None: + video_buffer = ctypes.string_at(data, size) + vision_frame = decode_video_frame(video_codec, video_decoder, video_buffer) + + if numpy.any(vision_frame) and video_queue.qsize() < video_queue.maxsize: + video_future = video_executor.submit(process_video_frame, vision_frame) + video_queue.put((time.monotonic(), video_future)) diff --git a/facefusion/codecs/__init__.py b/facefusion/codecs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/facefusion/codecs/aom_decoder.py b/facefusion/codecs/aom_decoder.py new file mode 100644 index 00000000..1c5ea56f --- /dev/null +++ b/facefusion/codecs/aom_decoder.py @@ -0,0 +1,70 @@ +import ctypes +import struct +from typing import Optional + +from facefusion.libraries import aom as aom_module +from facefusion.types import AomDecoder, AomPointer + + +def create(thread_count : int) -> Optional[AomDecoder]: + aom_library = aom_module.create_static_library() + + if aom_library: + aom_decoder = ctypes.create_string_buffer(128) + aom_codec = ctypes.c_void_p.in_dll(aom_library, 'aom_codec_av1_dx_algo') + config_buffer = ctypes.create_string_buffer(128) + + struct.pack_into('I', config_buffer, 0, thread_count) + struct.pack_into('I', config_buffer, 12, 1) + + if aom_library.aom_codec_dec_init_ver(aom_decoder, ctypes.byref(aom_codec), config_buffer, 0, 22) == 0: + return aom_decoder + + return None + + +def decode(aom_decoder : AomDecoder, input_buffer : bytes) -> Optional[AomPointer]: + aom_library = aom_module.create_static_library() + + if aom_library and input_buffer: + input_total = len(input_buffer) + temp_buffer = ctypes.create_string_buffer(input_buffer) + + if aom_library.aom_codec_decode(aom_decoder, temp_buffer, input_total, None) == 0: + address = aom_library.aom_codec_get_frame(aom_decoder, ctypes.byref(ctypes.c_void_p(0))) + + if address: + frame_width = ctypes.c_uint.from_address(address + 28).value & ~1 + frame_height = ctypes.c_uint.from_address(address + 32).value & ~1 + + return AomPointer( + buffer = collect(address, frame_width, frame_height), + resolution = (frame_width, frame_height) + ) + + return None + + +def collect(address : int, frame_width : int, frame_height : int) -> bytes: + output_parts = [] + + for index in range(3): + plane_pointer = ctypes.c_void_p.from_address(address + 64 + index * 8).value + stride = ctypes.c_int.from_address(address + 88 + index * 4).value + plane_width = frame_width >> (index > 0) + plane_height = frame_height >> (index > 0) + + if stride == plane_width: + output_parts.append(ctypes.string_at(plane_pointer, plane_width * plane_height)) + else: + for row in range(plane_height): + output_parts.append(ctypes.string_at(plane_pointer + row * stride, plane_width)) + + return bytes().join(output_parts) + + +def destroy(aom_decoder : AomDecoder) -> None: + aom_library = aom_module.create_static_library() + + if aom_library: + aom_library.aom_codec_destroy(aom_decoder) diff --git a/facefusion/codecs/aom_encoder.py b/facefusion/codecs/aom_encoder.py new file mode 100644 index 00000000..e696bf66 --- /dev/null +++ b/facefusion/codecs/aom_encoder.py @@ -0,0 +1,94 @@ +import ctypes +import struct +from typing import Optional + +from facefusion.libraries import aom as aom_module +from facefusion.types import AomEncoder, BitRate, Resolution + + +def create(frame_resolution : Resolution, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[AomEncoder]: + aom_library = aom_module.create_static_library() + + if aom_library: + aom_encoder = ctypes.create_string_buffer(1152) + aom_codec = ctypes.c_void_p.in_dll(aom_library, 'aom_codec_av1_cx_algo') + + config_buffer = ctypes.create_string_buffer(1024) + + if aom_library.aom_codec_enc_config_default(ctypes.byref(aom_codec), config_buffer, 1) == 0: + struct.pack_into('I', config_buffer, 4, thread_count) + struct.pack_into('I', config_buffer, 12, frame_resolution[0]) + struct.pack_into('I', config_buffer, 16, frame_resolution[1]) + struct.pack_into('I', config_buffer, 136, bitrate) + struct.pack_into('I', config_buffer, 192, 30) + + if aom_library.aom_codec_enc_init_ver(aom_encoder, ctypes.byref(aom_codec), config_buffer, 0, 25) == 0: + aom_library.aom_codec_control(aom_encoder, 13, ctypes.c_int(cpu_count)) + aom_library.aom_codec_control(aom_encoder, 75, ctypes.c_int(1)) + aom_library.aom_codec_control(aom_encoder, 106, ctypes.c_int(0)) + aom_library.aom_codec_control(aom_encoder, 122, ctypes.c_int(0)) + aom_library.aom_codec_control(aom_encoder, 123, ctypes.c_int(0)) + ctypes.memmove(ctypes.addressof(aom_encoder) + 128, config_buffer, 1024) + return aom_encoder + + return None + + +def encode(aom_encoder : AomEncoder, input_buffer : bytes, frame_resolution : Resolution, frame_index : int) -> bytes: + aom_library = aom_module.create_static_library() + output_buffer = bytes() + + if aom_library: + temp_buffer = ctypes.create_string_buffer(256) + encode_buffer = ctypes.create_string_buffer(input_buffer) + + if aom_library.aom_img_wrap(temp_buffer, 0x102, frame_resolution[0], frame_resolution[1], 1, encode_buffer) and aom_library.aom_codec_encode(aom_encoder, temp_buffer, frame_index, 1, 0, 1) == 0: + output_buffer = collect(aom_encoder) + + return output_buffer + + +def collect(aom_encoder : AomEncoder) -> bytes: + aom_library = aom_module.create_static_library() + output_parts = [] + + packet_cursor = ctypes.c_void_p(0) + packet = aom_library.aom_codec_get_cx_data(aom_encoder, ctypes.byref(packet_cursor)) + + while packet: + if ctypes.c_int.from_address(packet).value == 0: + buffer_pointer = ctypes.c_void_p.from_address(packet + 8).value + buffer_size = ctypes.c_size_t.from_address(packet + 16).value + output_parts.append(ctypes.string_at(buffer_pointer, buffer_size)) + + packet = aom_library.aom_codec_get_cx_data(aom_encoder, ctypes.byref(packet_cursor)) + + return bytes().join(output_parts) + + +def update_resolution(aom_encoder : AomEncoder, frame_resolution : Resolution) -> bool: + aom_library = aom_module.create_static_library() + + if aom_library: + struct.pack_into('I', aom_encoder, 128 + 12, frame_resolution[0]) + struct.pack_into('I', aom_encoder, 128 + 16, frame_resolution[1]) + return aom_library.aom_codec_enc_config_set(aom_encoder, ctypes.cast(ctypes.addressof(aom_encoder) + 128, ctypes.c_void_p)) == 0 + + return False + + +def update_bitrate(aom_encoder : AomEncoder, bitrate : BitRate) -> bool: + aom_library = aom_module.create_static_library() + + if aom_library: + struct.pack_into('I', aom_encoder, 128 + 136, bitrate) + return aom_library.aom_codec_enc_config_set(aom_encoder, ctypes.cast(ctypes.addressof(aom_encoder) + 128, ctypes.c_void_p)) == 0 + + return False + + +def destroy(aom_encoder : AomEncoder) -> None: + aom_library = aom_module.create_static_library() + + if aom_library: + aom_library.aom_codec_destroy(aom_encoder) diff --git a/facefusion/codecs/opus_decoder.py b/facefusion/codecs/opus_decoder.py new file mode 100644 index 00000000..3691f436 --- /dev/null +++ b/facefusion/codecs/opus_decoder.py @@ -0,0 +1,36 @@ +import ctypes +from typing import Optional + +from facefusion.libraries import opus as opus_module +from facefusion.types import OpusDecoder + + +def create(sample_rate : int, channel_total : int) -> Optional[OpusDecoder]: + opus_library = opus_module.create_static_library() + + if opus_library: + return opus_library.opus_decoder_create(sample_rate, channel_total, ctypes.byref(ctypes.c_int(0))) + + return None + + +def decode(opus_decoder : OpusDecoder, input_buffer : bytes, frame_size : int, channel_total : int) -> bytes: + opus_library = opus_module.create_static_library() + output_buffer = bytes() + + if opus_library: + input_total = len(input_buffer) + decode_buffer = (ctypes.c_float * (frame_size * channel_total))() + decode_length = opus_library.opus_decode_float(opus_decoder, input_buffer, input_total, decode_buffer, frame_size, 0) + + if decode_length: + output_buffer = ctypes.string_at(ctypes.addressof(decode_buffer), decode_length * channel_total * ctypes.sizeof(ctypes.c_float)) + + return output_buffer + + +def destroy(opus_decoder : OpusDecoder) -> None: + opus_library = opus_module.create_static_library() + + if opus_library: + opus_library.opus_decoder_destroy(opus_decoder) diff --git a/facefusion/codecs/opus_encoder.py b/facefusion/codecs/opus_encoder.py new file mode 100644 index 00000000..3e431db7 --- /dev/null +++ b/facefusion/codecs/opus_encoder.py @@ -0,0 +1,36 @@ +import ctypes +from typing import Optional + +from facefusion.libraries import opus as opus_module +from facefusion.types import OpusEncoder + + +def create(sample_rate : int, channel_total : int) -> Optional[OpusEncoder]: + opus_library = opus_module.create_static_library() + + if opus_library: + return opus_library.opus_encoder_create(sample_rate, channel_total, 2049, ctypes.byref(ctypes.c_int(0))) + + return None + + +def encode(opus_encoder : OpusEncoder, input_buffer : bytes, frame_size : int) -> bytes: + opus_library = opus_module.create_static_library() + output_buffer = bytes() + + if opus_library: + temp_buffer = ctypes.create_string_buffer(2048) + encode_buffer = ctypes.cast(ctypes.create_string_buffer(input_buffer), ctypes.POINTER(ctypes.c_float)) + encode_length = opus_library.opus_encode_float(opus_encoder, encode_buffer, frame_size, temp_buffer, 2048) + + if encode_length: + output_buffer = temp_buffer.raw[:encode_length] + + return output_buffer + + +def destroy(opus_encoder : OpusEncoder) -> None: + opus_library = opus_module.create_static_library() + + if opus_library: + opus_library.opus_encoder_destroy(opus_encoder) diff --git a/facefusion/codecs/vpx_decoder.py b/facefusion/codecs/vpx_decoder.py new file mode 100644 index 00000000..c6305f71 --- /dev/null +++ b/facefusion/codecs/vpx_decoder.py @@ -0,0 +1,74 @@ +import ctypes +import struct +from typing import Optional + +from facefusion.libraries import vpx as vpx_module +from facefusion.types import VpxDecoder, VpxPointer, VxpVideoCodec + + +def create(video_codec : VxpVideoCodec, thread_count : int) -> Optional[VpxDecoder]: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + vpx_decoder = ctypes.create_string_buffer(64) + vpx_algo = 'vpx_codec_vp8_dx_algo' + + if video_codec == 'vp9': + vpx_algo = 'vpx_codec_vp9_dx_algo' + + vpx_codec = ctypes.c_void_p.in_dll(vpx_library, vpx_algo) + config_buffer = ctypes.create_string_buffer(128) + + struct.pack_into('I', config_buffer, 0, thread_count) + + if vpx_library.vpx_codec_dec_init_ver(vpx_decoder, ctypes.byref(vpx_codec), config_buffer, 0, 12) == 0: + return vpx_decoder + + return None + + +def decode(vpx_decoder : VpxDecoder, input_buffer : bytes) -> Optional[VpxPointer]: + vpx_library = vpx_module.create_static_library() + + if vpx_library and input_buffer: + input_total = len(input_buffer) + temp_buffer = ctypes.create_string_buffer(input_buffer) + + if vpx_library.vpx_codec_decode(vpx_decoder, temp_buffer, input_total, None, 0) == 0: + address = vpx_library.vpx_codec_get_frame(vpx_decoder, ctypes.byref(ctypes.c_void_p(0))) + + if address: + frame_width = ctypes.c_uint.from_address(address + 24).value & ~1 + frame_height = ctypes.c_uint.from_address(address + 28).value & ~1 + + return VpxPointer( + buffer = collect(address, frame_width, frame_height), + resolution = (frame_width, frame_height) + ) + + return None + + +def collect(address : int, frame_width : int, frame_height : int) -> bytes: + output_parts = [] + + for index in range(3): + plane_pointer = ctypes.c_void_p.from_address(address + 48 + index * 8).value + stride = ctypes.c_int.from_address(address + 80 + index * 4).value + plane_width = frame_width >> (index > 0) + plane_height = frame_height >> (index > 0) + + if stride == plane_width: + output_parts.append(ctypes.string_at(plane_pointer, plane_width * plane_height)) + else: + for row in range(plane_height): + output_parts.append(ctypes.string_at(plane_pointer + row * stride, plane_width)) + + return bytes().join(output_parts) + + +def destroy(vpx_decoder : VpxDecoder) -> None: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + vpx_library.vpx_codec_destroy(vpx_decoder) diff --git a/facefusion/codecs/vpx_encoder.py b/facefusion/codecs/vpx_encoder.py new file mode 100644 index 00000000..ef3da620 --- /dev/null +++ b/facefusion/codecs/vpx_encoder.py @@ -0,0 +1,104 @@ +import ctypes +import struct +from typing import Optional + +from facefusion.libraries import vpx as vpx_module +from facefusion.types import BitRate, Resolution, VpxEncoder, VxpVideoCodec + + +def create(video_codec : VxpVideoCodec, frame_resolution : Resolution, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[VpxEncoder]: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + vpx_encoder = ctypes.create_string_buffer(640) + vpx_algo = 'vpx_codec_vp8_cx_algo' + + if video_codec == 'vp9': + vpx_algo = 'vpx_codec_vp9_cx_algo' + + vpx_codec = ctypes.c_void_p.in_dll(vpx_library, vpx_algo) + + config_buffer = ctypes.create_string_buffer(512) + + if vpx_library.vpx_codec_enc_config_default(ctypes.byref(vpx_codec), config_buffer, 0) == 0: + struct.pack_into('I', config_buffer, 4, thread_count) + struct.pack_into('I', config_buffer, 12, frame_resolution[0]) + struct.pack_into('I', config_buffer, 16, frame_resolution[1]) + struct.pack_into('I', config_buffer, 28, 1) + struct.pack_into('I', config_buffer, 36, 0) + struct.pack_into('I', config_buffer, 44, 0) + struct.pack_into('I', config_buffer, 72, 0) + struct.pack_into('I', config_buffer, 112, bitrate) + struct.pack_into('I', config_buffer, 116, 2) + struct.pack_into('I', config_buffer, 120, 50) + struct.pack_into('I', config_buffer, 124, 50) + struct.pack_into('I', config_buffer, 128, 50) + + if vpx_library.vpx_codec_enc_init_ver(vpx_encoder, ctypes.byref(vpx_codec), config_buffer, 0, 39) == 0: + vpx_library.vpx_codec_control_(vpx_encoder, 13, ctypes.c_int(cpu_count)) + vpx_library.vpx_codec_control_(vpx_encoder, 12, ctypes.c_int(0)) + vpx_library.vpx_codec_control_(vpx_encoder, 27, ctypes.c_int(10)) + ctypes.memmove(ctypes.addressof(vpx_encoder) + 64, config_buffer, 512) + return vpx_encoder + + return None + + +def encode(vpx_encoder : VpxEncoder, input_buffer : bytes, frame_resolution : Resolution, frame_index : int) -> bytes: + vpx_library = vpx_module.create_static_library() + output_buffer = bytes() + + if vpx_library: + temp_buffer = ctypes.create_string_buffer(256) + encode_buffer = ctypes.create_string_buffer(input_buffer) + + if vpx_library.vpx_img_wrap(temp_buffer, 0x102, frame_resolution[0], frame_resolution[1], 1, encode_buffer) and vpx_library.vpx_codec_encode(vpx_encoder, temp_buffer, frame_index, 1, 0, 1) == 0: + output_buffer = collect(vpx_encoder) + + return output_buffer + + +def collect(vpx_encoder : VpxEncoder) -> bytes: + vpx_library = vpx_module.create_static_library() + output_parts = [] + + packet_cursor = ctypes.c_void_p(0) + packet = vpx_library.vpx_codec_get_cx_data(vpx_encoder, ctypes.byref(packet_cursor)) + + while packet: + if ctypes.c_int.from_address(packet).value == 0: + buffer_pointer = ctypes.c_void_p.from_address(packet + 8).value + buffer_size = ctypes.c_size_t.from_address(packet + 16).value + output_parts.append(ctypes.string_at(buffer_pointer, buffer_size)) + + packet = vpx_library.vpx_codec_get_cx_data(vpx_encoder, ctypes.byref(packet_cursor)) + + return bytes().join(output_parts) + + +def update_resolution(vpx_encoder : VpxEncoder, frame_resolution : Resolution) -> bool: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + struct.pack_into('I', vpx_encoder, 64 + 12, frame_resolution[0]) + struct.pack_into('I', vpx_encoder, 64 + 16, frame_resolution[1]) + return vpx_library.vpx_codec_enc_config_set(vpx_encoder, ctypes.cast(ctypes.addressof(vpx_encoder) + 64, ctypes.c_void_p)) == 0 + + return False + + +def update_bitrate(vpx_encoder : VpxEncoder, bitrate : BitRate) -> bool: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + struct.pack_into('I', vpx_encoder, 64 + 112, bitrate) + return vpx_library.vpx_codec_enc_config_set(vpx_encoder, ctypes.cast(ctypes.addressof(vpx_encoder) + 64, ctypes.c_void_p)) == 0 + + return False + + +def destroy(vpx_encoder : VpxEncoder) -> None: + vpx_library = vpx_module.create_static_library() + + if vpx_library: + vpx_library.vpx_codec_destroy(vpx_encoder) diff --git a/facefusion/core.py b/facefusion/core.py index 10d8127b..8f70cff6 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -16,6 +16,7 @@ from facefusion.filesystem import get_file_extension, get_file_name, resolve_fil from facefusion.filesystem import has_audio, has_image, has_video from facefusion.jobs import job_helper, job_manager, job_runner from facefusion.jobs.job_list import compose_job_list +from facefusion.libraries import aom as aom_module, datachannel as datachannel_module, opus as opus_module, vpx as vpx_module from facefusion.processors.core import get_processors_modules from facefusion.program import create_program from facefusion.program_helper import validate_args @@ -101,13 +102,17 @@ def pre_check() -> bool: def common_pre_check() -> bool: common_modules =\ [ + aom_module, + datachannel_module, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, - voice_extractor + opus_module, + voice_extractor, + vpx_module ] content_analyser_content = inspect.getsource(content_analyser).encode() diff --git a/facefusion/libraries/aom.py b/facefusion/libraries/aom.py new file mode 100644 index 00000000..663ee7c2 --- /dev/null +++ b/facefusion/libraries/aom.py @@ -0,0 +1,134 @@ +import ctypes +from functools import lru_cache +from typing import Optional + +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.types import LibrarySet + + +@lru_cache +def create_static_library_set() -> Optional[LibrarySet]: + if is_linux(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libaom.hash'), + 'path': resolve_relative_path('../.libraries/libaom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libaom.so'), + 'path': resolve_relative_path('../.libraries/libaom.so') + } + } + } + if is_macos(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libaom.hash'), + 'path': resolve_relative_path('../.libraries/libaom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libaom.dylib'), + 'path': resolve_relative_path('../.libraries/libaom.dylib') + } + } + } + if is_windows(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/aom.hash'), + 'path': resolve_relative_path('../.libraries/aom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/aom.dll'), + 'path': resolve_relative_path('../.libraries/aom.dll') + } + } + } + + return None + + +def pre_check() -> bool: + library_hash_set = create_static_library_set().get('hashes') + library_source_set = create_static_library_set().get('sources') + + return conditional_download_hashes(library_hash_set) and conditional_download_sources(library_source_set) + + +@lru_cache +def create_static_library() -> Optional[ctypes.CDLL]: + library_path = create_static_library_set().get('sources').get('aom').get('path') + + if library_path: + if is_windows(): + library = ctypes.CDLL(library_path, winmode = 0) + else: + library = ctypes.CDLL(library_path) + + if library: + return init_ctypes(library) + + return None + + +def init_ctypes(library : ctypes.CDLL) -> ctypes.CDLL: + library.aom_codec_enc_config_default.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint ] + library.aom_codec_enc_config_default.restype = ctypes.c_int + + library.aom_codec_enc_init_ver.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long, ctypes.c_int ] + library.aom_codec_enc_init_ver.restype = ctypes.c_int + + library.aom_codec_encode.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int64, ctypes.c_ulong, ctypes.c_long, ctypes.c_ulong ] + library.aom_codec_encode.restype = ctypes.c_int + + library.aom_codec_get_cx_data.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p) ] + library.aom_codec_get_cx_data.restype = ctypes.c_void_p + + library.aom_codec_enc_config_set.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ] + library.aom_codec_enc_config_set.restype = ctypes.c_int + + library.aom_codec_destroy.argtypes = [ ctypes.c_void_p ] + library.aom_codec_destroy.restype = ctypes.c_int + + library.aom_img_wrap.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint, ctypes.c_void_p ] + library.aom_img_wrap.restype = ctypes.c_void_p + + library.aom_codec_control.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_int ] + library.aom_codec_control.restype = ctypes.c_int + + library.aom_codec_dec_init_ver.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long, ctypes.c_int ] + library.aom_codec_dec_init_ver.restype = ctypes.c_int + + library.aom_codec_decode.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p ] + library.aom_codec_decode.restype = ctypes.c_int + + library.aom_codec_get_frame.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p) ] + library.aom_codec_get_frame.restype = ctypes.c_void_p + + return library diff --git a/facefusion/libraries/datachannel.py b/facefusion/libraries/datachannel.py new file mode 100644 index 00000000..15b6dc39 --- /dev/null +++ b/facefusion/libraries/datachannel.py @@ -0,0 +1,294 @@ +import ctypes +from functools import lru_cache +from typing import Optional + +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.types import LibrarySet + + +@lru_cache +def create_static_library_set() -> Optional[LibrarySet]: + if is_linux(): + return\ + { + 'hashes': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libcrypto.hash'), + 'path': resolve_relative_path('../.libraries/libcrypto.hash') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libdatachannel_next.hash'), + 'path': resolve_relative_path('../.libraries/libdatachannel_next.hash') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libssl.hash'), + 'path': resolve_relative_path('../.libraries/libssl.hash') + } + }, + 'sources': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libcrypto.so'), + 'path': resolve_relative_path('../.libraries/libcrypto.so') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libdatachannel_next.so'), + 'path': resolve_relative_path('../.libraries/libdatachannel_next.so') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libssl.so'), + 'path': resolve_relative_path('../.libraries/libssl.so') + } + } + } + if is_macos(): + return\ + { + 'hashes': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libcrypto.hash'), + 'path': resolve_relative_path('../.libraries/libcrypto.hash') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libdatachannel_next.hash'), + 'path': resolve_relative_path('../.libraries/libdatachannel_next.hash') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libssl.hash'), + 'path': resolve_relative_path('../.libraries/libssl.hash') + } + }, + 'sources': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libcrypto.dylib'), + 'path': resolve_relative_path('../.libraries/libcrypto.dylib') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libdatachannel_next.dylib'), + 'path': resolve_relative_path('../.libraries/libdatachannel_next.dylib') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libssl.dylib'), + 'path': resolve_relative_path('../.libraries/libssl.dylib') + } + } + } + if is_windows(): + return\ + { + 'hashes': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/crypto.hash'), + 'path': resolve_relative_path('../.libraries/crypto.hash') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/datachannel_next.hash'), + 'path': resolve_relative_path('../.libraries/datachannel_next.hash') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/ssl.hash'), + 'path': resolve_relative_path('../.libraries/ssl.hash') + } + }, + 'sources': + { + 'crypto': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/crypto.dll'), + 'path': resolve_relative_path('../.libraries/crypto.dll') + }, + 'datachannel': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/datachannel_next.dll'), + 'path': resolve_relative_path('../.libraries/datachannel_next.dll') + }, + 'ssl': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/ssl.dll'), + 'path': resolve_relative_path('../.libraries/ssl.dll') + } + } + } + + return None + + +def pre_check() -> bool: + library_hash_set = create_static_library_set().get('hashes') + library_source_set = create_static_library_set().get('sources') + + return conditional_download_hashes(library_hash_set) and conditional_download_sources(library_source_set) + + +@lru_cache +def create_static_library() -> Optional[ctypes.CDLL]: + crypto_source_path = create_static_library_set().get('sources').get('crypto').get('path') + datachannel_source_path = create_static_library_set().get('sources').get('datachannel').get('path') + ssl_source_path = create_static_library_set().get('sources').get('ssl').get('path') + + if crypto_source_path and datachannel_source_path and ssl_source_path: + if is_windows(): + ctypes.CDLL(crypto_source_path, winmode = 0) + ctypes.CDLL(ssl_source_path, winmode = 0) + library = ctypes.CDLL(datachannel_source_path, winmode = 0) + else: + ctypes.CDLL(crypto_source_path) + ctypes.CDLL(ssl_source_path) + library = ctypes.CDLL(datachannel_source_path) + + if library: + return init_ctypes(library) + + return None + + +def init_ctypes(library : ctypes.CDLL) -> ctypes.CDLL: + library.rtcInitLogger.argtypes = [ ctypes.c_int, ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p) ] + library.rtcInitLogger.restype = None + library.rtcInitLogger(2, ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p)(0)) + + library.rtcCreatePeerConnection.restype = ctypes.c_int + + library.rtcDeletePeerConnection.argtypes = [ ctypes.c_int ] + library.rtcDeletePeerConnection.restype = ctypes.c_int + + library.rtcSetLocalDescription.argtypes = [ ctypes.c_int, ctypes.c_char_p ] + library.rtcSetLocalDescription.restype = ctypes.c_int + + library.rtcSetRemoteDescription.argtypes = [ ctypes.c_int, ctypes.c_char_p, ctypes.c_char_p ] + library.rtcSetRemoteDescription.restype = ctypes.c_int + + library.rtcAddTrackEx.restype = ctypes.c_int + + library.rtcSendMessage.argtypes = [ ctypes.c_int, ctypes.c_void_p, ctypes.c_int ] + library.rtcSendMessage.restype = ctypes.c_int + + library.rtcSetAV1Packetizer.restype = ctypes.c_int + library.rtcSetVP8Packetizer.restype = ctypes.c_int + library.rtcSetVP9Packetizer.restype = ctypes.c_int + + library.rtcChainRtcpSrReporter.argtypes = [ ctypes.c_int ] + library.rtcChainRtcpSrReporter.restype = ctypes.c_int + + library.rtcSetTrackRtpTimestamp.argtypes = [ ctypes.c_int, ctypes.c_uint32 ] + library.rtcSetTrackRtpTimestamp.restype = ctypes.c_int + + library.rtcIsOpen.argtypes = [ ctypes.c_int ] + library.rtcIsOpen.restype = ctypes.c_bool + + library.rtcChainRtcpNackResponder.argtypes = [ ctypes.c_int, ctypes.c_uint ] + library.rtcChainRtcpNackResponder.restype = ctypes.c_int + + library.rtcGetLocalDescription.argtypes = [ ctypes.c_int, ctypes.c_char_p, ctypes.c_int ] + library.rtcGetLocalDescription.restype = ctypes.c_int + + library.rtcSetOpusPacketizer.restype = ctypes.c_int + + library.rtcGetPayloadTypesForCodec.argtypes = [ ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_int), ctypes.c_int ] + library.rtcGetPayloadTypesForCodec.restype = ctypes.c_int + + library.rtcSetAV1Depacketizer.argtypes = [ ctypes.c_int, ctypes.c_int ] + library.rtcSetAV1Depacketizer.restype = ctypes.c_int + library.rtcSetVP8Depacketizer.restype = ctypes.c_int + library.rtcSetVP9Depacketizer.restype = ctypes.c_int + library.rtcSetOpusDepacketizer.restype = ctypes.c_int + + library.rtcChainRtcpReceivingSession.argtypes = [ ctypes.c_int ] + library.rtcChainRtcpReceivingSession.restype = ctypes.c_int + + library.rtcSetFrameCallback.argtypes = [ ctypes.c_int, ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p) ] + library.rtcSetFrameCallback.restype = ctypes.c_int + + library.rtcSetUserPointer.argtypes = [ ctypes.c_int, ctypes.c_void_p ] + library.rtcSetUserPointer.restype = None + + library.rtcChainRembHandler.argtypes = [ ctypes.c_int, ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_uint, ctypes.c_void_p) ] + library.rtcChainRembHandler.restype = ctypes.c_int + + library.rtcRequestBitrate.argtypes = [ ctypes.c_int, ctypes.c_uint ] + library.rtcRequestBitrate.restype = ctypes.c_int + + library.rtcSetClosedCallback.argtypes = [ ctypes.c_int, ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p) ] + library.rtcSetClosedCallback.restype = ctypes.c_int + + return library + + +def define_rtc_configuration() -> ctypes.Structure: + return type('RTC_CONFIGURATION', (ctypes.Structure,), + { + '_fields_': + [ + ('iceServers', ctypes.POINTER(ctypes.c_char_p)), + ('iceServersCount', ctypes.c_int), + ('proxyServer', ctypes.c_char_p), + ('bindAddress', ctypes.c_char_p), + ('certificateType', ctypes.c_int), + ('iceTransportPolicy', ctypes.c_int), + ('enableIceTcp', ctypes.c_bool), + ('enableIceUdpMux', ctypes.c_bool), + ('disableAutoNegotiation', ctypes.c_bool), + ('forceMediaTransport', ctypes.c_bool), + ('portRangeBegin', ctypes.c_ushort), + ('portRangeEnd', ctypes.c_ushort), + ('mtu', ctypes.c_int), + ('maxMessageSize', ctypes.c_int) + ] + })() + + +def define_rtc_track_init() -> ctypes.Structure: + return type('RTC_TRACK_INIT', (ctypes.Structure,), + { + '_fields_': + [ + ('direction', ctypes.c_int), + ('codec', ctypes.c_int), + ('payloadType', ctypes.c_int), + ('ssrc', ctypes.c_uint32), + ('mid', ctypes.c_char_p), + ('name', ctypes.c_char_p), + ('msid', ctypes.c_char_p), + ('trackId', ctypes.c_char_p), + ('profile', ctypes.c_char_p) + ] + })() + + +def define_rtc_packetizer_init() -> ctypes.Structure: + return type('RTC_PACKETIZER_INIT', (ctypes.Structure,), + { + '_fields_': + [ + ('ssrc', ctypes.c_uint32), + ('cname', ctypes.c_char_p), + ('payloadType', ctypes.c_uint8), + ('clockRate', ctypes.c_uint32), + ('sequenceNumber', ctypes.c_uint16), + ('timestamp', ctypes.c_uint32), + ('maxFragmentSize', ctypes.c_uint16), + ('nalSeparator', ctypes.c_int), + ('obuPacketization', ctypes.c_int) + ] + })() diff --git a/facefusion/libraries/opus.py b/facefusion/libraries/opus.py new file mode 100644 index 00000000..ccb2052d --- /dev/null +++ b/facefusion/libraries/opus.py @@ -0,0 +1,121 @@ +import ctypes +from functools import lru_cache +from typing import Optional + +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.types import LibrarySet + + +@lru_cache +def create_static_library_set() -> Optional[LibrarySet]: + if is_linux(): + return\ + { + 'hashes': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libopus.hash'), + 'path': resolve_relative_path('../.libraries/libopus.hash') + } + }, + 'sources': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libopus.so'), + 'path': resolve_relative_path('../.libraries/libopus.so') + } + } + } + + if is_macos(): + return\ + { + 'hashes': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libopus.hash'), + 'path': resolve_relative_path('../.libraries/libopus.hash') + } + }, + 'sources': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libopus.dylib'), + 'path': resolve_relative_path('../.libraries/libopus.dylib') + } + } + } + + if is_windows(): + return\ + { + 'hashes': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/opus.hash'), + 'path': resolve_relative_path('../.libraries/opus.hash') + } + }, + 'sources': + { + 'opus': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/opus.dll'), + 'path': resolve_relative_path('../.libraries/opus.dll') + } + } + } + + return None + + +def pre_check() -> bool: + library_hash_set = create_static_library_set().get('hashes') + library_source_set = create_static_library_set().get('sources') + + return conditional_download_hashes(library_hash_set) and conditional_download_sources(library_source_set) + + +@lru_cache +def create_static_library() -> Optional[ctypes.CDLL]: + library_path = create_static_library_set().get('sources').get('opus').get('path') + + if library_path: + if is_windows(): + library = ctypes.CDLL(library_path, winmode = 0) + else: + library = ctypes.CDLL(library_path) + + if library: + return init_ctypes(library) + + return None + + +def init_ctypes(library : ctypes.CDLL) -> ctypes.CDLL: + library.opus_encoder_create.argtypes = [ ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int) ] + library.opus_encoder_create.restype = ctypes.c_void_p + + library.opus_encode_float.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_float), ctypes.c_int, ctypes.c_char_p, ctypes.c_int ] + library.opus_encode_float.restype = ctypes.c_int + + library.opus_encoder_destroy.argtypes = [ ctypes.c_void_p ] + library.opus_encoder_destroy.restype = None + + library.opus_decoder_create.argtypes = [ ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int) ] + library.opus_decoder_create.restype = ctypes.c_void_p + + library.opus_decode_float.argtypes = [ ctypes.c_void_p, ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_float), ctypes.c_int, ctypes.c_int ] + library.opus_decode_float.restype = ctypes.c_int + + library.opus_decoder_destroy.argtypes = [ ctypes.c_void_p ] + library.opus_decoder_destroy.restype = None + + return library diff --git a/facefusion/libraries/vpx.py b/facefusion/libraries/vpx.py new file mode 100644 index 00000000..b128643d --- /dev/null +++ b/facefusion/libraries/vpx.py @@ -0,0 +1,134 @@ +import ctypes +from functools import lru_cache +from typing import Optional + +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.types import LibrarySet + + +@lru_cache +def create_static_library_set() -> Optional[LibrarySet]: + if is_linux(): + return\ + { + 'hashes': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libvpx.hash'), + 'path': resolve_relative_path('../.libraries/libvpx.hash') + } + }, + 'sources': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libvpx.so'), + 'path': resolve_relative_path('../.libraries/libvpx.so') + } + } + } + if is_macos(): + return\ + { + 'hashes': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libvpx.hash'), + 'path': resolve_relative_path('../.libraries/libvpx.hash') + } + }, + 'sources': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libvpx.dylib'), + 'path': resolve_relative_path('../.libraries/libvpx.dylib') + } + } + } + if is_windows(): + return\ + { + 'hashes': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/vpx.hash'), + 'path': resolve_relative_path('../.libraries/vpx.hash') + } + }, + 'sources': + { + 'vpx': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/vpx.dll'), + 'path': resolve_relative_path('../.libraries/vpx.dll') + } + } + } + + return None + + +def pre_check() -> bool: + library_hash_set = create_static_library_set().get('hashes') + library_source_set = create_static_library_set().get('sources') + + return conditional_download_hashes(library_hash_set) and conditional_download_sources(library_source_set) + + +@lru_cache +def create_static_library() -> Optional[ctypes.CDLL]: + library_path = create_static_library_set().get('sources').get('vpx').get('path') + + if library_path: + if is_windows(): + library = ctypes.CDLL(library_path, winmode = 0) + else: + library = ctypes.CDLL(library_path) + + if library: + return init_ctypes(library) + + return None + + +def init_ctypes(library : ctypes.CDLL) -> ctypes.CDLL: + library.vpx_codec_enc_config_default.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint ] + library.vpx_codec_enc_config_default.restype = ctypes.c_int + + library.vpx_codec_enc_init_ver.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long, ctypes.c_int ] + library.vpx_codec_enc_init_ver.restype = ctypes.c_int + + library.vpx_codec_encode.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int64, ctypes.c_ulong, ctypes.c_long, ctypes.c_ulong ] + library.vpx_codec_encode.restype = ctypes.c_int + + library.vpx_codec_get_cx_data.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p) ] + library.vpx_codec_get_cx_data.restype = ctypes.c_void_p + + library.vpx_codec_enc_config_set.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ] + library.vpx_codec_enc_config_set.restype = ctypes.c_int + + library.vpx_codec_destroy.argtypes = [ ctypes.c_void_p ] + library.vpx_codec_destroy.restype = ctypes.c_int + + library.vpx_img_wrap.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint, ctypes.c_void_p ] + library.vpx_img_wrap.restype = ctypes.c_void_p + + library.vpx_codec_control_.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_int ] + library.vpx_codec_control_.restype = ctypes.c_int + + library.vpx_codec_dec_init_ver.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long, ctypes.c_int ] + library.vpx_codec_dec_init_ver.restype = ctypes.c_int + + library.vpx_codec_decode.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p, ctypes.c_long ] + library.vpx_codec_decode.restype = ctypes.c_int + + library.vpx_codec_get_frame.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p) ] + library.vpx_codec_get_frame.restype = ctypes.c_void_p + + return library diff --git a/facefusion/rtc.py b/facefusion/rtc.py new file mode 100644 index 00000000..0097c46e --- /dev/null +++ b/facefusion/rtc.py @@ -0,0 +1,263 @@ +import ctypes +from typing import List, Optional + +from facefusion.libraries import datachannel as datachannel_module +from facefusion.types import AudioCodec, BitRate, MediaDirection, PeerConnection, RtcAudioTrack, RtcPeer, RtcTrackInit, RtcVideoTrack, SdpAnswer, SdpOffer, VideoCodec + + +def create_peer_connection() -> PeerConnection: + datachannel_library = datachannel_module.create_static_library() + rtc_configuration = datachannel_module.define_rtc_configuration() + + rtc_configuration.enableIceUdpMux = True + rtc_configuration.forceMediaTransport = True + rtc_configuration.disableAutoNegotiation = True + + return datachannel_library.rtcCreatePeerConnection(ctypes.byref(rtc_configuration)) + + +def create_sdp_offer(peer_connection : PeerConnection) -> Optional[SdpOffer]: + datachannel_library = datachannel_module.create_static_library() + datachannel_library.rtcSetLocalDescription(peer_connection, b'offer') + + sdp_buffer = ctypes.create_string_buffer(8192) + + if datachannel_library.rtcGetLocalDescription(peer_connection, sdp_buffer, 8192): + return sdp_buffer.value.decode() + + return None + + +def create_sdp_answer(peer_connection : PeerConnection) -> Optional[SdpAnswer]: + datachannel_library = datachannel_module.create_static_library() + datachannel_library.rtcSetLocalDescription(peer_connection, b'answer') + + sdp_buffer = ctypes.create_string_buffer(8192) + + if datachannel_library.rtcGetLocalDescription(peer_connection, sdp_buffer, 8192): + return sdp_buffer.value.decode() + + return None + + +def set_remote_description(peer_connection : PeerConnection, sdp_offer : SdpOffer) -> None: + datachannel_library = datachannel_module.create_static_library() + datachannel_library.rtcSetRemoteDescription(peer_connection, sdp_offer.encode(), b'offer') + + return None + + +def send_video(rtc_peer : RtcPeer, video_buffer : bytes, video_timestamp : int) -> None: + datachannel_library = datachannel_module.create_static_library() + + if rtc_peer.get('video'): + video_track = rtc_peer.get('video').get('sender_track') + + if datachannel_library.rtcIsOpen(video_track): + video_total = len(video_buffer) + datachannel_library.rtcSetTrackRtpTimestamp(video_track, video_timestamp) + datachannel_library.rtcSendMessage(video_track, video_buffer, video_total) + + return None + + +def send_audio(rtc_peer : RtcPeer, audio_buffer : bytes, audio_timestamp : int) -> None: + datachannel_library = datachannel_module.create_static_library() + + if rtc_peer.get('audio'): + audio_track = rtc_peer.get('audio').get('sender_track') + + if datachannel_library.rtcIsOpen(audio_track): + audio_total = len(audio_buffer) + datachannel_library.rtcSetTrackRtpTimestamp(audio_track, audio_timestamp) + datachannel_library.rtcSendMessage(audio_track, audio_buffer, audio_total) + + return None + + +def delete_peers(rtc_peers : List[RtcPeer]) -> None: + datachannel_library = datachannel_module.create_static_library() + + for rtc_peer in rtc_peers: + peer_connection = rtc_peer.get('peer_connection') + + if peer_connection: + datachannel_library.rtcDeletePeerConnection(peer_connection) + + 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() + audio_track_init = create_audio_track_init(media_direction, audio_codec, payload_type) + audio_track = datachannel_library.rtcAddTrackEx(peer_connection, audio_track_init) + + if media_direction == 'sendonly': + 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) + + if media_direction == 'recvonly': + audio_depacketizer = datachannel_module.define_rtc_packetizer_init() + audio_depacketizer.ssrc = 0 + audio_depacketizer.cname = b'audio' + audio_depacketizer.payloadType = payload_type + audio_depacketizer.clockRate = 48000 + + if audio_codec == 'opus': + datachannel_library.rtcSetOpusDepacketizer(audio_track, ctypes.byref(audio_depacketizer)) + + datachannel_library.rtcChainRtcpReceivingSession(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() + video_track_init = create_video_track_init(media_direction, video_codec, payload_type) + video_track = datachannel_library.rtcAddTrackEx(peer_connection, video_track_init) + + if media_direction == 'sendonly': + 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)) + + if video_codec == 'vp9': + datachannel_library.rtcSetVP9Packetizer(video_track, ctypes.byref(video_packetizer)) + + datachannel_library.rtcChainRtcpSrReporter(video_track) + datachannel_library.rtcChainRtcpNackResponder(video_track, 512) + + if media_direction == 'recvonly': + if video_codec == 'av1': + datachannel_library.rtcSetAV1Depacketizer(video_track, 1) + + if video_codec == 'vp8': + video_depacketizer = datachannel_module.define_rtc_packetizer_init() + video_depacketizer.ssrc = 0 + video_depacketizer.cname = b'video' + video_depacketizer.payloadType = payload_type + video_depacketizer.clockRate = 90000 + datachannel_library.rtcSetVP8Depacketizer(video_track, ctypes.byref(video_depacketizer)) + + if video_codec == 'vp9': + video_depacketizer = datachannel_module.define_rtc_packetizer_init() + video_depacketizer.ssrc = 0 + video_depacketizer.cname = b'video' + video_depacketizer.payloadType = payload_type + video_depacketizer.clockRate = 90000 + datachannel_library.rtcSetVP9Depacketizer(video_track, ctypes.byref(video_depacketizer)) + + datachannel_library.rtcChainRtcpReceivingSession(video_track) + + return video_track + + +def create_audio_track_init(media_direction : MediaDirection, audio_codec : AudioCodec, payload_type : int) -> RtcTrackInit: + track_init = datachannel_module.define_rtc_track_init() + track_init.name = b'audio' + track_init.payloadType = payload_type + + if media_direction == 'sendonly': + track_init.direction = 1 + track_init.mid = b'3' + track_init.ssrc = 43 + + if media_direction == 'recvonly': + track_init.direction = 2 + track_init.mid = b'2' + track_init.ssrc = 45 + + if media_direction == 'sendrecv': + track_init.direction = 3 + track_init.mid = b'1' + track_init.ssrc = 43 + + if audio_codec == 'opus': + track_init.codec = 128 + + return ctypes.byref(track_init) + + +def create_video_track_init(media_direction : MediaDirection, video_codec : VideoCodec, payload_type : int) -> RtcTrackInit: + track_init = datachannel_module.define_rtc_track_init() + track_init.name = b'video' + track_init.payloadType = payload_type + + if media_direction == 'sendonly': + track_init.direction = 1 + track_init.mid = b'1' + track_init.ssrc = 42 + + if media_direction == 'recvonly': + track_init.direction = 2 + track_init.mid = b'0' + track_init.ssrc = 44 + + if media_direction == 'sendrecv': + track_init.direction = 3 + track_init.mid = b'0' + track_init.ssrc = 42 + + if video_codec == 'av1': + track_init.codec = 4 + + if video_codec == 'vp8': + track_init.codec = 1 + + if video_codec == 'vp9': + track_init.codec = 2 + + return ctypes.byref(track_init) + + +def get_payload_type(sdp_offer : SdpOffer, codec : AudioCodec | VideoCodec) -> int: + datachannel_library = datachannel_module.create_static_library() + payload_type_buffer = (ctypes.c_int * 16)() + payload_type_total = datachannel_library.rtcGetPayloadTypesForCodec(sdp_offer.encode(), codec.lower().encode(), payload_type_buffer, 16) + + if payload_type_total: + return payload_type_buffer[0] + + return 0 + + +@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_uint, ctypes.c_void_p) +def handle_sender_bitrate(_ : int, bitrate : BitRate, pointer : int) -> None: + ctypes.cast(pointer, ctypes.POINTER(ctypes.c_uint)).contents.value = bitrate // 1000 + + +def wire_sender_bitrate(video_track : RtcVideoTrack, bitrate : ctypes.c_uint) -> None: + datachannel_library = datachannel_module.create_static_library() + datachannel_library.rtcSetUserPointer(video_track, ctypes.cast(ctypes.byref(bitrate), ctypes.c_void_p)) + datachannel_library.rtcChainRembHandler(video_track, handle_sender_bitrate) + + +def adapt_receiver_bitrate(rtc_peer : RtcPeer, bitrate : BitRate) -> None: + datachannel_library = datachannel_module.create_static_library() + receiver_track = rtc_peer.get('video').get('receiver_track') + + rtc_peer.get('receiver_bitrate').value = bitrate + datachannel_library.rtcRequestBitrate(receiver_track, bitrate * 1000) + + +def clear_bitrate(rtc_peer : RtcPeer) -> None: + rtc_peer.get('sender_bitrate').value = 0 + rtc_peer.get('receiver_bitrate').value = 0 diff --git a/facefusion/rtc_store.py b/facefusion/rtc_store.py new file mode 100644 index 00000000..65ec6725 --- /dev/null +++ b/facefusion/rtc_store.py @@ -0,0 +1,33 @@ +from typing import List + +from facefusion import rtc +from facefusion.types import RtcPeer, RtcStore, SessionId + +RTC_STORE : RtcStore = {} + + +def init_peers(session_id : SessionId) -> None: + RTC_STORE[session_id] = [] + + +def has_peers(session_id : SessionId) -> bool: + return bool(RTC_STORE.get(session_id)) + + +def get_peers(session_id : SessionId) -> List[RtcPeer]: + return RTC_STORE.get(session_id) + + +def delete_peers(session_id : SessionId) -> None: + if session_id in RTC_STORE: + rtc_peers = get_peers(session_id) + + if rtc_peers: + rtc.delete_peers(rtc_peers) + del RTC_STORE[session_id] + + return None + + +def clear() -> None: + RTC_STORE.clear() diff --git a/facefusion/types.py b/facefusion/types.py index 404fad0f..4a7b88fd 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -1,6 +1,7 @@ +import ctypes from collections import namedtuple from datetime import datetime -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, TypeAlias, TypedDict +from typing import Any, Callable, Dict, List, Literal, NotRequired, Optional, Tuple, TypeAlias, TypedDict import cv2 import numpy @@ -555,3 +556,68 @@ State = TypedDict('State', }) ApplyStateItem : TypeAlias = Callable[[Any, Any], None] StateSet : TypeAlias = Dict[AppContext, State] + +BitRate : TypeAlias = int + +AudioCodec : TypeAlias = Literal['opus'] +VideoCodec : TypeAlias = Literal['av1', 'vp8', 'vp9'] + +AomVideoCodec : TypeAlias = Literal['av1'] +VxpVideoCodec : TypeAlias = Literal['vp8', 'vp9'] + +FrameHandler : TypeAlias = Callable[..., None] + +AomEncoder : TypeAlias = ctypes.Array[ctypes.c_char] +AomDecoder : TypeAlias = ctypes.Array[ctypes.c_char] +OpusEncoder : TypeAlias = ctypes.c_void_p +OpusDecoder : TypeAlias = ctypes.c_void_p +VpxEncoder : TypeAlias = ctypes.Array[ctypes.c_char] +VpxDecoder : TypeAlias = ctypes.Array[ctypes.c_char] + +AomPointer = TypedDict('AomPointer', +{ + 'buffer' : bytes, + 'resolution' : Resolution +}) +VpxPointer = TypedDict('VpxPointer', +{ + 'buffer' : bytes, + 'resolution' : Resolution +}) + +PeerConnection : TypeAlias = int +SdpOffer : TypeAlias = str +SdpAnswer : TypeAlias = str +MediaDirection : TypeAlias = Literal['sendonly', 'recvonly', 'sendrecv'] + +RtcTrackInit : TypeAlias = Any + +RtcVideoTrack : TypeAlias = int +RtcAudioTrack : TypeAlias = int + +RtcPeerAudio = TypedDict('RtcPeerAudio', +{ + 'sender_track': RtcAudioTrack, + 'receiver_track': RtcAudioTrack, + 'codec': AudioCodec, +}) + +RtcPeerVideo = TypedDict('RtcPeerVideo', +{ + 'sender_track': RtcVideoTrack, + 'receiver_track': RtcVideoTrack, + 'codec': VideoCodec, +}) + +RtcPeer = TypedDict('RtcPeer', +{ + 'peer_connection': PeerConnection, + 'audio': NotRequired[RtcPeerAudio], + 'video': RtcPeerVideo, + 'sender_bitrate': ctypes.c_uint, + 'receiver_bitrate': ctypes.c_uint +}) +RtcStore : TypeAlias = Dict[SessionId, List[RtcPeer]] + +LibraryOptions : TypeAlias = Dict[str, Any] +LibrarySet : TypeAlias = Dict[str, LibraryOptions]