Improve encoder tests with hash assertion (#1103)

* improve test_encode_opus_buffer

* try different hash per os

* fix lint

* add windows check

* update windows hash

* fix test and lint

* update windows hash

* update CI for test_video_encoder.py

* update hash for macos

* update method to use single cpu

* update mac hash

* update windows hash

* cleanup

* restore ci.yml

* remove argument defaults

* selected CI tests

* selected CI tests

* restore ci.yml

---------

Co-authored-by: henryruhs <info@henryruhs.com>

facefusion/apis/stream_helper.py

@@ -44,7 +44,7 @@ async def receive_vision_frames(websocket : WebSocket) -> AsyncIterator[VisionFr
 
 # TODO: move to facefusion/vpx_encoder.py, throttle loop to avoid spinning on same frame
 def run_video_encode_loop(vision_frame_deque : deque[VisionFrame], session_id : SessionId, initial_resolution : Resolution, keyframe_interval : int) -> None:
-	vpx_encoder = create_vpx_encoder(initial_resolution[0], initial_resolution[1], 4500)
+	vpx_encoder = create_vpx_encoder(initial_resolution[0], initial_resolution[1], 4500, 8, 16)
 	current_resolution = initial_resolution
 	pts = 0
 
@@ -59,7 +59,7 @@ def run_video_encode_loop(vision_frame_deque : deque[VisionFrame], session_id :
 				destroy_vpx_encoder(vpx_encoder)
 
 			current_resolution = frame_resolution
-			vpx_encoder = create_vpx_encoder(current_resolution[0], current_resolution[1], 4500)
+			vpx_encoder = create_vpx_encoder(current_resolution[0], current_resolution[1], 4500, 8, 16)
 			pts = 0
 
 		if vpx_encoder:

facefusion/video_encoder.py

@@ -6,7 +6,7 @@ from facefusion.libraries import vpx as vpx_module
 from facefusion.types import BitRate, VpxEncoder
 
 
-def create_vpx_encoder(width : int, height : int, bitrate : BitRate) -> Optional[VpxEncoder]:
+def create_vpx_encoder(width : int, height : int, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[VpxEncoder]:
 	vpx_library = vpx_module.create_static_library()
 
 	if vpx_library:
@@ -16,7 +16,7 @@ def create_vpx_encoder(width : int, height : int, bitrate : BitRate) -> Optional
 		config_buffer = ctypes.create_string_buffer(4096)
 
 		if vpx_library.vpx_codec_enc_config_default(ctypes.byref(vp8_codec), config_buffer, 0) == 0:
-			struct.pack_into('I', config_buffer, 4, 8)
+			struct.pack_into('I', config_buffer, 4, thread_count)
 			struct.pack_into('I', config_buffer, 12, width)
 			struct.pack_into('I', config_buffer, 16, height)
 			struct.pack_into('I', config_buffer, 28, 1)
@@ -29,7 +29,7 @@ def create_vpx_encoder(width : int, height : int, bitrate : BitRate) -> Optional
 			struct.pack_into('I', config_buffer, 128, 50)
 
 			if vpx_library.vpx_codec_enc_init_ver(vpx_encoder, ctypes.byref(vp8_codec), config_buffer, 0, 39) == 0:
-				vpx_library.vpx_codec_control_(vpx_encoder, 13, ctypes.c_int(16))
+				vpx_library.vpx_codec_control_(vpx_encoder, 13, ctypes.c_int(cpu_count))
 				vpx_library.vpx_codec_control_(vpx_encoder, 12, ctypes.c_int(3))
 				vpx_library.vpx_codec_control_(vpx_encoder, 27, ctypes.c_int(10))
 				return vpx_encoder

tests/test_audio_encoder.py

@@ -7,8 +7,10 @@ from tests.assert_helper import get_test_example_file, get_test_examples_directo
 
 from facefusion import state_manager
 from facefusion.audio_encoder import create_opus_encoder, destroy_opus_encoder, encode_opus_buffer
+from facefusion.common_helper import is_linux, is_macos, is_windows
 from facefusion.download import conditional_download
 from facefusion.ffmpeg import read_audio_buffer
+from facefusion.hash_helper import create_hash
 from facefusion.libraries import opus as opus_module
 
 
@@ -26,15 +28,17 @@ def test_create_opus_encoder() -> None:
 	assert create_opus_encoder(0, 0) is None
 
 
-#TODO: rename to test_encode_opus_buffer
 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_buffer(opus_encoder, pcm_pointer, 960)
-	assert encode_opus_buffer(opus_encoder, pcm_pointer, 0) == b''
+	if is_linux() or is_windows():
+		assert create_hash(encode_opus_buffer(opus_encoder, pcm_pointer, 960)) == '8abe71cf'
+
+	if is_macos():
+		assert create_hash(encode_opus_buffer(opus_encoder, pcm_pointer, 960)) == '8ecd1108'
 
 
 def test_destroy_opus_encoder() -> None:

tests/test_video_encoder.py

@@ -5,7 +5,9 @@ import pytest
 from tests.assert_helper import get_test_example_file, get_test_examples_directory
 
 from facefusion import state_manager
+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 vpx as vpx_module
 from facefusion.video_encoder import create_vpx_encoder, destroy_vpx_encoder, encode_vpx_buffer
 from facefusion.vision import read_video_frame
@@ -21,24 +23,29 @@ def before_all() -> None:
 
 
 def test_create_vpx_encoder() -> None:
-	assert create_vpx_encoder(320, 240, 1000)
-	assert create_vpx_encoder(0, 0, 0) is None
+	assert create_vpx_encoder(320, 240, 1000, 8, 16)
+	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'))
 	height, width = vision_frame.shape[:2]
-	vpx_encoder = create_vpx_encoder(width, height, 1000)
+	vpx_encoder = create_vpx_encoder(width, height, 1000, 1, 0)
 
 	buffer_valid = cv2.cvtColor(vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes()
 	buffer_invalid = bytes(0)
 
-	assert encode_vpx_buffer(vpx_encoder, buffer_valid, width, height, 3, 1)
+	if is_linux() or is_windows():
+		assert create_hash(encode_vpx_buffer(vpx_encoder, buffer_valid, width, height, 3, 1)) == 'ce133a1f'
+
+	if is_macos():
+		assert create_hash(encode_vpx_buffer(vpx_encoder, buffer_valid, width, height, 3, 1)) == '21c36925'
+
 	assert encode_vpx_buffer(vpx_encoder, buffer_invalid, width, height, 0, 0) == b''
 
 
 def test_destroy_vpx_encoder() -> None:
-	vpx_encoder = create_vpx_encoder(320, 240, 1000)
+	vpx_encoder = create_vpx_encoder(320, 240, 1000, 8, 16)
 
 	with patch.object(vpx_module.create_static_library(), 'vpx_codec_destroy') as mock:
 		destroy_vpx_encoder(vpx_encoder)