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

Combine encode loop methods (#1119)

Browse Source 
* combine run_aom_encode_loop and run_vpx_encode_loop to encode_video_loop

* run_opus_encode_loop -> encode_audio_loop

* use else instead of continue

* rename to video_codec
This commit is contained in:
Harisreedhar
2026-05-16 23:29:36 +05:30
committed by GitHub
parent dd1ded1408
commit c48c238f88
6 changed files with 64 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_codec_aom.py
+4 -4
View File
@@ -23,15 +23,15 @@ def before_all() -> None:
def test_create_aom_encoder() -> None:
assert create_aom_encoder((320, 240), 1000, 8, 16)
assert create_aom_encoder((0, 0), 0, 0, 0) is None
assert create_aom_encoder(1000, 8, 16, (320, 240))
assert create_aom_encoder(0, 0, 0, (0, 0)) is None
def test_encode_aom_buffer() -> None:
vision_frame = read_video_frame(get_test_example_file('target-240p.mp4'))
video_buffer = cv2.cvtColor(vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes()
video_resolution = (vision_frame.shape[1], vision_frame.shape[0])
aom_encoder = create_aom_encoder(video_resolution, 1000, 1, 0)
aom_encoder = create_aom_encoder(1000, 1, 0, video_resolution)
if is_linux() or is_windows():
assert create_hash(encode_aom_buffer(aom_encoder, video_buffer, video_resolution, 3)) == '3ab6cc31'
@@ -41,7 +41,7 @@ def test_encode_aom_buffer() -> None:
def test_destroy_aom_encoder() -> None:
aom_encoder = create_aom_encoder((320, 240), 1000, 8, 16)
aom_encoder = create_aom_encoder(1000, 8, 16, (320, 240))
with patch.object(aom_module.create_static_library(), 'aom_codec_destroy') as mock:
destroy_aom_encoder(aom_encoder)
tests/test_codec_vpx.py
+4 -4
View File
@@ -23,15 +23,15 @@ def before_all() -> None:
def test_create_vpx_encoder() -> None:
assert create_vpx_encoder((320, 240), 1000, 8, 16)
assert create_vpx_encoder((0, 0), 0, 0, 0) is None
assert create_vpx_encoder(1000, 8, 16, (320, 240))
assert create_vpx_encoder(0, 0, 0, (0, 0)) is None
def test_encode_vpx_buffer() -> None:
vision_frame = read_video_frame(get_test_example_file('target-240p.mp4'))
video_buffer = cv2.cvtColor(vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes()
video_resolution = (vision_frame.shape[1], vision_frame.shape[0])
vpx_encoder = create_vpx_encoder(video_resolution, 1000, 1, 0)
vpx_encoder = create_vpx_encoder(1000, 1, 0, video_resolution)
if is_linux() or is_windows():
assert create_hash(encode_vpx_buffer(vpx_encoder, video_buffer, video_resolution, 3)) == 'ce133a1f'
@@ -41,7 +41,7 @@ def test_encode_vpx_buffer() -> None:
def test_destroy_vpx_encoder() -> None:
vpx_encoder = create_vpx_encoder((320, 240), 1000, 8, 16)
vpx_encoder = create_vpx_encoder(1000, 8, 16, (320, 240))
with patch.object(vpx_module.create_static_library(), 'vpx_codec_destroy') as mock:
destroy_vpx_encoder(vpx_encoder)
tests/test_stream_helper.py
+26 -28
View File
@@ -9,9 +9,9 @@ import pytest
from numpy.typing import NDArray
from starlette.websockets import WebSocketState
from facefusion.apis.stream_helper import handle_video_stream, run_aom_encode_loop, run_opus_encode_loop, run_vp8_encode_loop
from facefusion.apis.stream_helper import encode_audio_loop, encode_video_loop, handle_video_stream
from facefusion.hash_helper import create_hash
from facefusion.types import VisionFrame
from facefusion.types import VideoCodec, VisionFrame
def _make_handler_websocket(events : list[Any]) -> MagicMock:
@@ -35,7 +35,7 @@ def _make_audio_packet(samples : NDArray[Any]) -> bytes:
@pytest.mark.parametrize('video_codec', [ 'av1', 'vp8' ])
def test_run_video_encode_loop(video_codec : str) -> None:
def test_encode_video_loop(video_codec : VideoCodec) -> None:
frame = numpy.full((64, 64, 3), 128, dtype = numpy.uint8)
small_frame = numpy.full((64, 64, 3), 128, dtype = numpy.uint8)
large_frame = numpy.full((128, 128, 3), 128, dtype = numpy.uint8)
@@ -45,13 +45,11 @@ def test_run_video_encode_loop(video_codec : str) -> None:
create_name = prefix + 'create_aom_encoder'
encode_name = prefix + 'encode_aom_buffer'
destroy_name = prefix + 'destroy_aom_encoder'
run_loop = run_aom_encode_loop
if video_codec == 'vp8':
create_name = prefix + 'create_vpx_encoder'
encode_name = prefix + 'encode_vpx_buffer'
destroy_name = prefix + 'destroy_vpx_encoder'
run_loop = run_vp8_encode_loop
vision_frame_queue : queue.Queue[Optional[VisionFrame]] = queue.Queue()
vision_frame_queue.put(frame)
@@ -62,8 +60,8 @@ def test_run_video_encode_loop(video_codec : str) -> None:
patch(destroy_name), \
patch(prefix + 'rtc_store') as mock_rtc_store, \
patch(prefix + 'rtc') as mock_rtc:
mock_rtc_store.get_rtc_peers.return_value = [ MagicMock() ]
run_loop(vision_frame_queue, 'session-1', (64, 64))
mock_rtc_store.get_peers.return_value = [ MagicMock() ]
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
mock_rtc.send_video_to_peers.assert_called_once()
vision_frame_queue = queue.Queue()
@@ -77,8 +75,8 @@ def test_run_video_encode_loop(video_codec : str) -> None:
patch(destroy_name), \
patch(prefix + 'rtc_store') as mock_rtc_store, \
patch(prefix + 'rtc') as mock_rtc:
mock_rtc_store.get_rtc_peers.return_value = [ MagicMock() ]
run_loop(vision_frame_queue, 'session-1', (64, 64))
mock_rtc_store.get_peers.return_value = [ MagicMock() ]
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
assert mock_rtc.send_video_to_peers.call_count == 3
vision_frame_queue = queue.Queue()
@@ -86,7 +84,7 @@ def test_run_video_encode_loop(video_codec : str) -> None:
with patch(create_name, return_value = MagicMock()), \
patch(destroy_name) as mock_destroy, \
patch(prefix + 'rtc') as mock_rtc:
run_loop(vision_frame_queue, 'session-1', (64, 64))
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
mock_rtc.send_video_to_peers.assert_not_called()
mock_destroy.assert_called_once()
@@ -99,7 +97,7 @@ def test_run_video_encode_loop(video_codec : str) -> None:
patch(destroy_name), \
patch(prefix + 'rtc_store'), \
patch(prefix + 'rtc') as mock_rtc:
run_loop(vision_frame_queue, 'session-1', (64, 64))
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
mock_rtc.send_video_to_peers.assert_not_called()
vision_frame_queue = queue.Queue()
@@ -111,8 +109,8 @@ def test_run_video_encode_loop(video_codec : str) -> None:
patch(destroy_name) as mock_destroy, \
patch(prefix + 'rtc_store') as mock_rtc_store, \
patch(prefix + 'rtc') as mock_rtc:
mock_rtc_store.get_rtc_peers.return_value = [ MagicMock() ]
run_loop(vision_frame_queue, 'session-1', (64, 64))
mock_rtc_store.get_peers.return_value = [ MagicMock() ]
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
assert mock_create.call_count == 2
assert mock_destroy.call_count == 2
mock_rtc.send_video_to_peers.assert_called_once()
@@ -122,12 +120,12 @@ def test_run_video_encode_loop(video_codec : str) -> None:
vision_frame_queue.put(None)
with patch(create_name, return_value = None), \
patch(prefix + 'rtc') as mock_rtc:
run_loop(vision_frame_queue, 'session-1', (64, 64))
encode_video_loop(video_codec, vision_frame_queue, 'session-1', (64, 64))
mock_rtc.send_video_to_peers.assert_not_called()
# TODO: refine test
def test_run_opus_encode_loop() -> None:
def test_encode_audio_loop() -> None:
audio_chunk = numpy.zeros(1920, dtype = numpy.float32).tobytes()
audio_chunk_queue : queue.Queue[Optional[bytes]] = queue.Queue()
@@ -138,8 +136,8 @@ def test_run_opus_encode_loop() -> None:
patch('facefusion.apis.stream_helper.destroy_opus_encoder'), \
patch('facefusion.apis.stream_helper.rtc_store') as mock_rtc_store, \
patch('facefusion.apis.stream_helper.rtc') as mock_rtc:
mock_rtc_store.get_rtc_peers.return_value = [ MagicMock() ]
run_opus_encode_loop(audio_chunk_queue, 'session-1')
mock_rtc_store.get_peers.return_value = [ MagicMock() ]
encode_audio_loop('opus', audio_chunk_queue, 'session-1')
mock_rtc.send_audio_to_peers.assert_called_once()
assert mock_rtc.send_audio_to_peers.call_args[0][2] == 0
@@ -152,8 +150,8 @@ def test_run_opus_encode_loop() -> None:
patch('facefusion.apis.stream_helper.destroy_opus_encoder'), \
patch('facefusion.apis.stream_helper.rtc_store') as mock_rtc_store, \
patch('facefusion.apis.stream_helper.rtc') as mock_rtc:
mock_rtc_store.get_rtc_peers.return_value = [ MagicMock() ]
run_opus_encode_loop(audio_chunk_queue, 'session-1')
mock_rtc_store.get_peers.return_value = [ MagicMock() ]
encode_audio_loop('opus', audio_chunk_queue, 'session-1')
assert mock_rtc.send_audio_to_peers.call_count == 2
assert mock_rtc.send_audio_to_peers.call_args_list[0][0][2] == 0
assert mock_rtc.send_audio_to_peers.call_args_list[1][0][2] == 960
@@ -166,7 +164,7 @@ def test_run_opus_encode_loop() -> None:
patch('facefusion.apis.stream_helper.destroy_opus_encoder'), \
patch('facefusion.apis.stream_helper.rtc_store'), \
patch('facefusion.apis.stream_helper.rtc') as mock_rtc:
run_opus_encode_loop(audio_chunk_queue, 'session-1')
encode_audio_loop('opus', audio_chunk_queue, 'session-1')
mock_rtc.send_audio_to_peers.assert_not_called()
audio_chunk_queue = queue.Queue()
@@ -177,7 +175,7 @@ def test_run_opus_encode_loop() -> None:
patch('facefusion.apis.stream_helper.destroy_opus_encoder') as mock_destroy, \
patch('facefusion.apis.stream_helper.rtc_store'), \
patch('facefusion.apis.stream_helper.rtc'):
run_opus_encode_loop(audio_chunk_queue, 'session-1')
encode_audio_loop('opus', audio_chunk_queue, 'session-1')
mock_destroy.assert_called_once()
audio_chunk_queue = queue.Queue()
@@ -185,7 +183,7 @@ def test_run_opus_encode_loop() -> None:
with patch('facefusion.apis.stream_helper.create_opus_encoder', return_value = MagicMock()), \
patch('facefusion.apis.stream_helper.destroy_opus_encoder') as mock_destroy, \
patch('facefusion.apis.stream_helper.rtc') as mock_rtc:
run_opus_encode_loop(audio_chunk_queue, 'session-1')
encode_audio_loop('opus', audio_chunk_queue, 'session-1')
mock_rtc.send_audio_to_peers.assert_not_called()
mock_destroy.assert_called_once()
@@ -202,8 +200,8 @@ def test_handle_video_stream() -> None:
patch('facefusion.apis.stream_helper.session_manager.find_session_id', return_value = 'session-1'), \
patch('facefusion.apis.stream_helper.session_context.set_session_id'), \
patch('facefusion.apis.stream_helper.state_manager.get_item', return_value = 30), \
patch('facefusion.apis.stream_helper.run_aom_encode_loop') as mock_loop, \
patch('facefusion.apis.stream_helper.run_opus_encode_loop'), \
patch('facefusion.apis.stream_helper.encode_video_loop') as mock_loop, \
patch('facefusion.apis.stream_helper.encode_audio_loop'), \
patch('facefusion.apis.stream_helper.rtc_store') as mock_rtc:
asyncio.run(handle_video_stream(websocket))
websocket.accept.assert_called_once_with(subprotocol = 'proto')
@@ -211,7 +209,7 @@ def test_handle_video_stream() -> None:
websocket.close.assert_called_once()
mock_rtc.init_peers.assert_called_once_with('session-1')
mock_rtc.delete_peers.assert_called_once_with('session-1')
_, loop_session_id, loop_resolution = mock_loop.call_args[0]
_, _, loop_session_id, loop_resolution = mock_loop.call_args[0]
assert loop_session_id == 'session-1'
assert loop_resolution == (64, 64)
@@ -232,9 +230,9 @@ def test_handle_video_stream() -> None:
patch('facefusion.apis.stream_helper.session_manager.find_session_id', return_value = 'session-1'), \
patch('facefusion.apis.stream_helper.session_context.set_session_id'), \
patch('facefusion.apis.stream_helper.state_manager.get_item', return_value = 30), \
patch('facefusion.apis.stream_helper.run_aom_encode_loop'), \
patch('facefusion.apis.stream_helper.run_opus_encode_loop') as mock_audio_loop, \
patch('facefusion.apis.stream_helper.encode_video_loop'), \
patch('facefusion.apis.stream_helper.encode_audio_loop') as mock_audio_loop, \
patch('facefusion.apis.stream_helper.rtc_store'):
asyncio.run(handle_video_stream(websocket))
audio_queue = mock_audio_loop.call_args[0][0]
audio_queue = mock_audio_loop.call_args[0][1]
assert create_hash(audio_queue.get_nowait()) == '6d72f0fc'