From 7affeb7e3be9bf436f65e3a0c7ebe64d4d0e289c Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Thu, 19 Feb 2026 17:25:01 +0530 Subject: [PATCH] process image --- facefusion/apis/core.py | 4 +- facefusion/apis/endpoints/process.py | 81 +++++++++++++++++++ tests/test_api_process.py | 113 +++++++++++++++++++++++++++ 3 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 facefusion/apis/endpoints/process.py create mode 100644 tests/test_api_process.py diff --git a/facefusion/apis/core.py b/facefusion/apis/core.py index c4d1d5a5..8a64a59a 100644 --- a/facefusion/apis/core.py +++ b/facefusion/apis/core.py @@ -7,6 +7,7 @@ from facefusion.apis.endpoints.assets import delete_assets, get_asset, get_asset from facefusion.apis.endpoints.capabilities import get_capabilities from facefusion.apis.endpoints.metrics import get_metrics, websocket_metrics from facefusion.apis.endpoints.ping import websocket_ping +from facefusion.apis.endpoints.process import websocket_process_image from facefusion.apis.endpoints.session import create_session, create_session_guard, destroy_session, get_session, refresh_session from facefusion.apis.endpoints.state import get_state, set_state @@ -28,7 +29,8 @@ def create_api() -> Starlette: Route('/capabilities', get_capabilities, methods = [ 'GET' ]), Route('/metrics', get_metrics, methods = [ 'GET' ], middleware = [ session_guard ]), WebSocketRoute('/metrics', websocket_metrics, middleware = [ session_guard ]), - WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ]) + WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ]), + WebSocketRoute('/process/image', websocket_process_image, middleware = [ session_guard ]) ] api = Starlette(routes = routes) diff --git a/facefusion/apis/endpoints/process.py b/facefusion/apis/endpoints/process.py new file mode 100644 index 00000000..7dc07937 --- /dev/null +++ b/facefusion/apis/endpoints/process.py @@ -0,0 +1,81 @@ +import os +import uuid +from time import time +from typing import Optional + +from starlette.websockets import WebSocket + +from facefusion import ffmpeg, process_manager, session_context, session_manager, state_manager +from facefusion.apis.api_helper import get_sec_websocket_protocol +from facefusion.apis.endpoints.session import extract_access_token +from facefusion.filesystem import create_directory, remove_file +from facefusion.workflows import image_to_image + + +def save_image_bytes(image_bytes : bytes, temp_path : str, file_extension : str) -> str: + raw_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension) + + with open(raw_path, 'wb') as raw_file: + raw_file.write(image_bytes) + + return raw_path + + +def sanitize_upload(raw_path : str, temp_path : str, file_extension : str) -> Optional[str]: + target_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension) + + process_manager.start() + sanitized = ffmpeg.sanitize_image(raw_path, target_path) + process_manager.end() + remove_file(raw_path) + + if sanitized: + return target_path + return None + + +def process_target(target_path : str, output_path : str) -> bool: + state_manager.set_item('target_path', target_path) + state_manager.set_item('output_path', output_path) + state_manager.set_item('workflow', 'image-to-image') + + error_code = image_to_image.process(time()) + remove_file(target_path) + + return error_code == 0 + + +async def websocket_process_image(websocket : WebSocket) -> None: + subprotocol = get_sec_websocket_protocol(websocket.scope) + access_token = extract_access_token(websocket.scope) + session_id = session_manager.find_session_id(access_token) + + session_context.set_session_id(session_id) + source_paths = state_manager.get_item('source_paths') + + await websocket.accept(subprotocol = subprotocol) + + if source_paths: + temp_path = state_manager.get_temp_path() + create_directory(temp_path) + + try: + while True: + image_bytes = await websocket.receive_bytes() + raw_path = save_image_bytes(image_bytes, temp_path, '.jpg') + target_path = sanitize_upload(raw_path, temp_path, '.jpg') + + if target_path: + output_path = os.path.join(temp_path, str(uuid.uuid4()) + '.jpg') + + if process_target(target_path, output_path): + with open(output_path, 'rb') as output_file: + await websocket.send_bytes(output_file.read()) + + remove_file(output_path) + + except Exception: + pass + return + + await websocket.close(code = 1008) diff --git a/tests/test_api_process.py b/tests/test_api_process.py new file mode 100644 index 00000000..43418e01 --- /dev/null +++ b/tests/test_api_process.py @@ -0,0 +1,113 @@ +import shutil +import tempfile +from typing import Iterator +from unittest.mock import patch + +import pytest +from starlette.testclient import TestClient + +from facefusion import metadata, session_manager, state_manager +from facefusion.apis import asset_store +from facefusion.apis.core import create_api +from facefusion.download import conditional_download +from .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/source.jpg' + ]) + + +@pytest.fixture(scope = 'module') +def test_client() -> Iterator[TestClient]: + with TestClient(create_api()) as test_client: + yield test_client + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + state_manager.init_item('temp_path', tempfile.gettempdir()) + state_manager.init_item('temp_frame_format', 'png') + state_manager.init_item('source_paths', None) + session_manager.SESSIONS.clear() + asset_store.clear() + + +def mock_process_image(start_time : float) -> int: + output_path = state_manager.get_item('output_path') + target_path = state_manager.get_item('target_path') + shutil.copy(target_path, output_path) + return 0 + + +def test_websocket_process_image_without_auth(test_client : TestClient) -> None: + with pytest.raises(Exception): + with test_client.websocket_connect('/process/image') as _: + pass + + +def test_websocket_process_image_without_source(test_client : TestClient) -> None: + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + access_token = create_session_response.json().get('access_token') + + with test_client.websocket_connect('/process/image', subprotocols = + [ + 'access_token.' + access_token + ]) as websocket: + with pytest.raises(Exception): + websocket.receive_bytes() + + +def test_websocket_process_image_single(test_client : TestClient) -> None: + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + access_token = create_session_response.json().get('access_token') + source_path = get_test_example_file('source.jpg') + state_manager.init_item('source_paths', [ source_path ]) + + with open(source_path, 'rb') as source_file: + image_bytes = source_file.read() + + with patch('facefusion.apis.endpoints.process.image_to_image.process', side_effect = mock_process_image): + with test_client.websocket_connect('/process/image', subprotocols = + [ + 'access_token.' + access_token + ]) as websocket: + websocket.send_bytes(image_bytes) + result_bytes = websocket.receive_bytes() + + assert len(result_bytes) > 0 + + +def test_websocket_process_image_batch(test_client : TestClient) -> None: + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + access_token = create_session_response.json().get('access_token') + source_path = get_test_example_file('source.jpg') + state_manager.init_item('source_paths', [ source_path ]) + + with open(source_path, 'rb') as source_file: + image_bytes = source_file.read() + + with patch('facefusion.apis.endpoints.process.image_to_image.process', side_effect = mock_process_image): + with test_client.websocket_connect('/process/image', subprotocols = + [ + 'access_token.' + access_token + ]) as websocket: + websocket.send_bytes(image_bytes) + result_bytes_1 = websocket.receive_bytes() + websocket.send_bytes(image_bytes) + result_bytes_2 = websocket.receive_bytes() + + assert len(result_bytes_1) > 0 + assert len(result_bytes_2) > 0