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