Upload nodes.py

custom_nodes/comfyui-reactor-node/nodes.py

@@ -2,8 +2,7 @@ import os, glob, sys
 import logging
 
 import torch
-import torch.nn.functional as F
-import torchvision.transforms as T
+import torch.nn.functional as torchfn
 from torchvision.transforms.functional import normalize
 from torchvision.ops import masks_to_boxes
 
@@ -12,7 +11,6 @@ import cv2
 import math
 from typing import List
 from PIL import Image
-import io
 from scipy import stats
 from insightface.app.common import Face
 from segment_anything import sam_model_registry
@@ -23,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    # ← добавили (не работает - вызывает ошибку ImportError: cannot import name 'add_folder_path_and_extensions' from 'folder_paths' (/home/user/app/folder_paths.py))
 
 import scripts.reactor_version
 from r_chainner import model_loading
@@ -51,9 +50,7 @@ from reactor_utils import (
     prepare_cropped_face,
     normalize_cropped_face,
     add_folder_path_and_extensions,
-    rgba2rgb_tensor,
-    progress_bar,
-    progress_bar_reset
+    rgba2rgb_tensor
 )
 from reactor_patcher import apply_patch
 from r_facelib.utils.face_restoration_helper import FaceRestoreHelper
@@ -63,19 +60,39 @@ import scripts.r_masking.subcore as subcore
 import scripts.r_masking.core as core
 import scripts.r_masking.segs as masking_segs
 
-
+# … сервисные константы …
 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)
 
+# ——————————————————————————————————————————————————————————————
+# Настраиваем папку facerestore_models, чтобы туда «виделось» расширение .onnx
+# Вместо старого куска:
+# 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
@@ -155,8 +172,7 @@ class reactor:
             "hidden": {"faces_order": "FACES_ORDER"},
         }
 
-    RETURN_TYPES = ("IMAGE","FACE_MODEL","IMAGE")
-    RETURN_NAMES = ("SWAPPED_IMAGE","FACE_MODEL","ORIGINAL_IMAGE")
+    RETURN_TYPES = ("IMAGE","FACE_MODEL")
     FUNCTION = "execute"
     CATEGORY = "🌌 ReActor"
 
@@ -170,37 +186,24 @@ class reactor:
         self.interpolation = "Bicubic"
         self.boost_model_visibility = 1
         self.boost_cf_weight = 0.5
-
+        
+    # Начинаем restore face
     def restore_face(
-        self,
-        input_image,
-        face_restore_model,
-        face_restore_visibility,
-        codeformer_weight,
-        facedetection,
-        face_selection="all",
-        sort_by="area",
-        descending=True,
-        min_x_position=0.0,
-        max_x_position=1.0,
-        min_y_position=0.0,
-        max_y_position=1.0,
-        take_start=0,
-        take_count=1,
-        face_index=0,
-    ):
-
-        # >>>>> ПРИНУДИТЕЛЬНЫЙ ВЫВОД ОТЛАДКИ <<<<<
-    #   print(f"\n--- [ReActor Debug] Face selection: {face_selection}, Sort by: {sort_by}, Descending: {descending}")
-    
+            self,
+            input_image,
+            face_restore_model,
+            face_restore_visibility,
+            codeformer_weight,
+            facedetection,
+        ):
         result = input_image
 
         if face_restore_model != "none" and not model_management.processing_interrupted():
-
             global FACE_SIZE, FACE_HELPER
 
             self.face_helper = FACE_HELPER
             
+            # Подготавливаем размер лица
             faceSize = 512
             if "1024" in face_restore_model.lower():
                 faceSize = 1024
@@ -209,9 +212,30 @@ 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)
+                
+            # И вот здесь, если model_path всё ещё неправильный, мы сразу бросим понятную ошибку
+            if not os.path.isfile(model_path):
+                raise FileNotFoundError(f"ONNX-модель не найдена: {model_path}")
+                
             device = model_management.get_torch_device()
+            
+            # Единственный блок для .onnx — гарантируем, что сюда попадает только он
+            if face_restore_model.lower().endswith(".onnx"):
+                ort_session     = set_ort_session(model_path, providers=providers)
+                ort_session_inputs = {}
+                facerestore_model = ort_session            
+            
+            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():
 
@@ -243,19 +267,16 @@ class reactor:
                 FACE_SIZE = faceSize
                 FACE_HELPER = self.face_helper
 
-            # Copying Tensor to CPU (if it isn't) to convert torch.Tensor to np.ndarray
-            image_np = 255. * result.cpu().numpy()
+            image_np = 255. * result.numpy()
 
             total_images = image_np.shape[0]
 
             out_images = []
 
-            pbar = progress_bar(total_images)
-
             for i in range(total_images):
 
-                # if total_images > 1:
-                #     logger.status(f"Restoring {i}")
+                if total_images > 1:
+                    logger.status(f"Restoring {i+1}")
 
                 cur_image_np = image_np[i,:, :, ::-1]
 
@@ -269,110 +290,16 @@ class reactor:
                 self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
                 self.face_helper.align_warp_face()
 
-                # Фильтрация лиц
-                if face_selection != "all" and self.face_helper.cropped_faces:
-                    # Собираем информацию о лицах для фильтрации
-                    face_info = []
-                    img_height, img_width = cur_image_np.shape[0:2]
-                    
-                    for j, face in enumerate(self.face_helper.cropped_faces):
-                        # Используем центр лица вместо левого верхнего угла
-                        if hasattr(self.face_helper, 'det_faces') and len(self.face_helper.det_faces) > j:
-                            bbox = self.face_helper.det_faces[j]
-                            # Вычисляем центр лица для более точного позиционирования
-                            x1 = ((bbox[0] + bbox[2]) / 2) / img_width  # центр x
-                            y1 = ((bbox[1] + bbox[3]) / 2) / img_height  # центр y
-                            area = face.shape[0] * face.shape[1]
-                            confidence = bbox[4] if len(bbox) > 4 else 1.0
-                        else:
-                            # Если информация о bbox недоступна, используем приблизительные данные
-                            area = face.shape[0] * face.shape[1]
-                            x1, y1 = 0.5, 0.5  # центр изображения
-                            confidence = 1.0
-                            
-                        face_info.append({
-                            'index': j,
-                            'area': area,
-                            'x_position': x1,
-                            'y_position': y1,
-                            'detection_confidence': confidence
-                        })
-                    
-                    # Сначала сортируем все лица по выбранному критерию
-                    all_indices = list(range(len(self.face_helper.cropped_faces)))
-                    
-                    # Отладочный вывод перед сортировкой
-                #   print(f"--- [ReActor Debug] Sorting all faces by {sort_by}, descending={descending}")
-                    
-                    # Вывод для x_position и y_position
-                    if sort_by == "y_position":
-                        all_positions = [(idx, face_info[idx]['y_position']) for idx in all_indices]
-                    #   print(f"--- [ReActor Debug] All positions before sort: {all_positions}")
-                    elif sort_by == "x_position":
-                        all_positions = [(idx, face_info[idx]['x_position']) for idx in all_indices]
-                    #   print(f"--- [ReActor Debug] All X positions before sort: {all_positions}")
-                    
-                    # Сортировка по выбранному критерию
-                    sorted_indices = sorted(
-                        all_indices,
-                        key=lambda idx: face_info[idx][sort_by],
-                        reverse=descending
-                    )
-                    
-                    # Отладочный вывод после сортировки
-                    if sort_by == "y_position":
-                        sorted_positions = [(idx, face_info[idx]['y_position']) for idx in sorted_indices]
-                    #   print(f"--- [ReActor Debug] All positions after sort: {sorted_positions}")
-                    elif sort_by == "x_position":
-                        sorted_positions = [(idx, face_info[idx]['x_position']) for idx in sorted_indices]
-                    #   print(f"--- [ReActor Debug] All X positions after sort: {sorted_positions}")
-                    
-                    # Применяем фильтрацию в зависимости от режима
-                    if face_selection == "filter":
-                        # Фильтрация по координатам
-                        filtered_indices = [
-                            idx for idx in sorted_indices
-                            if min_x_position <= face_info[idx]['x_position'] <= max_x_position and
-                               min_y_position <= face_info[idx]['y_position'] <= max_y_position
-                        ]
-                        
-                    #   print(f"--- [ReActor Debug] Filtered faces: {len(filtered_indices)}")
-                        
-                        # Выборка по take_start и take_count
-                        selected_indices = filtered_indices[take_start:take_start + take_count]
-                    
-                    elif face_selection == "largest":
-                        # При выборе "largest" просто берем take_count лиц с наибольшей площадью, начиная с take_start
-                        selected_indices = sorted_indices[take_start:take_start + take_count]
-                    #   print(f"--- [ReActor Debug] Selected {take_count} face(s) starting from {take_start}: {selected_indices}")
-                    
-                    elif face_selection == "index":
-                        # В режиме "index" просто берем лица, начиная с take_start
-                        selected_indices = sorted_indices[take_start:take_start + take_count]
-                    #   print(f"--- [ReActor Debug] Selected {take_count} face(s) starting from {take_start}: {selected_indices}")
-                        
-                        # Дополнительная отладочная информация
-                    #   for i, idx in enumerate(selected_indices):
-                    #       if sort_by == "x_position":
-                    #           print(f"--- [ReActor Debug] Selected face {i}: X position = {face_info[idx]['x_position']}")
-                    #       elif sort_by == "y_position":
-                    #           print(f"--- [ReActor Debug] Selected face {i}: Y position = {face_info[idx]['y_position']}")
-                    
-                    # Применяем фильтрацию ко всем спискам
-                #   print(f"--- [ReActor Debug] Final selected indices: {selected_indices}")
-                    if selected_indices:
-                        self.face_helper.cropped_faces = [self.face_helper.cropped_faces[j] for j in selected_indices]
-                        if hasattr(self.face_helper, 'restored_faces'):
-                            self.face_helper.restored_faces = []
-                        if hasattr(self.face_helper, 'affine_matrices'):
-                            self.face_helper.affine_matrices = [self.face_helper.affine_matrices[j] for j in selected_indices]
-                        if hasattr(self.face_helper, 'det_faces'):
-                            self.face_helper.det_faces = [self.face_helper.det_faces[j] for j in selected_indices]
-
                 restored_face = None
 
+                # после подготовки self.face_helper и выравнивания лица начинается цикл по всем фрагментам:
                 for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
-
+                    
+                    # если у вас несколько лиц, логируем прогресс
+                    if len(self.face_helper.cropped_faces) > 1:
+                        logger.status(f"Restoring face {idx+1}")
+                        
+                    # превращаем PIL→tensor и нормализуем (для PTH-моделей):
                     # if ".pth" in face_restore_model:
                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
                     normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
@@ -381,33 +308,43 @@ class reactor:
                     try:
 
                         with torch.no_grad():
-
-                            if ".onnx" in face_restore_model: # ONNX models
-
-                                for ort_session_input in ort_session.get_inputs():
-                                    if ort_session_input.name == "input":
-                                        cropped_face_prep = prepare_cropped_face(cropped_face)
-                                        ort_session_inputs[ort_session_input.name] = cropped_face_prep
-                                    if ort_session_input.name == "weight":
-                                        weight = np.array([ 1 ], dtype = np.double)
-                                        ort_session_inputs[ort_session_input.name] = weight
-
+                            
+                             # ────── 🚩 МЕСТО ВСТАВКИ SNIPPET-а for ONNX ──────
+                            if face_restore_model.lower().endswith(".onnx"):
+                                # тут мы готовим входы для onnxruntime
+                                for inp in ort_session.get_inputs():
+                                    if inp.name == "input":
+                                        ort_session_inputs[inp.name] = prepare_cropped_face(cropped_face)
+                                    elif inp.name == "weight":
+                                        ort_session_inputs[inp.name] = np.array([1], dtype=np.double)
+
+                                # собственно, инференс
                                 output = ort_session.run(None, ort_session_inputs)[0][0]
+                                # возвращаем в изображение
                                 restored_face = normalize_cropped_face(output)
+                                
+                            else:
+                                # ────── ВЕТКА ДЛЯ PTH / CODEFORMER ──────
+                                # если CodeFormer
+                                if "codeformer" in face_restore_model.lower():
+                                    output = facerestore_model(cropped_face_t, w=codeformer_weight)[0]
+                                else:
+                                    # обычная PTH-модель (GFPGAN и т.п.)
+                                    output = facerestore_model(cropped_face_t)[0]
 
-                            else: # PTH models
-
-                                output = facerestore_model(cropped_face_t, w=codeformer_weight)[0] if "codeformer" in face_restore_model.lower() else facerestore_model(cropped_face_t)[0]
                                 restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
-
+                            # ─────────────────────────────────────────────────
+                            
+                        # освобождаем память    
                         del output
                         torch.cuda.empty_cache()
 
                     except Exception as error:
-
+                        # на случай краха инференса — возвращаем «сырой» кроп
                         print(f"\tFailed inference: {error}", file=sys.stderr)
                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
 
+                    # далее вы накладываете прозрачность и собираете кадр:
                     if face_restore_visibility < 1:
                         restored_face = cropped_face * (1 - face_restore_visibility) + restored_face * face_restore_visibility
 
@@ -430,24 +367,15 @@ class reactor:
                 if state.interrupted or model_management.processing_interrupted():
                     logger.status("Interrupted by User")
                     return input_image
-                
-                pbar.update(1)
 
             restored_img_np = np.array(out_images).astype(np.float32) / 255.0
             restored_img_tensor = torch.from_numpy(restored_img_np)
 
             result = restored_img_tensor
 
-            progress_bar_reset(pbar)
-
         return result
-        
-    def execute(self, enabled, input_image, swap_model, detect_gender_source, detect_gender_input, source_faces_index, input_faces_index, console_log_level, face_restore_model, face_restore_visibility, codeformer_weight, facedetection, source_image=None, face_model=None, faces_order=None, face_boost=None):
-
-        device = model_management.get_torch_device()
-
-        if isinstance(input_image, torch.Tensor) and input_image.device != device:
-            input_image = input_image.to(device)
+    
+    def execute(self, enabled, input_image, swap_model, detect_gender_source, detect_gender_input, source_faces_index, input_faces_index, console_log_level, face_restore_model,face_restore_visibility, codeformer_weight, facedetection, source_image=None, face_model=None, faces_order=None, face_boost=None):
 
         if face_boost is not None:
             self.face_boost_enabled = face_boost["enabled"]
@@ -473,35 +401,8 @@ class reactor:
         if face_model == "none":
             face_model = None
         
-        # Сохраняем параметры для последующего использования при restore
-        target_indices = []
-        if input_faces_index == "0" or input_faces_index == "":
-            target_indices = [0]
-        else:
-            try:
-                target_indices = [int(x.strip()) for x in input_faces_index.split(",") if x.strip()]
-            except:
-                target_indices = [0]
-        
-        # Определяем параметры сортировки
-        sort_by = "area"
-        descending = True
-        if faces_order is not None:
-            input_order = faces_order[0]
-            if input_order in ["left-right", "right-left"]:
-                sort_by = "x_position"
-                descending = (input_order == "right-left")
-            elif input_order in ["top-bottom", "bottom-top"]:
-                sort_by = "y_position"
-                descending = (input_order == "bottom-top")
-            elif input_order in ["small-large", "large-small"]:
-                sort_by = "area"
-                descending = (input_order == "large-small")
-        
-        # Выполняем face swap
         script = FaceSwapScript()
         pil_images = batch_tensor_to_pil(input_image)
-
         if source_image is not None:
             source = tensor_to_pil(source_image)
         else:
@@ -527,8 +428,7 @@ class reactor:
             codeformer_weight=self.boost_cf_weight,
             interpolation=self.interpolation,
         )
-        swapped_result = batched_pil_to_tensor(p.init_images)
-        original_image = batched_pil_to_tensor(pil_images)
+        result = batched_pil_to_tensor(p.init_images)
 
         if face_model is None:
             current_face_model = get_current_faces_model()
@@ -536,100 +436,11 @@ class reactor:
         else:
             face_model_to_provide = face_model
 
-        # Применяем restore face к результату face swap
         if self.restore or not self.face_boost_enabled:
-            # НОВЫЙ ПОДХОД: Анализируем лица ПОСЛЕ face swap
-            target_faces_coords = []
-            try:
-                # Преобразуем результат face swap в изображение для анализа
-                swapped_img_tensor = swapped_result[0].cpu()
-                swapped_img_np = (255 * swapped_img_tensor.numpy()).astype(np.uint8)
-                swapped_img_pil = Image.fromarray(swapped_img_np)
-                swapped_img_cv = cv2.cvtColor(np.array(swapped_img_pil), cv2.COLOR_RGB2BGR)
-                
-                # Определяем лица после face swap
-                face_analyser = get_current_faces_model()
-                detected_faces = analyze_faces(swapped_img_cv, (640, 640))
-                
-                if not detected_faces:
-                    # Пробуем с меньшим размером детекции
-                    detected_faces = analyze_faces(swapped_img_cv, (320, 320))
-                
-                if detected_faces:
-                    # Сортируем лица тем же способом, что и в face swap
-                    if sort_by == "x_position":
-                        detected_faces.sort(key=lambda x: (x.bbox[0] + x.bbox[2])/2, reverse=descending)
-                    elif sort_by == "y_position":
-                        detected_faces.sort(key=lambda x: (x.bbox[1] + x.bbox[3])/2, reverse=descending)
-                    elif sort_by == "area":
-                        detected_faces.sort(key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=descending)
-                    
-                    # Выбираем лица по тем же индексам, что и для face swap
-                    for idx in target_indices:
-                        if idx < len(detected_faces):
-                            face = detected_faces[idx]
-                            # Сохраняем координаты центра лица
-                            center_x = (face.bbox[0] + face.bbox[2]) / 2 / swapped_img_cv.shape[1]  # нормализуем
-                            center_y = (face.bbox[1] + face.bbox[3]) / 2 / swapped_img_cv.shape[0]  # нормализуем
-                            target_faces_coords.append((center_x, center_y))
-                    
-                #   print(f"--- [ReActor Debug] Detected face coordinates after swap: {target_faces_coords}")
-            #   else:
-                #   print("--- [ReActor Debug] No faces detected after face swap")
-            
-            except Exception as e:
-            #   print(f"--- [ReActor Debug] Error analyzing faces after swap: {str(e)}")
-                target_faces_coords = []
-
-            # Если определены координаты лиц, применяем restore только к ним
-            if target_faces_coords:
-                # Используем небольшой отступ вокруг каждого лица
-                margin = 0.15  # 15% от размера изображения
-                
-                restored_result = swapped_result
-                for center_x, center_y in target_faces_coords:
-                    min_x = max(0.0, center_x - margin)
-                    max_x = min(1.0, center_x + margin)
-                    min_y = max(0.0, center_y - margin)
-                    max_y = min(1.0, center_y + margin)
-                    
-                #   print(f"--- [ReActor Debug] Restoring faces in region: x={min_x:.2f}-{max_x:.2f}, y={min_y:.2f}-{max_y:.2f}")
-                    
-                    # Применяем restore_face к указанной области
-                    restored_result = reactor.restore_face(
-                        self,
-                        restored_result,
-                        face_restore_model,
-                        face_restore_visibility,
-                        codeformer_weight,
-                        facedetection,
-                        "filter",  # Используем filter для выбора по координатам
-                        sort_by,
-                        descending,
-                        min_x,
-                        max_x,
-                        min_y,
-                        max_y,
-                        0,  # take_start
-                        10  # take_count - берем больше лиц для надежности
-                    )
-                
-                return (restored_result, face_model_to_provide, original_image)
-            else:
-                # Если координаты не определены, восстанавливаем все лица
-            #   print("--- [ReActor Debug] Falling back to restoring all faces")
-                restored_result = reactor.restore_face(
-                    self,
-                    swapped_result,
-                    face_restore_model,
-                    face_restore_visibility,
-                    codeformer_weight,
-                    facedetection
-                )
-                return (restored_result, face_model_to_provide, original_image)
-        else:
-            # Если restore не требуется
-            return (swapped_result, face_model_to_provide, original_image)
+            result = reactor.restore_face(self,result,face_restore_model,face_restore_visibility,codeformer_weight,facedetection)
+
+        return (result,face_model_to_provide)
+
 
 class ReActorPlusOpt:
     @classmethod
@@ -652,8 +463,7 @@ class ReActorPlusOpt:
             }
         }
 
-    RETURN_TYPES = ("IMAGE","FACE_MODEL","IMAGE")
-    RETURN_NAMES = ("SWAPPED_IMAGE","FACE_MODEL","ORIGINAL_IMAGE")
+    RETURN_TYPES = ("IMAGE","FACE_MODEL")
     FUNCTION = "execute"
     CATEGORY = "🌌 ReActor"
 
@@ -1010,44 +820,40 @@ class RestoreFace:
                 "model": (get_model_names(get_restorers),),
                 "visibility": ("FLOAT", {"default": 1, "min": 0.0, "max": 1, "step": 0.05}),
                 "codeformer_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1, "step": 0.05}),
-                "face_selection": (["all", "filter", "largest"],{"default": "all"}), # ["all", "filter", "largest", "index"]
             },
-            "optional": {
-                "sort_by": (["area", "x_position", "y_position", "detection_confidence"],{"default": "area"}),
-                "descending": ("BOOLEAN", {"default": True}),
-            #   "min_x_position": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-            #   "max_x_position": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-            #   "min_y_position": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-            #   "max_y_position": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "take_start": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
-                "take_count": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
-            #   "face_index": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
-            }
         }
-    
+
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "execute"
     CATEGORY = "🌌 ReActor"
 
-    def execute(self, image, model, visibility, codeformer_weight, facedetection, face_selection="all", 
-                sort_by="area", descending=True, min_x_position=0.0, max_x_position=1.0, 
-                min_y_position=0.0, max_y_position=1.0, take_start=0, take_count=1, face_index=0):
-        result = reactor.restore_face(
-            self, image, model, visibility, codeformer_weight, facedetection, 
-            face_selection, sort_by, descending, min_x_position, max_x_position, 
-            min_y_position, max_y_position, take_start, take_count, face_index
-        )
+    # def __init__(self):
+    #     self.face_helper = None
+    #     self.face_size = 512
+
+    def execute(self, image, model, visibility, codeformer_weight, facedetection):
+        result = reactor.restore_face(self,image,model,visibility,codeformer_weight,facedetection)
         return (result,)
 
 
 class MaskHelper:
     def __init__(self):
+        # self.threshold = 0.5
+        # self.dilation = 10
+        # self.crop_factor = 3.0
+        # self.drop_size = 1
         self.labels = "all"
         self.detailer_hook = None
         self.device_mode = "AUTO"
         self.detection_hint = "center-1"
-        self._sam_cache = {}
-        self._bbox_cache = {}
+        # self.sam_dilation = 0
+        # self.sam_threshold = 0.93
+        # self.bbox_expansion = 0
+        # self.mask_hint_threshold = 0.7
+        # self.mask_hint_use_negative = "False"
+        # self.force_resize_width = 0
+        # self.force_resize_height = 0
+        # self.resize_behavior = "source_size"
     
     @classmethod
     def INPUT_TYPES(s):
@@ -1085,114 +891,82 @@ class MaskHelper:
     CATEGORY = "🌌 ReActor"
 
     def execute(self, image, swapped_image, bbox_model_name, bbox_threshold, bbox_dilation, bbox_crop_factor, bbox_drop_size, sam_model_name, sam_dilation, sam_threshold, bbox_expansion, mask_hint_threshold, mask_hint_use_negative, morphology_operation, morphology_distance, blur_radius, sigma_factor, mask_optional=None):
-        device = model_management.get_torch_device()
 
         # images = [image[i:i + 1, ...] for i in range(image.shape[0])]
-        # Оптимально перемещаем тензоры
-        if isinstance(image, torch.Tensor) and image.device != device:
-            image = image.to(device)
 
         images = image
 
-        if mask_optional is not None:
-            combined_mask = mask_optional
-        else:
-            # Load and cache BBox model
-            if bbox_model_name not in self._bbox_cache:
-                bbox_model_path = folder_paths.get_full_path("ultralytics", bbox_model_name)
-                model = subcore.load_yolo(bbox_model_path)
-                self._bbox_cache[bbox_model_name] = subcore.UltraBBoxDetector(model)
-            bbox_detector = self._bbox_cache[bbox_model_name]
-
-            segs_all, seg_labels = bbox_detector.detect(image, bbox_threshold, bbox_dilation, bbox_crop_factor, bbox_drop_size, self.detailer_hook)
-
-            if self.labels != 'all':
-                labels = self.labels.split(',') if isinstance(self.labels, str) else self.labels
-                segs_all, _ = masking_segs.filter(segs_all, labels)
-
-            # Load and cache SAM model
-            if sam_model_name not in self._sam_cache:
-                sam_model_path = folder_paths.get_full_path("sams", sam_model_name)
-                if 'vit_h' in sam_model_name:
-                    model_kind = 'vit_h'
-                elif 'vit_l' in sam_model_name:
-                    model_kind = 'vit_l'
-                else:
-                    model_kind = 'vit_b'
-                sam = sam_model_registry[model_kind](checkpoint=sam_model_path)
-                size = os.path.getsize(sam_model_path)
-                sam.safe_to = core.SafeToGPU(size)
-                sam.safe_to.to_device(sam, device)
-                sam.is_auto_mode = self.device_mode == "AUTO"
-                self._sam_cache[sam_model_name] = sam
-            else:
-                sam = self._sam_cache[sam_model_name]
-
-            # Handle batched input
-            if image.ndim == 4:
-                combined_masks = []
-                for i in range(image.shape[0]):
-                    segs_i = segs_all[i] if i < len(segs_all) else []
-                    segs_tuple = ([segs_i], seg_labels) if isinstance(segs_i, dict) else (segs_i, seg_labels)
-                    image_device = image.to(device) if image.device != device else image
-                    image_i = image_device[i]
-                    mask_i, _ = core.make_sam_mask_segmented(
-                        sam, segs_tuple, image_i, self.detection_hint,
-                        sam_dilation, sam_threshold, bbox_expansion,
-                        mask_hint_threshold, mask_hint_use_negative
-                    )
-                    combined_masks.append(mask_i)
-                combined_mask = torch.stack(combined_masks)
+        if mask_optional is None:
+
+            bbox_model_path = folder_paths.get_full_path("ultralytics", bbox_model_name)
+            bbox_model = subcore.load_yolo(bbox_model_path)
+            bbox_detector = subcore.UltraBBoxDetector(bbox_model)
+
+            segs = bbox_detector.detect(images, bbox_threshold, bbox_dilation, bbox_crop_factor, bbox_drop_size, self.detailer_hook)
+
+            if isinstance(self.labels, list):
+                self.labels = str(self.labels[0])
+
+            if self.labels is not None and self.labels != '':
+                self.labels = self.labels.split(',')
+                if len(self.labels) > 0:
+                    segs, _ = masking_segs.filter(segs, self.labels)
+            # segs, _ = masking_segs.filter(segs, "all")
+            
+            sam_modelname = folder_paths.get_full_path("sams", sam_model_name)
+
+            if 'vit_h' in sam_model_name:
+                model_kind = 'vit_h'
+            elif 'vit_l' in sam_model_name:
+                model_kind = 'vit_l'
             else:
-                image_device = image.to(device) if image.device != device else image
-                combined_mask, _ = core.make_sam_mask_segmented(
-                    sam, (segs_all, seg_labels), image_device, self.detection_hint,
-                    sam_dilation, sam_threshold, bbox_expansion,
-                    mask_hint_threshold, mask_hint_use_negative
-                )
+                model_kind = 'vit_b'
+
+            sam = sam_model_registry[model_kind](checkpoint=sam_modelname)
+            size = os.path.getsize(sam_modelname)
+            sam.safe_to = core.SafeToGPU(size)
+
+            device = model_management.get_torch_device()
 
-        # # *** MASK TO IMAGE ***:
+            sam.safe_to.to_device(sam, device)
+
+            sam.is_auto_mode = self.device_mode == "AUTO"
+
+            combined_mask, _ = core.make_sam_mask_segmented(sam, segs, images, self.detection_hint, sam_dilation, sam_threshold, bbox_expansion, mask_hint_threshold, mask_hint_use_negative)
         
-        # mask_image = combined_mask.reshape((-1, 1, combined_mask.shape[-2], combined_mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+        else:
+            combined_mask = mask_optional
+
+        # *** MASK TO IMAGE ***:
+        
+        mask_image = combined_mask.reshape((-1, 1, combined_mask.shape[-2], combined_mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
 
-        # # *** MASK MORPH ***:
+        # *** MASK MORPH ***:
 
-        # mask_image = core.tensor2mask(mask_image)
+        mask_image = core.tensor2mask(mask_image)
 
-        # Morph operations
         if morphology_operation == "dilate":
-            combined_mask = self.iterative_morphology(combined_mask, morphology_distance, op="dilate")
+            mask_image = self.dilate(mask_image, morphology_distance)
         elif morphology_operation == "erode":
-            combined_mask = self.iterative_morphology(combined_mask, morphology_distance, op="erode")
+            mask_image = self.erode(mask_image, morphology_distance)
         elif morphology_operation == "open":
-            combined_mask = self.iterative_morphology(self.iterative_morphology(combined_mask, morphology_distance, op="erode"), morphology_distance, op="dilate")
+            mask_image = self.erode(mask_image, morphology_distance)
+            mask_image = self.dilate(mask_image, morphology_distance)
         elif morphology_operation == "close":
-            combined_mask = self.iterative_morphology(self.iterative_morphology(combined_mask, morphology_distance, op="dilate"), morphology_distance, op="erode")
+            mask_image = self.dilate(mask_image, morphology_distance)
+            mask_image = self.erode(mask_image, morphology_distance)
         
-        # # *** MASK BLUR ***:
+        # *** MASK BLUR ***:
         
-        # if len(mask_image.size()) == 3:
-        #     mask_image = mask_image.unsqueeze(3)
-        # Gaussian blur
-
-        if blur_radius > 0:
-            blur = T.GaussianBlur(kernel_size=blur_radius * 2 + 1, sigma=sigma_factor)
-            mask_blurred = blur(combined_mask.unsqueeze(1)).squeeze(1)
-        else:
-            mask_blurred = combined_mask
+        if len(mask_image.size()) == 3:
+            mask_image = mask_image.unsqueeze(3)
         
-        # mask_image = mask_image.permute(0, 3, 1, 2)
-        # kernel_size = blur_radius * 2 + 1
-        # sigma = sigma_factor * (0.6 * blur_radius - 0.3)
-        # mask_image_final = self.gaussian_blur(mask_image, kernel_size, sigma).permute(0, 2, 3, 1)
-        # if mask_image_final.size()[3] == 1:
-        #     mask_image_final = mask_image_final[:, :, :, 0]
-
-        # Apply mask to swapped image (basic RGBA composite)
-        swapped_image = swapped_image.to(device) if swapped_image.device != device else swapped_image
-        swapped_rgba = core.tensor2rgba(swapped_image)
-
-        mask_image_final = mask_blurred
+        mask_image = mask_image.permute(0, 3, 1, 2)
+        kernel_size = blur_radius * 2 + 1
+        sigma = sigma_factor * (0.6 * blur_radius - 0.3)
+        mask_image_final = self.gaussian_blur(mask_image, kernel_size, sigma).permute(0, 2, 3, 1)
+        if mask_image_final.size()[3] == 1:
+            mask_image_final = mask_image_final[:, :, :, 0]
 
         # *** CUT BY MASK ***:
         
@@ -1231,11 +1005,15 @@ class MaskHelper:
         use_width = int(torch.max(width).item())
         use_height = int(torch.max(height).item())
 
+        # if self.force_resize_width > 0:
+        #     use_width = self.force_resize_width
+
+        # if self.force_resize_height > 0:
+        #     use_height = self.force_resize_height
+
         alpha_mask = torch.ones((B, H, W, 4))
         alpha_mask[:,:,:,3] = mask
 
-        alpha_mask = alpha_mask.to(device) if alpha_mask.device != device else alpha_mask
-
         swapped_image = swapped_image * alpha_mask
 
         cutted_image = torch.zeros((B, use_height, use_width, 4))
@@ -1257,7 +1035,7 @@ class MaskHelper:
 
         # *** PASTE BY MASK ***:
 
-        image_base = core.tensor2rgba(image)
+        image_base = core.tensor2rgba(images)
         image_to_paste = core.tensor2rgba(cutted_image)
         mask = core.tensor2mask(mask_image_final)
 
@@ -1307,18 +1085,35 @@ class MaskHelper:
                 continue
             else:
                 image_index = i
-                # source_size = image_to_paste.size()
+                source_size = image_to_paste.size()
                 SB, SH, SW, _ = image_to_paste.shape
 
                 # Figure out the desired size
                 width = int(target_width[i].item())
                 height = int(target_height[i].item())
+                # if self.resize_behavior == "keep_ratio_fill":
+                #     target_ratio = width / height
+                #     actual_ratio = SW / SH
+                #     if actual_ratio > target_ratio:
+                #         width = int(height * actual_ratio)
+                #     elif actual_ratio < target_ratio:
+                #         height = int(width / actual_ratio)
+                # elif self.resize_behavior == "keep_ratio_fit":
+                #     target_ratio = width / height
+                #     actual_ratio = SW / SH
+                #     if actual_ratio > target_ratio:
+                #         height = int(width / actual_ratio)
+                #     elif actual_ratio < target_ratio:
+                #         width = int(height * actual_ratio)
+                # elif self.resize_behavior == "source_size" or self.resize_behavior == "source_size_unmasked":
 
                 width = SW
                 height = SH
 
                 # Resize the image we're pasting if needed
                 resized_image = image_to_paste[i].unsqueeze(0)
+                # if SH != height or SW != width:
+                #     resized_image = torch.nn.functional.interpolate(resized_image.permute(0, 3, 1, 2), size=(height,width), mode='bicubic').permute(0, 2, 3, 1)
 
                 pasting = torch.ones([H, W, C])
                 ymid = float(mid_y[i].item())
@@ -1350,11 +1145,12 @@ class MaskHelper:
                 pasting_alpha = torch.zeros([H, W])
                 pasting_alpha[ymin:ymax, xmin:xmax] = resized_image[0, source_ymin:source_ymax, source_xmin:source_xmax, 3]
 
+                # if self.resize_behavior == "keep_ratio_fill" or self.resize_behavior == "source_size_unmasked":
+                #     # If we explicitly want to fill the area, we are ok with extending outside
+                #     paste_mask = pasting_alpha.unsqueeze(2).repeat(1, 1, 4)
+                # else:
+                #     paste_mask = torch.min(pasting_alpha, mask[i]).unsqueeze(2).repeat(1, 1, 4)
                 paste_mask = torch.min(pasting_alpha, mask[i]).unsqueeze(2).repeat(1, 1, 4)
-                
-                pasting = pasting.to(device) if pasting.device != device else pasting
-                paste_mask = paste_mask.to(device) if paste_mask.device != device else paste_mask
-
                 result[image_index] = pasting * paste_mask + result[image_index] * (1. - paste_mask)
 
                 face_segment = result
@@ -1363,24 +1159,44 @@ class MaskHelper:
 
                 result = rgba2rgb_tensor(result)
         
-        # return (result,combined_mask,mask_image_final,face_segment,)
-        try:
-            torch.cuda.empty_cache()
-        except:
-            pass
-
-        return (result, combined_mask, mask_blurred, face_segment)
+        return (result,combined_mask,mask_image_final,face_segment,)
+
+    def gaussian_blur(self, image, kernel_size, sigma):
+        kernel = torch.Tensor(kernel_size, kernel_size).to(device=image.device)
+        center = kernel_size // 2
+        variance = sigma**2
+        for i in range(kernel_size):
+            for j in range(kernel_size):
+                x = i - center
+                y = j - center
+                kernel[i, j] = math.exp(-(x**2 + y**2)/(2*variance))
+        kernel /= kernel.sum()
+
+        # Pad the input tensor
+        padding = (kernel_size - 1) // 2
+        input_pad = torch.nn.functional.pad(image, (padding, padding, padding, padding), mode='reflect')
+
+        # Reshape the padded input tensor for batched convolution
+        batch_size, num_channels, height, width = image.shape
+        input_reshaped = input_pad.reshape(batch_size*num_channels, 1, height+padding*2, width+padding*2)
+
+        # Perform batched convolution with the Gaussian kernel
+        output_reshaped = torch.nn.functional.conv2d(input_reshaped, kernel.unsqueeze(0).unsqueeze(0))
+
+        # Reshape the output tensor to its original shape
+        output_tensor = output_reshaped.reshape(batch_size, num_channels, height, width)
+
+        return output_tensor
     
-    def iterative_morphology(self, image, distance, op="dilate"):
-        if distance <= 0:
-            return image
-        image = image.unsqueeze(1)  # shape [B, 1, H, W] or [1, 1, H, W]
-        for _ in range(distance):
-            if op == "dilate":
-                image = F.max_pool2d(image, kernel_size=3, stride=1, padding=1)
-            elif op == "erode":
-                image = -F.max_pool2d(-image, kernel_size=3, stride=1, padding=1)
-        return image.squeeze(1)
+    def erode(self, image, distance):
+        return 1. - self.dilate(1. - image, distance)
+
+    def dilate(self, image, distance):
+        kernel_size = 1 + distance * 2
+        # Add the channels dimension
+        image = image.unsqueeze(1)
+        out = torchfn.max_pool2d(image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2).squeeze(1)
+        return out
 
 
 class ImageDublicator: