Upload nodes.py

custom_nodes/comfyui-reactor-node/nodes.py

@@ -21,6 +21,7 @@ from modules.shared import state
 import comfy.model_management as model_management
 import comfy.utils
 import folder_paths
+from folder_paths import add_folder_path_and_extensions    # ← добавили
 
 import scripts.reactor_version
 from r_chainner import model_loading
@@ -64,14 +65,32 @@ models_dir = folder_paths.models_dir
 REACTOR_MODELS_PATH = os.path.join(models_dir, "reactor")
 FACE_MODELS_PATH = os.path.join(REACTOR_MODELS_PATH, "faces")
 
+os.makedirs(REACTOR_MODELS_PATH, exist_ok=True)
+os.makedirs(FACE_MODELS_PATH, exist_ok=True)
+
 if not os.path.exists(REACTOR_MODELS_PATH):
     os.makedirs(REACTOR_MODELS_PATH)
     if not os.path.exists(FACE_MODELS_PATH):
         os.makedirs(FACE_MODELS_PATH)
 
+# ——————————————————————————————————————————————————————————————
+# Вместо старого куска:
+# dir_facerestore_models = os.path.join(models_dir, "facerestore_models")
+# os.makedirs(dir_facerestore_models, exist_ok=True)
+# folder_paths.folder_names_and_paths["facerestore_models"] = (
+#     [dir_facerestore_models],
+#     folder_paths.supported_pt_extensions
+# )
+#
+# Делаем так:
 dir_facerestore_models = os.path.join(models_dir, "facerestore_models")
 os.makedirs(dir_facerestore_models, exist_ok=True)
-folder_paths.folder_names_and_paths["facerestore_models"] = ([dir_facerestore_models], folder_paths.supported_pt_extensions)
+add_folder_path_and_extensions(
+    "facerestore_models",
+    [dir_facerestore_models],
+    [".pth", ".pt", ".onnx"]         # ← явно указываем .onnx
+)
+# ——————————————————————————————————————————————————————————————
 
 BLENDED_FACE_MODEL = None
 FACE_SIZE: int = 512
@@ -192,8 +211,17 @@ class reactor:
             logger.status(f"Restoring with {face_restore_model} | Face Size is set to {faceSize}")
 
             model_path = folder_paths.get_full_path("facerestore_models", face_restore_model)
-
+            # если по какой-то причине не нашли — склеиваем вручную
+            if model_path is None or not os.path.isfile(model_path):
+                model_path = os.path.join(dir_facerestore_models, face_restore_model)
+                
             device = model_management.get_torch_device()
+            
+            if ".onnx" in face_restore_model.lower():
+                # в этом месте model_path уже точно строка, а не None
+                ort_session = set_ort_session(model_path, providers=providers)
+                ort_session_inputs = {}
+                facerestore_model = ort_session
 
             if "codeformer" in face_restore_model.lower():