diff --git a/facefusion/apis/jobs_helper.py b/facefusion/apis/jobs_helper.py index 5420a11a..03f61649 100644 --- a/facefusion/apis/jobs_helper.py +++ b/facefusion/apis/jobs_helper.py @@ -1,9 +1,10 @@ from facefusion.apis import asset_store from facefusion.filesystem import is_file from facefusion.jobs import job_manager +from facefusion.types import JobId -def capture_output_asset(job_id : str) -> None: +def capture_output_asset(job_id : JobId) -> None: job = job_manager.read_job_file(job_id) if job and job.get('steps'): diff --git a/facefusion/core.py b/facefusion/core.py index 091ddcf7..44978265 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -19,7 +19,7 @@ from facefusion.jobs.job_list import compose_job_list from facefusion.processors.core import get_processors_modules from facefusion.program import create_program from facefusion.program_helper import validate_args -from facefusion.types import Args, ErrorCode +from facefusion.types import Args, ErrorCode, JobId, StepIndex from facefusion.workflows.core import detect_workflow_mode from facefusion.workflows.modules import audio_to_image, audio_to_image_as_frames, image_to_image, image_to_video, image_to_video_as_frames @@ -319,7 +319,7 @@ def process_batch(args : Args) -> ErrorCode: return 1 -def process_step(job_id : str, step_index : int, step_args : Args) -> bool: +def process_step(job_id : JobId, step_index : StepIndex, step_args : Args) -> bool: step_total = job_manager.count_step_total(job_id) fork_id = session_manager.fork_session() session_manager.observe_cli_session(fork_id) diff --git a/facefusion/jobs/job_helper.py b/facefusion/jobs/job_helper.py index 8e3e7688..9e0856b4 100644 --- a/facefusion/jobs/job_helper.py +++ b/facefusion/jobs/job_helper.py @@ -4,9 +4,10 @@ from datetime import datetime from typing import Optional from facefusion.filesystem import get_file_extension, get_file_name +from facefusion.types import JobId, StepIndex -def get_step_output_path(job_id : str, step_index : int, output_path : str) -> Optional[str]: +def get_step_output_path(job_id : JobId, step_index : StepIndex, output_path : str) -> Optional[str]: if output_path: output_directory_path, output_file_path = os.path.split(output_path) output_file_name = get_file_name(output_file_path) @@ -21,5 +22,5 @@ def get_step_output_path(job_id : str, step_index : int, output_path : str) -> O return None -def suggest_job_id(job_prefix : str = 'job') -> str: +def suggest_job_id(job_prefix : str = 'job') -> JobId: return job_prefix + '-' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S') + '-' + secrets.token_hex(4) diff --git a/facefusion/jobs/job_manager.py b/facefusion/jobs/job_manager.py index 681b3755..7cb6f057 100644 --- a/facefusion/jobs/job_manager.py +++ b/facefusion/jobs/job_manager.py @@ -9,7 +9,7 @@ from facefusion.jobs.job_helper import get_step_output_path from facefusion.json import read_json, write_json from facefusion.sanitizer import sanitize_job_id from facefusion.time_helper import get_current_date_time -from facefusion.types import Args, Content, Job, JobSet, JobStatus, JobStep, JobStepStatus +from facefusion.types import Args, Content, Job, JobId, JobSet, JobStatus, Step, StepIndex, StepStatus def init_jobs(jobs_path : str) -> bool: @@ -24,7 +24,7 @@ def clear_jobs(jobs_path : str) -> bool: return remove_directory(jobs_path) -def create_job(job_id : str) -> bool: +def create_job(job_id : JobId) -> bool: job : Job =\ { 'version': '1', @@ -36,7 +36,7 @@ def create_job(job_id : str) -> bool: return create_job_file(job_id, job) -def submit_job(job_id : str) -> bool: +def submit_job(job_id : JobId) -> bool: drafted_job_ids = find_job_ids('drafted') steps = get_steps(job_id) @@ -59,7 +59,7 @@ def submit_jobs(halt_on_error : bool) -> bool: return False -def delete_job(job_id : str) -> bool: +def delete_job(job_id : JobId) -> bool: return delete_job_file(job_id) @@ -86,7 +86,7 @@ def find_jobs(job_status : JobStatus) -> JobSet: return job_set -def find_job_ids(job_status : JobStatus) -> List[str]: +def find_job_ids(job_status : JobStatus) -> List[JobId]: jobs_path = state_manager.get_jobs_path() job_file_pattern = os.path.join(jobs_path, job_status, '*.json') job_file_paths = resolve_file_pattern(job_file_pattern) @@ -99,17 +99,17 @@ def find_job_ids(job_status : JobStatus) -> List[str]: return job_ids -def validate_job(job_id : str) -> bool: +def validate_job(job_id : JobId) -> bool: job = read_job_file(job_id) return bool(job and 'version' in job and 'date_created' in job and 'date_updated' in job and 'steps' in job) -def has_step(job_id : str, step_index : int) -> bool: +def has_step(job_id : JobId, step_index : StepIndex) -> bool: step_total = count_step_total(job_id) return step_index in range(step_total) -def add_step(job_id : str, step_args : Args) -> bool: +def add_step(job_id : JobId, step_args : Args) -> bool: job = read_job_file(job_id) if job: @@ -122,7 +122,7 @@ def add_step(job_id : str, step_args : Args) -> bool: return False -def remix_step(job_id : str, step_index : int, step_args : Args) -> bool: +def remix_step(job_id : JobId, step_index : StepIndex, step_args : Args) -> bool: steps = get_steps(job_id) step_args = copy(step_args) @@ -136,7 +136,7 @@ def remix_step(job_id : str, step_index : int, step_args : Args) -> bool: return False -def insert_step(job_id : str, step_index : int, step_args : Args) -> bool: +def insert_step(job_id : JobId, step_index : StepIndex, step_args : Args) -> bool: job = read_job_file(job_id) step_args = copy(step_args) @@ -153,7 +153,7 @@ def insert_step(job_id : str, step_index : int, step_args : Args) -> bool: return False -def remove_step(job_id : str, step_index : int) -> bool: +def remove_step(job_id : JobId, step_index : StepIndex) -> bool: job = read_job_file(job_id) if step_index and step_index < 0: @@ -165,7 +165,7 @@ def remove_step(job_id : str, step_index : int) -> bool: return False -def get_steps(job_id : str) -> List[JobStep]: +def get_steps(job_id : JobId) -> List[Step]: job = read_job_file(job_id) if job: @@ -173,7 +173,7 @@ def get_steps(job_id : str) -> List[JobStep]: return [] -def count_step_total(job_id : str) -> int: +def count_step_total(job_id : JobId) -> int: steps = get_steps(job_id) if steps: @@ -181,7 +181,7 @@ def count_step_total(job_id : str) -> int: return 0 -def set_step_status(job_id : str, step_index : int, step_status : JobStepStatus) -> bool: +def set_step_status(job_id : JobId, step_index : StepIndex, step_status : StepStatus) -> bool: job = read_job_file(job_id) if job: @@ -192,7 +192,7 @@ def set_step_status(job_id : str, step_index : int, step_status : JobStepStatus) return False -def set_steps_status(job_id : str, step_status : JobStepStatus) -> bool: +def set_steps_status(job_id : JobId, step_status : StepStatus) -> bool: job = read_job_file(job_id) if job: @@ -202,12 +202,12 @@ def set_steps_status(job_id : str, step_status : JobStepStatus) -> bool: return False -def read_job_file(job_id : str) -> Optional[Job]: +def read_job_file(job_id : JobId) -> Optional[Job]: job_path = find_job_path(job_id) return cast(Job, read_json(job_path)) -def create_job_file(job_id : str, job : Job) -> bool: +def create_job_file(job_id : JobId, job : Job) -> bool: job_path = find_job_path(job_id) if not is_file(job_path): @@ -216,7 +216,7 @@ def create_job_file(job_id : str, job : Job) -> bool: return False -def update_job_file(job_id : str, job : Job) -> bool: +def update_job_file(job_id : JobId, job : Job) -> bool: job_path = find_job_path(job_id) if is_file(job_path): @@ -225,18 +225,18 @@ def update_job_file(job_id : str, job : Job) -> bool: return False -def move_job_file(job_id : str, job_status : JobStatus) -> bool: +def move_job_file(job_id : JobId, job_status : JobStatus) -> bool: job_path = find_job_path(job_id) job_move_path = suggest_job_path(job_id, job_status) return move_file(job_path, job_move_path) -def delete_job_file(job_id : str) -> bool: +def delete_job_file(job_id : JobId) -> bool: job_path = find_job_path(job_id) return remove_file(job_path) -def suggest_job_path(job_id : str, job_status : JobStatus) -> Optional[str]: +def suggest_job_path(job_id : JobId, job_status : JobStatus) -> Optional[str]: jobs_path = state_manager.get_jobs_path() job_file_name = get_job_file_name(job_id) @@ -245,7 +245,7 @@ def suggest_job_path(job_id : str, job_status : JobStatus) -> Optional[str]: return None -def find_job_path(job_id : str) -> Optional[str]: +def find_job_path(job_id : JobId) -> Optional[str]: jobs_path = state_manager.get_jobs_path() job_file_name = get_job_file_name(job_id) @@ -259,7 +259,7 @@ def find_job_path(job_id : str) -> Optional[str]: return None -def get_job_file_name(job_id : str) -> Optional[str]: +def get_job_file_name(job_id : JobId) -> Optional[str]: if job_id: job_id = sanitize_job_id(job_id) return job_id + '.json' diff --git a/facefusion/jobs/job_runner.py b/facefusion/jobs/job_runner.py index 8a05d38a..3a7caf27 100644 --- a/facefusion/jobs/job_runner.py +++ b/facefusion/jobs/job_runner.py @@ -3,10 +3,10 @@ import os from facefusion.ffmpeg import concat_video from facefusion.filesystem import are_images, are_videos, copy_file, create_directory, is_directory, is_file, move_directory, move_file, remove_directory, remove_file, resolve_file_paths from facefusion.jobs import job_helper, job_manager -from facefusion.types import JobOutputSet, JobStep, ProcessStep +from facefusion.types import JobId, JobOutputSet, ProcessStep, Step, StepIndex -def run_job(job_id : str, process_step : ProcessStep) -> bool: +def run_job(job_id : JobId, process_step : ProcessStep) -> bool: queued_job_ids = job_manager.find_job_ids('queued') if job_id in queued_job_ids: @@ -32,7 +32,7 @@ def run_jobs(process_step : ProcessStep, halt_on_error : bool) -> bool: return False -def retry_job(job_id : str, process_step : ProcessStep) -> bool: +def retry_job(job_id : JobId, process_step : ProcessStep) -> bool: failed_job_ids = job_manager.find_job_ids('failed') if job_id in failed_job_ids: @@ -54,7 +54,7 @@ def retry_jobs(process_step : ProcessStep, halt_on_error : bool) -> bool: return False -def run_step(job_id : str, step_index : int, step : JobStep, process_step : ProcessStep) -> bool: +def run_step(job_id : JobId, step_index : StepIndex, step : Step, process_step : ProcessStep) -> bool: step_args = step.get('args') if job_manager.set_step_status(job_id, step_index, 'started') and process_step(job_id, step_index, step_args): @@ -68,7 +68,7 @@ def run_step(job_id : str, step_index : int, step : JobStep, process_step : Proc return False -def run_steps(job_id : str, process_step : ProcessStep) -> bool: +def run_steps(job_id : JobId, process_step : ProcessStep) -> bool: steps = job_manager.get_steps(job_id) if steps: @@ -79,7 +79,7 @@ def run_steps(job_id : str, process_step : ProcessStep) -> bool: return False -def finalize_steps(job_id : str) -> bool: +def finalize_steps(job_id : JobId) -> bool: output_set = collect_output_set(job_id) for output_path, temp_output_paths in output_set.items(): @@ -106,7 +106,7 @@ def finalize_steps(job_id : str) -> bool: return True -def clean_steps(job_id: str) -> bool: +def clean_steps(job_id : JobId) -> bool: output_set = collect_output_set(job_id) for temp_output_paths in output_set.values(): @@ -120,7 +120,7 @@ def clean_steps(job_id: str) -> bool: return True -def collect_output_set(job_id : str) -> JobOutputSet: +def collect_output_set(job_id : JobId) -> JobOutputSet: steps = job_manager.get_steps(job_id) job_output_set : JobOutputSet = {} diff --git a/facefusion/sanitizer.py b/facefusion/sanitizer.py index 3f85e532..5d34c92d 100644 --- a/facefusion/sanitizer.py +++ b/facefusion/sanitizer.py @@ -2,9 +2,10 @@ import hashlib from typing import Any, Sequence from facefusion.common_helper import cast_int +from facefusion.types import JobId -def sanitize_job_id(job_id : str) -> str: +def sanitize_job_id(job_id : JobId) -> JobId: __job_id__ = job_id.replace('-', '') if __job_id__.isalnum(): diff --git a/facefusion/types.py b/facefusion/types.py index d924ef06..436b635a 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -152,7 +152,6 @@ CapabilityStore = TypedDict('CapabilityStore', CapabilityGroup : TypeAlias = Dict[Group, CapabilitySet] ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] -ProcessStep : TypeAlias = Callable[[str, int, Args], bool] Content : TypeAlias = Dict[str, Any] @@ -517,22 +516,27 @@ ApiSecurityStrategy = Literal['strict', 'moderate'] InferencePool : TypeAlias = Dict[str, InferenceSession] -JobOutputSet : TypeAlias = Dict[str, List[str]] -JobStatus = Literal['drafted', 'queued', 'completed', 'failed'] -JobStepStatus = Literal['drafted', 'queued', 'started', 'completed', 'failed'] -JobStep = TypedDict('JobStep', +StepIndex : TypeAlias = int +StepStatus = Literal['drafted', 'queued', 'started', 'completed', 'failed'] +Step = TypedDict('Step', { 'args' : Args, - 'status' : JobStepStatus + 'status' : StepStatus }) + +JobId : TypeAlias = str +JobStatus = Literal['drafted', 'queued', 'completed', 'failed'] Job = TypedDict('Job', { 'version' : str, 'date_created' : str, 'date_updated' : Optional[str], - 'steps' : List[JobStep] + 'steps' : List[Step] }) -JobSet : TypeAlias = Dict[str, Job] +JobSet : TypeAlias = Dict[JobId, Job] + +JobOutputSet : TypeAlias = Dict[str, List[str]] +ProcessStep : TypeAlias = Callable[[JobId, StepIndex, Args], bool] StateValue : TypeAlias = Any StateType : TypeAlias = Any @@ -684,9 +688,9 @@ State = TypedDict('State', 'api_key' : str, 'api_security_strategy' : ApiSecurityStrategy, 'api_session_limit' : int, - 'job_id' : str, + 'job_id' : JobId, 'job_status' : JobStatus, - 'step_index' : int + 'step_index' : StepIndex }) ApplyStateItem : TypeAlias = Callable[[StateKey, StateValue], None] diff --git a/tests/jobs/test_job_runner.py b/tests/jobs/test_job_runner.py index afb767c5..fbeb7b22 100644 --- a/tests/jobs/test_job_runner.py +++ b/tests/jobs/test_job_runner.py @@ -7,7 +7,7 @@ from facefusion.download import conditional_download from facefusion.filesystem import copy_file, create_directory, get_file_extension from facefusion.jobs.job_manager import add_step, clear_jobs, create_job, init_jobs, move_job_file, submit_job, submit_jobs from facefusion.jobs.job_runner import clean_steps, collect_output_set, finalize_steps, retry_job, retry_jobs, run_job, run_jobs, run_steps -from facefusion.types import Args +from facefusion.types import Args, JobId, StepIndex from tests.assert_helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory @@ -42,7 +42,7 @@ def before_each() -> None: prepare_test_output_directory() -def process_step(job_id : str, step_index : int, step_args : Args) -> bool: +def process_step(job_id : JobId, step_index : StepIndex, step_args : Args) -> bool: output_path = step_args.get('output_path') target_path = step_args.get('target_path')