convert video manager (#1240)

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 content_store, inference_manager, process_manager, session_context, session_manager, state_manager, translator
+from facefusion import content_store, inference_manager, process_manager, session_context, session_manager, state_manager, translator, video_manager
 from facefusion.apis import asset_store
 from facefusion.apis.session_helper import validate_api_key
 from facefusion.apis.stream_manager import destroy_stream
@@ -23,6 +23,7 @@ async def create_session(request : Request) -> JSONResponse:
 		state_manager.init()
 		content_store.init()
 		inference_manager.init()
+		video_manager.init()
 		process_manager.init()
 
 		return JSONResponse(
@@ -88,6 +89,7 @@ async def destroy_session(request : Request) -> JSONResponse:
 	state_manager.clear()
 	content_store.clear()
 	inference_manager.clear()
+	video_manager.clear()
 	process_manager.clear()
 
 	return JSONResponse(

facefusion/core.py

@@ -8,7 +8,7 @@ from time import time
 import uvicorn
 
 import facefusion.apis.core
-from facefusion import args_helper, benchmarker, cli_helper, content_analyser, content_store, hash_helper, inference_manager, logger, process_manager, session_manager, state_manager, translator
+from facefusion import args_helper, benchmarker, cli_helper, content_analyser, content_store, hash_helper, inference_manager, logger, process_manager, session_manager, state_manager, translator, video_manager
 from facefusion.args_helper import apply_args
 from facefusion.download import conditional_download_hashes, conditional_download_sources
 from facefusion.exit_helper import hard_exit, signal_exit
@@ -41,6 +41,7 @@ def cli() -> None:
 				content_store.init()
 				inference_manager.init()
 				process_manager.init()
+				video_manager.init()
 
 				route(args)
 			else:
@@ -312,6 +313,8 @@ def process_step(job_id : str, step_index : int, step_args : Args) -> bool:
 	session_manager.fork_session()
 	state_manager.clone_state()
 
+	video_manager.init()
+
 	cli_args = args_helper.extract_cli_args(state_manager.get_state())
 	args = cli_args.copy()
 	args.update(step_args)

facefusion/processors/modules/age_modifier/core.py

@@ -181,7 +181,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/background_remover/core.py

@@ -559,7 +559,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/deep_swapper/core.py

@@ -335,7 +335,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/expression_restorer/core.py

@@ -168,7 +168,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/face_debugger/core.py

@@ -78,7 +78,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') == 'strict':
 		for common_module in get_common_modules():

facefusion/processors/modules/face_editor/core.py

@@ -303,7 +303,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/face_enhancer/core.py

@@ -359,7 +359,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/face_swapper/core.py

@@ -641,7 +641,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		get_static_model_initializer.cache_clear()

facefusion/processors/modules/frame_colorizer/core.py

@@ -251,7 +251,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/frame_enhancer/core.py

@@ -637,7 +637,7 @@ def pre_process(mode : ProcessMode) -> bool:
 def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/processors/modules/lip_syncer/core.py

@@ -187,7 +187,7 @@ def post_process() -> None:
 	read_static_image.cache_clear()
 	read_static_video_frame.cache_clear()
 	read_static_voice.cache_clear()
-	video_manager.clear_video_pool()
+	video_manager.clear()
 
 	if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]:
 		clear_inference_pool()

facefusion/video_manager.py

@@ -5,24 +5,32 @@ from typing import Optional, cast
 
 import numpy
 
-from facefusion import ffmpeg, ffprobe, frame_store, vision
+from facefusion import ffmpeg, ffprobe, frame_store, store_creator, vision
 from facefusion.common_helper import get_first, get_last
-from facefusion.types import Fps, Resolution, VideoPoolSet, VideoReader, VideoWriter, VisionFrame, VisionFrameSet
+from facefusion.session_context import get_session_id
+from facefusion.types import Fps, Resolution, Store, VideoReader, VideoWriter, VisionFrame, VisionFrameSet
 
-VIDEO_POOL_SET : VideoPoolSet =\
+VIDEO_POOL_STORE : Store = store_creator.create_store(
 {
 	'reader': {},
 	'writer': {}
-}
+})
+
+
+def init() -> None:
+	session_id = get_session_id()
+	store_creator.init_content(VIDEO_POOL_STORE, session_id)
 
 
 def get_reader(video_path : str, context : str) -> VideoReader:
-	reader_id = hashlib.sha1((video_path + '_' + context).encode()).hexdigest()
+	session_id = get_session_id()
+	reader_id = hashlib.sha1((video_path + '_' + context + '_' + session_id).encode()).hexdigest()
+	video_pool = store_creator.get_content(VIDEO_POOL_STORE, session_id)
 
-	if reader_id not in VIDEO_POOL_SET.get('reader'):
+	if reader_id not in video_pool.get('reader'):
 		video_metadata = ffprobe.extract_static_video_metadata(video_path)
 
-		VIDEO_POOL_SET['reader'][reader_id] =\
+		video_pool['reader'][reader_id] =\
 		{
 			'id': reader_id,
 			'file_path': video_path,
@@ -31,7 +39,7 @@ def get_reader(video_path : str, context : str) -> VideoReader:
 			'frame_index': 0
 		}
 
-	return VIDEO_POOL_SET.get('reader').get(reader_id)
+	return video_pool.get('reader').get(reader_id)
 
 
 def conditional_seek_video_reader(video_reader : VideoReader, frame_index : int = 0) -> None:
@@ -116,8 +124,11 @@ def close_video_reader(video_reader : VideoReader) -> None:
 
 
 def get_writer(video_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> VideoWriter:
-	if video_path not in VIDEO_POOL_SET.get('writer'):
+	session_id = get_session_id()
-		VIDEO_POOL_SET['writer'][video_path] =\
+	video_pool = store_creator.get_content(VIDEO_POOL_STORE, session_id)
+
+	if video_path not in video_pool.get('writer'):
+		video_pool['writer'][video_path] =\
 		{
 			'id': uuid.uuid4().hex,
 			'file_path': video_path,
@@ -129,7 +140,7 @@ def get_writer(video_path : str, temp_video_fps : Fps, temp_video_resolution : R
 			}
 		}
 
-	return VIDEO_POOL_SET.get('writer').get(video_path)
+	return video_pool.get('writer').get(video_path)
 
 
 def write_video_frame(video_writer : VideoWriter, vision_frame : VisionFrame) -> None:
@@ -143,13 +154,15 @@ def close_video_writer(video_writer : VideoWriter) -> bool:
 	return video_writer.get('process').returncode == 0
 
 
-def clear_video_pool() -> None:
+def clear() -> None:
-	for video_reader in VIDEO_POOL_SET.get('reader').values():
+	session_id = get_session_id()
+	video_pool = store_creator.get_content(VIDEO_POOL_STORE, session_id)
+
+	for video_reader in video_pool.get('reader').values():
 		close_video_reader(video_reader)
 		frame_store.clear_frames(video_reader.get('id'))
 
-	for video_writer in VIDEO_POOL_SET.get('writer').values():
+	for video_writer in video_pool.get('writer').values():
 		close_video_writer(video_writer)
 
-	VIDEO_POOL_SET['reader'].clear()
+	store_creator.init_content(VIDEO_POOL_STORE, session_id)
-	VIDEO_POOL_SET['writer'].clear()

facefusion/workflows/to_video.py

@@ -177,7 +177,7 @@ def finalize_video(start_time : float) -> ErrorCode:
 
 def conditional_clear_video_pool() -> None:
 	if state_manager.get_item('workflow_mode') == 'image-to-video':
-		video_manager.clear_video_pool()
+		video_manager.clear()
 
 
 def conditional_restrict_video_fps() -> Fps:

tests/test_codec_aom_decoder.py

@@ -4,7 +4,7 @@ import cv2
 import pytest
 from tests.assert_helper import get_test_example_file, get_test_examples_directory
 
-from facefusion import state_manager
+from facefusion import state_manager, video_manager
 from facefusion.codecs.aom_decoder import create, decode, destroy
 from facefusion.codecs.aom_encoder import create as create_encoder, encode
 from facefusion.common_helper import is_linux, is_macos, is_windows
@@ -19,6 +19,8 @@ def before_all() -> None:
 	state_manager.init()
 	state_manager.init_item('download_providers', [ 'github', 'huggingface' ])
 
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(), [ 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ])
 
 	aom_module.pre_check()

tests/test_codec_aom_encoder.py

@@ -4,7 +4,7 @@ import cv2
 import pytest
 from tests.assert_helper import get_test_example_file, get_test_examples_directory
 
-from facefusion import state_manager
+from facefusion import state_manager, video_manager
 from facefusion.codecs.aom_encoder import create, destroy, encode
 from facefusion.common_helper import is_linux, is_macos, is_windows
 from facefusion.download import conditional_download
@@ -18,6 +18,8 @@ def before_all() -> None:
 	state_manager.init()
 	state_manager.init_item('download_providers', [ 'github', 'huggingface' ])
 
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(), [ 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ])
 
 	aom_module.pre_check()

tests/test_codec_vpx_decoder.py

@@ -4,7 +4,7 @@ import cv2
 import pytest
 from tests.assert_helper import get_test_example_file, get_test_examples_directory
 
-from facefusion import state_manager
+from facefusion import state_manager, video_manager
 from facefusion.codecs.vpx_decoder import create, decode, destroy
 from facefusion.codecs.vpx_encoder import create as create_encoder, encode
 from facefusion.common_helper import is_linux, is_macos, is_windows
@@ -20,6 +20,8 @@ def before_all() -> None:
 	state_manager.init()
 	state_manager.init_item('download_providers', [ 'github', 'huggingface' ])
 
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(), [ 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ])
 
 	vpx_module.pre_check()

tests/test_codec_vpx_encoder.py

@@ -4,7 +4,7 @@ import cv2
 import pytest
 from tests.assert_helper import get_test_example_file, get_test_examples_directory
 
-from facefusion import state_manager
+from facefusion import state_manager, video_manager
 from facefusion.codecs.vpx_encoder import create, destroy, encode
 from facefusion.common_helper import is_linux, is_macos, is_windows
 from facefusion.download import conditional_download
@@ -19,6 +19,8 @@ def before_all() -> None:
 	state_manager.init()
 	state_manager.init_item('download_providers', [ 'github', 'huggingface' ])
 
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(), [ 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ])
 
 	vpx_module.pre_check()

tests/test_face_tracker.py

@@ -1,7 +1,7 @@
 import numpy
 import pytest
 
-from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, inference_manager, state_manager
+from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, inference_manager, state_manager, video_manager
 from facefusion.common_helper import get_first, get_last
 from facefusion.download import conditional_download
 from facefusion.face_creator import get_many_faces, get_one_face
@@ -16,6 +16,8 @@ def before_all() -> None:
 	state_manager.init()
 	inference_manager.init()
 
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(),
 	[
 		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'

tests/test_frame_store.py

@@ -1,6 +1,6 @@
 import pytest
 
-from facefusion import process_manager
+from facefusion import process_manager, video_manager
 from facefusion.download import conditional_download
 from facefusion.frame_store import clear_frames, get_frame_store, reduce_frames, select_frame_set, set_frame
 from facefusion.vision import read_video_frame
@@ -10,6 +10,9 @@ from .assert_helper import get_test_example_file, get_test_examples_directory
 @pytest.fixture(scope = 'module', autouse = True)
 def before_all() -> None:
 	process_manager.start()
+
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(),
 	[
 		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'

tests/test_video_manager.py

@@ -3,13 +3,13 @@ import tempfile
 import numpy
 import pytest
 
-from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager
+from facefusion import ffmpeg, ffmpeg_builder, process_manager, session_context, state_manager
 from facefusion.common_helper import is_linux, is_macos, is_windows
 from facefusion.download import conditional_download
 from facefusion.ffprobe import extract_video_metadata
 from facefusion.frame_store import get_frame_store
 from facefusion.temp_helper import create_temp_directory, get_temp_file_path
-from facefusion.video_manager import clear_video_pool, close_video_reader, close_video_writer, collect_video_frames, conditional_seek_video_reader, drain_video_reader, get_reader, get_writer, read_video_frame, read_video_frames, seek_video_reader, write_video_frame
+from facefusion.video_manager import clear, close_video_reader, close_video_writer, collect_video_frames, conditional_seek_video_reader, drain_video_reader, get_reader, get_writer, init, read_video_frame, read_video_frames, seek_video_reader, write_video_frame
 from .assert_helper import get_test_example_file, get_test_examples_directory
 
 
@@ -18,6 +18,9 @@ def before_all() -> None:
 	state_manager.init()
 
 	process_manager.start()
+
+	init()
+
 	conditional_download(get_test_examples_directory(),
 	[
 		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'
@@ -42,7 +45,23 @@ def before_all() -> None:
 
 @pytest.fixture(scope = 'function', autouse = True)
 def before_each() -> None:
-	clear_video_pool()
+	clear()
+
+
+def test_init() -> None:
+	local_id = session_context.resolve_local_id()
+	video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame')
+
+	session_context.set_session_id('session-a')
+	state_manager.init()
+	init()
+
+	assert not get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame').get('id') == video_reader.get('id')
+
+	clear()
+	session_context.set_session_id(local_id)
+
+	assert get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') is video_reader
 
 
 def test_get_reader() -> None:
@@ -186,7 +205,7 @@ def test_close_video_writer() -> None:
 	assert close_video_writer(video_writer) is True
 
 
-def test_clear_video_pool() -> None:
+def test_clear() -> None:
 	target_path = get_test_example_file('target-240p-25fps.mp4')
 	create_temp_directory(state_manager.get_temp_path(), target_path)
 	video_reader = get_reader(target_path, 'select_video_frames')
@@ -194,7 +213,7 @@ def test_clear_video_pool() -> None:
 
 	read_video_frames(video_reader, 0, 4)
 	write_video_frame(video_writer, read_video_frame(video_reader))
-	clear_video_pool()
+	clear()
 
 	if is_windows():
 		assert video_reader.get('process').returncode == 1

tests/test_vision.py

@@ -2,7 +2,7 @@
 import numpy
 import pytest
 
-from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager
+from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager, video_manager
 from facefusion.download import conditional_download
 from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_frame, resolve_extract_frame_index, resolve_target_frame_index, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, unpack_resolution, write_image
 from .assert_helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory
@@ -13,6 +13,9 @@ def before_all() -> None:
 	state_manager.init()
 
 	process_manager.start()
+
+	video_manager.init()
+
 	conditional_download(get_test_examples_directory(),
 	[
 		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg',