From fa813d80258cf051b23d36349811224034152668 Mon Sep 17 00:00:00 2001
From: Somdev Sangwan <s0md3v@gmail.com>
Date: Mon, 29 May 2023 19:04:00 +0530
Subject: [PATCH] handle when user has gpu but won't use it

---
 core/config.py    | 3 ++-
 core/processor.py | 3 ++-
 run.py            | 7 +++++++
 3 files changed, 11 insertions(+), 2 deletions(-)

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