CalvinBackup/facefusion
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion.git synced 2026-07-15 06:27:21 +02:00
Code Issues Packages Projects Releases Wiki Activity

Resolve final stream TODOs (#1155)

Browse Source 
* refactor a lot

* refactor a lot

* fix test

* fix test
This commit is contained in:
Henry Ruhs
2026-06-15 11:37:25 +02:00
committed by GitHub
parent d5271d21a1
commit 2854bf1db0
7 changed files with 64 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_api_stream_audio.py
+14 -14
View File
@@ -14,7 +14,7 @@ from facefusion.download import conditional_download
from facefusion.ffmpeg import read_audio_buffer
from facefusion.hash_helper import create_hash
from facefusion.libraries import datachannel as datachannel_module, opus as opus_module
from facefusion.types import AudioCodec, AudioFrame, FrameHandler, RtcPeer, RtcPeerAudio
from facefusion.types import AudioCodec, AudioFrame, Buffer, FrameHandler, RtcPeer, RtcPeerAudio
from .assert_helper import get_test_example_file, get_test_examples_directory
@@ -37,8 +37,11 @@ def before_each() -> None:
rtc_store.clear()
def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None:
def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event:
frame_handler(buffer, 0)
ready_event = threading.Event()
ready_event.set()
return ready_event
def test_run_audio_encode_loop() -> None:
@@ -48,6 +51,12 @@ def test_run_audio_encode_loop() -> None:
rtc_peer : RtcPeer =\
{
'peer_connection': peer_connection,
'audio':
{
'sender_track': 0,
'receiver_track': 0,
'codec': 'opus'
},
'video':
{
'sender_track': 0,
@@ -82,11 +91,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None:
audio_frame = numpy.frombuffer(audio_buffer, dtype = numpy.int16).astype(numpy.float32) / 32768.0
audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300)
datachannel_mock = MagicMock()
ready_event = threading.Event()
datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event)
with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock):
with patch('facefusion.apis.stream_audio.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))):
with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()):
rtc_peer_audio : RtcPeerAudio =\
{
@@ -94,12 +99,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None:
'receiver_track': 0,
'codec': audio_codec
}
audio_receiver_thread = threading.Thread(target = receive_audio_frames, args = (rtc_peer_audio, audio_queue), daemon = True)
audio_receiver_thread.start()
ready_event.wait(timeout = 5.0)
datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None)
datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None)
audio_receiver_thread.join(timeout = 5.0)
receive_audio_frames(rtc_peer_audio, audio_queue)
_, temp_audio_frame = audio_queue.get_nowait()
@@ -113,7 +113,7 @@ def test_handle_audio_frame() -> None:
audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300)
with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()):
handle_audio_frame('opus', audio_decoder_mock, audio_queue, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p())
handle_audio_frame('opus', audio_decoder_mock, audio_queue, bytes(), 0)
_, temp_audio_frame = audio_queue.get_nowait()
tests/test_api_stream_video.py
+15 -15
View File
@@ -5,7 +5,7 @@ from concurrent.futures import Future, ThreadPoolExecutor
from functools import partial
from queue import Queue
from typing import Tuple
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import cv2
import numpy
@@ -18,7 +18,7 @@ from facefusion.common_helper import is_linux, is_macos, is_windows
from facefusion.download import conditional_download
from facefusion.hash_helper import create_hash
from facefusion.libraries import aom as aom_module, datachannel as datachannel_module, vpx as vpx_module
from facefusion.types import BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec
from facefusion.types import Buffer, BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec
from facefusion.vision import read_video_frame
from .assert_helper import get_test_example_file, get_test_examples_directory
@@ -44,8 +44,11 @@ def before_each() -> None:
rtc_store.clear()
def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None:
def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event:
frame_handler(buffer, 0)
ready_event = threading.Event()
ready_event.set()
return ready_event
@pytest.mark.parametrize('video_codec, payload_type', [ ('av1', 35), ('vp8', 96), ('vp9', 98) ])
@@ -57,6 +60,12 @@ def test_run_video_encode_loop(video_codec : VideoCodec, payload_type : int) ->
rtc_peer : RtcPeer =\
{
'peer_connection': peer_connection,
'audio':
{
'sender_track': 0,
'receiver_track': 0,
'codec': 'opus'
},
'video':
{
'sender_track': video_sender_track,
@@ -101,12 +110,8 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None:
video_frame = read_video_frame(get_test_example_file('target-240p.mp4'))
video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = 30)
datachannel_mock = MagicMock()
ready_event = threading.Event()
datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event)
with ThreadPoolExecutor(max_workers = 1) as executor:
with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock):
with patch('facefusion.apis.stream_video.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))):
with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame):
with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))):
rtc_peer_video : RtcPeerVideo =\
@@ -115,12 +120,7 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None:
'receiver_track': 0,
'codec': video_codec
}
video_receiver_thread = threading.Thread(target = receive_video_frames, args = (rtc_peer_video, video_queue, executor), daemon = True)
video_receiver_thread.start()
ready_event.wait(timeout = 5.0)
datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None)
datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None)
video_receiver_thread.join(timeout = 5.0)
receive_video_frames(rtc_peer_video, video_queue, executor)
_, video_future = video_queue.get_nowait()
video_buffer = video_future.result().get('buffer')
@@ -267,7 +267,7 @@ def test_handle_video_frame(video_codec : VideoCodec) -> None:
with ThreadPoolExecutor(max_workers = 1) as executor:
with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame):
with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))):
handle_video_frame(video_codec, video_decoder, video_queue, executor, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p())
handle_video_frame(video_codec, video_decoder, video_queue, executor, bytes(), 0)
_, video_future = video_queue.get_nowait()
video_buffer = video_future.result().get('buffer')