DegMaTsu commited on
Commit
0b1b1be
·
verified ·
1 Parent(s): 1e06f23

Upload nodes.py

Browse files
custom_nodes/comfyui-reactor-node/nodes.py CHANGED
@@ -21,6 +21,7 @@ from modules.shared import state
21
  import comfy.model_management as model_management
22
  import comfy.utils
23
  import folder_paths
 
24
 
25
  import scripts.reactor_version
26
  from r_chainner import model_loading
@@ -64,14 +65,32 @@ models_dir = folder_paths.models_dir
64
  REACTOR_MODELS_PATH = os.path.join(models_dir, "reactor")
65
  FACE_MODELS_PATH = os.path.join(REACTOR_MODELS_PATH, "faces")
66
 
 
 
 
67
  if not os.path.exists(REACTOR_MODELS_PATH):
68
  os.makedirs(REACTOR_MODELS_PATH)
69
  if not os.path.exists(FACE_MODELS_PATH):
70
  os.makedirs(FACE_MODELS_PATH)
71
 
 
 
 
 
 
 
 
 
 
 
72
  dir_facerestore_models = os.path.join(models_dir, "facerestore_models")
73
  os.makedirs(dir_facerestore_models, exist_ok=True)
74
- folder_paths.folder_names_and_paths["facerestore_models"] = ([dir_facerestore_models], folder_paths.supported_pt_extensions)
 
 
 
 
 
75
 
76
  BLENDED_FACE_MODEL = None
77
  FACE_SIZE: int = 512
@@ -192,8 +211,17 @@ class reactor:
192
  logger.status(f"Restoring with {face_restore_model} | Face Size is set to {faceSize}")
193
 
194
  model_path = folder_paths.get_full_path("facerestore_models", face_restore_model)
195
-
 
 
 
196
  device = model_management.get_torch_device()
 
 
 
 
 
 
197
 
198
  if "codeformer" in face_restore_model.lower():
199
 
 
21
  import comfy.model_management as model_management
22
  import comfy.utils
23
  import folder_paths
24
+ from folder_paths import add_folder_path_and_extensions # ← добавили
25
 
26
  import scripts.reactor_version
27
  from r_chainner import model_loading
 
65
  REACTOR_MODELS_PATH = os.path.join(models_dir, "reactor")
66
  FACE_MODELS_PATH = os.path.join(REACTOR_MODELS_PATH, "faces")
67
 
68
+ os.makedirs(REACTOR_MODELS_PATH, exist_ok=True)
69
+ os.makedirs(FACE_MODELS_PATH, exist_ok=True)
70
+
71
  if not os.path.exists(REACTOR_MODELS_PATH):
72
  os.makedirs(REACTOR_MODELS_PATH)
73
  if not os.path.exists(FACE_MODELS_PATH):
74
  os.makedirs(FACE_MODELS_PATH)
75
 
76
+ # ——————————————————————————————————————————————————————————————
77
+ # Вместо старого куска:
78
+ # dir_facerestore_models = os.path.join(models_dir, "facerestore_models")
79
+ # os.makedirs(dir_facerestore_models, exist_ok=True)
80
+ # folder_paths.folder_names_and_paths["facerestore_models"] = (
81
+ # [dir_facerestore_models],
82
+ # folder_paths.supported_pt_extensions
83
+ # )
84
+ #
85
+ # Делаем так:
86
  dir_facerestore_models = os.path.join(models_dir, "facerestore_models")
87
  os.makedirs(dir_facerestore_models, exist_ok=True)
88
+ add_folder_path_and_extensions(
89
+ "facerestore_models",
90
+ [dir_facerestore_models],
91
+ [".pth", ".pt", ".onnx"] # ← явно указываем .onnx
92
+ )
93
+ # ——————————————————————————————————————————————————————————————
94
 
95
  BLENDED_FACE_MODEL = None
96
  FACE_SIZE: int = 512
 
211
  logger.status(f"Restoring with {face_restore_model} | Face Size is set to {faceSize}")
212
 
213
  model_path = folder_paths.get_full_path("facerestore_models", face_restore_model)
214
+ # если по какой-то причине не нашли — склеиваем вручную
215
+ if model_path is None or not os.path.isfile(model_path):
216
+ model_path = os.path.join(dir_facerestore_models, face_restore_model)
217
+
218
  device = model_management.get_torch_device()
219
+
220
+ if ".onnx" in face_restore_model.lower():
221
+ # в этом месте model_path уже точно строка, а не None
222
+ ort_session = set_ort_session(model_path, providers=providers)
223
+ ort_session_inputs = {}
224
+ facerestore_model = ort_session
225
 
226
  if "codeformer" in face_restore_model.lower():
227