import copy
import os
from dataclasses import dataclass
from pprint import pformat
import traceback
from typing import Any, Dict, Generator, List, Set, Tuple, Optional, Union
import tempfile
from tqdm import tqdm
import sys
from io import StringIO
from contextlib import contextmanager
import cv2
import insightface
import numpy as np
from insightface.app.common import Face as ISFace
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
from scripts.faceswaplab_swapping import upscaled_inswapper
from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
from scripts.faceswaplab_utils.imgutils import (
pil_to_cv2,
check_against_nsfw,
)
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
from scripts import faceswaplab_globals
from functools import lru_cache
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
from scripts.faceswaplab_postprocessing.postprocessing_options import (
PostProcessingOptions,
)
from scripts.faceswaplab_utils.models_utils import get_current_swap_model
from scripts.faceswaplab_utils.typing import CV2ImgU8, PILImage, Face
from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
from modules import shared
import onnxruntime
from scripts.faceswaplab_utils.sd_utils import get_sd_option
def use_gpu() -> bool:
return (
getattr(shared.cmd_opts, "faceswaplab_gpu", False)
or get_sd_option("faceswaplab_use_gpu", False)
) and sys.platform != "darwin"
@lru_cache
def force_install_gpu_providers() -> None:
# Ugly Ugly hack due to SDNEXT :
try:
from scripts.faceswaplab_utils.install_utils import check_install
logger.warning("Try to reinstall gpu dependencies")
check_install()
logger.warning("IF onnxruntime-gpu has been installed successfully, RESTART")
logger.warning(
"On SD.NEXT/vladmantic you will also need to check numpy>=1.24.2 and tensorflow>=2.13.0"
)
except:
logger.error(
"Reinstall has failed (which is normal on windows), please install requirements-gpu.txt manually to enable gpu."
)
def get_providers() -> List[str]:
providers = ["CPUExecutionProvider"]
if use_gpu():
if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
providers = ["CUDAExecutionProvider"]
else:
logger.error(
"CUDAExecutionProvider not found in onnxruntime.available_providers : %s, use CPU instead. Check onnxruntime-gpu is installed.",
onnxruntime.get_available_providers(),
)
force_install_gpu_providers()
return providers
def is_cpu_provider() -> bool:
return get_providers() == ["CPUExecutionProvider"]
def cosine_similarity_face(face1: Face, face2: Face) -> float:
"""
Calculates the cosine similarity between two face embeddings.
Args:
face1 (Face): The first face object containing an embedding.
face2 (Face): The second face object containing an embedding.
Returns:
float: The cosine similarity between the face embeddings.
Note:
The cosine similarity ranges from 0 to 1, where 1 indicates identical embeddings and 0 indicates completely
dissimilar embeddings. In this implementation, the similarity is clamped to a minimum value of 0 to ensure a
non-negative similarity score.
"""
# Reshape the face embeddings to have a shape of (1, -1)
assert face1.normed_embedding is not None and face2.normed_embedding is not None
vec1 = face1.normed_embedding.reshape(1, -1)
vec2 = face2.normed_embedding.reshape(1, -1)
# Calculate the cosine similarity between the reshaped embeddings
similarity = cosine_similarity(vec1, vec2)
# Return the maximum of 0 and the calculated similarity as the final similarity score
return max(0, similarity[0, 0])
def compare_faces(img1: PILImage, img2: PILImage) -> float:
"""
Compares the similarity between two faces extracted from images using cosine similarity.
Args:
img1: The first image containing a face.
img2: The second image containing a face.
Returns:
A float value representing the similarity between the two faces (0 to 1).
Returns -1 if one or both of the images do not contain any faces.
"""
# Extract faces from the images
face1 = get_or_default(get_faces(pil_to_cv2(img1)), 0, None)
face2 = get_or_default(get_faces(pil_to_cv2(img2)), 0, None)
# Check if both faces are detected
if face1 is not None and face2 is not None:
# Calculate the cosine similarity between the faces
return cosine_similarity_face(face1, face2)
# Return -1 if one or both of the images do not contain any faces
return -1
def batch_process(
src_images: List[Union[PILImage, str]], # image or filename
save_path: Optional[str],
units: List[FaceSwapUnitSettings],
postprocess_options: Optional[PostProcessingOptions],
) -> Optional[List[PILImage]]:
"""
Process a batch of images, apply face swapping according to the given settings, and optionally save the resulting images to a specified path.
Args:
src_images (List[Union[PILImage, str]]): List of source PIL Images to process or list of images file names
save_path (Optional[str]): Destination path where the processed images will be saved. If None, no images are saved.
units (List[FaceSwapUnitSettings]): List of FaceSwapUnitSettings to apply to the images.
postprocess_options (PostProcessingOptions): Post-processing settings to be applied to the images.
Returns:
Optional[List[PILImage]]: List of processed images, or None in case of an exception.
Raises:
Any exceptions raised by the underlying process will be logged and the function will return None.
"""
try:
if save_path:
os.makedirs(save_path, exist_ok=True)
units = [u for u in units if u.enable]
if src_images is not None and len(units) > 0:
result_images = []
for src_image in src_images:
path: str = ""
if isinstance(src_image, str):
if save_path:
path = os.path.join(
save_path, "swapped_" + os.path.basename(src_image)
)
src_image = Image.open(src_image)
elif save_path:
path = tempfile.NamedTemporaryFile(
delete=False, suffix=".png", dir=save_path
).name
assert isinstance(src_image, Image.Image)
current_images = []
swapped_images = process_images_units(
get_current_swap_model(), images=[(src_image, None)], units=units
)
if swapped_images and len(swapped_images) > 0:
current_images += [img for img, _ in swapped_images]
logger.info("%s images generated", len(current_images))
if postprocess_options:
for i, img in enumerate(current_images):
current_images[i] = enhance_image(img, postprocess_options)
if save_path:
for img in current_images:
img.save(path)
result_images += current_images
return result_images
except Exception as e:
logger.error("Batch Process error : %s", e)
import traceback
traceback.print_exc()
return None
def extract_faces(
images: List[PILImage],
extract_path: Optional[str],
postprocess_options: PostProcessingOptions,
) -> Optional[List[PILImage]]:
"""
Extracts faces from a list of image files.
Given a list of image file paths, this function opens each image, extracts the faces,
and saves them in a specified directory. Post-processing is applied to each extracted face,
and the processed faces are saved as separate PNG files.
Parameters:
files (Optional[List[Image]]): List of file paths to the images to extract faces from.
extract_path (Optional[str]): Path where the extracted faces will be saved.
If no path is provided, a temporary directory will be created.
postprocess_options (PostProcessingOptions): Post-processing settings to be applied to the images.
Returns:
Optional[List[img]]: List of face images
"""
try:
if extract_path:
os.makedirs(extract_path, exist_ok=True)
if images:
result_images: list[PILImage] = []
for img in images:
faces = get_faces(pil_to_cv2(img))
if faces:
face_images = []
for face in faces:
bbox = face.bbox.astype(int) # type: ignore
x_min, y_min, x_max, y_max = bbox
face_image = img.crop((x_min, y_min, x_max, y_max))
if postprocess_options and (
postprocess_options.face_restorer_name
or postprocess_options.restorer_visibility
):
postprocess_options.scale = (
1 if face_image.width > 512 else 512 // face_image.width
)
face_image = enhance_image(face_image, postprocess_options)
if extract_path:
path = tempfile.NamedTemporaryFile(
delete=False, suffix=".png", dir=extract_path
).name
face_image.save(path)
face_images.append(face_image)
result_images += face_images
return result_images
except Exception as e:
logger.error("Failed to extract : %s", e)
import traceback
traceback.print_exc()
return None
class FaceModelException(Exception):
"""Exception raised when an error is encountered in the face model."""
def __init__(self, message: str) -> None:
"""
Args:
message: A string containing the error description.
"""
self.message = message
super().__init__(self.message)
@contextmanager
def capture_stdout() -> Generator[StringIO, None, None]:
"""
Capture and yield the printed messages to stdout.
This context manager temporarily replaces sys.stdout with a StringIO object,
capturing all printed output. After the context block is exited, sys.stdout
is restored to its original value.
Example usage:
with capture_stdout() as captured:
print("Hello, World!")
output = captured.getvalue()
# output now contains "Hello, World!\n"
Returns:
A StringIO object containing the captured output.
"""
original_stdout = sys.stdout # Type: ignore
captured_stdout = StringIO()
sys.stdout = captured_stdout # Type: ignore
try:
yield captured_stdout
finally:
sys.stdout = original_stdout # Type: ignore
@lru_cache(maxsize=3)
def getAnalysisModel(
det_size: Tuple[int, int] = (640, 640),
det_thresh: float = 0.5,
use_gpu: bool = False,
) -> insightface.app.FaceAnalysis:
"""
Retrieves the analysis model for face analysis.
Returns:
insightface.app.FaceAnalysis: The analysis model for face analysis.
"""
try:
if not os.path.exists(faceswaplab_globals.ANALYZER_DIR):
os.makedirs(faceswaplab_globals.ANALYZER_DIR)
providers = get_providers()
logger.info(
f"Load analysis model det_size={det_size}, det_thresh={det_thresh}, providers = {providers}, will take some time. (> 30s)"
)
# Initialize the analysis model with the specified name and providers
with tqdm(
total=1,
desc=f"Loading {det_size} analysis model (first time is slow)",
unit="model",
) as pbar:
with capture_stdout() as captured:
model = insightface.app.FaceAnalysis(
name="buffalo_l",
providers=providers,
root=faceswaplab_globals.ANALYZER_DIR,
)
# Prepare the analysis model for face detection with the specified detection size
model.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size)
pbar.update(1)
logger.info("%s", pformat(captured.getvalue()))
return model
except Exception as e:
logger.error(
"Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
)
raise FaceModelException("Loading of analysis model failed")
@lru_cache(maxsize=1)
def getFaceSwapModel(
model_path: str, use_gpu: bool = False
) -> upscaled_inswapper.UpscaledINSwapper:
"""
Retrieves the face swap model and initializes it if necessary.
Args:
model_path (str): Path to the face swap model.
Returns:
insightface.model_zoo.FaceModel: The face swap model.
"""
try:
providers = get_providers()
with tqdm(total=1, desc="Loading swap model", unit="model") as pbar:
with capture_stdout() as captured:
model = upscaled_inswapper.UpscaledINSwapper(
insightface.model_zoo.get_model(model_path, providers=providers) # type: ignore
)
pbar.update(1)
logger.info("%s", pformat(captured.getvalue()))
return model
except Exception as e:
logger.error(
"Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
)
traceback.print_exc()
raise FaceModelException("Loading of swapping model failed")
def get_faces(
img_data: CV2ImgU8,
det_thresh: Optional[float] = None,
det_size: Tuple[int, int] = (640, 640),
) -> List[Face]:
"""
Detects and retrieves faces from an image using an analysis model.
Args:
img_data (CV2ImgU8): The image data as a NumPy array.
det_size (tuple): The desired detection size (width, height). Defaults to (640, 640).
sort_by_face_size (bool) : Will sort the faces by their size from larger to smaller face
Returns:
list: A list of detected faces, sorted by their x-coordinate of the bounding box.
"""
if det_thresh is None:
det_thresh = get_sd_option("faceswaplab_detection_threshold", 0.5)
auto_det_size = get_sd_option("faceswaplab_auto_det_size", True)
if not auto_det_size:
x = get_sd_option("faceswaplab_det_size", 640)
det_size = (x, x)
face_analyser = getAnalysisModel(
det_size=det_size, det_thresh=det_thresh, use_gpu=not is_cpu_provider()
)
# Get the detected faces from the image using the analysis model
faces = face_analyser.get(img_data)
# If no faces are detected and the detection size is larger than 320x320,
# recursively call the function with a smaller detection size
if len(faces) == 0:
if auto_det_size:
if det_size[0] > 320 and det_size[1] > 320:
det_size_half = (det_size[0] // 2, det_size[1] // 2)
return get_faces(
img_data, det_size=det_size_half, det_thresh=det_thresh
)
# If no faces are detected print a warning to user about change in detection
else:
if det_size[0] > 320:
logger.warning(
"No faces detected, you might want to play with det_size by reducing it (in sd global settings). Lower (320) means more detection but less precise. Or activate auto-det-size."
)
try:
# Sort the detected faces based on their x-coordinate of the bounding box
return sorted(faces, key=lambda x: x.bbox[0]) # type: ignore
except Exception as e:
logger.error("Failed to get faces %s", e)
traceback.print_exc()
return []
@dataclass
class FaceFilteringOptions:
faces_index: Set[int]
source_gender: Optional[int] = None # if none will not use same gender
sort_by_face_size: bool = False
def filter_faces(
all_faces: List[Face], filtering_options: FaceFilteringOptions
) -> List[Face]:
"""
Sorts and filters a list of faces based on specified criteria.
This function takes a list of Face objects and can sort them by face size and filter them by gender.
Sorting by face size is performed if sort_by_face_size is set to True, and filtering by gender is
performed if source_gender is provided.
:param faces: A list of Face objects representing the faces to be sorted and filtered.
:param faces_index: A set of faces index
:param source_gender: An optional integer representing the gender by which to filter the faces.
If provided, only faces with the specified gender will be included in the result.
:param sort_by_face_size: A boolean indicating whether to sort the faces by size. If True, faces are
sorted in descending order by size, calculated as the area of the bounding box.
:return: A list of Face objects sorted and filtered according to the specified criteria.
"""
filtered_faces = copy.copy(all_faces)
if filtering_options.sort_by_face_size:
filtered_faces = sorted(
all_faces,
reverse=True,
key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), # type: ignore
)
if filtering_options.source_gender is not None:
filtered_faces = [
face
for face in filtered_faces
if face["gender"] == filtering_options.source_gender
]
return [
face
for i, face in enumerate(filtered_faces)
if i in filtering_options.faces_index
]
@dataclass
class ImageResult:
"""
Represents the result of an image swap operation
"""
image: PILImage
"""
The image object with the swapped face
"""
similarity: Dict[int, float]
"""
A dictionary mapping face indices to their similarity scores.
The similarity scores are represented as floating-point values between 0 and 1.
"""
ref_similarity: Dict[int, float]
"""
A dictionary mapping face indices to their similarity scores compared to a reference image.
The similarity scores are represented as floating-point values between 0 and 1.
"""
def get_or_default(l: List[Any], index: int, default: Any) -> Any:
"""
Retrieve the value at the specified index from the given list.
If the index is out of bounds, return the default value instead.
Args:
l (list): The input list.
index (int): The index to retrieve the value from.
default: The default value to return if the index is out of bounds.
Returns:
The value at the specified index if it exists, otherwise the default value.
"""
return l[index] if index < len(l) else default
def get_faces_from_img_files(images: List[PILImage]) -> List[Face]:
"""
Extracts faces from a list of image files.
Args:
images (list): A list of PILImage objects representing image files.
Returns:
list: A list of detected faces.
"""
faces: List[Face] = []
if len(images) > 0:
for img in images:
face = get_or_default(
get_faces(pil_to_cv2(img)), 0, None
) # Extract faces from the image
if face is not None:
faces.append(face) # Add the detected face to the list of faces
return faces
def blend_faces(faces: List[Face]) -> Optional[Face]:
"""
Blends the embeddings of multiple faces into a single face.
Args:
faces (List[Face]): List of Face objects.
Returns:
Face: The blended Face object with the averaged embedding.
Returns None if the input list is empty.
Raises:
ValueError: If the embeddings have different shapes.
"""
embeddings: list[Any] = [face.embedding for face in faces]
if len(embeddings) > 0:
embedding_shape = embeddings[0].shape
# Check if all embeddings have the same shape
for embedding in embeddings:
if embedding.shape != embedding_shape:
raise ValueError("embedding shape mismatch")
# Compute the mean of all embeddings
blended_embedding = np.mean(embeddings, axis=0)
# Create a new Face object using the properties of the first face in the list
# Assign the blended embedding to the blended Face object
blended = ISFace(
embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age
)
return blended
# Return None if the input list is empty
return None
def swap_face(
source_face: Face,
target_img: PILImage,
target_faces: List[Face],
model: str,
swapping_options: Optional[InswappperOptions],
) -> ImageResult:
"""
Swaps faces in the target image with the source face.
Args:
source_face (CV2ImgU8): The source face to be swapped.
target_img (PILImage): The target image to swap faces in.
model (str): Path to the face swap model.
Returns:
ImageResult: An object containing the swapped image and similarity scores.
"""
return_result = ImageResult(target_img, {}, {})
target_img_cv2: CV2ImgU8 = cv2.cvtColor(
np.array(target_img), cv2.COLOR_RGB2BGR
).astype("uint8")
try:
gender = source_face["gender"]
logger.info("Source Gender %s", gender)
if source_face is not None:
result: CV2ImgU8 = target_img_cv2
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
face_swapper = getFaceSwapModel(model_path, use_gpu=not is_cpu_provider())
logger.info("Target faces count : %s", len(target_faces))
for i, swapped_face in enumerate(target_faces):
logger.info(f"swap face {i}")
result = face_swapper.get(
img=result,
target_face=swapped_face,
source_face=source_face,
options=swapping_options,
) # type: ignore
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
return_result.image = result_image
except Exception as e:
logger.error("Conversion failed %s", e)
raise e
return return_result
def compute_similarity(
reference_face: Face,
source_face: Face,
swapped_image: PILImage,
filtering: FaceFilteringOptions,
) -> Tuple[Dict[int, float], Dict[int, float]]:
similarity: Dict[int, float] = {}
ref_similarity: Dict[int, float] = {}
try:
swapped_image_cv2: CV2ImgU8 = cv2.cvtColor(
np.array(swapped_image), cv2.COLOR_RGB2BGR
)
new_faces = filter_faces(get_faces(swapped_image_cv2), filtering)
if len(new_faces) == 0:
logger.error("compute_similarity : No faces to compare with !")
return None
for i, swapped_face in enumerate(new_faces):
logger.info(f"compare face {i}")
similarity[i] = cosine_similarity_face(source_face, swapped_face)
ref_similarity[i] = cosine_similarity_face(reference_face, swapped_face)
logger.info(f"similarity {similarity}")
logger.info(f"ref similarity {ref_similarity}")
return (similarity, ref_similarity)
except Exception as e:
logger.error("Similarity processing failed %s", e)
raise e
return None
def process_image_unit(
model: str,
unit: FaceSwapUnitSettings,
image: PILImage,
info: Optional[str] = None,
force_blend: bool = False,
) -> List[Tuple[PILImage, Optional[str]]]:
"""Process one image and return a List of (image, info) (one if blended, many if not).
Args:
unit : the current unit
image : the image where to apply swapping
info : The info
Returns:
List of tuple of (image, info) where image is the image where swapping has been applied and info is the image info with similarity infos.
"""
results = []
if unit.enable:
faces = get_faces(pil_to_cv2(image))
if check_against_nsfw(image):
return [(image, info)]
if not unit.blend_faces and not force_blend:
src_faces = unit.faces
logger.info(f"will generate {len(src_faces)} images")
else:
logger.info("blend all faces together")
src_faces = [unit.blended_faces]
for i, src_face in enumerate(src_faces):
current_image = image
logger.info(f"Process face {i}")
if unit.reference_face is not None:
reference_face = unit.reference_face
else:
logger.info("Use source face as reference face")
reference_face = src_face
face_filtering_options = FaceFilteringOptions(
faces_index=unit.faces_index,
source_gender=src_face["gender"] if unit.same_gender else None,
sort_by_face_size=unit.sort_by_size,
)
target_faces: List[Face] = filter_faces(
all_faces=faces, filtering_options=face_filtering_options
)
# Apply pre-inpainting to image
if unit.pre_inpainting.inpainting_denoising_strengh > 0:
current_image = img2img_diffusion(
img=current_image, faces=target_faces, options=unit.pre_inpainting
)
save_img_debug(image, "Before swap")
result: ImageResult = swap_face(
source_face=src_face,
target_img=current_image,
target_faces=target_faces,
model=model,
swapping_options=unit.swapping_options,
)
# Apply post-inpainting to image
if unit.post_inpainting.inpainting_denoising_strengh > 0:
result.image = img2img_diffusion(
img=result.image, faces=target_faces, options=unit.post_inpainting
)
save_img_debug(result.image, "After swap")
if unit.compute_similarity:
similarities = compute_similarity(
reference_face=reference_face,
source_face=src_face,
swapped_image=result.image,
filtering=face_filtering_options,
)
if similarities:
(result.similarity, result.ref_similarity) = similarities
else:
logger.error("Failed to compute similarity")
if result.image is None:
logger.error("Result image is None")
if (
(not unit.check_similarity)
or result.similarity
and all(
[result.similarity.values() != 0]
+ [x >= unit.min_sim for x in result.similarity.values()]
)
and all(
[result.ref_similarity.values() != 0]
+ [x >= unit.min_ref_sim for x in result.ref_similarity.values()]
)
):
results.append(
(
result.image,
f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}",
)
)
else:
logger.warning(
f"skip, similarity to low, sim = {result.similarity} (target {unit.min_sim}) ref sim = {result.ref_similarity} (target = {unit.min_ref_sim})"
)
logger.debug("process_image_unit : Unit produced %s results", len(results))
return results
def process_images_units(
model: str,
units: List[FaceSwapUnitSettings],
images: List[Tuple[Optional[PILImage], Optional[str]]],
force_blend: bool = False,
) -> Optional[List[Tuple[PILImage, str]]]:
"""
Process a list of images using a specified model and unit settings for face swapping.
Args:
model (str): The name of the model to use for processing.
units (List[FaceSwapUnitSettings]): A list of settings for face swap units to apply on each image.
images (List[Tuple[Optional[PILImage], Optional[str]]]): A list of tuples, each containing
an image and its associated info string. If an image or info string is not available,
its value can be None.
upscaled_swapper (bool, optional): If True, uses an upscaled version of the face swapper.
Defaults to False.
force_blend (bool, optional): If True, forces the blending of the swapped face on the original
image. Defaults to False.
Returns:
Optional[List[Tuple[PILImage, str]]]: A list of tuples, each containing a processed image
and its associated info string. If no units are provided for processing, returns None.
"""
if len(units) == 0:
logger.info("Finished processing image, return %s images", len(images))
return None
logger.debug("%s more units", len(units))
processed_images = []
for i, (image, info) in enumerate(images):
logger.debug("Processing image %s", i)
swapped = process_image_unit(model, units[0], image, info, force_blend)
logger.debug("Image %s -> %s images", i, len(swapped))
nexts = process_images_units(model, units[1:], swapped, force_blend)
if nexts:
processed_images.extend(nexts)
else:
processed_images.extend(swapped)
return processed_images