From c7531b333ccfb95f66508bd2cd5f499296e9219e Mon Sep 17 00:00:00 2001 From: henryruhs Date: Sat, 7 Nov 2026 09:53:03 +0100 Subject: [PATCH] return session and fork via /process --- facefusion/apis/endpoints/process.py | 19 +++--- facefusion/process_manager.py | 8 +++ facefusion/types.py | 5 ++ tests/apis/endpoints/test_process.py | 90 +++++++++++++++++++++++----- tests/test_process_manager.py | 17 +++++- 5 files changed, 116 insertions(+), 23 deletions(-) diff --git a/facefusion/apis/endpoints/process.py b/facefusion/apis/endpoints/process.py index 9c64e00a..672c3343 100644 --- a/facefusion/apis/endpoints/process.py +++ b/facefusion/apis/endpoints/process.py @@ -5,6 +5,8 @@ from starlette.websockets import WebSocket from facefusion import process_manager from facefusion.apis import websocket_manager from facefusion.apis.websocket_helper import get_sec_websocket_protocol +from facefusion.session_manager import find_fork_session_id, resolve_owner_id +from facefusion.types import ProcessSet async def websocket_process(websocket : WebSocket) -> None: @@ -16,14 +18,17 @@ async def websocket_process(websocket : WebSocket) -> None: receive_task = asyncio.create_task(websocket.receive()) while websocket_manager.has_connection('process') and not receive_task.done(): - process_state = process_manager.get_state() - stages = process_manager.get_stages() - - await websocket.send_json( + owner_id = resolve_owner_id() + fork_id = find_fork_session_id(owner_id) + process_set : ProcessSet =\ { - 'status': process_state, - 'stages': stages - }) + 'session': process_manager.get_process(owner_id) + } + + if fork_id and process_manager.has_process(fork_id): + process_set['fork'] = process_manager.get_process(fork_id) + + await websocket.send_json(process_set) await asyncio.wait([ receive_task ], timeout = 0.5) receive_task.cancel() diff --git a/facefusion/process_manager.py b/facefusion/process_manager.py index 14887207..b46e61e6 100644 --- a/facefusion/process_manager.py +++ b/facefusion/process_manager.py @@ -18,6 +18,14 @@ def init() -> None: store_creator.init_content(PROCESS_STORE, session_id) +def has_process(session_id : SessionId) -> bool: + return store_creator.has_content(PROCESS_STORE, session_id) + + +def get_process(session_id : SessionId) -> Process: + return store_creator.get_content(PROCESS_STORE, session_id) + + def get_status(session_id : SessionId) -> Optional[ProcessState]: if store_creator.has_content(PROCESS_STORE, session_id): return store_creator.get_content(PROCESS_STORE, session_id).get('status') diff --git a/facefusion/types.py b/facefusion/types.py index 9820697b..c9eba83f 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -234,6 +234,11 @@ Process = TypedDict('Process', 'status' : ProcessState, 'stages' : List[Stage] }) +ProcessSet = TypedDict('ProcessSet', +{ + 'session' : Process, + 'fork' : NotRequired[Process] +}) TableHeader : TypeAlias = str TableContent : TypeAlias = Any diff --git a/tests/apis/endpoints/test_process.py b/tests/apis/endpoints/test_process.py index ad608685..91d84057 100644 --- a/tests/apis/endpoints/test_process.py +++ b/tests/apis/endpoints/test_process.py @@ -52,8 +52,11 @@ def test_websocket_process(test_client : TestClient) -> None: ]) as websocket: assert websocket.receive_json() ==\ { - 'status': 'pending', - 'stages': [] + 'session': + { + 'status': 'pending', + 'stages': [] + } } session_context.set_session_id(session_id) @@ -68,23 +71,80 @@ def test_websocket_process(test_client : TestClient) -> None: assert websocket.receive_json() ==\ { - 'status': 'processing', - 'stages': - [ - { - 'type': 'stream', - 'title': 'streaming', - 'progress': + 'session': + { + 'status': 'processing', + 'stages': + [ { - 'unit': 'frame', - 'current': 0, - 'started_at': 10.0, - 'updated_at': 10.0 + 'type': 'stream', + 'title': 'streaming', + 'progress': + { + 'unit': 'frame', + 'current': 0, + 'started_at': 10.0, + 'updated_at': 10.0 + } } - } - ] + ] + } } + session_manager.fork_session() + process_manager.start() + process_manager.stage('infer', 'processing', + { + 'unit': 'frame', + 'current': 0, + 'total': 100, + 'started_at': 20.0, + 'updated_at': 20.0 + }) + + assert websocket.receive_json() ==\ + { + 'session': + { + 'status': 'processing', + 'stages': + [ + { + 'type': 'stream', + 'title': 'streaming', + 'progress': + { + 'unit': 'frame', + 'current': 0, + 'started_at': 10.0, + 'updated_at': 10.0 + } + } + ] + }, + 'fork': + { + 'status': 'processing', + 'stages': + [ + { + 'type': 'infer', + 'title': 'processing', + 'progress': + { + 'unit': 'frame', + 'current': 0, + 'total': 100, + 'started_at': 20.0, + 'updated_at': 20.0 + } + } + ] + } + } + + process_manager.end() + session_manager.join_session() process_manager.end() websocket_manager.destroy(session_id) diff --git a/tests/test_process_manager.py b/tests/test_process_manager.py index ba07b2bd..25e3a05c 100644 --- a/tests/test_process_manager.py +++ b/tests/test_process_manager.py @@ -4,7 +4,7 @@ from unittest.mock import patch import pytest from facefusion import session_context, session_manager, store_creator -from facefusion.process_manager import PROCESS_STORE, clear, count, describe, destroy, end, get_stages, get_state, get_status, init, is_pending, is_processing, is_stopping, seek, set_state, stage, start, stop, update +from facefusion.process_manager import PROCESS_STORE, clear, count, describe, destroy, end, get_process, get_stages, get_state, get_status, has_process, init, is_pending, is_processing, is_stopping, seek, set_state, stage, start, stop, update @pytest.fixture(scope = 'function', autouse = True) @@ -31,6 +31,21 @@ def test_init() -> None: assert get_state() == 'pending' +def test_has_process() -> None: + assert has_process('session-a') is True + assert has_process('session-c') is False + + +def test_get_process() -> None: + set_state('processing') + + assert get_process('session-a') ==\ + { + 'status': 'processing', + 'stages': [] + } + + def test_get_status() -> None: set_state('processing')