import os
import shutil
from typing import List, Union
import cv2
import numpy as np
from PIL import Image
import insightface
from insightface.app.common import Face
# try:
# import torch.cuda as cuda
# except:
# cuda = None
import torch
import folder_paths
import comfy.model_management as model_management
from modules.shared import state
from scripts.reactor_logger import logger
from reactor_utils import (
move_path,
get_image_md5hash,
)
from scripts.r_faceboost import swapper, restorer
import warnings
np.warnings = warnings
np.warnings.filterwarnings('ignore')
# PROVIDERS
try:
if torch.cuda.is_available():
providers = ["CUDAExecutionProvider"]
elif torch.backends.mps.is_available():
providers = ["CoreMLExecutionProvider"]
elif hasattr(torch,'dml') or hasattr(torch,'privateuseone'):
providers = ["ROCMExecutionProvider"]
else:
providers = ["CPUExecutionProvider"]
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")
insightface_models_path_old = os.path.join(insightface_path_old, "models")
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")
if os.path.exists(models_path_old):
move_path(insightface_models_path_old, insightface_models_path)
move_path(insightface_path_old, insightface_path)
move_path(models_path_old, models_path)
if os.path.exists(insightface_path) and os.path.exists(insightface_path_old):
shutil.rmtree(insightface_path_old)
shutil.rmtree(models_path_old)
FS_MODEL = None
CURRENT_FS_MODEL_PATH = None
ANALYSIS_MODELS = {
"640": None,
"320": None,
}
SOURCE_FACES = None
SOURCE_IMAGE_HASH = None
TARGET_FACES = None
TARGET_IMAGE_HASH = None
TARGET_FACES_LIST = []
TARGET_IMAGE_LIST_HASH = []
def unload_model(model):
if model is not None:
# check if model has unload method
# if "unload" in model:
# model.unload()
# if "model_unload" in model:
# model.model_unload()
del model
return None
def unload_all_models():
global FS_MODEL, CURRENT_FS_MODEL_PATH
FS_MODEL = unload_model(FS_MODEL)
ANALYSIS_MODELS["320"] = unload_model(ANALYSIS_MODELS["320"])
ANALYSIS_MODELS["640"] = unload_model(ANALYSIS_MODELS["640"])
def get_current_faces_model():
global SOURCE_FACES
return SOURCE_FACES
def getAnalysisModel(det_size = (640, 640)):
global ANALYSIS_MODELS
ANALYSIS_MODEL = ANALYSIS_MODELS[str(det_size[0])]
if ANALYSIS_MODEL is None:
ANALYSIS_MODEL = insightface.app.FaceAnalysis(
name="buffalo_l", providers=providers, root=insightface_path
)
ANALYSIS_MODEL.prepare(ctx_id=0, det_size=det_size)
ANALYSIS_MODELS[str(det_size[0])] = ANALYSIS_MODEL
return ANALYSIS_MODEL
def getFaceSwapModel(model_path: str):
global FS_MODEL, CURRENT_FS_MODEL_PATH
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)
return FS_MODEL
def sort_by_order(face, order: str):
if order == "left-right":
return sorted(face, key=lambda x: x.bbox[0])
if order == "right-left":
return sorted(face, key=lambda x: x.bbox[0], reverse = True)
if order == "top-bottom":
return sorted(face, key=lambda x: x.bbox[1])
if order == "bottom-top":
return sorted(face, key=lambda x: x.bbox[1], reverse = True)
if order == "small-large":
return sorted(face, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
# if order == "large-small":
# return sorted(face, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse = True)
# by default "large-small":
return sorted(face, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse = True)
def get_face_gender(
face,
face_index,
gender_condition,
operated: str,
order: str,
):
gender = [
x.sex
for x in face
]
gender.reverse()
# If index is outside of bounds, return None, avoid exception
if face_index >= len(gender):
logger.status("Requested face index (%s) is out of bounds (max available index is %s)", face_index, len(gender))
return None, 0
face_gender = gender[face_index]
logger.status("%s Face %s: Detected Gender -%s-", operated, face_index, face_gender)
if (gender_condition == 1 and face_gender == "F") or (gender_condition == 2 and face_gender == "M"):
logger.status("OK - Detected Gender matches Condition")
try:
faces_sorted = sort_by_order(face, order)
return faces_sorted[face_index], 0
# return sorted(face, key=lambda x: x.bbox[0])[face_index], 0
except IndexError:
return None, 0
else:
logger.status("WRONG - Detected Gender doesn't match Condition")
faces_sorted = sort_by_order(face, order)
return faces_sorted[face_index], 1
# return sorted(face, key=lambda x: x.bbox[0])[face_index], 1
def half_det_size(det_size):
logger.status("Trying to halve 'det_size' parameter")
return (det_size[0] // 2, det_size[1] // 2)
def analyze_faces(img_data: np.ndarray, det_size=(640, 640)):
face_analyser = getAnalysisModel(det_size)
faces = face_analyser.get(img_data)
# Try halving det_size if no faces are found
if len(faces) == 0 and det_size[0] > 320 and det_size[1] > 320:
det_size_half = half_det_size(det_size)
return analyze_faces(img_data, det_size_half)
return faces
def get_face_single(img_data: np.ndarray, face, face_index=0, det_size=(640, 640), gender_source=0, gender_target=0, order="large-small"):
buffalo_path = os.path.join(insightface_models_path, "buffalo_l.zip")
if os.path.exists(buffalo_path):
os.remove(buffalo_path)
if gender_source != 0:
if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
det_size_half = half_det_size(det_size)
return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target, order)
return get_face_gender(face,face_index,gender_source,"Source", order)
if gender_target != 0:
if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
det_size_half = half_det_size(det_size)
return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target, order)
return get_face_gender(face,face_index,gender_target,"Target", order)
if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
det_size_half = half_det_size(det_size)
return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target, order)
try:
faces_sorted = sort_by_order(face, order)
return faces_sorted[face_index], 0
# return sorted(face, key=lambda x: x.bbox[0])[face_index], 0
except IndexError:
return None, 0
def swap_face(
source_img: Union[Image.Image, None],
target_img: Image.Image,
model: Union[str, None] = None,
source_faces_index: List[int] = [0],
faces_index: List[int] = [0],
gender_source: int = 0,
gender_target: int = 0,
face_model: Union[Face, None] = None,
faces_order: List = ["large-small", "large-small"],
face_boost_enabled: bool = False,
face_restore_model = None,
face_restore_visibility: int = 1,
codeformer_weight: float = 0.5,
interpolation: str = "Bicubic",
):
global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH
result_image = target_img
if model is not None:
if isinstance(source_img, str): # source_img is a base64 string
import base64, io
if 'base64,' in source_img: # check if the base64 string has a data URL scheme
# split the base64 string to get the actual base64 encoded image data
base64_data = source_img.split('base64,')[-1]
# decode base64 string to bytes
img_bytes = base64.b64decode(base64_data)
else:
# if no data URL scheme, just decode
img_bytes = base64.b64decode(source_img)
source_img = Image.open(io.BytesIO(img_bytes))
target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
if source_img is not None:
source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
source_image_md5hash = get_image_md5hash(source_img)
if SOURCE_IMAGE_HASH is None:
SOURCE_IMAGE_HASH = source_image_md5hash
source_image_same = False
else:
source_image_same = True if SOURCE_IMAGE_HASH == source_image_md5hash else False
if not source_image_same:
SOURCE_IMAGE_HASH = source_image_md5hash
logger.info("Source Image MD5 Hash = %s", SOURCE_IMAGE_HASH)
logger.info("Source Image the Same? %s", source_image_same)
if SOURCE_FACES is None or not source_image_same:
logger.status("Analyzing Source Image...")
source_faces = analyze_faces(source_img)
SOURCE_FACES = source_faces
elif source_image_same:
logger.status("Using Hashed Source Face(s) Model...")
source_faces = SOURCE_FACES
elif face_model is not None:
source_faces_index = [0]
logger.status("Using Loaded Source Face Model...")
source_face_model = [face_model]
source_faces = source_face_model
else:
logger.error("Cannot detect any Source")
if source_faces is not None:
target_image_md5hash = get_image_md5hash(target_img)
if TARGET_IMAGE_HASH is None:
TARGET_IMAGE_HASH = target_image_md5hash
target_image_same = False
else:
target_image_same = True if TARGET_IMAGE_HASH == target_image_md5hash else False
if not target_image_same:
TARGET_IMAGE_HASH = target_image_md5hash
logger.info("Target Image MD5 Hash = %s", TARGET_IMAGE_HASH)
logger.info("Target Image the Same? %s", target_image_same)
if TARGET_FACES is None or not target_image_same:
logger.status("Analyzing Target Image...")
target_faces = analyze_faces(target_img)
TARGET_FACES = target_faces
elif target_image_same:
logger.status("Using Hashed Target Face(s) Model...")
target_faces = TARGET_FACES
# No use in trying to swap faces if no faces are found, enhancement
if len(target_faces) == 0:
logger.status("Cannot detect any Target, skipping swapping...")
return result_image
if source_img is not None:
# separated management of wrong_gender between source and target, enhancement
source_face, src_wrong_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source, order=faces_order[1])
else:
# source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]]
source_face = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse = True)[source_faces_index[0]]
src_wrong_gender = 0
if len(source_faces_index) != 0 and len(source_faces_index) != 1 and len(source_faces_index) != len(faces_index):
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:
result = target_img
if "inswapper" in model:
model_path = model_path = os.path.join(insightface_path, model)
elif "reswapper" in model:
model_path = model_path = os.path.join(reswapper_path, model)
face_swapper = getFaceSwapModel(model_path)
source_face_idx = 0
for face_num in faces_index:
# No use in trying to swap faces if no further faces are found, enhancement
if face_num >= len(target_faces):
logger.status("Checked all existing target faces, skipping swapping...")
break
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
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 face_boost_enabled:
logger.status(f"Face Boost is enabled")
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
result = swapper.in_swap(target_img, bgr_fake, M)
else:
# logger.status(f"Swapping as-is")
result = face_swapper.get(result, target_face, source_face)
elif wrong_gender == 1:
wrong_gender = 0
# Keep searching for other faces if wrong gender is detected, enhancement
#if source_face_idx == len(source_faces_index):
# result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
# return result_image
logger.status("Wrong target gender detected")
continue
else:
logger.status(f"No target face found for {face_num}")
elif src_wrong_gender == 1:
src_wrong_gender = 0
# Keep searching for other faces if wrong gender is detected, enhancement
#if source_face_idx == len(source_faces_index):
# result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
# return result_image
logger.status("Wrong source gender detected")
continue
else:
logger.status(f"No source face found for face number {source_face_idx}.")
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
else:
logger.status("No source face(s) in the provided Index")
else:
logger.status("No source face(s) found")
return result_image
def swap_face_many(
source_img: Union[Image.Image, None],
target_imgs: List[Image.Image],
model: Union[str, None] = None,
source_faces_index: List[int] = [0],
faces_index: List[int] = [0],
gender_source: int = 0,
gender_target: int = 0,
face_model: Union[Face, None] = None,
faces_order: List = ["large-small", "large-small"],
face_boost_enabled: bool = False,
face_restore_model = None,
face_restore_visibility: int = 1,
codeformer_weight: float = 0.5,
interpolation: str = "Bicubic",
):
global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, TARGET_FACES_LIST, TARGET_IMAGE_LIST_HASH
result_images = target_imgs
if model is not None:
if isinstance(source_img, str): # source_img is a base64 string
import base64, io
if 'base64,' in source_img: # check if the base64 string has a data URL scheme
# split the base64 string to get the actual base64 encoded image data
base64_data = source_img.split('base64,')[-1]
# decode base64 string to bytes
img_bytes = base64.b64decode(base64_data)
else:
# if no data URL scheme, just decode
img_bytes = base64.b64decode(source_img)
source_img = Image.open(io.BytesIO(img_bytes))
target_imgs = [cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) for target_img in target_imgs]
if source_img is not None:
source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
source_image_md5hash = get_image_md5hash(source_img)
if SOURCE_IMAGE_HASH is None:
SOURCE_IMAGE_HASH = source_image_md5hash
source_image_same = False
else:
source_image_same = True if SOURCE_IMAGE_HASH == source_image_md5hash else False
if not source_image_same:
SOURCE_IMAGE_HASH = source_image_md5hash
logger.info("Source Image MD5 Hash = %s", SOURCE_IMAGE_HASH)
logger.info("Source Image the Same? %s", source_image_same)
if SOURCE_FACES is None or not source_image_same:
logger.status("Analyzing Source Image...")
source_faces = analyze_faces(source_img)
SOURCE_FACES = source_faces
elif source_image_same:
logger.status("Using Hashed Source Face(s) Model...")
source_faces = SOURCE_FACES
elif face_model is not None:
source_faces_index = [0]
logger.status("Using Loaded Source Face Model...")
source_face_model = [face_model]
source_faces = source_face_model
else:
logger.error("Cannot detect any Source")
if source_faces is not None:
target_faces = []
for i, target_img in enumerate(target_imgs):
if state.interrupted or model_management.processing_interrupted():
logger.status("Interrupted by User")
break
target_image_md5hash = get_image_md5hash(target_img)
if len(TARGET_IMAGE_LIST_HASH) == 0:
TARGET_IMAGE_LIST_HASH = [target_image_md5hash]
target_image_same = False
elif len(TARGET_IMAGE_LIST_HASH) == i:
TARGET_IMAGE_LIST_HASH.append(target_image_md5hash)
target_image_same = False
else:
target_image_same = True if TARGET_IMAGE_LIST_HASH[i] == target_image_md5hash else False
if not target_image_same:
TARGET_IMAGE_LIST_HASH[i] = target_image_md5hash
logger.info("(Image %s) Target Image MD5 Hash = %s", i, TARGET_IMAGE_LIST_HASH[i])
logger.info("(Image %s) Target Image the Same? %s", i, target_image_same)
if len(TARGET_FACES_LIST) == 0:
logger.status(f"Analyzing Target Image {i}...")
target_face = analyze_faces(target_img)
TARGET_FACES_LIST = [target_face]
elif len(TARGET_FACES_LIST) == i and not target_image_same:
logger.status(f"Analyzing Target Image {i}...")
target_face = analyze_faces(target_img)
TARGET_FACES_LIST.append(target_face)
elif len(TARGET_FACES_LIST) != i and not target_image_same:
logger.status(f"Analyzing Target Image {i}...")
target_face = analyze_faces(target_img)
TARGET_FACES_LIST[i] = target_face
elif target_image_same:
logger.status("(Image %s) Using Hashed Target Face(s) Model...", i)
target_face = TARGET_FACES_LIST[i]
# logger.status(f"Analyzing Target Image {i}...")
# target_face = analyze_faces(target_img)
if target_face is not None:
target_faces.append(target_face)
# No use in trying to swap faces if no faces are found, enhancement
if len(target_faces) == 0:
logger.status("Cannot detect any Target, skipping swapping...")
return result_images
if source_img is not None:
# separated management of wrong_gender between source and target, enhancement
source_face, src_wrong_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source, order=faces_order[1])
else:
# source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]]
source_face = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse = True)[source_faces_index[0]]
src_wrong_gender = 0
if len(source_faces_index) != 0 and len(source_faces_index) != 1 and len(source_faces_index) != len(faces_index):
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)
face_swapper = getFaceSwapModel(model_path)
source_face_idx = 0
for face_num in faces_index:
# No use in trying to swap faces if no further faces are found, enhancement
if face_num >= len(target_faces):
logger.status("Checked all existing target faces, skipping swapping...")
break
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
if source_face is not None and src_wrong_gender == 0:
# Reading results to make current face swap on a previous face result
for i, (target_img, target_face) in enumerate(zip(results, target_faces)):
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:
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")
result = face_swapper.get(target_img, target_face_single, source_face)
results[i] = result
elif wrong_gender == 1:
wrong_gender = 0
logger.status("Wrong target gender detected")
continue
else:
logger.status(f"No target face found for {face_num}")
elif src_wrong_gender == 1:
src_wrong_gender = 0
logger.status("Wrong source gender detected")
continue
else:
logger.status(f"No source face found for face number {source_face_idx}.")
result_images = [Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) for result in results]
else:
logger.status("No source face(s) in the provided Index")
else:
logger.status("No source face(s) found")
return result_images