diff --git a/roop/core.py b/roop/core.py index d1ce7dc..e6c740e 100755 --- a/roop/core.py +++ b/roop/core.py @@ -126,8 +126,8 @@ def update_status(message: str, scope: str = 'ROOP.CORE') -> None: def start() -> None: - for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): - if not frame_processor.pre_start(): + for frame_processor_module in get_frame_processors_modules(roop.globals.frame_processors): + if not frame_processor_module.pre_start(): return # process image to image if has_image_extension(roop.globals.target_path): @@ -135,10 +135,10 @@ def start() -> None: destroy() shutil.copy2(roop.globals.target_path, roop.globals.output_path) # process frame - for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): - update_status('Progressing...', frame_processor.NAME) - frame_processor.process_image(roop.globals.source_path, roop.globals.output_path, roop.globals.output_path) - frame_processor.post_process() + for frame_processor_module in get_frame_processors_modules(roop.globals.frame_processors): + update_status('Progressing...', frame_processor_module.NAME) + frame_processor_module.process_image(roop.globals.source_path, roop.globals.output_path, roop.globals.output_path) + frame_processor_module.post_process() # validate image if is_image(roop.globals.target_path): update_status('Processing to image succeed!') @@ -161,10 +161,10 @@ def start() -> None: # process frame temp_frame_paths = get_temp_frame_paths(roop.globals.target_path) if temp_frame_paths: - for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): - update_status('Progressing...', frame_processor.NAME) - frame_processor.process_video(roop.globals.source_path, temp_frame_paths) - frame_processor.post_process() + for frame_processor_module in get_frame_processors_modules(roop.globals.frame_processors): + update_status('Progressing...', frame_processor_module.NAME) + frame_processor_module.process_video(roop.globals.source_path, temp_frame_paths) + frame_processor_module.post_process() else: update_status('Frames not found...') return diff --git a/roop/processors/frame/core.py b/roop/processors/frame/core.py index 498169d..68a9e94 100644 --- a/roop/processors/frame/core.py +++ b/roop/processors/frame/core.py @@ -45,6 +45,12 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType return FRAME_PROCESSORS_MODULES +def clear_frame_processors_modules() -> None: + global FRAME_PROCESSORS_MODULES + + FRAME_PROCESSORS_MODULES = [] + + def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None: with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor: futures = [] diff --git a/roop/uis/core.py b/roop/uis/core.py index 804a8a0..f269c46 100644 --- a/roop/uis/core.py +++ b/roop/uis/core.py @@ -11,7 +11,7 @@ def init() -> None: with gradio.Blocks(theme=get_theme()) as ui: with gradio.Row(): settings.render() - with gradio.Row(): + with gradio.Column(): source.render() target.render() preview.render() diff --git a/roop/uis/preview.py b/roop/uis/preview.py index b19b0cf..58853ce 100644 --- a/roop/uis/preview.py +++ b/roop/uis/preview.py @@ -9,7 +9,7 @@ from roop.core import destroy from roop.face_analyser import get_one_face from roop.face_reference import get_face_reference, set_face_reference from roop.predictor import predict_frame -from roop.processors.frame.core import get_frame_processors_modules +from roop.processors.frame.core import get_frame_processors_modules, load_frame_processor_module from roop.typing import Frame from roop.uis import core as ui from roop.utilities import is_video, is_image @@ -21,7 +21,6 @@ def render() -> None: with gradio.Box(): preview_image_args: Dict[str, Any] = { 'label': 'preview_image', - 'height': 400, 'visible': False } preview_slider_args: Dict[str, Any] = { @@ -47,10 +46,13 @@ def render() -> None: preview_slider.change(update, inputs=preview_slider, outputs=[preview_image, preview_slider], show_progress=False) source_file = ui.get_component('source_file') target_file = ui.get_component('target_file') + frame_processors_checkbox_group = ui.get_component('frame_processors_checkbox_group') if source_file: source_file.change(update, outputs=[preview_image, preview_slider]) if target_file: target_file.change(update, outputs=[preview_image, preview_slider]) + if frame_processors_checkbox_group: + frame_processors_checkbox_group.change(update, inputs=preview_slider, outputs=[preview_image, preview_slider]) def update(frame_number: int = 0) -> Tuple[Dict[Any, Any], Dict[Any, Any]]: @@ -76,9 +78,9 @@ def get_preview_frame(temp_frame: Frame) -> Frame: reference_face = get_one_face(reference_frame, roop.globals.reference_face_position) set_face_reference(reference_face) reference_face = get_face_reference() if not roop.globals.many_faces else None - for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): - if frame_processor.pre_start(): - temp_frame = frame_processor.process_frame( + for frame_processor_module in get_frame_processors_modules(roop.globals.frame_processors): + if frame_processor_module.pre_start(): + temp_frame = frame_processor_module.process_frame( source_face, reference_face, temp_frame diff --git a/roop/uis/settings.py b/roop/uis/settings.py index 9f22c93..f6ac3f9 100644 --- a/roop/uis/settings.py +++ b/roop/uis/settings.py @@ -3,28 +3,45 @@ import gradio import onnxruntime import roop.globals +from roop.processors.frame.core import clear_frame_processors_modules +from roop.uis import core as ui NAME = 'ROOP.UIS.OUTPUT' def render() -> None: - with gradio.Row(): - execution_providers_checkbox_group = gradio.CheckboxGroup( - label='execution_providers', - choices=onnxruntime.get_available_providers(), - value=roop.globals.execution_providers - ) - execution_threads_slider = gradio.Slider( - label='execution_threads', - value=roop.globals.execution_threads, - step=1, - minimum=1, - maximum=64 - ) + with gradio.Column(): + with gradio.Box(): + frame_processors_checkbox_group = gradio.CheckboxGroup( + label='frame_processors', + choices=['face_swapper', 'face_enhancer'], + value=roop.globals.frame_processors + ) + ui.register_component('frame_processors_checkbox_group', frame_processors_checkbox_group) + with gradio.Box(): + execution_providers_checkbox_group = gradio.CheckboxGroup( + label='execution_providers', + choices=onnxruntime.get_available_providers(), + value=roop.globals.execution_providers + ) + execution_threads_slider = gradio.Slider( + label='execution_threads', + value=roop.globals.execution_threads, + step=1, + minimum=1, + maximum=64 + ) + frame_processors_checkbox_group.change(update_frame_processors, inputs=frame_processors_checkbox_group, outputs=frame_processors_checkbox_group) execution_providers_checkbox_group.change(update_execution_providers, inputs=execution_providers_checkbox_group, outputs=execution_providers_checkbox_group) execution_threads_slider.change(update_execution_threads, inputs=execution_threads_slider, outputs=execution_threads_slider) +def update_frame_processors(frame_processors: List[str]) -> Dict[Any, Any]: + clear_frame_processors_modules() + roop.globals.frame_processors = frame_processors + return gradio.update(value=frame_processors) + + def update_execution_providers(execution_providers: List[str]) -> Dict[Any, Any]: roop.globals.execution_providers = execution_providers return gradio.update(value=execution_providers) diff --git a/roop/uis/source.py b/roop/uis/source.py index 90e1fa9..c3c21dc 100644 --- a/roop/uis/source.py +++ b/roop/uis/source.py @@ -10,7 +10,7 @@ NAME = 'ROOP.UIS.SOURCE' def render() -> None: - with gradio.Column(): + with gradio.Box(): is_source_image = is_image(roop.globals.source_path) source_file = gradio.File( file_count='single', @@ -22,7 +22,6 @@ def render() -> None: source_image = gradio.Image( label='source_image', value=source_file.value['name'] if is_source_image else None, - height=200, visible=is_source_image ) source_file.change(update, inputs=source_file, outputs=source_image) diff --git a/roop/uis/target.py b/roop/uis/target.py index e2ba114..4632e1f 100644 --- a/roop/uis/target.py +++ b/roop/uis/target.py @@ -10,7 +10,7 @@ NAME = 'ROOP.UIS.TARGET' def render() -> None: - with gradio.Column(): + with gradio.Box(): is_target_image = is_image(roop.globals.target_path) is_target_video = is_video(roop.globals.target_path) target_file = gradio.File( @@ -23,19 +23,18 @@ def render() -> None: target_image = gradio.Image( label='target_image', value=target_file.value['name'] if is_target_image else None, - height=200, visible=is_target_image ) target_video = gradio.Video( label='target_video', value=target_file.value['name'] if is_target_video else None, - height=200, visible=is_target_video ) target_file.change(update, inputs=target_file, outputs=[target_image, target_video]) def update(file: IO[Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]: + clear_face_reference() if file and is_image(file.name): roop.globals.target_path = file.name return gradio.update(value=file.name, visible=True), gradio.update(value=None, visible=False) @@ -43,5 +42,4 @@ def update(file: IO[Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]: roop.globals.target_path = file.name return gradio.update(value=None, visible=False), gradio.update(value=file.name, visible=True) roop.globals.target_path = None - clear_face_reference() return gradio.update(value=None, visible=False), gradio.update(value=None, visible=False) diff --git a/roop/uis/typing.py b/roop/uis/typing.py index e61500d..ab9b0be 100644 --- a/roop/uis/typing.py +++ b/roop/uis/typing.py @@ -2,4 +2,4 @@ from typing import Literal import gradio Component = gradio.File or gradio.Image or gradio.Video -ComponentName = Literal['source_file', 'target_file'] +ComponentName = Literal['source_file', 'target_file', 'frame_processors_checkbox_group']