diff --git a/facefusion/apis/stream_helper.py b/facefusion/apis/stream_helper.py index cce1af18..da34e38b 100644 --- a/facefusion/apis/stream_helper.py +++ b/facefusion/apis/stream_helper.py @@ -11,10 +11,10 @@ from starlette.websockets import WebSocket, WebSocketState from facefusion import rtc_store, session_context, session_manager, state_manager from facefusion.apis.api_helper import get_sec_websocket_protocol from facefusion.apis.session_helper import extract_access_token -from facefusion.audio_encoder import create_opus_encoder, destroy_opus_encoder, encode_opus +from facefusion.audio_encoder import create_opus_encoder, destroy_opus_encoder, encode_opus_buffer from facefusion.streamer import process_vision_frame from facefusion.types import Resolution, SessionId, VisionFrame -from facefusion.video_encoder import create_vpx_encoder, destroy_vpx_encoder, encode_vpx +from facefusion.video_encoder import create_vpx_encoder, destroy_vpx_encoder, encode_vpx_buffer async def receive_stream_frames(websocket : WebSocket) -> AsyncIterator[Tuple[int, bytes]]: @@ -69,7 +69,7 @@ def run_video_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : if pts % keyframe_interval == 0: vpx_flags = 1 - frame_buffer = encode_vpx(vpx_encoder, yuv_frame.tobytes(), width, height, pts, vpx_flags) + frame_buffer = encode_vpx_buffer(vpx_encoder, yuv_frame.tobytes(), width, height, pts, vpx_flags) if frame_buffer: rtc_store.send_rtc_video(session_id, frame_buffer) @@ -151,7 +151,7 @@ async def handle_video_stream(websocket : WebSocket) -> None: audio_chunk = audio_temp[:1920] audio_temp = audio_temp[1920:] pcm_pointer = audio_chunk.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) - audio_buffer = encode_opus(opus_encoder, pcm_pointer, 960) + audio_buffer = encode_opus_buffer(opus_encoder, pcm_pointer, 960) if audio_buffer: rtc_store.send_rtc_audio(session_id, audio_buffer, audio_timestamp) diff --git a/facefusion/audio_encoder.py b/facefusion/audio_encoder.py index d6700c31..a69a6988 100644 --- a/facefusion/audio_encoder.py +++ b/facefusion/audio_encoder.py @@ -14,8 +14,8 @@ def create_opus_encoder(sample_rate : int, channel_total : int) -> Optional[Opus return None -# TODO this method needs refinement - rename to encode_opus_buffer -def encode_opus(opus_encoder : OpusEncoder, pcm_pointer : ctypes.c_void_p, frame_size : int) -> bytes: +# TODO this method needs refinement +def encode_opus_buffer(opus_encoder : OpusEncoder, pcm_pointer : ctypes.c_void_p, frame_size : int) -> bytes: opus_library = opus_module.create_static_library() audio_buffer = b'' diff --git a/facefusion/video_encoder.py b/facefusion/video_encoder.py index 1d4fa9ce..9a34f22f 100644 --- a/facefusion/video_encoder.py +++ b/facefusion/video_encoder.py @@ -37,8 +37,8 @@ def create_vpx_encoder(width : int, height : int, bitrate : BitRate) -> Optional return None -# TODO this method needs refinement - rename to encode_vpx_buffer -def encode_vpx(vpx_encoder : VpxEncoder, yuv_buffer : bytes, width : int, height : int, presentation_timestamp : int, flags : int) -> bytes: +# TODO this method needs refinement +def encode_vpx_buffer(vpx_encoder : VpxEncoder, yuv_buffer : bytes, width : int, height : int, presentation_timestamp : int, flags : int) -> bytes: vpx_library = vpx_module.create_static_library() frame_buffer = b'' diff --git a/tests/test_audio_encoder.py b/tests/test_audio_encoder.py index d294d33c..a1ceb538 100644 --- a/tests/test_audio_encoder.py +++ b/tests/test_audio_encoder.py @@ -6,7 +6,7 @@ import pytest from tests.assert_helper import get_test_example_file, get_test_examples_directory from facefusion import state_manager -from facefusion.audio_encoder import create_opus_encoder, destroy_opus_encoder, encode_opus +from facefusion.audio_encoder import create_opus_encoder, destroy_opus_encoder, encode_opus_buffer from facefusion.download import conditional_download from facefusion.ffmpeg import read_audio_buffer from facefusion.libraries import opus as opus_module @@ -27,14 +27,14 @@ def test_create_opus_encoder() -> None: #TODO: rename to test_encode_opus_buffer -def test_encode_opus() -> None: +def test_encode_opus_buffer() -> None: audio_buffer = read_audio_buffer(get_test_example_file('source.mp3'), 48000, 16, 2) pcm_samples = numpy.frombuffer(audio_buffer, dtype = numpy.int16).astype(numpy.float32) / 32768.0 pcm_pointer = pcm_samples[:1920].ctypes.data_as(ctypes.POINTER(ctypes.c_float)) opus_encoder = create_opus_encoder(48000, 2) - assert encode_opus(opus_encoder, pcm_pointer, 960) - assert encode_opus(opus_encoder, pcm_pointer, 0) == b'' + assert encode_opus_buffer(opus_encoder, pcm_pointer, 960) + assert encode_opus_buffer(opus_encoder, pcm_pointer, 0) == b'' def test_destroy_opus_encoder() -> None: diff --git a/tests/test_video_encoder.py b/tests/test_video_encoder.py index da3a3e66..a8e88440 100644 --- a/tests/test_video_encoder.py +++ b/tests/test_video_encoder.py @@ -7,7 +7,7 @@ from tests.assert_helper import get_test_example_file, get_test_examples_directo from facefusion import state_manager from facefusion.download import conditional_download from facefusion.libraries import vpx as vpx_module -from facefusion.video_encoder import create_vpx_encoder, destroy_vpx_encoder, encode_vpx +from facefusion.video_encoder import create_vpx_encoder, destroy_vpx_encoder, encode_vpx_buffer from facefusion.vision import read_video_frame @@ -25,8 +25,7 @@ def test_create_vpx_encoder() -> None: assert create_vpx_encoder(0, 0, 0) is None -# TODO: rename to test_encode_vpx_buffer -def test_encode_vpx() -> None: +def test_encode_vpx_buffer() -> None: vision_frame = read_video_frame(get_test_example_file('target-240p.mp4')) height, width = vision_frame.shape[:2] vpx_encoder = create_vpx_encoder(width, height, 1000) @@ -34,8 +33,8 @@ def test_encode_vpx() -> None: buffer_valid = cv2.cvtColor(vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes() buffer_invalid = bytes(0) - assert encode_vpx(vpx_encoder, buffer_valid, width, height, 3, 1) - assert encode_vpx(vpx_encoder, buffer_invalid, width, height, 0, 0) == b'' + assert encode_vpx_buffer(vpx_encoder, buffer_valid, width, height, 3, 1) + assert encode_vpx_buffer(vpx_encoder, buffer_invalid, width, height, 0, 0) == b'' def test_destroy_vpx_encoder() -> None: