Update app.py

app.py

@@ -115,9 +115,11 @@ if USE_CUDA:
     else:
         USE_CUDA = False
         print("\n********** CUDA unavailable running on CPU **********\n")
+        PROVIDER = ["CPUExecutionProvider"]
 else:
     USE_CUDA = False
     print("\n********** Running on CPU **********\n")
+    PROVIDER = ["CPUExecutionProvider"]
 
 device = "cuda" if USE_CUDA else "cpu"
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
@@ -146,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=device)
+        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"):
@@ -155,7 +157,10 @@ def load_nsfw_detector_model(path="./assets/pretrained_models/open-nsfw.onnx"):
         NSFW_DETECTOR = NSFWChecker(model_path=path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
 
 
+@spaces.GPU(enable_queue=True)
 load_face_analyser_model()
+
+@spaces.GPU(enable_queue=True)
 load_face_swapper_model()
 
 ## ------------------------------ MAIN PROCESS ------------------------------
@@ -198,9 +203,22 @@ def process(
     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"]
+    if USE_CUDA:
+        available_providers = onnxruntime.get_available_providers()
+        if "CUDAExecutionProvider" in available_providers:
+            print("\n********** Running on CUDA **********\n")
+            PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        else:
+            USE_CUDA = False
+            print("\n********** CUDA unavailable running on CPU **********\n")
+            PROVIDER = ["CPUExecutionProvider"]
+    else:
+        USE_CUDA = False
+        print("\n********** Running on CPU **********\n")
+        PROVIDER = ["CPUExecutionProvider"]
+
+    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
     FACE_ANALYSER = None
@@ -968,6 +986,7 @@ if __name__ == "__main__":
     if USE_COLAB:
         print("Running in colab mode")
 
+    interface.queue()
     interface.launch()