diff --git a/facefusion/apis/core.py b/facefusion/apis/core.py index 4a1ddc0a..bee0e3ba 100644 --- a/facefusion/apis/core.py +++ b/facefusion/apis/core.py @@ -8,7 +8,7 @@ from starlette.routing import Route, WebSocketRoute from facefusion.apis.endpoints.assets import delete_assets, get_asset, get_assets, upload_asset from facefusion.apis.endpoints.capabilities import get_capabilities -from facefusion.apis.endpoints.jobs import create_job, delete_job, delete_jobs, get_job, get_jobs, update_job, update_jobs +from facefusion.apis.endpoints.jobs import create_job, create_step, delete_job, delete_jobs, delete_step, get_job, get_jobs, update_job, update_jobs from facefusion.apis.endpoints.metrics import get_metrics, websocket_metrics from facefusion.apis.endpoints.ping import websocket_ping from facefusion.apis.endpoints.session import create_session, destroy_session, get_session, refresh_session @@ -53,8 +53,11 @@ def create_api() -> Starlette: Route('/jobs', update_jobs, methods = [ 'PATCH' ], middleware = [ session_guard ]), Route('/jobs', delete_jobs, methods = [ 'DELETE' ], middleware = [ session_guard ]), Route('/jobs/{job_id}', get_job, methods = [ 'GET' ], middleware = [ session_guard ]), + Route('/jobs/{job_id}', create_step, methods = [ 'POST' ], middleware = [ session_guard ]), Route('/jobs/{job_id}', update_job, methods = [ 'PATCH' ], middleware = [ session_guard ]), Route('/jobs/{job_id}', delete_job, methods = [ 'DELETE' ], middleware = [ session_guard ]), + Route('/jobs/{job_id}/{step_index:int}', create_step, methods = [ 'POST' ], middleware = [ session_guard ]), + Route('/jobs/{job_id}/{step_index:int}', delete_step, methods = [ 'DELETE' ], middleware = [ session_guard ]), WebSocketRoute('/metrics', websocket_metrics, middleware = [ session_guard ]), WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ]), WebSocketRoute('/stream', websocket_stream, middleware = [ session_guard ]) diff --git a/facefusion/apis/endpoints/jobs.py b/facefusion/apis/endpoints/jobs.py index 88d57a02..24cff175 100644 --- a/facefusion/apis/endpoints/jobs.py +++ b/facefusion/apis/endpoints/jobs.py @@ -183,3 +183,67 @@ async def delete_job(request : Request) -> JSONResponse: { 'message': translator.get('job_not_deleted', 'facefusion.apis') }, status_code = HTTP_404_NOT_FOUND) + + +async def create_step(request : Request) -> JSONResponse: + job_id = request.path_params.get('job_id') + step_index = request.path_params.get('step_index') + action = request.query_params.get('action') + step_args = await request.json() + + if action == 'add': + if job_manager.add_step(job_id, step_args): + return JSONResponse( + { + 'message': translator.get('ok', 'facefusion.apis') + }, status_code = HTTP_201_CREATED) + + return JSONResponse( + { + 'message': translator.get('job_step_not_added', 'facefusion.apis') + }, status_code = HTTP_400_BAD_REQUEST) + + if action == 'insert': + if job_manager.insert_step(job_id, step_index, step_args): + return JSONResponse( + { + 'message': translator.get('ok', 'facefusion.apis') + }, status_code = HTTP_201_CREATED) + + return JSONResponse( + { + 'message': translator.get('job_step_not_inserted', 'facefusion.apis') + }, status_code = HTTP_400_BAD_REQUEST) + + if action == 'remix': + if job_manager.remix_step(job_id, step_index, step_args): + return JSONResponse( + { + 'message': translator.get('ok', 'facefusion.apis') + }, status_code = HTTP_201_CREATED) + + return JSONResponse( + { + 'message': translator.get('job_step_not_remixed', 'facefusion.apis') + }, status_code = HTTP_400_BAD_REQUEST) + + return JSONResponse( + { + 'message': translator.get('invalid_job_action', 'facefusion.apis') + }, status_code = HTTP_400_BAD_REQUEST) + + +async def delete_step(request : Request) -> JSONResponse: + job_id = request.path_params.get('job_id') + step_index = request.path_params.get('step_index') + + if job_manager.remove_step(job_id, step_index): + return JSONResponse( + { + 'message': translator.get('ok', 'facefusion.apis') + }, status_code = HTTP_200_OK) + + return JSONResponse( + { + 'message': translator.get('job_step_not_removed', 'facefusion.apis') + }, status_code = HTTP_404_NOT_FOUND) diff --git a/facefusion/apis/locales.py b/facefusion/apis/locales.py index f0d90315..2c01b9c7 100644 --- a/facefusion/apis/locales.py +++ b/facefusion/apis/locales.py @@ -22,6 +22,10 @@ LOCALES : Locales =\ 'job_all_not_submitted': 'jobs not submitted', 'job_all_not_deleted': 'jobs not deleted', 'job_all_not_run': 'jobs not run', - 'job_all_not_retried': 'jobs not retried' + 'job_all_not_retried': 'jobs not retried', + 'job_step_not_added': 'step not added', + 'job_step_not_inserted': 'step not inserted', + 'job_step_not_remixed': 'step not remixed', + 'job_step_not_removed': 'step not removed' } } diff --git a/tests/test_api_jobs.py b/tests/test_api_jobs.py index 216dd543..c6a6d66c 100644 --- a/tests/test_api_jobs.py +++ b/tests/test_api_jobs.py @@ -6,7 +6,7 @@ from starlette.testclient import TestClient from facefusion import metadata, session_manager from facefusion.apis.core import create_api -from facefusion.jobs.job_manager import clear_jobs, create_job, find_job_ids, init_jobs +from facefusion.jobs.job_manager import clear_jobs, count_step_total, create_job, find_job_ids, init_jobs from .assert_helper import get_test_jobs_directory @@ -415,3 +415,108 @@ def test_delete_job(test_client : TestClient) -> None: assert find_job_ids('drafted') == [] assert delete_job_response.status_code == 200 + + +def test_create_step(test_client : TestClient) -> None: + create_step_response = test_client.post('/jobs/job-test-create-step?action=add', json = + { + 'processors': [ 'face_swapper' ] + }) + + assert create_step_response.status_code == 401 + + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + create_session_body = create_session_response.json() + access_token = create_session_body.get('access_token') + + create_job('job-test-create-step') + + create_step_response = test_client.post('/jobs/job-test-create-step?action=add', headers = + { + 'Authorization': 'Bearer ' + access_token + }, json = + { + 'processors': [ 'face_swapper' ] + }) + create_step_body = create_step_response.json() + + assert create_step_body.get('message') == 'ok' + assert count_step_total('job-test-create-step') == 1 + assert create_step_response.status_code == 201 + + create_step_response = test_client.post('/jobs/job-test-create-step/0?action=insert', headers = + { + 'Authorization': 'Bearer ' + access_token + }, json = + { + 'processors': [ 'face_swapper' ] + }) + + assert count_step_total('job-test-create-step') == 2 + assert create_step_response.status_code == 201 + + create_step_response = test_client.post('/jobs/job-test-create-step/0?action=remix', headers = + { + 'Authorization': 'Bearer ' + access_token + }, json = + { + 'processors': [ 'face_swapper' ] + }) + + assert count_step_total('job-test-create-step') == 3 + assert create_step_response.status_code == 201 + + create_step_response = test_client.post('/jobs/job-test-create-step?action=invalid', headers = + { + 'Authorization': 'Bearer ' + access_token + }, json = + { + 'processors': [ 'face_swapper' ] + }) + create_step_body = create_step_response.json() + + assert create_step_body.get('message') == 'invalid job action' + assert create_step_response.status_code == 400 + + +def test_delete_step(test_client : TestClient) -> None: + delete_step_response = test_client.delete('/jobs/job-test-delete-step/0') + + assert delete_step_response.status_code == 401 + + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + create_session_body = create_session_response.json() + access_token = create_session_body.get('access_token') + + create_job('job-test-delete-step') + + delete_step_response = test_client.delete('/jobs/job-test-delete-step/0', headers = + { + 'Authorization': 'Bearer ' + access_token + }) + delete_step_body = delete_step_response.json() + + assert delete_step_body.get('message') == 'step not removed' + assert delete_step_response.status_code == 404 + + test_client.post('/jobs/job-test-delete-step?action=add', headers = + { + 'Authorization': 'Bearer ' + access_token + }, json = + { + 'processors': [ 'face_swapper' ] + }) + + delete_step_response = test_client.delete('/jobs/job-test-delete-step/0', headers = + { + 'Authorization': 'Bearer ' + access_token + }) + + assert count_step_total('job-test-delete-step') == 0 + assert delete_step_response.status_code == 200