Update app.py

app.py

@@ -121,7 +121,7 @@ else:
     print("\n********** Running on CPU **********\n")
     PROVIDER = ["CPUExecutionProvider"]
 
-device = "cuda" if USE_CUDA else "cpu"
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
 ## ------------------------------ LOAD MODELS ------------------------------
@@ -148,7 +148,7 @@ def load_face_swapper_model(path="./assets/pretrained_models/inswapper_128.onnx"
 def load_face_parser_model(path="./assets/pretrained_models/79999_iter.pth"):
     global FACE_PARSER
     if FACE_PARSER is None:
-        FACE_PARSER = init_parsing_model(path, device="cuda")
+        FACE_PARSER = init_parsing_model(path, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
 
 @spaces.GPU(enable_queue=True)
 def load_nsfw_detector_model(path="./assets/pretrained_models/open-nsfw.onnx"):
@@ -215,7 +215,7 @@ def process(
         print("\n********** Running on CPU **********\n")
         PROVIDER = ["CPUExecutionProvider"]
 
-    device = "cuda" if USE_CUDA else "cpu"
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
     # Reset models to None to reload them with GPU
@@ -269,7 +269,7 @@ def process(
     if face_enhancer_name != "NONE":
         if face_enhancer_name not in cv2_interpolations:
             yield f"### \n ⌛ Loading {face_enhancer_name} model...", *ui_before()
-        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
+        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
     else:
         FACE_ENHANCER = None
 
@@ -351,7 +351,7 @@ def process(
             yield "### \n ⌛ Face-parsing mask...", *ui_before()
             masks = []
             count = 0
-            for batch_mask in get_parsed_mask(FACE_PARSER, preds, classes=includes, device=device, batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
+            for batch_mask in get_parsed_mask(FACE_PARSER, preds, classes=includes, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
                 masks.append(batch_mask)
                 EMPTY_CACHE()
                 count += 1