Update app.py

app.py

@@ -30,7 +30,7 @@ from utils import trim_video, StreamerThread, ProcessBar, open_directory, split_
 
 parser = argparse.ArgumentParser(description="Swap Face Swapper")
 parser.add_argument("--out_dir", help="Default Output directory", default=os.getcwd())
-parser.add_argument("--batch_size", help="Gpu batch size", default=32)
+parser.add_argument("--batch_size", help="Gpu batch size", default=1)
 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
@@ -76,14 +76,30 @@ FACE_ENHANCER_LIST = ["NONE"]
 FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
 FACE_ENHANCER_LIST.extend(cv2_interpolations)
 
+## ------------------------------ SET EXECUTION PROVIDER ------------------------------
+
 ## ------------------------------ SET EXECUTION PROVIDER ------------------------------
 # Note: Non CUDA users may change settings here
 
-PROVIDER = ["CPUExecutionProvider"]  # Default to CPU provider
-device = "cpu"
+PROVIDER = ["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")
+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 ------------------------------
+
 
 ## ------------------------------ LOAD MODELS ------------------------------