Update app.py

app.py

@@ -1,6 +1,3 @@
-
-#### APP.PY CODE END ###
-
 import os
 import spaces
 import cv2
@@ -36,7 +33,7 @@ parser.add_argument("--out_dir", help="Default Output directory", default=os.get
 parser.add_argument("--batch_size", help="Gpu batch size", default=32)
 parser.add_argument("--cuda", action="store_true", help="Enable cuda", default=False)
 parser.add_argument("--colab", action="store_true", help="Enable colab mode", default=False)
-parser.add_argument("--device", default="cuda:0", type=str)
+parser.add_argument("--device", default="cuda", type=str)
 user_args = parser.parse_args()
 
 @spaces.GPU
@@ -121,7 +118,7 @@ else:
     print("\n********** Running on CPU **********\n")
     PROVIDER = ["CPUExecutionProvider"]
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = "cuda" if USE_CUDA else "cpu"
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
 ## ------------------------------ LOAD MODELS ------------------------------
@@ -148,7 +145,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=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+        FACE_PARSER = init_parsing_model(path, device="cuda")
 
 @spaces.GPU(enable_queue=True)
 def load_nsfw_detector_model(path="./assets/pretrained_models/open-nsfw.onnx"):
@@ -215,7 +212,7 @@ def process(
         print("\n********** Running on CPU **********\n")
         PROVIDER = ["CPUExecutionProvider"]
 
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = "cuda" if USE_CUDA 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 +266,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=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
     else:
         FACE_ENHANCER = None
 
@@ -351,7 +348,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=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
+            for batch_mask in get_parsed_mask(FACE_PARSER, preds, classes=includes, device=device, batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
                 masks.append(batch_mask)
                 EMPTY_CACHE()
                 count += 1
@@ -987,5 +984,3 @@ if __name__ == "__main__":
     interface.queue()
     interface.launch()
 
-
-#### APP.PY CODE END ###