diff --git a/facefusion/args.py b/facefusion/args.py index 98c67f43..8ceab01c 100644 --- a/facefusion/args.py +++ b/facefusion/args.py @@ -105,7 +105,7 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('output_video_fps', output_video_fps) apply_state_item('skip_audio', args.get('skip_audio')) # processors - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] apply_state_item('processors', args.get('processors')) for processor_module in get_processors_modules(available_processors): processor_module.apply_args(args, apply_state_item) diff --git a/facefusion/core.py b/facefusion/core.py index 5775ef6a..25eb7d8f 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -133,7 +133,7 @@ def force_download() -> ErrorCode: face_recognizer, voice_extractor ] - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] processor_modules = get_processors_modules(available_processors) for module in common_modules + processor_modules: diff --git a/facefusion/filesystem.py b/facefusion/filesystem.py index 75ffa9f6..09974f93 100644 --- a/facefusion/filesystem.py +++ b/facefusion/filesystem.py @@ -7,6 +7,7 @@ from typing import List, Optional import filetype from facefusion.common_helper import is_windows +from facefusion.typing import File if is_windows(): import ctypes @@ -126,11 +127,23 @@ def create_directory(directory_path : str) -> bool: return False -def list_directory(directory_path : str) -> Optional[List[str]]: +def list_directory(directory_path : str) -> Optional[List[File]]: if is_directory(directory_path): - file_paths = os.listdir(directory_path) - file_paths = [ Path(file_path).stem for file_path in file_paths if not Path(file_path).stem.startswith(('.', '__')) ] - return sorted(file_paths) + file_paths = sorted(os.listdir(directory_path)) + files: List[File] = [] + + for file_path in file_paths: + file_name, file_extension = os.path.splitext(file_path) + + if not file_name.startswith(('.', '__')): + files.append( + { + 'name': file_name, + 'extension': file_extension, + 'path': os.path.join(directory_path, file_path) + }) + + return files return None diff --git a/facefusion/program.py b/facefusion/program.py index d6b4901c..b95705ec 100755 --- a/facefusion/program.py +++ b/facefusion/program.py @@ -169,7 +169,7 @@ def create_output_creation_program() -> ArgumentParser: def create_processors_program() -> ArgumentParser: program = ArgumentParser(add_help = False) - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] group_processors = program.add_argument_group('processors') group_processors.add_argument('--processors', help = wording.get('help.processors').format(choices = ', '.join(available_processors)), default = config.get_str_list('processors.processors', 'face_swapper'), nargs = '+') job_store.register_step_keys([ 'processors' ]) @@ -180,7 +180,7 @@ def create_processors_program() -> ArgumentParser: def create_uis_program() -> ArgumentParser: program = ArgumentParser(add_help = False) - available_ui_layouts = list_directory('facefusion/uis/layouts') + available_ui_layouts = [ file.get('name') for file in list_directory('facefusion/uis/layouts') ] group_uis = program.add_argument_group('uis') group_uis.add_argument('--open-browser', help = wording.get('help.open_browser'), action = 'store_true', default = config.get_bool_value('uis.open_browser')) group_uis.add_argument('--ui-layouts', help = wording.get('help.ui_layouts').format(choices = ', '.join(available_ui_layouts)), default = config.get_str_list('uis.ui_layouts', 'default'), nargs = '+') diff --git a/facefusion/typing.py b/facefusion/typing.py index cd1cec2e..0b9dc089 100755 --- a/facefusion/typing.py +++ b/facefusion/typing.py @@ -167,6 +167,13 @@ DownloadSet = Dict[str, Download] VideoMemoryStrategy = Literal['strict', 'moderate', 'tolerant'] +File = TypedDict('File', +{ + 'name' : str, + 'extension' : str, + 'path': str +}) + AppContext = Literal['cli', 'ui'] InferencePool = Dict[str, InferenceSession] diff --git a/facefusion/uis/components/download.py b/facefusion/uis/components/download.py index 2c34b0a8..f08cbbf8 100644 --- a/facefusion/uis/components/download.py +++ b/facefusion/uis/components/download.py @@ -36,7 +36,7 @@ def update_download_providers(download_providers : List[DownloadProviderKey]) -> face_masker, voice_extractor ] - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] processor_modules = get_processors_modules(available_processors) for module in common_modules + processor_modules: diff --git a/facefusion/uis/components/execution.py b/facefusion/uis/components/execution.py index e0c0e209..961419ee 100644 --- a/facefusion/uis/components/execution.py +++ b/facefusion/uis/components/execution.py @@ -36,7 +36,7 @@ def update_execution_providers(execution_providers : List[ExecutionProviderKey]) face_recognizer, voice_extractor ] - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] processor_modules = get_processors_modules(available_processors) for module in common_modules + processor_modules: diff --git a/facefusion/uis/components/processors.py b/facefusion/uis/components/processors.py index 3e0c6d6a..990006ff 100644 --- a/facefusion/uis/components/processors.py +++ b/facefusion/uis/components/processors.py @@ -39,5 +39,5 @@ def update_processors(processors : List[str]) -> gradio.CheckboxGroup: def sort_processors(processors : List[str]) -> List[str]: - available_processors = list_directory('facefusion/processors/modules') + available_processors = [ file.get('name') for file in list_directory('facefusion/processors/modules') ] return sorted(available_processors, key = lambda processor : processors.index(processor) if processor in processors else len(processors)) diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py index fedac41c..137bbcf6 100644 --- a/tests/test_filesystem.py +++ b/tests/test_filesystem.py @@ -105,8 +105,11 @@ def test_create_directory() -> None: def test_list_directory() -> None: - assert list_directory(get_test_examples_directory()) - assert list_directory(get_test_example_file('source.jpg')) is None + files = list_directory(get_test_examples_directory()) + + for file in files: + assert file.get('path') == get_test_example_file(file.get('name') + file.get('extension')) + assert list_directory('invalid') is None