Fix broken inference pools part2

This commit is contained in:
henryruhs
2025-02-09 01:43:58 +01:00
parent b438b899e8
commit 1bdc02014c
18 changed files with 92 additions and 65 deletions

View File

@@ -16,16 +16,17 @@ def before_all() -> None:
def test_get_inference_pool() -> None:
model_names = [ 'yolo_nsfw' ]
model_sources = content_analyser.get_model_options().get('sources')
with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'):
get_inference_pool('test', model_sources)
get_inference_pool('facefusion.content_analyser', model_names, model_sources)
assert isinstance(INFERENCE_POOLS.get('cli').get('test.content_analyser.0.cpu').get('content_analyser'), InferenceSession)
assert isinstance(INFERENCE_POOLS.get('cli').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser'), InferenceSession)
with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'):
get_inference_pool('test', model_sources)
get_inference_pool('facefusion.content_analyser', model_names, model_sources)
assert isinstance(INFERENCE_POOLS.get('ui').get('test.content_analyser.0.cpu').get('content_analyser'), InferenceSession)
assert isinstance(INFERENCE_POOLS.get('ui').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser'), InferenceSession)
assert INFERENCE_POOLS.get('cli').get('test.content_analyser.0.cpu').get('content_analyser') == INFERENCE_POOLS.get('ui').get('test.content_analyser.0.cpu').get('content_analyser')
assert INFERENCE_POOLS.get('cli').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser') == INFERENCE_POOLS.get('ui').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser')