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

@@ -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')

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)

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)

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:

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)

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)
 

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
 

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)

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)

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):

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'

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: