From 8c92608569fdb378db7436e51d5d2cf709421aef Mon Sep 17 00:00:00 2001 From: Harisreedhar <46858047+harisreedhar@users.noreply.github.com> Date: Fri, 24 Jul 2026 16:23:54 +0530 Subject: [PATCH] Add frame_store module (#1194) * add frame_store module * rename and change tests * rename and update tests --- facefusion/frame_store.py | 34 +++++++++++++++++++ facefusion/types.py | 1 + tests/test_frame_store.py | 70 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+) create mode 100644 facefusion/frame_store.py create mode 100644 tests/test_frame_store.py diff --git a/facefusion/frame_store.py b/facefusion/frame_store.py new file mode 100644 index 00000000..6ce448b6 --- /dev/null +++ b/facefusion/frame_store.py @@ -0,0 +1,34 @@ +from facefusion.types import FrameStoreSet, VisionFrame, VisionFrameSet + +FRAME_STORE_SET : FrameStoreSet = {} + + +def get_frame_store(id : str) -> VisionFrameSet: + if id not in FRAME_STORE_SET: + FRAME_STORE_SET[id] = {} + + return FRAME_STORE_SET.get(id) + + +def set_frame(id : str, frame_number : int, vision_frame : VisionFrame) -> None: + frame_store = get_frame_store(id) + frame_store[frame_number] = vision_frame + + +def select_frame_set(id : str, frame_start : int, frame_end : int) -> VisionFrameSet: + frame_store = get_frame_store(id) + frame_set = {} + + for frame_number in range(frame_start, frame_end + 1): + if frame_number in frame_store: + frame_set[frame_number] = frame_store.get(frame_number) + + return frame_set + + +def reduce_frames(id : str, frame_min : int, frame_max : int) -> None: + FRAME_STORE_SET[id] = select_frame_set(id, frame_min, frame_max) + + +def clear_frames() -> None: + FRAME_STORE_SET.clear() diff --git a/facefusion/types.py b/facefusion/types.py index 62427886..5e946f58 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -146,6 +146,7 @@ VideoPoolSet = TypedDict('VideoPoolSet', 'reader' : VideoReaderSet, 'writer' : VideoWriterSet }) +FrameStoreSet : TypeAlias = Dict[str, VisionFrameSet] ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] Args : TypeAlias = Dict[str, Any] diff --git a/tests/test_frame_store.py b/tests/test_frame_store.py new file mode 100644 index 00000000..057a121f --- /dev/null +++ b/tests/test_frame_store.py @@ -0,0 +1,70 @@ +import pytest + +from facefusion import process_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 +from .helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + process_manager.start() + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear_frames() + + +def test_get_frame_store() -> None: + frame_store = get_frame_store('reader-1') + + assert frame_store == {} + assert get_frame_store('reader-1') is frame_store + + +def test_set_frame() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + set_frame('reader-1', 5, target_frame) + + assert get_frame_store('reader-1').get(5) is target_frame + + +def test_select_frame_set() -> None: + first_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + fifth_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 5) + + set_frame('reader-1', 2, first_frame) + set_frame('reader-1', 5, fifth_frame) + + assert sorted(select_frame_set('reader-1', 0, 4)) == [ 2 ] + assert select_frame_set('reader-1', 0, 4).get(2) is first_frame + assert sorted(select_frame_set('reader-1', 2, 6)) == [ 2, 5 ] + assert select_frame_set('reader-1', 2, 6).get(5) is fifth_frame + assert select_frame_set('reader-1', 8, 12) == {} + + +def test_reduce_frames() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + for frame_number in range(0, 10): + set_frame('reader-1', frame_number, target_frame) + + reduce_frames('reader-1', 4, 6) + + assert sorted(get_frame_store('reader-1')) == [ 4, 5, 6 ] + + +def test_clear_frames() -> None: + target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) + + set_frame('reader-1', 0, target_frame) + clear_frames() + + assert get_frame_store('reader-1') == {}