Merge branch 'next' of https://github.com/K1llMan/roop into next

roop/core.py

@@ -37,7 +37,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument('-f', '--face', help='use this face', dest='source_img')
 parser.add_argument('-t', '--target', help='replace this face', dest='target_path')
 parser.add_argument('-o', '--output', help='save output to this file', dest='output_file')
-parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False)
+parser.add_argument('--gpu', help='choice your gpu vendor', dest='gpu', choices=['amd', 'nvidia'])
 parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
 parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
 parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', type=int)
@@ -47,7 +47,10 @@ parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_fac
 for name, value in vars(parser.parse_args()).items():
     args[name] = value
 
-if '--all-faces' in sys.argv or '-a' in sys.argv:
+if 'gpu' in args:
+    roop.globals.gpu = args['gpu']
+
+if 'all-faces' in args:
     roop.globals.all_faces = True
 
 sep = "/"
@@ -75,32 +78,31 @@ def pre_check():
     model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx')
     if not os.path.isfile(model_path):
         quit('File "inswapper_128.onnx" does not exist!')
-    if '--gpu' in sys.argv:
-        NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
-        if len(list(set(roop.globals.providers) - set(NVIDIA_PROVIDERS))) == 1:
-            CUDA_VERSION = torch.version.cuda
-            CUDNN_VERSION = torch.backends.cudnn.version()
-            if not torch.cuda.is_available() or not CUDA_VERSION:
-                quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
-            if CUDA_VERSION > '11.8':
-                quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8")
-            if CUDA_VERSION < '11.4':
-                quit(f"CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8")
-            if CUDNN_VERSION < 8220:
-                quit(f"CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1")
-            if CUDNN_VERSION > 8910:
-                quit(f"CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1")
+    if roop.globals.gpu == 'amd':
+        if 'ROCMExecutionProvider' not in roop.globals.providers:
+            quit("You are using --gpu=amd flag but ROCM isn't available or properly installed on your system.")
+    if roop.globals.gpu == 'nvidia':
+        CUDA_VERSION = torch.version.cuda
+        CUDNN_VERSION = torch.backends.cudnn.version()
+        if not torch.cuda.is_available() or not CUDA_VERSION:
+            quit("You are using --gpu=nvidia flag but CUDA isn't available or properly installed on your system.")
+        if CUDA_VERSION > '11.8':
+            quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8")
+        if CUDA_VERSION < '11.4':
+            quit(f"CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8")
+        if CUDNN_VERSION < 8220:
+            quit(f"CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1")
+        if CUDNN_VERSION > 8910:
+            quit(f"CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1")
     else:
         roop.globals.providers = ['CPUExecutionProvider']
-    if '--all-faces' in sys.argv or '-a' in sys.argv:
-        roop.globals.all_faces = True
 
 
 def start_processing():
     frame_paths = args["frame_paths"]
     n = len(frame_paths) // (args['cores_count'])
     # single thread
-    if args['gpu'] or n < 2:
+    if roop.globals.gpu == 'amd' or roop.globals.gpu == 'nvidia' or n < 2:
         process_video(args['source_img'], args["frame_paths"], preview.update)
         return
     # multithread if total frames to cpu cores ratio is greater than 2
@@ -345,4 +347,4 @@ def run():
     status_label = tk.Label(window, width=580, justify="center", text="Status: waiting for input...", fg="#2ecc71", bg="#2d3436")
     status_label.place(x=10,y=640,width=580,height=30)
 
-    window.mainloop()
+    window.mainloop()

roop/globals.py

@@ -1,7 +1,8 @@
 import onnxruntime
 
-use_gpu = False
+gpu = None
 all_faces = False
+log_level = 'error'
 providers = onnxruntime.get_available_providers()
 
 if 'TensorrtExecutionProvider' in providers:

roop/utils.py

@@ -1,5 +1,6 @@
 import os
 import shutil
+import roop.globals
 
 sep = "/"
 if os.name == "nt":
@@ -29,26 +30,35 @@ def detect_fps(input_path):
     return 30, 30
 
 
+def run_ffmpeg(args):
+
+    log_level = f'-loglevel {roop.globals.log_level}'
+
+    os.system(f'ffmpeg {log_level} {args}')
+
+
 def set_fps(input_path, output_path, fps):
     input_path, output_path = path(input_path), path(output_path)
-    os.system(f'ffmpeg -i "{input_path}" -filter:v fps=fps={fps} "{output_path}" -loglevel error')
+    run_ffmpeg(f'-i "{input_path}" -filter:v fps=fps={fps} "{output_path}"')
 
 
 def create_video(video_name, fps, output_dir):
+    hwaccel_option = '-hwaccel cuda' if roop.globals.gpu == 'nvidia' else ''
     output_dir = path(output_dir)
-    os.system(f'ffmpeg -framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4" -loglevel error')
+    run_ffmpeg(f'{hwaccel_option} -framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"')
 
 
 def extract_frames(input_path, output_dir):
+    hwaccel_option = '-hwaccel cuda' if roop.globals.gpu == 'nvidia' else ''
     input_path, output_dir = path(input_path), path(output_dir)
-    os.system(f'ffmpeg -i "{input_path}" "{output_dir}{sep}%04d.png" -loglevel error')
+    run_ffmpeg(f' {hwaccel_option} -i "{input_path}" "{output_dir}{sep}%04d.png"')
 
 
 def add_audio(output_dir, target_path, video, keep_frames, output_file):
     video_name = os.path.splitext(video)[0]
     save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4"
     save_to_ff, output_dir_ff = path(save_to), path(output_dir)
-    os.system(f'ffmpeg -i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}" -loglevel error')
+    run_ffmpeg(f'-i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"')
     if not os.path.isfile(save_to):
         shutil.move(output_dir + "/output.mp4", save_to)
     if not keep_frames: