diff --git a/scripts/reactor_faceswap.py b/scripts/reactor_faceswap.py index b969f0c..cf27d4c 100644 --- a/scripts/reactor_faceswap.py +++ b/scripts/reactor_faceswap.py @@ -26,7 +26,8 @@ import comfy.model_management as model_management def get_models(): swappers = [ "insightface", - "reswapper" + "reswapper", + "hyperswap" ] models_list = [] for folder in swappers: diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py index fc317dd..82245ab 100644 --- a/scripts/reactor_swapper.py +++ b/scripts/reactor_swapper.py @@ -6,12 +6,10 @@ import cv2 import numpy as np from PIL import Image +import onnxruntime as ort + import insightface from insightface.app.common import Face -# try: -# import torch.cuda as cuda -# except: -# cuda = None import torch import folder_paths @@ -45,13 +43,6 @@ try: except Exception as e: logger.debug(f"ExecutionProviderError: {e}.\nEP is set to CPU.") providers = ["CPUExecutionProvider"] -# if cuda is not None: -# if cuda.is_available(): -# providers = ["CUDAExecutionProvider"] -# else: -# providers = ["CPUExecutionProvider"] -# else: -# providers = ["CPUExecutionProvider"] models_path_old = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models") insightface_path_old = os.path.join(models_path_old, "insightface") @@ -61,6 +52,7 @@ models_path = folder_paths.models_dir insightface_path = os.path.join(models_path, "insightface") insightface_models_path = os.path.join(insightface_path, "models") reswapper_path = os.path.join(models_path, "reswapper") +hyperswap_path = os.path.join(models_path, "hyperswap") if os.path.exists(models_path_old): move_path(insightface_models_path_old, insightface_models_path) @@ -122,11 +114,165 @@ 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) - FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers) + + model_filename = os.path.basename(model_path) + if "hyperswap" in model_filename.lower(): + model_path = os.path.join(folder_paths.models_dir, "hyperswap", model_filename) + FS_MODEL = ort.InferenceSession(model_path, providers=providers) + elif "reswapper" in model_filename.lower(): + model_path = os.path.join(folder_paths.models_dir, "reswapper", model_filename) + FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers) + else: + FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers) return FS_MODEL +# Функция для получения 5 ключевых точек из объекта Face +def get_landmarks_5(face): + if hasattr(face, 'landmark_5') and face.landmark_5 is not None: + return face.landmark_5 + elif hasattr(face, 'kps') and face.kps is not None: + return face.kps + elif hasattr(face, 'landmark') and face.landmark is not None: + if face.landmark.shape[0] >= 68: + idxs = [36, 45, 30, 48, 54] + return face.landmark[idxs] + return None + +# Функция для вычисления аффинного преобразования +def get_affine_transform(src_pts, dst_pts): + M, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts) + return M + +# Создаём градиентную маску овальной формы без обрезки +def create_gradient_mask(crop_size=256): + # 1. Создаём пустую маску (все пиксели = 0) + mask = np.zeros((crop_size, crop_size), dtype=np.float32) + + # 2. Определяем центр и размеры эллипса + center = (crop_size // 2, crop_size // 2) + axes = (int(crop_size * 0.35), int(crop_size * 0.4)) + + # 3. Рисуем эллипс (заполняем белым цветом, значение=1.0) + cv2.ellipse( + mask, # Массив для рисования + center, # Центр эллипса + axes, # Полуоси (ширина, высота) + angle=0, # Угол поворота + startAngle=0, # Начальный угол дуги + endAngle=360, # Конечный угол дуги (360 = полный эллипс) + color=1.0, # Значение для заполнения (белый = 1.0) + thickness=-1 # -1 = заполнить всю область эллипса + ) + + # 4. Применяем размытие для плавных краёв + blur_ksize = 15 # Нечётное число, чтобы ядро было симметричным + mask = cv2.GaussianBlur(mask, (blur_ksize, blur_ksize), 0) + + # 5. Ограничим значения в диапазоне [0, 1] + mask = np.clip(mask, 0, 1) + + return mask + +def paste_back(target_img, swapped_face, M, crop_size=256): + + # 1. Создание мягкой маски (Эрозия + Размытие) + mask = create_gradient_mask(crop_size) + + # Преобразуем в трехканальную маску + 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, + (w, h), + flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP, + borderMode=cv2.BORDER_TRANSPARENT + ) + + # Для маски (INTER_CUBIC — плавные границы) + inv_mask = cv2.warpAffine( + mask_3c, + M, + (w, h), + flags=cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP, + borderMode=cv2.BORDER_TRANSPARENT + ) + + # 4. Ограничение значений маски [0, 1] + inv_mask = np.clip(inv_mask, 0, 1) + + # 5. Дополнительное размытие для устранения артефактов + inv_mask = cv2.GaussianBlur(inv_mask, (3, 3), 0) + + # 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) + + return result + +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) + +# Итоговая функция run_hyperswap с аффинным преобразованием +def run_hyperswap(session, source_face, target_face, target_img): + # 1. Подготовка эмбеддинга + source_embedding = source_face.normed_embedding.reshape(1, -1).astype(np.float32) + + # 2. Получаем 5 точек target + target_landmarks_5 = get_landmarks_5(target_face) + visualize_points(target_img, target_landmarks_5, (0, 255, 0)) # Зеленые точки + + if target_landmarks_5 is None: + return None, None + + # 3. Определение эталонных точек для выравнивания 256x256 (FFHQ Alignment) + std_landmarks_256 = np.array([ + [ 84.87, 105.94], # Левый глаз + [171.13, 105.94], # Правый глаз + [128.00, 146.66], # Кончик носа + [ 96.95, 188.64], # Левый уголок рта + [159.05, 188.64] # Правый уголок рта + ], dtype=np.float32) + + # Вычисляем аффинную матрицу + M = get_affine_transform(target_landmarks_5.astype(np.float32), std_landmarks_256) + + # Применяем аффинное преобразование с новой матрицей M + crop = cv2.warpAffine(target_img, M, (256, 256), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REFLECT) + + # 4. Преобразуем crop для модели + crop_input = crop[:, :, ::-1].astype(np.float32) / 255.0 # RGB -> [0,1] + crop_input = (crop_input - 0.5) / 0.5 # Нормализация + crop_input = crop_input.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32) + + # 5. Инференс + try: + output = session.run(None, {'source': source_embedding, 'target': crop_input})[0][0] + except: + return target_img + + # 6. Обратная нормализация + output = (output * 0.5 + 0.5) * 255.0 # [-1..1] -> [0..255] + output = np.clip(output, 0, 255).astype(np.uint8) + output = output.transpose(1, 2, 0) # CHW -> HWC + output = output[:, :, ::-1] # BGR -> RGB + + return output, M # Возвращаем лицо (256x256) и матрицу M + + def sort_by_order(face, order: str): if order == "left-right": return sorted(face, key=lambda x: x.bbox[0]) @@ -346,6 +492,9 @@ def swap_face( 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) + face_swapper = getFaceSwapModel(model_path) source_face_idx = 0 @@ -364,8 +513,12 @@ def swap_face( target_face, wrong_gender, target_face_index = 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 face_boost_enabled: - logger.status(f"Face Boost is enabled") + if "hyperswap" in model: + 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) + elif face_boost_enabled: + logger.status(f"Face Boost is enabled (inswapper/reswapper only)") bgr_fake, M = face_swapper.get(result, target_face, 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 @@ -541,7 +694,13 @@ 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) + face_swapper = getFaceSwapModel(model_path) source_face_idx = 0 @@ -566,14 +725,17 @@ def swap_face_many( target_face_single, wrong_gender, target_face_index = 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 face_boost_enabled: - logger.status(f"Face Boost is enabled") + if "hyperswap" in model: + 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 (inswapper/reswapper only)") 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") result = face_swapper.get(target_img, target_face_single, source_face) results[i] = result bbox.append(tuple(map(float, target_face_single.bbox)))