From f79770aa13e3a824988e6fc88e6669dee4ad9d27 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 10 Jun 2025 20:03:27 +0200 Subject: [PATCH] Polish content analyser --- tests/test_inference_manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_inference_manager.py b/tests/test_inference_manager.py index 28bf43f..2b4d61e 100644 --- a/tests/test_inference_manager.py +++ b/tests/test_inference_manager.py @@ -16,17 +16,17 @@ def before_all() -> None: def test_get_inference_pool() -> None: - model_names = [ 'yolo_11m', 'marqo', 'freepik' ] + model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ] _, model_source_set = content_analyser.collect_model_downloads() with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m'), InferenceSession) + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1'), InferenceSession) with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m'), InferenceSession) + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1'), InferenceSession) - assert INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m') == INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m') + assert INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1') == INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.nsfw_1.nsfw_2.nsfw_3.0.cpu').get('nsfw_1')