Update app.py

app.py

@@ -1,4 +1,4 @@
-import os
+
 import cv2
 import glob
 import time
@@ -30,9 +30,7 @@ parser = argparse.ArgumentParser(description="Free 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("--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("--colab", action="store_true", help="Enable colab mode", default=False)
 user_args = parser.parse_args()
 
 ## ------------------------------ DEFAULTS ------------------------------
@@ -74,23 +72,16 @@ FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
 FACE_ENHANCER_LIST.extend(cv2_interpolations)
 
 ## ------------------------------ SET EXECUTION PROVIDER ------------------------------
-# Note: Non CUDA users may change settings here
 
-PROVIDER = ["CPUExecutionProvider"]
+# Использование CUDA (если доступно)
+PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"] if torch.cuda.is_available() else ["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")
+if "CUDAExecutionProvider" in PROVIDER:
+    print("\n********** Running on CUDA **********\n")
 else:
-    USE_CUDA = False
     print("\n********** Running on CPU **********\n")
 
-device = "cuda" if USE_CUDA else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
 ## ------------------------------ LOAD MODELS ------------------------------
@@ -100,7 +91,9 @@ def load_face_analyser_model(name="buffalo_l"):
     if FACE_ANALYSER is None:
         FACE_ANALYSER = insightface.app.FaceAnalysis(name=name, providers=PROVIDER)
         FACE_ANALYSER.prepare(
-            ctx_id=0, det_size=(DETECT_SIZE, DETECT_SIZE), det_thresh=DETECT_THRESH
+            ctx_id=0 if "CUDAExecutionProvider" in PROVIDER else -1,
+            det_size=(DETECT_SIZE, DETECT_SIZE),
+            det_thresh=DETECT_THRESH,
         )
 
 
@@ -122,7 +115,6 @@ load_face_swapper_model()
 
 ## ------------------------------ MAIN PROCESS ------------------------------
 
-
 def process(
     input_type,
     image_path,
@@ -155,6 +147,25 @@ def process(
     global PREVIEW
     WORKSPACE, OUTPUT_FILE, PREVIEW = None, None, None
 
+    ## ------------------------------ LOADING MODELS ------------------------------
+    yield "### \n 🌀 Loading face analyser model...", gr.update(visible=True)
+    load_face_analyser_model()
+
+    yield "### \n ⚙️ Loading face swapper model...", gr.update(visible=True)
+    load_face_swapper_model()
+
+    if face_enhancer_name != "NONE":
+        if face_enhancer_name not in cv2_interpolations:
+            yield f"### \n 💡 Loading {face_enhancer_name} model...", gr.update(visible=True)
+        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
+    else:
+        FACE_ENHANCER = None
+
+    if enable_face_parser:
+        yield "### \n 📀 Loading face parsing model...", gr.update(visible=True)
+        load_face_parser_model()
+
+    # Прочие шаги процесса остаются без изменений
     ## ------------------------------ GUI UPDATE FUNC ------------------------------
 
     def ui_before():