From 54f800dd0e1dcec145601b49e7f3931c910f1166 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 30 May 2023 00:40:02 +0200 Subject: [PATCH 1/4] Add GPU support, Quit on missing model, Remove globals (sorry) --- core/config.py | 3 +-- core/globals.py | 4 ---- core/processor.py | 9 +++++++-- core/utils.py | 2 ++ requirements.txt | 3 ++- run.py | 10 +++++----- 6 files changed, 17 insertions(+), 14 deletions(-) delete mode 100644 core/globals.py mode change 100644 => 100755 run.py diff --git a/core/config.py b/core/config.py index bbaad5e..df1c3ca 100644 --- a/core/config.py +++ b/core/config.py @@ -1,8 +1,7 @@ import insightface import onnxruntime -import core.globals -face_analyser = insightface.app.FaceAnalysis(name='buffalo_l', providers=core.globals.providers) +face_analyser = insightface.app.FaceAnalysis(name='buffalo_l', providers=onnxruntime.get_available_providers()) face_analyser.prepare(ctx_id=0, det_size=(640, 640)) diff --git a/core/globals.py b/core/globals.py deleted file mode 100644 index 83d7277..0000000 --- a/core/globals.py +++ /dev/null @@ -1,4 +0,0 @@ -import onnxruntime - -use_gpu = False -providers = onnxruntime.get_available_providers() diff --git a/core/processor.py b/core/processor.py index eb75975..29dc082 100644 --- a/core/processor.py +++ b/core/processor.py @@ -1,11 +1,15 @@ +import os import cv2 import insightface import onnxruntime -import core.globals from core.config import get_face from core.utils import rreplace -face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=core.globals.providers) +if os.path.isfile('inswapper_128.onnx'): + face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=onnxruntime.get_available_providers()) +else: + quit('File "inswapper_128.onnx" does not exist!') + def process_video(source_img, frame_paths): @@ -25,6 +29,7 @@ def process_video(source_img, frame_paths): pass print(flush=True) + def process_img(source_img, target_path): frame = cv2.imread(target_path) face = get_face(frame) diff --git a/core/utils.py b/core/utils.py index 9c5cd75..228d83b 100644 --- a/core/utils.py +++ b/core/utils.py @@ -17,6 +17,7 @@ def run_command(command, mode="silent"): return os.system(command) return os.popen(command).read() + def detect_fps(input_path): input_path = path(input_path) output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read() @@ -58,6 +59,7 @@ def add_audio(output_dir, target_path, keep_frames, output_file): def is_img(path): return path.lower().endswith(("png", "jpg", "jpeg", "bmp")) + def rreplace(s, old, new, occurrence): li = s.rsplit(old, occurrence) return new.join(li) diff --git a/requirements.txt b/requirements.txt index b4d8baa..96153a2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ numpy==1.24.3 opencv-python==4.7.0.72 onnx==1.14.0 -onnxruntime==1.15.0 +onnxruntime-gpu==1.15.0 insightface==0.7.3 psutil==5.9.5 tk==0.1.0 +torch==2.0.1 \ No newline at end of file diff --git a/run.py b/run.py old mode 100644 new mode 100755 index 45badee..a47916d --- a/run.py +++ b/run.py @@ -1,12 +1,11 @@ -import sys +#!/usr/bin/env python3 + +import torch import shutil -import core.globals if not shutil.which('ffmpeg'): print('ffmpeg is not installed. Read the docs: https://github.com/s0md3v/roop#installation.\n' * 10) quit() -if '--gpu' not in sys.argv: - core.globals.providers = ['CPUExecutionProvider'] import glob import argparse @@ -31,7 +30,8 @@ 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('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) -parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) +if torch.cuda.is_available(): + parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) for name, value in vars(parser.parse_args()).items(): From dc0653ad39d8f6e928651f6983e277c26834baf8 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 30 May 2023 01:03:43 +0200 Subject: [PATCH 2/4] Restore globals, add process time for better comparison --- core/config.py | 4 ++-- core/globals.py | 4 ++++ core/processor.py | 10 +++++++--- run.py | 5 ++++- 4 files changed, 17 insertions(+), 6 deletions(-) create mode 100644 core/globals.py diff --git a/core/config.py b/core/config.py index df1c3ca..151ad33 100644 --- a/core/config.py +++ b/core/config.py @@ -1,7 +1,7 @@ import insightface -import onnxruntime +import core.globals -face_analyser = insightface.app.FaceAnalysis(name='buffalo_l', providers=onnxruntime.get_available_providers()) +face_analyser = insightface.app.FaceAnalysis(name='buffalo_l', providers=core.globals.providers) face_analyser.prepare(ctx_id=0, det_size=(640, 640)) diff --git a/core/globals.py b/core/globals.py new file mode 100644 index 0000000..cbd26c2 --- /dev/null +++ b/core/globals.py @@ -0,0 +1,4 @@ +import onnxruntime + +use_gpu = False +providers = onnxruntime.get_available_providers() diff --git a/core/processor.py b/core/processor.py index 29dc082..631feab 100644 --- a/core/processor.py +++ b/core/processor.py @@ -1,18 +1,20 @@ import os +import time + import cv2 import insightface -import onnxruntime +import core.globals from core.config import get_face from core.utils import rreplace if os.path.isfile('inswapper_128.onnx'): - face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=onnxruntime.get_available_providers()) + face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=core.globals.providers) else: quit('File "inswapper_128.onnx" does not exist!') - def process_video(source_img, frame_paths): + start_time = time.time() source_face = get_face(cv2.imread(source_img)) for frame_path in frame_paths: frame = cv2.imread(frame_path) @@ -28,6 +30,8 @@ def process_video(source_img, frame_paths): print('E', end='', flush=True) pass print(flush=True) + end_time = time.time() + print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) def process_img(source_img, target_path): diff --git a/run.py b/run.py index a47916d..7e3c92a 100755 --- a/run.py +++ b/run.py @@ -1,11 +1,14 @@ #!/usr/bin/env python3 - +import sys import torch import shutil +import core.globals if not shutil.which('ffmpeg'): print('ffmpeg is not installed. Read the docs: https://github.com/s0md3v/roop#installation.\n' * 10) quit() +if '--gpu' not in sys.argv: + core.globals.providers = ['CPUExecutionProvider'] import glob import argparse From 2b14613a467dcbaca421a3df3de0081c935112b7 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 30 May 2023 01:28:16 +0200 Subject: [PATCH 3/4] Move process time to the correct place --- core/processor.py | 6 ------ run.py | 8 ++++++++ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/core/processor.py b/core/processor.py index 631feab..d67b1cf 100644 --- a/core/processor.py +++ b/core/processor.py @@ -1,6 +1,4 @@ import os -import time - import cv2 import insightface import core.globals @@ -14,7 +12,6 @@ else: def process_video(source_img, frame_paths): - start_time = time.time() source_face = get_face(cv2.imread(source_img)) for frame_path in frame_paths: frame = cv2.imread(frame_path) @@ -29,9 +26,6 @@ def process_video(source_img, frame_paths): except Exception as e: print('E', end='', flush=True) pass - print(flush=True) - end_time = time.time() - print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) def process_img(source_img, target_path): diff --git a/run.py b/run.py index 7e3c92a..7390c9e 100755 --- a/run.py +++ b/run.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import sys +import time import torch import shutil import core.globals @@ -47,8 +48,12 @@ if os.name == "nt": def start_processing(): + start_time = time.time() if args['gpu']: process_video(args['source_img'], args["frame_paths"]) + print(flush=True) + end_time = time.time() + print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) return frame_paths = args["frame_paths"] n = len(frame_paths)//(psutil.cpu_count()-1) @@ -60,6 +65,9 @@ def start_processing(): p.get() pool.close() pool.join() + print(flush=True) + end_time = time.time() + print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) def select_face(): From 0f782ff516ad2db61816787c6291d3d104f75747 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Tue, 30 May 2023 01:28:42 +0200 Subject: [PATCH 4/4] Move process time to the correct place --- run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/run.py b/run.py index 7390c9e..d2075dd 100755 --- a/run.py +++ b/run.py @@ -51,8 +51,8 @@ def start_processing(): start_time = time.time() if args['gpu']: process_video(args['source_img'], args["frame_paths"]) - print(flush=True) end_time = time.time() + print(flush=True) print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True) return frame_paths = args["frame_paths"] @@ -65,8 +65,8 @@ def start_processing(): p.get() pool.close() pool.join() - print(flush=True) end_time = time.time() + print(flush=True) print(f"Processing time: {end_time - start_time:.2f} seconds", flush=True)