Upload reactor_swapper.py

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

@@ -200,71 +200,71 @@ def create_gradient_mask(crop_size=256):
 #### 1. Используйте `cv2.BORDER_TRANSPARENT` (OpenCV ≥ 4.5)  
 # Этот флаг позволяет **не заполнять** области за пределами маски никаким цветом (пиксели остаются `0` или "прозрачные").    
 def paste_back(target_img, swapped_face, M, crop_size=256):
+    # Улучшенная функция paste_back с идеальной овальной маской и исправлениями артефактов
     
-    # Улучшенная функция paste_back с идеальной овальной маской
-#    Стандартная функция InsightFace для вставки лица обратно с использованием WARP_INVERSE_MAP.   
-#    target_img: Исходное изображение (BGR, numpy)
-#    swapped_face: Результат работы модели (256x256, BGR, numpy)
-#    M: Матрица аффинного преобразования (Target -> Crop)
-#    crop_size: Размер кропа (для HyperSwap это 256)
+    # target_img: Исходное изображение (BGR, numpy, uint8)
+    # swapped_face: Результат работы модели (256x256, BGR, uint8)
+    # M: Матрица аффинного преобразования (Target -> Crop), но здесь используется M_inv из run_hyperswap
+    # crop_size: Размер кропа (для HyperSwap это 256)
     
     # 1. Создание мягкой маски (Эрозия + Размытие)
     mask = create_gradient_mask(crop_size)
-
-    # Сохраняем для отладки
-    cv2.imwrite("debug_original_mask.png", (mask * 255).astype(np.uint8))
-
+    
     # Преобразуем в трехканальную маску
     mask_3c = np.stack([mask] * 3, axis=2)
-
+    
     # 2. Получаем размеры целевого изображения
     h, w = target_img.shape[:2]
-
-    # 3. Обратное преобразование (WARP_INVERSE_MAP) для лица И маски
-    # Для лица (INTER_LANCZOS4 — высококачественная интерполяция)
-    inv_face = cv2.warpAffine(
-        swapped_face.astype(np.float32),
-        M,
+    
+    # 3. Нормализация swapped_face к float32 [0,1] для warp
+    swapped_face_norm = swapped_face.astype(np.float32) / 255.0
+    mask_norm = mask_3c.astype(np.float32)  # Маска уже [0,1]
+    
+    # 4. Обратное преобразование (WARP_INVERSE_MAP) для лица И маски
+    # Используем BORDER_CONSTANT с borderValue=0.5 (серый, чтобы избежать синих/зеленых артефактов)
+    warped_face = cv2.warpAffine(
+        swapped_face_norm,
+        M,  # Это M_inv из run_hyperswap
         (w, h),
         flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
-        borderMode=cv2.BORDER_TRANSPARENT  # <<< ИЗМЕНЕНИЕ: TRANSPARENT вместо CONSTANT
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=0.5  # Серый фон вместо черного/белого
     )
-
-    # Для маски (INTER_CUBIC — плавные границы)
-    inv_mask = cv2.warpAffine(
-        mask_3c,
-        M,
+    
+    warped_mask = cv2.warpAffine(
+        mask_norm,
+        M,  # Это M_inv из run_hyperswap
         (w, h),
         flags=cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP,
-        borderMode=cv2.BORDER_TRANSPARENT  # <<< ИЗМЕНЕНИЕ: TRANSPARENT вместо CONSTANT
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=0.0  # Маска: 0 за пределами
     )
-
-    # 4. Ограничение значений маски [0, 1]
-    inv_mask = np.clip(inv_mask, 0, 1)
-
-    # 5. Дополнительное размытие для устранения артефактов
-    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())
-
-    # Визуализация
-  # 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)
-    inv_face_float = inv_face.astype(np.float32)
-    result = target_img_float * (1.0 - inv_mask) + inv_face_float * inv_mask
-
-    # 7. Ограничение результата [0, 255]
-    result = np.clip(result, 0, 255).astype(np.uint8)
-
-    # Сохраняем результат
-    cv2.imwrite("debug_result.png", result)
-
+    
+    # 5. Обработка после warp: Clip, NaN fix
+    warped_face = np.clip(warped_face, 0, 1)  # Убираем отрицательные
+    warped_face = np.nan_to_num(warped_face, nan=0.5)  # NaN -> серый
+    
+    warped_mask = np.clip(warped_mask, 0, 1)
+    warped_mask = np.nan_to_num(warped_mask, nan=0.0)
+    
+    # 6. Дополнительное размытие для устранения артефактов (опционально, но помогает)
+    warped_mask = cv2.GaussianBlur(warped_mask, (3, 3), 0)
+    
+    # Отладочные логи (добавьте после warp)
+    logger.debug("Warped face shape: %s | Min: %s | Max: %s | NaN count: %s", 
+                 warped_face.shape, warped_face.min(), warped_face.max(), np.isnan(warped_face).sum())
+    logger.debug("Warped mask shape: %s | Min: %s | Max: %s | NaN count: %s", 
+                 warped_mask.shape, warped_mask.min(), warped_mask.max(), np.isnan(warped_mask).sum())
+    
+    # 7. Плавное наложение в float32
+    target_float = target_img.astype(np.float32) / 255.0
+    result_float = target_float * (1.0 - warped_mask) + warped_face * warped_mask
+    
+    # 8. Обратная нормализация к uint8
+    result = (result_float * 255).clip(0, 255).astype(np.uint8)
+    
+    logger.debug("Final result: shape %s | Min: %s | Max: %s", result.shape, result.min(), result.max())
+    
     return result
     
 #### Что проверить:
@@ -338,7 +338,6 @@ def run_hyperswap(session, source_face, target_face, target_img):
 
     # --- CPU FLOAT NORMALIZATION FIX ---
     # предотвращает появление "синей кожи" и "шума" при работе на CPU
-    # (адаптировано из патча patch_cpu_fix.diff)
     if isinstance(output, np.ndarray):
         # устранение NaN и бесконечностей
         output = np.nan_to_num(output, nan=0.0, posinf=255.0, neginf=0.0)
@@ -360,7 +359,7 @@ def run_hyperswap(session, source_face, target_face, target_img):
     # (ваш код без изменений, но без старой денормализации)
     output = output.transpose(1, 2, 0)  # CHW -> HWC
     output = output[:, :, ::-1]  # RGB -> BGR (Убедитесь, что это BGR, если вход был BGR)
-    logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
+    logger.debug("Output after denormalization: Min: %s | Max: %s | NaN count: %s", output.min(), output.max(), np.isnan(output).sum())
     
     # Визуализация после денормализации
     #### Что проверить: