diff --git a/facefusion/apis/endpoints/metrics.py b/facefusion/apis/endpoints/metrics.py index 6d8c5e8a..cb51de46 100644 --- a/facefusion/apis/endpoints/metrics.py +++ b/facefusion/apis/endpoints/metrics.py @@ -5,6 +5,7 @@ from starlette.responses import JSONResponse, Response from starlette.status import HTTP_404_NOT_FOUND from starlette.websockets import WebSocket +from facefusion.apis import websocket_store from facefusion.apis.api_helper import get_sec_websocket_protocol from facefusion.system import get_metrics_set @@ -21,6 +22,7 @@ async def get_metrics(request : Request) -> Response: async def websocket_metrics(websocket : WebSocket) -> None: subprotocol = get_sec_websocket_protocol(websocket.scope) await websocket.accept(subprotocol = subprotocol) + websocket_store.set_websocket(websocket) try: while True: @@ -30,3 +32,5 @@ async def websocket_metrics(websocket : WebSocket) -> None: except Exception: pass + + websocket_store.delete_websocket(websocket) diff --git a/facefusion/apis/endpoints/ping.py b/facefusion/apis/endpoints/ping.py index f508c851..53362bbc 100644 --- a/facefusion/apis/endpoints/ping.py +++ b/facefusion/apis/endpoints/ping.py @@ -1,5 +1,6 @@ from starlette.websockets import WebSocket +from facefusion.apis import websocket_store from facefusion.apis.api_helper import get_sec_websocket_protocol @@ -7,6 +8,7 @@ async def websocket_ping(websocket : WebSocket) -> None: subprotocol = get_sec_websocket_protocol(websocket.scope) await websocket.accept(subprotocol = subprotocol) + websocket_store.set_websocket(websocket) try: while True: @@ -14,3 +16,5 @@ async def websocket_ping(websocket : WebSocket) -> None: except Exception: pass + + websocket_store.delete_websocket(websocket) diff --git a/facefusion/apis/endpoints/session.py b/facefusion/apis/endpoints/session.py index cf31d6fe..6ad8afa8 100644 --- a/facefusion/apis/endpoints/session.py +++ b/facefusion/apis/endpoints/session.py @@ -1,11 +1,12 @@ import secrets +from starlette.concurrency import run_in_threadpool 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, HTTP_429_TOO_MANY_REQUESTS from facefusion import content_store, face_store, inference_manager, process_manager, rtc_store, session_context, session_manager, state_manager, store_manager, translator, video_manager -from facefusion.apis import asset_store +from facefusion.apis import asset_store, websocket_store from facefusion.apis.session_helper import validate_api_key from facefusion.filesystem import is_directory, remove_directory from facefusion.jobs import job_manager @@ -31,6 +32,7 @@ async def create_session(request : Request) -> JSONResponse: video_manager.init() process_manager.init() rtc_store.init() + websocket_store.init() store_manager.observe_api_session(session_id) @@ -105,7 +107,7 @@ async def destroy_session(request : Request) -> JSONResponse: }, status_code = HTTP_404_NOT_FOUND) session_manager.clear_api_session(session_id) - store_manager.destroy(session_id) + await run_in_threadpool(store_manager.destroy, session_id) return JSONResponse( { diff --git a/facefusion/apis/endpoints/stream.py b/facefusion/apis/endpoints/stream.py index 4a81a3a5..f773aceb 100644 --- a/facefusion/apis/endpoints/stream.py +++ b/facefusion/apis/endpoints/stream.py @@ -4,6 +4,7 @@ from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_404_NOT_FOUND, from starlette.websockets import WebSocket, WebSocketState from facefusion import rtc_store +from facefusion.apis import websocket_store from facefusion.apis.api_helper import get_sec_websocket_protocol from facefusion.apis.stream_manager import destroy_stream, process_image, process_video @@ -12,7 +13,9 @@ async def websocket_stream(websocket : WebSocket) -> None: subprotocol = get_sec_websocket_protocol(websocket.scope) await websocket.accept(subprotocol = subprotocol) + websocket_store.set_websocket(websocket) await process_image(websocket) + websocket_store.delete_websocket(websocket) if websocket.client_state == WebSocketState.CONNECTED: await websocket.close() diff --git a/facefusion/apis/websocket_store.py b/facefusion/apis/websocket_store.py new file mode 100644 index 00000000..a0f223e9 --- /dev/null +++ b/facefusion/apis/websocket_store.py @@ -0,0 +1,41 @@ +from anyio.from_thread import run +from anyio.lowlevel import current_token + +from facefusion import store_creator +from facefusion.session_manager import resolve_owner_id +from facefusion.types import SessionId, SessionWebsocketSet, Store, Websocket + +WEBSOCKET_STORE : Store = store_creator.create_store({}) + + +def init() -> None: + owner_id = resolve_owner_id() + store_creator.init_content(WEBSOCKET_STORE, owner_id) + + +def set_websocket(websocket : Websocket) -> None: + owner_id = resolve_owner_id() + + if store_creator.has_content(WEBSOCKET_STORE, owner_id): + store_creator.get_content(WEBSOCKET_STORE, owner_id)[id(websocket)] =\ + { + 'websocket': websocket, + 'event_loop_token': current_token() + } + + +def delete_websocket(websocket : Websocket) -> None: + owner_id = resolve_owner_id() + session_websocket_set : SessionWebsocketSet = store_creator.get_content(WEBSOCKET_STORE, owner_id) + + if session_websocket_set and id(websocket) in session_websocket_set: + del session_websocket_set[id(websocket)] + + +def destroy(session_id : SessionId) -> None: + session_websocket_set : SessionWebsocketSet = store_creator.get_content(WEBSOCKET_STORE, session_id) + store_creator.delete_content(WEBSOCKET_STORE, session_id) + + if session_websocket_set: + for session_websocket in session_websocket_set.values(): + run(session_websocket.get('websocket').close, token = session_websocket.get('event_loop_token')) diff --git a/facefusion/store_manager.py b/facefusion/store_manager.py index 736067b3..83ef6bb7 100644 --- a/facefusion/store_manager.py +++ b/facefusion/store_manager.py @@ -5,12 +5,12 @@ from types import ModuleType from typing import List from facefusion import content_store, face_store, inference_manager, process_manager, rtc_store, session_manager, state_manager, video_manager -from facefusion.apis import asset_store +from facefusion.apis import asset_store, websocket_store from facefusion.types import SessionId def get_stores() -> List[ModuleType]: - return [ state_manager, asset_store, content_store, face_store, inference_manager, process_manager, rtc_store, video_manager ] + return [ state_manager, asset_store, content_store, face_store, inference_manager, process_manager, rtc_store, video_manager, websocket_store ] def observe_api_session(session_id : SessionId) -> None: diff --git a/facefusion/types.py b/facefusion/types.py index 4c2326f5..814a6445 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -156,6 +156,15 @@ Content : TypeAlias = Dict[str, Any] Token : TypeAlias = str SessionId : TypeAlias = str +Websocket : TypeAlias = Any +WebsocketId : TypeAlias = int +EventLoopToken : TypeAlias = Any +SessionWebsocket = TypedDict('SessionWebsocket', +{ + 'websocket' : Websocket, + 'event_loop_token' : EventLoopToken +}) +SessionWebsocketSet : TypeAlias = Dict[WebsocketId, SessionWebsocket] ApiSession = TypedDict('ApiSession', { 'access_token' : Token, diff --git a/tests/test_api_metrics.py b/tests/test_api_metrics.py index 18113ddb..bc91c50e 100644 --- a/tests/test_api_metrics.py +++ b/tests/test_api_metrics.py @@ -5,6 +5,7 @@ from pytest_mock import MockerFixture from starlette.testclient import TestClient from facefusion import metadata, session_manager, state_manager +from facefusion.apis import websocket_store from facefusion.apis.core import create_api from .assert_helper import get_test_jobs_directory @@ -247,3 +248,7 @@ def test_websocket_metrics(test_client : TestClient) -> None: assert metrics_set.get('processor').get('cores').get('value') == 16 assert metrics_set.get('processor').get('frequency').get('value') == 3500 assert metrics_set.get('processor').get('utilization').get('value') == 25 + + websocket_store.destroy(session_manager.find_api_session_id(create_session_body.get('access_token'))) + + assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' } diff --git a/tests/test_api_ping.py b/tests/test_api_ping.py index df597445..fb8497ef 100644 --- a/tests/test_api_ping.py +++ b/tests/test_api_ping.py @@ -4,6 +4,7 @@ import pytest from starlette.testclient import TestClient from facefusion import metadata, session_manager, state_manager +from facefusion.apis import websocket_store from facefusion.apis.core import create_api from .assert_helper import get_test_jobs_directory @@ -38,3 +39,7 @@ def test_ping(test_client : TestClient) -> None: 'access_token.' + create_session_body.get('access_token') ]) as websocket: assert websocket + + websocket_store.destroy(session_manager.find_api_session_id(create_session_body.get('access_token'))) + + assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' } diff --git a/tests/test_api_session.py b/tests/test_api_session.py index e7000c21..601a366f 100644 --- a/tests/test_api_session.py +++ b/tests/test_api_session.py @@ -8,7 +8,7 @@ import pytest from starlette.testclient import TestClient from facefusion import content_store, face_store, inference_manager, metadata, process_manager, rtc_store, session_context, session_manager, state_manager, store_creator, store_manager, video_manager -from facefusion.apis import asset_store +from facefusion.apis import asset_store, websocket_store from facefusion.apis.core import create_api from facefusion.download import conditional_download from facefusion.libraries import datachannel as datachannel_module @@ -300,7 +300,8 @@ def test_destroy_session_content(test_client : TestClient) -> None: (inference_manager, inference_manager.INFERENCE_POOL_STORE), (process_manager, process_manager.PROCESS_STORE), (rtc_store, rtc_store.RTC_STORE), - (video_manager, video_manager.VIDEO_POOL_STORE) + (video_manager, video_manager.VIDEO_POOL_STORE), + (websocket_store, websocket_store.WEBSOCKET_STORE) ] session_manager.set_api_session(session_id, diff --git a/tests/test_api_stream.py b/tests/test_api_stream.py index d1539cb7..f6078012 100644 --- a/tests/test_api_stream.py +++ b/tests/test_api_stream.py @@ -91,7 +91,14 @@ def test_stream_image(test_client : TestClient) -> None: websocket.send_bytes(source_content) output_buffer = websocket.receive_bytes() - assert create_hash(output_buffer) == '0142782f' + assert create_hash(output_buffer) == '0142782f' + + test_client.delete('/session', headers = + { + 'Authorization': 'Bearer ' + access_token + }) + + assert websocket.receive() == { 'type': 'websocket.close', 'code': 1000, 'reason': '' } @pytest.mark.parametrize('video_codec', [ 'av1', 'vp8' ]) diff --git a/tests/test_store_manager.py b/tests/test_store_manager.py new file mode 100644 index 00000000..71368b14 --- /dev/null +++ b/tests/test_store_manager.py @@ -0,0 +1,48 @@ +from typing import Iterator + +import pytest + +from facefusion import state_manager, store_creator, video_manager +from facefusion.download import conditional_download +from facefusion.session_context import resolve_local_id, set_session_id +from facefusion.session_manager import clear_api_session, clear_cli_session, create_api_session, fork_session, set_api_session +from facefusion.store_manager import conditional_cli_destroy +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + state_manager.init() + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> Iterator[None]: + local_id = resolve_local_id() + + set_session_id(local_id) + + yield + + set_session_id(local_id) + + +def test_conditional_cli_destroy() -> None: + set_api_session('session-a', create_api_session()) + set_session_id('session-a') + + fork_id = fork_session() + video_manager.init() + video_reader = video_manager.get_reader(get_test_example_file('target-240p.mp4'), 'read_video_frame') + + clear_api_session('session-a') + conditional_cli_destroy(fork_id) + + assert video_reader.get('process').poll() == -9 + assert store_creator.has_content(video_manager.VIDEO_POOL_STORE, fork_id) is False + + clear_cli_session(fork_id) diff --git a/tests/test_websocket_store.py b/tests/test_websocket_store.py new file mode 100644 index 00000000..fd58678c --- /dev/null +++ b/tests/test_websocket_store.py @@ -0,0 +1,72 @@ +from typing import Iterator +from unittest.mock import AsyncMock + +import anyio +import pytest + +from facefusion import store_creator +from facefusion.apis.websocket_store import WEBSOCKET_STORE, delete_websocket, destroy, init, set_websocket +from facefusion.session_context import resolve_local_id, set_session_id +from facefusion.types import SessionId, Websocket + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> Iterator[None]: + local_id = resolve_local_id() + + set_session_id(local_id) + init() + + yield + + set_session_id(local_id) + + +async def set_and_destroy_websocket(websocket : Websocket, session_id : SessionId) -> None: + set_websocket(websocket) + await anyio.to_thread.run_sync(destroy, session_id) + + +async def set_and_delete_websocket(websocket : Websocket) -> None: + set_websocket(websocket) + delete_websocket(websocket) + + +async def set_websocket_in_event_loop(websocket : Websocket) -> None: + set_websocket(websocket) + + +def test_init() -> None: + set_session_id('session-a') + init() + + assert store_creator.get_content(WEBSOCKET_STORE, 'session-a') == {} + + destroy('session-a') + + +def test_set_websocket() -> None: + websocket_mock = AsyncMock() + + anyio.run(set_websocket_in_event_loop, websocket_mock) + + assert store_creator.get_content(WEBSOCKET_STORE, resolve_local_id()).get(id(websocket_mock)).get('websocket') is websocket_mock + + +def test_delete_websocket() -> None: + websocket_mock = AsyncMock() + + anyio.run(set_and_delete_websocket, websocket_mock) + + assert store_creator.get_content(WEBSOCKET_STORE, resolve_local_id()) == {} + + +def test_destroy() -> None: + websocket_mock = AsyncMock() + + set_session_id('session-a') + init() + anyio.run(set_and_destroy_websocket, websocket_mock, 'session-a') + + websocket_mock.close.assert_awaited_once() + assert store_creator.has_content(WEBSOCKET_STORE, 'session-a') is False