v0.5

app.py

@@ -43,7 +43,7 @@ user_args = parser.parse_args()
 ## ------------------------------ DEFAULTS ------------------------------
 
 USE_COLAB = user_args.colab
-USE_CUDA = user_args.cuda
+USE_CUDA = False
 DEF_OUTPUT_PATH = user_args.out_dir
 BATCH_SIZE = int(user_args.batch_size)
 WORKSPACE = None
@@ -63,7 +63,8 @@ MASK_INCLUDE = [
     "Nose",
     "Mouth",
     "L-Lip",
-    "U-Lip"
+    "U-Lip",
+    "Hair"
 ]
 MASK_SOFT_KERNEL = 17
 MASK_SOFT_ITERATIONS = 10
@@ -82,22 +83,9 @@ FACE_ENHANCER_LIST.extend(cv2_interpolations)
 ## ------------------------------ SET EXECUTION PROVIDER ------------------------------
 # Note: Non CUDA users may change settings here
 
-PROVIDER = ["CPUExecutionProvider"]
+PROVIDER = ["CPUExecutionProvider"]  # Default to CPU provider
+device = "cpu"
 
-if USE_CUDA:
-    available_providers = onnxruntime.get_available_providers()
-    if "CUDAExecutionProvider" in available_providers:
-        print("\n********** Running on CUDA **********\n")
-        PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        cv2.setNumThreads(32)
-    else:
-        USE_CUDA = False
-        print("\n********** CUDA unavailable running on CPU **********\n")
-else:
-    USE_CUDA = False
-    print("\n********** Running on CPU **********\n")
-
-device = "cuda" if USE_CUDA else "cpu"
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
 ## ------------------------------ LOAD MODELS ------------------------------
@@ -166,8 +154,23 @@ def process(
     global WORKSPACE
     global OUTPUT_FILE
     global PREVIEW
-    WORKSPACE, OUTPUT_FILE, PREVIEW = None, None, None
+    global USE_CUDA  # Access global variables
+    global device
+    global PROVIDER
+    global FACE_ANALYSER, FACE_SWAPPER, FACE_ENHANCER, FACE_PARSER, NSFW_DETECTOR
 
+    WORKSPACE, OUTPUT_FILE, PREVIEW = None, None, None
+    USE_CUDA = True
+    device = "cuda"
+    PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+    # Reset models to None to reload them with GPU
+    FACE_ANALYSER = None
+    FACE_SWAPPER = None
+    FACE_ENHANCER = None
+    FACE_PARSER = None
+    NSFW_DETECTOR = None
+    
     ## ------------------------------ GUI UPDATE FUNC ------------------------------
 
     def ui_before():