UPD: HyperSwap models support

FR #143
Thanks @Buumcode for contribution
This commit is contained in:
Gourieff | 古仁
2025-10-06 22:15:27 +07:00
parent d8e3b90c82
commit 95ed7b248f
2 changed files with 182 additions and 19 deletions
+2 -1
View File
@@ -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:
+180 -18
View File
@@ -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)))