Introduce pre_start() and move globals check to pre_check()

roop/core.py

@@ -17,13 +17,11 @@ import torch
 import onnxruntime
 import tensorflow
 from opennsfw2 import predict_video_frames, predict_image
-import cv2
 
 import roop.globals
 import roop.ui as ui
 from roop.processors.frame.core import get_frame_processors_modules
 from roop.utilities import has_image_extension, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clean_temp, normalize_output_path
-from roop.face_analyser import get_one_face
 
 if 'ROCMExecutionProvider' in roop.globals.execution_providers:
     del torch
@@ -34,19 +32,19 @@ warnings.simplefilter(action='ignore', category=FutureWarning)
 def parse_args() -> None:
     signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
     parser = argparse.ArgumentParser()
-    parser.add_argument('-f', '--face', help='use a face image', dest='source_path')
-    parser.add_argument('-t', '--target', help='replace image or video with face', dest='target_path')
-    parser.add_argument('-o', '--output', help='save output to this file', dest='output_path')
+    parser.add_argument('-s', '--source', help='select an source image', dest='source_path')
+    parser.add_argument('-t', '--target', help='select an target image or video', dest='target_path')
+    parser.add_argument('-o', '--output', help='select output file or directory', dest='output_path')
     parser.add_argument('--frame-processor', help='list of frame processors to run', dest='frame_processor', default=['face_swapper'], choices=['face_swapper', 'face_enhancer'], nargs='+')
-    parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
-    parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True)
-    parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
-    parser.add_argument('--many-faces', help='swap every face in the frame', dest='many_faces', action='store_true', default=False)
+    parser.add_argument('--keep-fps', help='keep original fps', dest='keep_fps', action='store_true', default=False)
+    parser.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True)
+    parser.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False)
+    parser.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False)
     parser.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
     parser.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18)
-    parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int, default=suggest_max_memory())
+    parser.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
     parser.add_argument('--execution-provider', help='execution provider', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+')
-    parser.add_argument('--execution-threads', help='number of threads to be use for the GPU', dest='execution_threads', type=int, default=suggest_execution_threads())
+    parser.add_argument('--execution-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
 
     args = parser.parse_known_args()[0]
 
@@ -135,17 +133,10 @@ def update_status(message: str) -> None:
 
 
 def start() -> None:
-    # validate paths
-    if not is_image(roop.globals.source_path):
-        update_status('Select an image for source path.')
-        return
-    elif not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
-        update_status('Select an image or video for target path.')
-        return
-    has_face = get_one_face(cv2.imread(roop.globals.source_path))
-    if not has_face:
-        update_status('No face detected in source path.')
-        return
+    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
+        update_status(f'{frame_processor.NAME} is starting...')
+        frame_processor.pre_start()
+        release_resources()
     # process image to image
     if has_image_extension(roop.globals.target_path):
         if predict_image(roop.globals.target_path) > 0.85:

roop/processors/frame/core.py

@@ -9,6 +9,7 @@ import roop
 FRAME_PROCESSORS_MODULES = None
 FRAME_PROCESSORS_INTERFACE = [
     'pre_check',
+    'pre_start',
     'process_frame',
     'process_image',
     'process_video'

roop/processors/frame/face_enhancer.py

@@ -9,7 +9,7 @@ from codeformer.basicsr.utils import img2tensor, tensor2img
 
 import roop.globals
 import roop.processors.frame.core
-from roop.utilities import conditional_download, resolve_relative_path
+from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
 
 if 'ROCMExecutionProvider' in roop.globals.execution_providers:
     del torch
@@ -20,6 +20,11 @@ NAME = 'Face Enhancer'
 
 
 def pre_check() -> None:
+    if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
+        quit('Select an image or video for target path.')
+
+
+def pre_start() -> None:
     download_directory_path = resolve_relative_path('../models')
     conditional_download(download_directory_path, ['https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'])
 
@@ -45,14 +50,14 @@ def get_code_former():
 def get_face_enhancer(FACE_ENHANCER):
     if FACE_ENHANCER is None:
         FACE_ENHANCER = FaceRestoreHelper(
-        upscale_factor = int(2),
-        face_size=512,
-        crop_ratio=(1, 1),
-        det_model='retinaface_resnet50',
-        save_ext='png',
-        use_parse=True,
-        device='cuda',
-    )
+            upscale_factor = int(2),
+            face_size=512,
+            crop_ratio=(1, 1),
+            det_model='retinaface_resnet50',
+            save_ext='png',
+            use_parse=True,
+            device='cuda'
+        )
     return FACE_ENHANCER
 
 

roop/processors/frame/face_swapper.py

@@ -6,7 +6,7 @@ import threading
 import roop.globals
 import roop.processors.frame.core
 from roop.face_analyser import get_one_face, get_many_faces
-from roop.utilities import conditional_download, resolve_relative_path
+from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
 
 FACE_SWAPPER = None
 THREAD_LOCK = threading.Lock()
@@ -14,6 +14,15 @@ NAME = 'Face Swapper'
 
 
 def pre_check() -> None:
+    if not is_image(roop.globals.source_path):
+        quit('Select an image for source path.')
+    if not get_one_face(cv2.imread(roop.globals.source_path)):
+        quit('No face in source path detected.')
+    if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
+        quit('Select an image or video for target path.')
+
+
+def pre_start() -> None:
     download_directory_path = resolve_relative_path('../models')
     conditional_download(download_directory_path, ['https://huggingface.co/deepinsight/inswapper/resolve/main/inswapper_128.onnx'])
 

roop/ui.py

@@ -111,7 +111,7 @@ def select_source_path() -> None:
     global RECENT_DIRECTORY_SOURCE
 
     PREVIEW.withdraw()
-    source_path = ctk.filedialog.askopenfilename(title='Select an face image', initialdir=RECENT_DIRECTORY_SOURCE)
+    source_path = ctk.filedialog.askopenfilename(title='select an source image', initialdir=RECENT_DIRECTORY_SOURCE)
     if is_image(source_path):
         roop.globals.source_path = source_path
         RECENT_DIRECTORY_SOURCE = os.path.dirname(roop.globals.source_path)
@@ -128,7 +128,7 @@ def select_target_path() -> None:
     global RECENT_DIRECTORY_TARGET
 
     PREVIEW.withdraw()
-    target_path = ctk.filedialog.askopenfilename(title='Select an image or video target', initialdir=RECENT_DIRECTORY_TARGET)
+    target_path = ctk.filedialog.askopenfilename(title='select an target image or video', initialdir=RECENT_DIRECTORY_TARGET)
     if is_image(target_path):
         roop.globals.target_path = target_path
         RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path)
@@ -151,9 +151,9 @@ def select_output_path(start):
     global RECENT_DIRECTORY_OUTPUT
 
     if is_image(roop.globals.target_path):
-        output_path = ctk.filedialog.asksaveasfilename(title='Save image output', initialfile='output.png', initialdir=RECENT_DIRECTORY_OUTPUT)
+        output_path = ctk.filedialog.asksaveasfilename(title='save image output file', initialfile='output.png', initialdir=RECENT_DIRECTORY_OUTPUT)
     elif is_video(roop.globals.target_path):
-        output_path = ctk.filedialog.asksaveasfilename(title='Save video output', initialfile='output.mp4', initialdir=RECENT_DIRECTORY_OUTPUT)
+        output_path = ctk.filedialog.asksaveasfilename(title='save video output file', initialfile='output.mp4', initialdir=RECENT_DIRECTORY_OUTPUT)
     else:
         output_path = None
     if output_path: