Upload reactor_swapper.py

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

@@ -14,14 +14,13 @@ 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,
@@ -127,21 +126,9 @@ 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)
-        
-        # Извлекаем имя файла модели из пути
-        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)
+        if "hyperswap" in model_path.lower():
+            FS_MODEL = ort.InferenceSession(model_path, providers=providers)
         else:
-            # Для моделей insightface используем оригинальный путь
             FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
     return FS_MODEL
     
@@ -150,10 +137,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-точечная разметка, берём нужные индексы
@@ -161,9 +148,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
     
 #### Что проверить:
@@ -191,14 +178,14 @@ def create_gradient_mask(crop_size=256):
     
     # 3. Рисуем эллипс (заполняем белым цветом, значение=1.0)
     cv2.ellipse(
-        mask,          # Массив для рисования
-        center,        # Центр эллипса
-        axes,          # Полуоси (ширина, высота)
-        angle=0,       # Угол поворота
-        startAngle=0,  # Начальный угол дуги
-        endAngle=360,  # Конечный угол дуги (360 = полный эллипс)
-        color=1.0,     # Значение для заполнения (белый = 1.0)
-        thickness=-1   # -1 = заполнить всю область эллипса   
+        mask,          
+        center,        
+        axes,          
+        angle=0,       
+        startAngle=0,  
+        endAngle=360,  
+        color=1.0,     
+        thickness=-1   
     )
     
     # 4. Применяем размытие для плавных краёв
@@ -225,7 +212,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)
@@ -236,7 +223,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).copy(),
+        swapped_face.astype(np.float32),
         M,
         (w, h),
         flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
@@ -259,8 +246,8 @@ 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))
@@ -276,7 +263,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
     
@@ -302,7 +289,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
 
@@ -317,7 +304,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` или бесконечности.
@@ -330,7 +317,7 @@ 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())
+    logger.debug("Crop shape: %s | Min: %s | Max: %s", crop.shape, crop.min(), crop.max())
   # cv2.imshow("Crop Before Inference", crop)
   # cv2.waitKey(1)  # Отображает изображение
 
@@ -346,9 +333,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. Обратная нормализация
@@ -366,7 +353,7 @@ 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())
+    logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
   # cv2.imshow("Output After Denormalization", output)
   # cv2.waitKey(1)
     
@@ -374,13 +361,7 @@ def run_hyperswap(session, source_face, target_face, target_img):
     # Мы больше не вызываем warp_and_paste_face отсюда.
     # return result
     
-    # --- Безопасное преобразование перед возвратом ---
-    # Убедимся, что результат всегда 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
+    return output, M # Возвращаем лицо (256x256) и матрицу M    
 
 def sort_by_order(face, order: str):
     if order == "left-right":
@@ -475,13 +456,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
@@ -503,14 +484,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
@@ -626,19 +607,13 @@ 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)
@@ -824,17 +799,7 @@ 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
-                
-                # Определяем путь к модели в зависимости от типа
-                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)
-                    
+                model_path = model_path = os.path.join(insightface_path, model)
                 face_swapper = getFaceSwapModel(model_path)
 
                 source_face_idx = 0
@@ -858,23 +823,16 @@ 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
-                                
-                                # Обработка в зависимости от типа модели
-                                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"Swapping {i}...")
+                                if 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:
-                                    # Для остальных моделей используем метод get()
+                                    # logger.status(f"Swapping as-is")
                                     result = face_swapper.get(target_img, target_face_single, source_face)
-                                    
                                 results[i] = result
                                 pbar.update(1)
                             elif wrong_gender == 1: