From d96580b34c0c498d8dde21591c27f4ca12c625f9 Mon Sep 17 00:00:00 2001 From: Henry Ruhs Date: Fri, 11 Sep 2026 21:17:22 +0200 Subject: [PATCH] convert content store (#1236) * convert content store * add local_id variables * add local_id variables * fix misconcept of clear() in stores, pass session id to threads --- facefusion/apis/endpoints/session.py | 10 +++-- facefusion/apis/stream_manager.py | 7 ++-- facefusion/content_store.py | 39 +++++++++++++------ facefusion/core.py | 7 +++- facefusion/process_manager.py | 43 +++++++++++---------- tests/assert_helper.py | 3 +- tests/test_api_stream_manager.py | 14 ++++--- tests/test_cli_job_manager.py | 3 +- tests/test_cli_job_runner.py | 3 +- tests/test_content_store.py | 24 ++++++++++-- tests/test_process_manager.py | 57 ++++++++++++++-------------- tests/test_session_context.py | 3 +- 12 files changed, 134 insertions(+), 79 deletions(-) diff --git a/facefusion/apis/endpoints/session.py b/facefusion/apis/endpoints/session.py index 8ab5dc27..7ee4ff2d 100644 --- a/facefusion/apis/endpoints/session.py +++ b/facefusion/apis/endpoints/session.py @@ -4,7 +4,7 @@ from starlette.requests import Request from starlette.responses import JSONResponse from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_401_UNAUTHORIZED, HTTP_404_NOT_FOUND -from facefusion import process_manager, session_context, session_manager, state_manager, translator +from facefusion import content_store, process_manager, session_context, session_manager, state_manager, translator from facefusion.apis import asset_store from facefusion.apis.session_helper import extract_access_token, validate_api_key from facefusion.filesystem import is_directory, remove_directory @@ -18,7 +18,9 @@ async def create_session(request : Request) -> JSONResponse: session = session_manager.create_session() session_context.set_session_id(session_id) session_manager.set_session(session_id, session) - process_manager.init_process_state() + + content_store.init() + process_manager.init() return JSONResponse( { @@ -81,9 +83,11 @@ async def destroy_session(request : Request) -> JSONResponse: }, status_code = HTTP_404_NOT_FOUND) asset_store.delete_assets(session_id) - process_manager.clear_process_state() session_manager.clear_session(session_id) + content_store.clear() + process_manager.clear() + return JSONResponse( { 'message': translator.get('ok', 'facefusion.apis') diff --git a/facefusion/apis/stream_manager.py b/facefusion/apis/stream_manager.py index 3a2157fc..3ac9101a 100644 --- a/facefusion/apis/stream_manager.py +++ b/facefusion/apis/stream_manager.py @@ -7,7 +7,7 @@ from typing import Optional, Tuple from starlette.websockets import WebSocket -from facefusion import content_store, rtc, rtc_store, state_manager, streamer +from facefusion import content_store, rtc, rtc_store, session_context, state_manager, streamer from facefusion.apis.stream_audio import receive_audio_frames, run_audio_encode_loop from facefusion.apis.stream_video import receive_video_frames, run_video_encode_loop from facefusion.content_analyser import analyse_frame @@ -91,8 +91,8 @@ def process_video(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpA codec = audio_codec ) - rtc_store.set_peer(session_id, rtc_peer) content_store.clear() + rtc_store.set_peer(session_id, rtc_peer) threading.Thread(target = run_peer_loop, args = (session_id, rtc_peer), daemon = True).start() @@ -104,10 +104,11 @@ def process_video(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpA def run_peer_loop(session_id : SessionId, rtc_peer : RtcPeer) -> None: + session_context.set_session_id(session_id) execution_thread_count = state_manager.get_item('execution_thread_count') video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = execution_thread_count) audio_queue : Queue[Tuple[Time, AudioFrame]] = Queue(maxsize = execution_thread_count * 10) - video_executor = ThreadPoolExecutor(max_workers = execution_thread_count) + video_executor = ThreadPoolExecutor(max_workers = execution_thread_count, initializer = session_context.set_session_id, initargs = (session_id,)) video_receiver_thread = threading.Thread(target = receive_video_frames, args = (rtc_peer.get('video'), video_queue, video_executor), daemon = True) video_encoder_thread = threading.Thread(target = run_video_encode_loop, args = (rtc_peer, video_queue), daemon = True) diff --git a/facefusion/content_store.py b/facefusion/content_store.py index c0b5b2f9..297e8470 100644 --- a/facefusion/content_store.py +++ b/facefusion/content_store.py @@ -1,32 +1,49 @@ -from facefusion.types import ContentSet +from facefusion import store_creator +from facefusion.session_context import get_session_id +from facefusion.types import Store -CONTENT_STORE : ContentSet =\ +CONTENT_STORE : Store = store_creator.create_store( { 'hit': 0, 'total': 0 -} +}) + + +def init() -> None: + session_id = get_session_id() + store_creator.init_content(CONTENT_STORE, session_id) def tick(step : int = 30) -> bool: - CONTENT_STORE['total'] += 1 + session_id = get_session_id() + content_set = store_creator.get_content(CONTENT_STORE, session_id) + content_set['total'] += 1 - return CONTENT_STORE.get('total') % step == 0 + return content_set.get('total') % step == 0 def get_hit() -> int: - return CONTENT_STORE.get('hit') + session_id = get_session_id() + content_set = store_creator.get_content(CONTENT_STORE, session_id) + + return content_set.get('hit') def set_hit() -> None: - CONTENT_STORE['hit'] += 1 + session_id = get_session_id() + content_set = store_creator.get_content(CONTENT_STORE, session_id) + content_set['hit'] += 1 def calculate_rate(step : int = 30) -> float: - if CONTENT_STORE.get('hit') and CONTENT_STORE.get('total'): - return CONTENT_STORE.get('hit') / CONTENT_STORE.get('total') * step * 100 + session_id = get_session_id() + content_set = store_creator.get_content(CONTENT_STORE, session_id) + + if content_set.get('hit') and content_set.get('total'): + return content_set.get('hit') / content_set.get('total') * step * 100 return 0.0 def clear() -> None: - CONTENT_STORE['hit'] = 0 - CONTENT_STORE['total'] = 0 + session_id = get_session_id() + store_creator.init_content(CONTENT_STORE, session_id) diff --git a/facefusion/core.py b/facefusion/core.py index 09249e27..329ed5fd 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -35,7 +35,10 @@ def cli() -> None: if state_manager.get_item('command'): logger.init(state_manager.get_item('log_level')) - process_manager.init_process_state() + + content_store.init() + process_manager.init() + route(args) else: program.print_help() @@ -105,7 +108,7 @@ def pre_check() -> bool: def common_pre_check() -> bool: content_analyser_content = inspect.getsource(content_analyser).encode() content_store_content = inspect.getsource(content_store).encode() - return hash_helper.create_hash(content_analyser_content) == 'c5d70a17' and hash_helper.create_hash(content_store_content) == '9cd9f029' + return hash_helper.create_hash(content_analyser_content) == 'c5d70a17' and hash_helper.create_hash(content_store_content) == '83656ea4' def processors_pre_check() -> bool: diff --git a/facefusion/process_manager.py b/facefusion/process_manager.py index ce897fad..0daac344 100644 --- a/facefusion/process_manager.py +++ b/facefusion/process_manager.py @@ -5,49 +5,54 @@ from facefusion.types import ProcessState, Store PROCESS_STORE : Store = store_creator.create_store('pending') -def init_process_state() -> None: - store_creator.init_content(PROCESS_STORE, get_session_id()) +def init() -> None: + session_id = get_session_id() + store_creator.init_content(PROCESS_STORE, session_id) -def get_process_state() -> ProcessState: - return store_creator.get_content(PROCESS_STORE, get_session_id()) +def get_state() -> ProcessState: + session_id = get_session_id() - -def set_process_state(process_state : ProcessState) -> None: - store_creator.set_content(PROCESS_STORE, get_session_id(), process_state) - - -def clear_process_state() -> None: - store_creator.delete_content(PROCESS_STORE, get_session_id()) + return store_creator.get_content(PROCESS_STORE, session_id) def is_checking() -> bool: - return get_process_state() == 'checking' + return get_state() == 'checking' def is_processing() -> bool: - return get_process_state() == 'processing' + return get_state() == 'processing' def is_stopping() -> bool: - return get_process_state() == 'stopping' + return get_state() == 'stopping' def is_pending() -> bool: - return get_process_state() == 'pending' + return get_state() == 'pending' + + +def set_state(process_state : ProcessState) -> None: + session_id = get_session_id() + store_creator.set_content(PROCESS_STORE, session_id, process_state) def check() -> None: - set_process_state('checking') + set_state('checking') def start() -> None: - set_process_state('processing') + set_state('processing') def stop() -> None: - set_process_state('stopping') + set_state('stopping') def end() -> None: - set_process_state('pending') + set_state('pending') + + +def clear() -> None: + session_id = get_session_id() + store_creator.init_content(PROCESS_STORE, session_id) diff --git a/tests/assert_helper.py b/tests/assert_helper.py index 64588e6f..5f925059 100644 --- a/tests/assert_helper.py +++ b/tests/assert_helper.py @@ -11,7 +11,8 @@ def is_test_job_file(file_path : str, job_status : JobStatus) -> bool: def get_test_job_file(file_path : str, job_status : JobStatus) -> str: - jobs_path = os.path.join(get_test_jobs_directory(), resolve_local_id()) + local_id = resolve_local_id() + jobs_path = os.path.join(get_test_jobs_directory(), local_id) return os.path.join(jobs_path, job_status, file_path) diff --git a/tests/test_api_stream_manager.py b/tests/test_api_stream_manager.py index e90be621..cb3afebb 100644 --- a/tests/test_api_stream_manager.py +++ b/tests/test_api_stream_manager.py @@ -10,6 +10,7 @@ from facefusion.common_helper import is_linux, is_windows from facefusion.download import conditional_download from facefusion.hash_helper import create_hash from facefusion.libraries import datachannel as datachannel_module +from facefusion.session_context import set_session_id from facefusion.types import RtcPeer, SessionId, VideoCodec from .assert_helper import get_test_example_file, get_test_examples_directory @@ -150,11 +151,14 @@ def test_run_peer_loop(video_codec : VideoCodec, payload_type : int, session_id assert rtc_store.has_peer(session_id) is True - with patch('facefusion.apis.stream_manager.receive_video_frames'): - with patch('facefusion.apis.stream_manager.run_video_encode_loop'): - thread = threading.Thread(target = run_peer_loop, args = (session_id, rtc_peer), daemon = True) - thread.start() - thread.join(timeout = 5.0) + with patch('facefusion.apis.stream_manager.ThreadPoolExecutor') as thread_pool_executor_mock: + with patch('facefusion.apis.stream_manager.receive_video_frames'): + with patch('facefusion.apis.stream_manager.run_video_encode_loop'): + thread = threading.Thread(target = run_peer_loop, args = (session_id, rtc_peer), daemon = True) + thread.start() + thread.join(timeout = 5.0) + + thread_pool_executor_mock.assert_called_once_with(max_workers = 8, initializer = set_session_id, initargs = (session_id,)) assert rtc_store.has_peer(session_id) is False diff --git a/tests/test_cli_job_manager.py b/tests/test_cli_job_manager.py index 91a0d1e0..a52ab3d3 100644 --- a/tests/test_cli_job_manager.py +++ b/tests/test_cli_job_manager.py @@ -34,7 +34,8 @@ def before_all() -> None: @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - jobs_path = os.path.join(get_test_jobs_directory(), resolve_local_id()) + local_id = resolve_local_id() + jobs_path = os.path.join(get_test_jobs_directory(), local_id) clear_jobs(get_test_jobs_directory()) init_jobs(jobs_path) diff --git a/tests/test_cli_job_runner.py b/tests/test_cli_job_runner.py index 9bafc311..4d422cd0 100644 --- a/tests/test_cli_job_runner.py +++ b/tests/test_cli_job_runner.py @@ -34,7 +34,8 @@ def before_all() -> None: @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - jobs_path = os.path.join(get_test_jobs_directory(), resolve_local_id()) + local_id = resolve_local_id() + jobs_path = os.path.join(get_test_jobs_directory(), local_id) clear_jobs(get_test_jobs_directory()) init_jobs(jobs_path) diff --git a/tests/test_content_store.py b/tests/test_content_store.py index 6ce0f7af..e4b8407f 100644 --- a/tests/test_content_store.py +++ b/tests/test_content_store.py @@ -1,11 +1,29 @@ import pytest -from facefusion.content_store import calculate_rate, clear, get_hit, set_hit, tick +from facefusion import session_context +from facefusion.content_store import calculate_rate, clear, get_hit, init, set_hit, tick @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - clear() + local_id = session_context.resolve_local_id() + + session_context.set_session_id(local_id) + init() + + +def test_init() -> None: + local_id = session_context.resolve_local_id() + + session_context.set_session_id('session-a') + init() + set_hit() + + assert get_hit() == 1 + + session_context.set_session_id(local_id) + + assert get_hit() == 0 def test_get_hit() -> None: @@ -23,7 +41,7 @@ def test_set_hit() -> None: assert get_hit() == 2 -def test_get_rate() -> None: +def test_calculate_rate() -> None: assert calculate_rate() == 0.0 for _ in range(100): diff --git a/tests/test_process_manager.py b/tests/test_process_manager.py index b013487a..c3ae86d3 100644 --- a/tests/test_process_manager.py +++ b/tests/test_process_manager.py @@ -1,61 +1,60 @@ import pytest -from facefusion.process_manager import clear_process_state, end, get_process_state, init_process_state, is_pending, is_processing, is_stopping, set_process_state, start, stop -from facefusion.session_context import set_session_id +from facefusion import session_context +from facefusion.process_manager import clear, end, get_state, init, is_pending, is_processing, is_stopping, set_state, start, stop @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - set_session_id('session-a') - clear_process_state() - set_session_id('session-b') - clear_process_state() - set_session_id('session-a') + session_context.set_session_id('session-a') + clear() + session_context.set_session_id('session-b') + clear() + session_context.set_session_id('session-a') -def test_init_process_state() -> None: - assert get_process_state() is None +def test_init() -> None: + set_state('processing') + init() - init_process_state() - - assert get_process_state() == 'pending' + assert get_state() == 'pending' -def test_get_process_state() -> None: - set_process_state('processing') - set_session_id('session-b') - set_process_state('stopping') +def test_get_state() -> None: + set_state('processing') + session_context.set_session_id('session-b') + set_state('stopping') - assert get_process_state() == 'stopping' + assert get_state() == 'stopping' - set_session_id('session-a') + session_context.set_session_id('session-a') - assert get_process_state() == 'processing' - - -def test_clear_process_state() -> None: - set_process_state('processing') - clear_process_state() - - assert get_process_state() is None + assert get_state() == 'processing' def test_start() -> None: - set_process_state('pending') + set_state('pending') start() assert is_processing() def test_stop() -> None: - set_process_state('processing') + set_state('processing') stop() assert is_stopping() def test_end() -> None: - set_process_state('processing') + set_state('processing') end() assert is_pending() + + +def test_clear() -> None: + set_state('processing') + clear() + + assert get_state() == 'pending' diff --git a/tests/test_session_context.py b/tests/test_session_context.py index ae3877d5..ba4ee854 100644 --- a/tests/test_session_context.py +++ b/tests/test_session_context.py @@ -5,7 +5,8 @@ from facefusion.session_context import get_session_id, resolve_local_id, set_ses @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - set_session_id(resolve_local_id()) + local_id = resolve_local_id() + set_session_id(local_id) def test_get_session_id() -> None: