expose the output via api (#1223)

* expose the output via api

* remove mocks from job testing

facefusion/apis/endpoints/jobs.py

@@ -1,13 +1,17 @@
+import os
 from functools import partial
 
-from starlette.background import BackgroundTask
+from starlette.background import BackgroundTask, BackgroundTasks
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_202_ACCEPTED, HTTP_400_BAD_REQUEST, HTTP_404_NOT_FOUND
 
 import facefusion.choices
 import facefusion.core
-from facefusion import args_helper, state_manager, translator
+from facefusion import args_helper, session_context, session_manager, state_manager, translator
+from facefusion.apis import jobs_helper
+from facefusion.apis.session_helper import extract_access_token
+from facefusion.filesystem import create_directory, get_file_extension, is_directory
 from facefusion.jobs import job_helper, job_manager, job_runner
 
 
@@ -15,18 +19,18 @@ async def get_jobs(request : Request) -> JSONResponse:
 	job_status = request.query_params.get('status')
 
 	if job_status in facefusion.choices.job_statuses:
+		__job_set__ = {}
 		job_set = job_manager.find_jobs(job_status)
-		job_summaries = {}
 
 		for job_id, job in job_set.items():
-			job_summaries[job_id] =\
+			__job_set__[job_id] =\
 			{
 				'version': job.get('version'),
 				'date_created': job.get('date_created'),
 				'date_updated': job.get('date_updated')
 			}
 
-		return JSONResponse(job_summaries, status_code = HTTP_200_OK)
+		return JSONResponse(__job_set__, status_code = HTTP_200_OK)
 
 	return JSONResponse(
 	{
@@ -80,6 +84,7 @@ async def update_jobs(request : Request) -> JSONResponse:
 	if action == 'run':
 		if job_manager.find_job_ids('queued'):
 			run_jobs_task = BackgroundTask(partial(job_runner.run_jobs, facefusion.core.process_step, state_manager.get_item('halt_on_error')))
+
 			return JSONResponse(
 			{
 				'message': translator.get('ok', 'facefusion.apis')
@@ -93,6 +98,7 @@ async def update_jobs(request : Request) -> JSONResponse:
 	if action == 'retry':
 		if job_manager.find_job_ids('failed'):
 			retry_jobs_task = BackgroundTask(partial(job_runner.retry_jobs, facefusion.core.process_step, state_manager.get_item('halt_on_error')))
+
 			return JSONResponse(
 			{
 				'message': translator.get('ok', 'facefusion.apis')
@@ -112,6 +118,8 @@ async def update_jobs(request : Request) -> JSONResponse:
 async def update_job(request : Request) -> JSONResponse:
 	job_id = request.path_params.get('job_id')
 	action = request.query_params.get('action')
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
 
 	if action == 'submit':
 		if job_manager.submit_job(job_id):
@@ -127,11 +135,14 @@ async def update_job(request : Request) -> JSONResponse:
 
 	if action == 'run':
 		if job_id in job_manager.find_job_ids('queued'):
-			run_job_task = BackgroundTask(partial(job_runner.run_job, job_id, facefusion.core.process_step))
+			run_job_tasks = BackgroundTasks()
+			run_job_tasks.add_task(partial(job_runner.run_job, job_id, facefusion.core.process_step))
+			run_job_tasks.add_task(partial(jobs_helper.capture_output_asset, job_id, session_id))
+
 			return JSONResponse(
 			{
 				'message': translator.get('ok', 'facefusion.apis')
-			}, status_code = HTTP_202_ACCEPTED, background = run_job_task)
+			}, status_code = HTTP_202_ACCEPTED, background = run_job_tasks)
 
 		return JSONResponse(
 		{
@@ -140,11 +151,14 @@ async def update_job(request : Request) -> JSONResponse:
 
 	if action == 'retry':
 		if job_id in job_manager.find_job_ids('failed'):
-			retry_job_task = BackgroundTask(partial(job_runner.retry_job, job_id, facefusion.core.process_step))
+			retry_job_tasks = BackgroundTasks()
+			retry_job_tasks.add_task(partial(job_runner.retry_job, job_id, facefusion.core.process_step))
+			retry_job_tasks.add_task(partial(jobs_helper.capture_output_asset, job_id, session_id))
+
 			return JSONResponse(
 			{
 				'message': translator.get('ok', 'facefusion.apis')
-			}, status_code = HTTP_202_ACCEPTED, background = retry_job_task)
+			}, status_code = HTTP_202_ACCEPTED, background = retry_job_tasks)
 
 		return JSONResponse(
 		{
@@ -189,7 +203,23 @@ 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 = args_helper.filter_api_step_args(await request.json())
+
+	step_args = await request.json()
+	step_args = args_helper.filter_api_step_args(step_args)
+
+	if state_manager.get_item('source_paths'):
+		step_args['source_paths'] = state_manager.get_item('source_paths')
+
+	if state_manager.get_item('target_path'):
+		access_token = extract_access_token(request.scope)
+		session_id = session_manager.find_session_id(access_token)
+		session_context.set_session_id(session_id)
+		temp_path = state_manager.get_temp_path()
+
+		step_args['target_path'] = state_manager.get_item('target_path')
+
+		if is_directory(temp_path) or create_directory(temp_path):
+			step_args['output_path'] = os.path.join(temp_path, job_id + get_file_extension(state_manager.get_item('target_path')))
 
 	if action == 'add':
 		if job_manager.add_step(job_id, step_args):

facefusion/apis/jobs_helper.py

@@ -0,0 +1,14 @@
+from facefusion.apis import asset_store
+from facefusion.filesystem import is_file
+from facefusion.jobs import job_manager
+from facefusion.types import SessionId
+
+
+def capture_output_asset(job_id : str, session_id : SessionId) -> None:
+	job = job_manager.read_job_file(job_id)
+
+	if job and job.get('steps'):
+		output_path = job.get('steps')[-1].get('args').get('output_path')
+
+		if output_path and is_file(output_path):
+			asset_store.create_asset(session_id, 'output', output_path)

facefusion/jobs/job_helper.py

@@ -13,8 +13,10 @@ def get_step_output_path(job_id : str, step_index : int, output_path : str) -> O
 
 		if output_file_name and output_file_extension:
 			return os.path.join(output_directory_path, output_file_name + '-' + job_id + '-' + str(step_index) + output_file_extension)
+
 		if output_file_path and output_directory_path:
 			return os.path.join(output_directory_path, output_file_path + '-' + job_id + '-' + str(step_index))
+
 	return None
 
 

facefusion/types.py

@@ -222,7 +222,7 @@ EncoderSet = TypedDict('EncoderSet',
 VideoPreset = Literal['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow']
 
 AssetId : TypeAlias = str
-AssetType = Literal['source', 'target']
+AssetType = Literal['source', 'target', 'output']
 MediaType = Literal['image', 'video', 'audio']
 AudioMetadata = TypedDict('AudioMetadata',
 {

tests/test_api_jobs.py

@@ -1,24 +1,37 @@
+import os
 from typing import Iterator
 from unittest.mock import patch
 
 import pytest
 from starlette.testclient import TestClient
 
-from facefusion import metadata, session_manager
+from facefusion import metadata, session_manager, state_manager
+from facefusion.apis import asset_store
 from facefusion.apis.core import create_api
-from facefusion.jobs.job_manager import clear_jobs, count_step_total, create_job, find_job_ids, init_jobs
+from facefusion.download import conditional_download
+from facefusion.jobs.job_manager import clear_jobs, count_step_total, create_job, find_job_ids, init_jobs, move_job_file, set_steps_status
 from facefusion.program import create_program
-from .assert_helper import get_test_jobs_directory
+from .assert_helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory
 
 
 @pytest.fixture(scope = 'module', autouse = True)
 def before_all() -> None:
 	create_program()
 
+	conditional_download(get_test_examples_directory(),
+	[
+		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg',
+		'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4'
+	])
+
 
 @pytest.fixture(scope = 'function', autouse = True)
 def before_each() -> None:
 	session_manager.SESSIONS.clear()
+	asset_store.clear()
+	state_manager.init_item('source_paths', [ get_test_example_file('source.jpg') ])
+	state_manager.init_item('target_path', get_test_example_file('target-240p.mp4'))
+	state_manager.init_item('temp_path', get_test_jobs_directory())
 	clear_jobs(get_test_jobs_directory())
 	init_jobs(get_test_jobs_directory())
 
@@ -166,15 +179,23 @@ def test_submit_jobs(test_client : TestClient) -> None:
 	assert submit_jobs_body.get('message') == 'jobs not submitted'
 	assert submit_jobs_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.submit_jobs', return_value = True):
-		submit_jobs_response = test_client.patch('/jobs?action=submit', headers =
-		{
-			'Authorization': 'Bearer ' + access_token
-		})
-		submit_jobs_body = submit_jobs_response.json()
+	test_client.post('/jobs/job-test-submit-jobs?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
 
-		assert submit_jobs_body.get('message') == 'ok'
-		assert submit_jobs_response.status_code == 200
+	submit_jobs_response = test_client.patch('/jobs?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
+	submit_jobs_body = submit_jobs_response.json()
+
+	assert submit_jobs_body.get('message') == 'ok'
+	assert find_job_ids('queued') == [ 'job-test-submit-jobs' ]
+	assert submit_jobs_response.status_code == 200
 
 
 def test_submit_job(test_client : TestClient) -> None:
@@ -209,15 +230,23 @@ def test_submit_job(test_client : TestClient) -> None:
 	assert submit_job_body.get('message') == 'job not submitted'
 	assert submit_job_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.submit_job', return_value = True):
-		submit_job_response = test_client.patch('/jobs/job-test-submit-job?action=submit', headers =
-		{
-			'Authorization': 'Bearer ' + access_token
-		})
-		submit_job_body = submit_job_response.json()
+	test_client.post('/jobs/job-test-submit-job?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
 
-		assert submit_job_body.get('message') == 'ok'
-		assert submit_job_response.status_code == 200
+	submit_job_response = test_client.patch('/jobs/job-test-submit-job?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
+	submit_job_body = submit_job_response.json()
+
+	assert submit_job_body.get('message') == 'ok'
+	assert find_job_ids('queued') == [ 'job-test-submit-job' ]
+	assert submit_job_response.status_code == 200
 
 
 def test_run_jobs(test_client : TestClient) -> None:
@@ -241,17 +270,29 @@ def test_run_jobs(test_client : TestClient) -> None:
 	assert run_jobs_body.get('message') == 'jobs not run'
 	assert run_jobs_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.find_job_ids', return_value = [ 'job-test-run-jobs' ]):
-		with patch('facefusion.jobs.job_runner.run_jobs', return_value = True) as run_jobs_mock:
-			run_jobs_response = test_client.patch('/jobs?action=run', headers =
-			{
-				'Authorization': 'Bearer ' + access_token
-			})
-			run_jobs_body = run_jobs_response.json()
+	create_job('job-test-run-jobs')
+	test_client.post('/jobs/job-test-run-jobs?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
+	test_client.patch('/jobs/job-test-run-jobs?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
 
-			assert run_jobs_body.get('message') == 'ok'
-			assert run_jobs_response.status_code == 202
-			assert run_jobs_mock.called is True
+	with patch('facefusion.jobs.job_runner.run_jobs', return_value = True) as run_jobs_mock:
+		run_jobs_response = test_client.patch('/jobs?action=run', headers =
+		{
+			'Authorization': 'Bearer ' + access_token
+		})
+		run_jobs_body = run_jobs_response.json()
+
+		assert run_jobs_body.get('message') == 'ok'
+		assert run_jobs_response.status_code == 202
+		assert run_jobs_mock.called is True
 
 
 def test_run_job(test_client : TestClient) -> None:
@@ -277,17 +318,28 @@ def test_run_job(test_client : TestClient) -> None:
 	assert run_job_body.get('message') == 'job not run'
 	assert run_job_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.find_job_ids', return_value = [ 'job-test-run-job' ]):
-		with patch('facefusion.jobs.job_runner.run_job', return_value = True) as run_job_mock:
-			run_job_response = test_client.patch('/jobs/job-test-run-job?action=run', headers =
-			{
-				'Authorization': 'Bearer ' + access_token
-			})
-			run_job_body = run_job_response.json()
+	test_client.post('/jobs/job-test-run-job?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
+	test_client.patch('/jobs/job-test-run-job?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
 
-			assert run_job_body.get('message') == 'ok'
-			assert run_job_response.status_code == 202
-			assert run_job_mock.called is True
+	with patch('facefusion.jobs.job_runner.run_job', return_value = True) as run_job_mock:
+		run_job_response = test_client.patch('/jobs/job-test-run-job?action=run', headers =
+		{
+			'Authorization': 'Bearer ' + access_token
+		})
+		run_job_body = run_job_response.json()
+
+		assert run_job_body.get('message') == 'ok'
+		assert run_job_response.status_code == 202
+		assert run_job_mock.called is True
 
 
 def test_retry_jobs(test_client : TestClient) -> None:
@@ -311,17 +363,31 @@ def test_retry_jobs(test_client : TestClient) -> None:
 	assert retry_jobs_body.get('message') == 'jobs not retried'
 	assert retry_jobs_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.find_job_ids', return_value = [ 'job-test-retry-jobs' ]):
-		with patch('facefusion.jobs.job_runner.retry_jobs', return_value = True) as retry_jobs_mock:
-			retry_jobs_response = test_client.patch('/jobs?action=retry', headers =
-			{
-				'Authorization': 'Bearer ' + access_token
-			})
-			retry_jobs_body = retry_jobs_response.json()
+	create_job('job-test-retry-jobs')
+	test_client.post('/jobs/job-test-retry-jobs?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
+	test_client.patch('/jobs/job-test-retry-jobs?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
+	set_steps_status('job-test-retry-jobs', 'failed')
+	move_job_file('job-test-retry-jobs', 'failed')
 
-			assert retry_jobs_body.get('message') == 'ok'
-			assert retry_jobs_response.status_code == 202
-			assert retry_jobs_mock.called is True
+	with patch('facefusion.jobs.job_runner.retry_jobs', return_value = True) as retry_jobs_mock:
+		retry_jobs_response = test_client.patch('/jobs?action=retry', headers =
+		{
+			'Authorization': 'Bearer ' + access_token
+		})
+		retry_jobs_body = retry_jobs_response.json()
+
+		assert retry_jobs_body.get('message') == 'ok'
+		assert retry_jobs_response.status_code == 202
+		assert retry_jobs_mock.called is True
 
 
 def test_retry_job(test_client : TestClient) -> None:
@@ -345,17 +411,31 @@ def test_retry_job(test_client : TestClient) -> None:
 	assert retry_job_body.get('message') == 'job not retried'
 	assert retry_job_response.status_code == 400
 
-	with patch('facefusion.jobs.job_manager.find_job_ids', return_value = [ 'job-test-retry-job' ]):
-		with patch('facefusion.jobs.job_runner.retry_job', return_value = True) as retry_job_mock:
-			retry_job_response = test_client.patch('/jobs/job-test-retry-job?action=retry', headers =
-			{
-				'Authorization': 'Bearer ' + access_token
-			})
-			retry_job_body = retry_job_response.json()
+	create_job('job-test-retry-job')
+	test_client.post('/jobs/job-test-retry-job?action=add', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, json =
+	{
+		'processors': [ 'face_swapper' ]
+	})
+	test_client.patch('/jobs/job-test-retry-job?action=submit', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
+	set_steps_status('job-test-retry-job', 'failed')
+	move_job_file('job-test-retry-job', 'failed')
 
-			assert retry_job_body.get('message') == 'ok'
-			assert retry_job_response.status_code == 202
-			assert retry_job_mock.called is True
+	with patch('facefusion.jobs.job_runner.retry_job', return_value = True) as retry_job_mock:
+		retry_job_response = test_client.patch('/jobs/job-test-retry-job?action=retry', headers =
+		{
+			'Authorization': 'Bearer ' + access_token
+		})
+		retry_job_body = retry_job_response.json()
+
+		assert retry_job_body.get('message') == 'ok'
+		assert retry_job_response.status_code == 202
+		assert retry_job_mock.called is True
 
 
 def test_delete_jobs(test_client : TestClient) -> None:
@@ -445,6 +525,7 @@ def test_create_step(test_client : TestClient) -> None:
 		'Authorization': 'Bearer ' + access_token
 	}, json =
 	{
+		'target_path': 'invalid.mp4',
 		'processors': [ 'face_swapper' ]
 	})
 	create_step_body = create_step_response.json()
@@ -458,8 +539,17 @@ def test_create_step(test_client : TestClient) -> None:
 		'Authorization': 'Bearer ' + access_token
 	})
 	get_job_body = get_job_response.json()
+	step_args = get_job_body.get('steps')[0].get('args')
 
-	assert get_job_body.get('steps')[0].get('args') == { 'processors': [ 'face_swapper' ] }
+	access_session_id = session_manager.find_session_id(access_token)
+
+	assert step_args ==\
+	{
+		'processors': [ 'face_swapper' ],
+		'source_paths': [ get_test_example_file('source.jpg') ],
+		'target_path': get_test_example_file('target-240p.mp4'),
+		'output_path': os.path.join(get_test_jobs_directory(), access_session_id, 'job-test-create-step.mp4')
+	}
 
 	create_step_response = test_client.post('/jobs/job-test-create-step/0?action=insert', headers =
 	{