return session and fork via /process

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

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

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

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)
 

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