Upload reactor_swapper.py

custom_nodes/comfyui-reactor-node/scripts/reactor_swapper.py

@@ -14,13 +14,14 @@ from insightface.app.common import Face
 # except:
 #     cuda = None
 import torch
+import gc # для очистки кэша
 
 import folder_paths
 import comfy.model_management as model_management
 from modules.shared import state
 
 # 1. Добавьте импорт logging наверху файла, если его там нет:
-import logging
+# import logging - закомментирован 
 from scripts.reactor_logger import logger
 from reactor_utils import (
     move_path,
@@ -126,9 +127,21 @@ def getFaceSwapModel(model_path: str):
     if FS_MODEL is None or CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path:
         CURRENT_FS_MODEL_PATH = model_path
         FS_MODEL = unload_model(FS_MODEL)
-        if "hyperswap" in model_path.lower():
-            FS_MODEL = ort.InferenceSession(model_path, providers=providers)
+        
+        # Извлекаем имя файла модели из пути
+        model_filename = os.path.basename(model_path)
+        
+        # Определяем правильный путь в зависимости от типа модели
+        if "hyperswap" in model_filename.lower():
+            # Ищем в директории hyperswap
+            correct_path = os.path.join(folder_paths.models_dir, "hyperswap", model_filename)
+            FS_MODEL = ort.InferenceSession(correct_path, providers=providers)
+        elif "reswapper" in model_filename.lower():
+            # Ищем в директории reswapper
+            correct_path = os.path.join(folder_paths.models_dir, "reswapper", model_filename)
+            FS_MODEL = insightface.model_zoo.get_model(correct_path, providers=providers)
         else:
+            # Для моделей insightface используем оригинальный путь
             FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
     return FS_MODEL
     
@@ -137,10 +150,10 @@ def get_landmarks_5(face):
     # face.landmark_5: np.ndarray shape (5,2)
     # Если нет, попробуй face.kps или face.landmark
     if hasattr(face, 'landmark_5') and face.landmark_5 is not None:
-        logger.debug("landmark_5: %s", face.landmark_5)
+     #  logger.debug("landmark_5: %s", face.landmark_5)
         return face.landmark_5
     elif hasattr(face, 'kps') and face.kps is not None:
-        logger.debug("kps: %s", face.kps)
+     #  logger.debug("kps: %s", face.kps)
         return face.kps
     elif hasattr(face, 'landmark') and face.landmark is not None:
         # 68-точечная разметка, берём нужные индексы
@@ -148,9 +161,9 @@ def get_landmarks_5(face):
         # Пример: [36, 45, 30, 48, 54] — левый/правый глаз, нос, левый/правый рот
         if face.landmark.shape[0] >= 68:
             idxs = [36, 45, 30, 48, 54]
-            logger.debug("landmark (68 точек): %s", face.landmark[idxs])
+         #  logger.debug("landmark (68 точек): %s", face.landmark[idxs])
             return face.landmark[idxs]
-    logger.warning("Нет подходящих точек в объекте Face. Доступные атрибуты: %s", dir(face))
+ #  logger.warning("Нет подходящих точек в объекте Face. Доступные атрибуты: %s", dir(face))
     return None
     
 #### Что проверить:
@@ -178,14 +191,14 @@ def create_gradient_mask(crop_size=256):
     
     # 3. Рисуем эллипс (заполняем белым цветом, значение=1.0)
     cv2.ellipse(
-        mask,          
-        center,        
-        axes,          
-        angle=0,       
-        startAngle=0,  
-        endAngle=360,  
-        color=1.0,     
-        thickness=-1   
+        mask,          # Массив для рисования
+        center,        # Центр эллипса
+        axes,          # Полуоси (ширина, высота)
+        angle=0,       # Угол поворота
+        startAngle=0,  # Начальный угол дуги
+        endAngle=360,  # Конечный угол дуги (360 = полный эллипс)
+        color=1.0,     # Значение для заполнения (белый = 1.0)
+        thickness=-1   # -1 = заполнить всю область эллипса   
     )
     
     # 4. Применяем размытие для плавных краёв
@@ -212,7 +225,7 @@ def paste_back(target_img, swapped_face, M, crop_size=256):
     mask = create_gradient_mask(crop_size)
 
     # Сохраняем для отладки
-    cv2.imwrite("debug_original_mask.png", (mask * 255).astype(np.uint8))
+ #  cv2.imwrite("debug_original_mask.png", (mask * 255).astype(np.uint8))
 
     # Преобразуем в трехканальную маску
     mask_3c = np.stack([mask] * 3, axis=2)
@@ -223,7 +236,7 @@ def paste_back(target_img, swapped_face, M, crop_size=256):
     # 3. Обратное преобразование (WARP_INVERSE_MAP) для лица И маски
     # Для лица (INTER_LANCZOS4 — высококачественная интерполяция)
     inv_face = cv2.warpAffine(
-        swapped_face.astype(np.float32),
+        swapped_face.astype(np.float32).copy(),
         M,
         (w, h),
         flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
@@ -246,13 +259,13 @@ def paste_back(target_img, swapped_face, M, crop_size=256):
     inv_mask = cv2.GaussianBlur(inv_mask, (3, 3), 0)
 
     # Отладка
-    logger.debug("Warped face shape: %s | Min: %s | Max: %s", inv_face.shape, inv_face.min(), inv_face.max())
-    logger.debug("Warped mask shape: %s | Min: %s | Max: %s", inv_mask.shape, inv_mask.min(), inv_mask.max())
+  # logger.debug("Warped face shape: %s | Min: %s | Max: %s", inv_face.shape, inv_face.min(), inv_face.max())
+  # logger.debug("Warped mask shape: %s | Min: %s | Max: %s", inv_mask.shape, inv_mask.min(), inv_mask.max())
 
     # Визуализация
-#   cv2.imshow("Warped Face", inv_face.astype(np.uint8))
-#   cv2.imshow("Warped Mask", (inv_mask * 255).astype(np.uint8))
-#   cv2.waitKey(1)
+  # cv2.imshow("Warped Face", inv_face.astype(np.uint8))
+  # cv2.imshow("Warped Mask", (inv_mask * 255).astype(np.uint8))
+  # cv2.waitKey(1)
 
     # 6. Плавное наложение
     target_img_float = target_img.astype(np.float32)
@@ -263,7 +276,7 @@ def paste_back(target_img, swapped_face, M, crop_size=256):
     result = np.clip(result, 0, 255).astype(np.uint8)
 
     # Сохраняем результат
-    cv2.imwrite("debug_result.png", result)
+  # cv2.imwrite("debug_result.png", result)
 
     return result
     
@@ -276,8 +289,8 @@ def visualize_points(img, points, color=(0, 255, 0)):
     img = img.copy()
     for p in points:
         cv2.circle(img, tuple(p.astype(int)), 3, color, -1)
-#   cv2.imshow("Face Points", img)
-#   cv2.waitKey(1)    
+  # cv2.imshow("Face Points", img)
+  # cv2.waitKey(1)    
 
 # Итоговая функция run_hyperswap с аффинным преобразованием
 def run_hyperswap(session, source_face, target_face, target_img):
@@ -289,7 +302,7 @@ def run_hyperswap(session, source_face, target_face, target_img):
     visualize_points(target_img, target_landmarks_5, (0, 255, 0))  # Зеленые точки
     
     if target_landmarks_5 is None:
-        logger.error("Не удалось получить 5 точек для целевого лица")
+      # logger.error("Не удалось получить 5 точек для целевого лица")
         # Важно: Если ошибка, возвращаем None и исходную матрицу (или обрабатываем ошибку иначе)
         return None, None
 
@@ -304,7 +317,7 @@ def run_hyperswap(session, source_face, target_face, target_img):
 
     # Вычисляем аффинную матрицу
     M = get_affine_transform(target_landmarks_5.astype(np.float32), std_landmarks_256)
-    logger.debug("Affine Matrix M (used for cropping):\n%s", M)
+  # logger.debug("Affine Matrix M (used for cropping):\n%s", M)
 
 #### Что проверить:
 # Матрица `M` не должна содержать `NaN` или бесконечности.
@@ -317,9 +330,9 @@ def run_hyperswap(session, source_face, target_face, target_img):
     #### Что проверить:
 # Окно `"Crop Before Inference"` должно показывать лицо, вырезанное по аффинному преобразованию.
 # Если изображение черное — проблема в `M` или `target_landmarks_5`.
-    logger.debug("Crop shape: %s | Min: %s | Max: %s", crop.shape, crop.min(), crop.max())
-#   cv2.imshow("Crop Before Inference", crop)
-#   cv2.waitKey(1)  # Отображает изображение
+  # logger.debug("Crop shape: %s | Min: %s | Max: %s", crop.shape, crop.min(), crop.max())
+  # cv2.imshow("Crop Before Inference", crop)
+  # cv2.waitKey(1)  # Отображает изображение
 
     # 4. Преобразуем crop для модели
     # crop_input = crop[:, :, ::-1] / 255.0
@@ -333,9 +346,9 @@ def run_hyperswap(session, source_face, target_face, target_img):
     # 5. Инференс
     try:
         output = session.run(None, {'source': source_embedding, 'target': crop_input})[0][0]
-        logger.debug("Model output shape: %s | Min: %s | Max: %s", output.shape, output.min(), output.max())
+      # logger.debug("Model output shape: %s | Min: %s | Max: %s", output.shape, output.min(), output.max())
     except Exception as e:
-        logger.error("Ошибка выполнения модели: %s", e)
+      # logger.error("Ошибка выполнения модели: %s", e)
         return target_img
 
     # 6. Обратная нормализация
@@ -353,15 +366,21 @@ def run_hyperswap(session, source_face, target_face, target_img):
 #### Что проверить:
 # `output` должен быть в диапазоне `[0..255]` и содержать лицо.
 # Если `output` черный — проблема в нормализации/денормализации или в самой модели.
-    logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
-#   cv2.imshow("Output After Denormalization", output)
-#   cv2.waitKey(1)
+  # logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
+  # cv2.imshow("Output After Denormalization", output)
+  # cv2.waitKey(1)
     
     # 7. ИЗМЕНЕНИЕ: Возвращаем результат 256x256 И матрицу M
     # Мы больше не вызываем warp_and_paste_face отсюда.
     # return result
     
-    return output, M # Возвращаем лицо (256x256) и матрицу M    
+    # --- Безопасное преобразование перед возвратом ---
+    # Убедимся, что результат всегда float32 в диапазоне [0..255]
+    output = output.astype(np.float32)
+    output = np.clip(output, 0, 255)
+    torch.cuda.empty_cache()
+    gc.collect()
+    return output, M # Возвращаем лицо (256x256) и матрицу M
 
 def sort_by_order(face, order: str):
     if order == "left-right":
@@ -456,13 +475,13 @@ def get_face_single(img_data: np.ndarray, face, face_index=0, det_size=(640, 640
     try:
         faces_sorted = sort_by_order(face, order)
         selected_face = faces_sorted[face_index]
-        logger.debug("Выбрано лицо: bbox=%s, landmark_5=%s, kps=%s, landmark=%s",
-                     selected_face.bbox,
-                     hasattr(selected_face, "landmark_5"),
-                     hasattr(selected_face, "kps"),
-                     hasattr(selected_face, "landmark"))
+      # logger.debug("Выбрано лицо: bbox=%s, landmark_5=%s, kps=%s, landmark=%s",
+      #          selected_face.bbox,
+      #          hasattr(selected_face, "landmark_5"),
+      #          hasattr(selected_face, "kps"),
+      #          hasattr(selected_face, "landmark"))
         return selected_face, 0
-        return faces_sorted[face_index], 0
+        # return faces_sorted[face_index], 0
         # return sorted(face, key=lambda x: x.bbox[0])[face_index], 0
     except IndexError:
         return None, 0
@@ -484,14 +503,14 @@ def swap_face(
     codeformer_weight: float = 0.5,
     interpolation: str = "Bicubic",
 ):
-    # >>>>> РЕШЕНИЕ: Принудительная установка уровня DEBUG <<<<<
-    if logger.getEffectiveLevel() != logging.DEBUG:
-        print("\n--- [ReActor Debug] Принудительная установка уровня логирования на DEBUG (10) в swap_face ---")
-        logger.setLevel(logging.DEBUG)
+  #   >>>>> РЕШЕНИЕ: Принудительная установка уровня DEBUG <<<<<
+  # if logger.getEffectiveLevel() != logging.DEBUG:
+  #     print("\n--- [ReActor Debug] Принудительная установка уровня логирования на DEBUG (10) в swap_face ---")
+  #     logger.setLevel(logging.DEBUG)
     
-    # Проверочное сообщение (теперь оно должно появиться)
-    logger.debug("--- ТЕСТ: swap_face запущена, уровень логирования DEBUG активен. ---")
-    # >>>>> КОНЕЦ РЕШЕНИЯ <<<<<
+  #   Проверочное сообщение (теперь оно должно появиться)
+  # logger.debug("--- ТЕСТ: swap_face запущена, уровень логирования DEBUG активен. ---")
+  #   >>>>> КОНЕЦ РЕШЕНИЯ <<<<<
     
     global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH
     result_image = target_img
@@ -607,13 +626,19 @@ def swap_face(
                     if len(source_faces_index) > 1 and source_face_idx > 0:
                         source_face, src_wrong_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[source_face_idx], gender_source=gender_source, order=faces_order[1])
                     source_face_idx += 1
+                    
+                    # очистка кэша и памяти перед каждым Hyperswap вызовом
+                    torch.cuda.empty_cache()
+                    gc.collect()
+                    SOURCE_FACES = None
+                    TARGET_FACES = None
 
                     if source_face is not None and src_wrong_gender == 0:
                         target_face, wrong_gender = get_face_single(target_img, target_faces, face_index=face_num, gender_target=gender_target, order=faces_order[0])
                         if target_face is not None and wrong_gender == 0:
                             logger.status(f"Swapping...")
                             if "hyperswap" in model:
-                                logger.status(f"Swapping with Hyperswap...")
+                              # logger.status(f"Swapping with Hyperswap...")
                                 swapped_face_256, M = run_hyperswap(face_swapper, source_face, target_face, result)
                                 if swapped_face_256 is not None:
                                     result = paste_back(result, swapped_face_256, M, crop_size=256)
@@ -799,7 +824,17 @@ def swap_face_many(
                 logger.status(f'Source Faces must have no entries (default=0), one entry, or same number of entries as target faces.')
             elif source_face is not None:
                 results = target_imgs
-                model_path = model_path = os.path.join(insightface_path, model)
+                
+                # Определяем путь к модели в зависимости от типа
+                if "inswapper" in model:
+                    model_path = os.path.join(insightface_path, model)
+                elif "reswapper" in model:
+                    model_path = os.path.join(reswapper_path, model)
+                elif "hyperswap" in model:
+                    model_path = os.path.join(hyperswap_path, model)
+                else:
+                    model_path = os.path.join(insightface_path, model)
+                    
                 face_swapper = getFaceSwapModel(model_path)
 
                 source_face_idx = 0
@@ -823,16 +858,23 @@ def swap_face_many(
                             target_face_single, wrong_gender = get_face_single(target_img, target_face, face_index=face_num, gender_target=gender_target, order=faces_order[0])
                             if target_face_single is not None and wrong_gender == 0:
                                 result = target_img
-                                # logger.status(f"Swapping {i}...")
-                                if face_boost_enabled:
+                                
+                                # Обработка в зависимости от типа модели
+                                if "hyperswap" in model:
+                                    # Для Hyperswap используем специальную функцию
+                                    swapped_face_256, M = run_hyperswap(face_swapper, source_face, target_face_single, result)
+                                    if swapped_face_256 is not None:
+                                        result = paste_back(result, swapped_face_256, M, crop_size=256)
+                                elif face_boost_enabled:
                                     logger.status(f"Face Boost is enabled")
                                     bgr_fake, M = face_swapper.get(target_img, target_face_single, source_face, paste_back=False)
                                     bgr_fake, scale = restorer.get_restored_face(bgr_fake, face_restore_model, face_restore_visibility, codeformer_weight, interpolation)
                                     M *= scale
                                     result = swapper.in_swap(target_img, bgr_fake, M)
                                 else:
-                                    # logger.status(f"Swapping as-is")
+                                    # Для остальных моделей используем метод get()
                                     result = face_swapper.get(target_img, target_face_single, source_face)
+                                    
                                 results[i] = result
                                 pbar.update(1)
                             elif wrong_gender == 1: