diff --git a/core/config.py b/core/config.py index df1c3ca..bbaad5e 100644 --- a/core/config.py +++ b/core/config.py @@ -1,7 +1,8 @@ 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/processor.py b/core/processor.py index 535cc17..112fd71 100644 --- a/core/processor.py +++ b/core/processor.py @@ -1,10 +1,11 @@ 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=onnxruntime.get_available_providers()) +face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=core.globals.providers) def process_video(source_img, frame_paths): diff --git a/run.py b/run.py index c44eced..aeb16dc 100644 --- a/run.py +++ b/run.py @@ -1,3 +1,10 @@ +import sys +import onnxruntime +import core.globals + +if '--gpu' in sys.argv: + core.globals.providers = onnxruntime.get_available_providers() + import glob import argparse import multiprocessing as mp