CalvinBackup/sd-webui-faceswaplab
1
0
Fork 0
Code Issues 56 Pull Requests 4 Actions Packages Projects Releases 8 Wiki Activity

wip, add nsfw option due to perf, still some mypy warnings

Browse Source
This commit is contained in:
Tran Xen
2023-08-16 01:11:02 +02:00
parent f9fc0bbff1
commit afcfc7d255
20 changed files with 123 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/faceswaplab_swapping/face_checkpoints.py
+239
View File
@@ -0,0 +1,239 @@
import glob
import os
from typing import *
from insightface.app.common import Face
from safetensors.torch import save_file, safe_open
import torch
import modules.scripts as scripts
from modules import scripts
from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
from scripts.faceswaplab_utils.faceswaplab_logging import logger
from scripts.faceswaplab_utils.typing import *
from scripts.faceswaplab_utils import imgutils
from scripts.faceswaplab_utils.models_utils import get_swap_models
import traceback
import dill as pickle # will be removed in future versions
from scripts.faceswaplab_swapping import swapper
from pprint import pformat
import re
from client_api import api_utils
import tempfile
def sanitize_name(name: str) -> str:
"""
Sanitize the input name by removing special characters and replacing spaces with underscores.
Parameters:
name (str): The input name to be sanitized.
Returns:
str: The sanitized name with special characters removed and spaces replaced by underscores.
"""
name = re.sub("[^A-Za-z0-9_. ]+", "", name)
name = name.replace(" ", "_")
return name[:255]
def build_face_checkpoint_and_save(
images: List[PILImage], name: str, overwrite: bool = False, path: str = None
) -> PILImage:
"""
Builds a face checkpoint using the provided image files, performs face swapping,
and saves the result to a file. If a blended face is successfully obtained and the face swapping
process succeeds, the resulting image is returned. Otherwise, None is returned.
Args:
batch_files (list): List of image file paths used to create the face checkpoint.
name (str): The name assigned to the face checkpoint.
Returns:
PIL.PILImage or None: The resulting swapped face image if the process is successful; None otherwise.
"""
try:
name = sanitize_name(name)
images = images or []
logger.info("Build %s with %s images", name, len(images))
faces = swapper.get_faces_from_img_files(images)
blended_face = swapper.blend_faces(faces)
preview_path = os.path.join(
scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
)
reference_preview_img: PILImage
if blended_face:
if blended_face["gender"] == 0:
reference_preview_img = Image.open(
os.path.join(preview_path, "woman.png")
)
else:
reference_preview_img = Image.open(
os.path.join(preview_path, "man.png")
)
if name == "":
name = "default_name"
logger.debug("Face %s", pformat(blended_face))
target_face = swapper.get_or_default(
swapper.get_faces(imgutils.pil_to_cv2(reference_preview_img)), 0, None
)
if target_face is None:
logger.error(
"Failed to open reference image, cannot create preview : That should not happen unless you deleted the references folder or change the detection threshold."
)
else:
result = swapper.swap_face(
target_faces=[target_face],
source_face=blended_face,
target_img=reference_preview_img,
model=get_swap_models()[0],
swapping_options=InswappperOptions(face_restorer_name="Codeformer"),
)
preview_image = result.image
if path:
file_path = path
else:
file_path = os.path.join(get_checkpoint_path(), f"{name}.safetensors")
if not overwrite:
file_number = 1
while os.path.exists(file_path):
file_path = os.path.join(
get_checkpoint_path(), f"{name}_{file_number}.safetensors"
)
file_number += 1
save_face(filename=file_path, face=blended_face)
preview_image.save(file_path + ".png")
try:
data = load_face(file_path)
logger.debug(data)
except Exception as e:
logger.error("Error loading checkpoint, after creation %s", e)
traceback.print_exc()
return preview_image
else:
logger.error("No face found")
return None
except Exception as e:
logger.error("Failed to build checkpoint %s", e)
traceback.print_exc()
return None
def save_face(face: Face, filename: str) -> None:
try:
tensors = {
"embedding": torch.tensor(face["embedding"]),
"gender": torch.tensor(face["gender"]),
"age": torch.tensor(face["age"]),
}
save_file(tensors, filename)
except Exception as e:
traceback.print_exc
logger.error("Failed to save checkpoint %s", e)
raise e
def load_face(name: str) -> Face:
if name.startswith("data:application/face;base64,"):
with tempfile.NamedTemporaryFile(delete=True) as temp_file:
api_utils.base64_to_safetensors(name, temp_file.name)
face = {}
with safe_open(temp_file.name, framework="pt", device="cpu") as f:
for k in f.keys():
logger.debug("load key %s", k)
face[k] = f.get_tensor(k).numpy()
return Face(face)
filename = matching_checkpoint(name)
if filename is None:
return None
if filename.endswith(".pkl"):
logger.warning(
"Pkl files for faces are deprecated to enhance safety, they will be unsupported in future versions."
)
logger.warning("The file will be converted to .safetensors")
logger.warning(
"You can also use this script https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d"
)
with open(filename, "rb") as file:
logger.info("Load pkl")
face = Face(pickle.load(file))
logger.warning(
"Convert to safetensors, you can remove the pkl version once you have ensured that the safetensor is working"
)
save_face(face, filename.replace(".pkl", ".safetensors"))
return face
elif filename.endswith(".safetensors"):
face = {}
with safe_open(filename, framework="pt", device="cpu") as f:
for k in f.keys():
logger.debug("load key %s", k)
face[k] = f.get_tensor(k).numpy()
return Face(face)
raise NotImplementedError("Unknown file type, face extraction not implemented")
def get_checkpoint_path() -> str:
checkpoint_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")
os.makedirs(checkpoint_path, exist_ok=True)
return checkpoint_path
def matching_checkpoint(name: str) -> Optional[str]:
"""
Retrieve the full path of a checkpoint file matching the given name.
If the name already includes a path separator, it is returned as-is. Otherwise, the function looks for a matching
file with the extensions ".safetensors" or ".pkl" in the checkpoint directory.
Args:
name (str): The name or path of the checkpoint file.
Returns:
Optional[str]: The full path of the matching checkpoint file, or None if no match is found.
"""
# If the name already includes a path separator, return it as is
if os.path.sep in name:
return name
# If the name doesn't end with the specified extensions, look for a matching file
if not (name.endswith(".safetensors") or name.endswith(".pkl")):
# Try appending each extension and check if the file exists in the checkpoint path
for ext in [".safetensors", ".pkl"]:
full_path = os.path.join(get_checkpoint_path(), name + ext)
if os.path.exists(full_path):
return full_path
# If no matching file is found, return None
return None
# If the name already ends with the specified extensions, simply complete the path
return os.path.join(get_checkpoint_path(), name)
def get_face_checkpoints() -> List[str]:
"""
Retrieve a list of face checkpoint paths.
This function searches for face files with the extension ".safetensors" in the specified directory and returns a list
containing the paths of those files.
Returns:
list: A list of face paths, including the string "None" as the first element.
"""
faces_path = os.path.join(get_checkpoint_path(), "*.safetensors")
faces = glob.glob(faces_path)
faces_path = os.path.join(get_checkpoint_path(), "*.pkl")
faces += glob.glob(faces_path)
return ["None"] + [os.path.basename(face) for face in sorted(faces)]
scripts/faceswaplab_swapping/facemask.py
+1 -1
View File
@@ -83,7 +83,7 @@ def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarr
convert_bgr_to_rgb=True,
use_float32=True,
)
normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) # type: ignore
assert isinstance(face_input, torch.Tensor)
face_input = torch.unsqueeze(face_input, 0).to(device)
scripts/faceswaplab_swapping/swapper.py
+17 -17
View File
@@ -26,7 +26,6 @@ from scripts.faceswaplab_utils.imgutils import (
)
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
from scripts import faceswaplab_globals
from modules.shared import opts
from functools import lru_cache
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
@@ -38,12 +37,13 @@ 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 opts.data.get("faceswaplab_use_gpu", False)
or get_sd_option("faceswaplab_use_gpu", False)
) and sys.platform != "darwin"
@@ -166,6 +166,7 @@ def batch_process(
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(
@@ -182,7 +183,7 @@ def batch_process(
swapped_images = process_images_units(
get_current_swap_model(), images=[(src_image, None)], units=units
)
if len(swapped_images) > 0:
if swapped_images and len(swapped_images) > 0:
current_images += [img for img, _ in swapped_images]
logger.info("%s images generated", len(current_images))
@@ -209,7 +210,7 @@ def extract_faces(
images: List[PILImage],
extract_path: Optional[str],
postprocess_options: PostProcessingOptions,
) -> Optional[List[str]]:
) -> Optional[List[PILImage]]:
"""
Extracts faces from a list of image files.
@@ -232,14 +233,14 @@ def extract_faces(
os.makedirs(extract_path, exist_ok=True)
if images:
result_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)
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))
@@ -370,7 +371,7 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
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)
insightface.model_zoo.get_model(model_path, providers=providers) # type: ignore
)
pbar.update(1)
logger.info("%s", pformat(captured.getvalue()))
@@ -402,11 +403,11 @@ def get_faces(
"""
if det_thresh is None:
det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5)
det_thresh = get_sd_option("faceswaplab_detection_threshold", 0.5)
auto_det_size = opts.data.get("faceswaplab_auto_det_size", True)
auto_det_size = get_sd_option("faceswaplab_auto_det_size", True)
if not auto_det_size:
x = opts.data.get("faceswaplab_det_size", 640)
x = get_sd_option("faceswaplab_det_size", 640)
det_size = (x, x)
face_analyser = getAnalysisModel(det_size, det_thresh)
@@ -433,7 +434,7 @@ def get_faces(
try:
# Sort the detected faces based on their x-coordinate of the bounding box
return sorted(faces, key=lambda x: x.bbox[0])
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()
@@ -470,7 +471,7 @@ def filter_faces(
filtered_faces = sorted(
all_faces,
reverse=True,
key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),
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:
@@ -566,7 +567,7 @@ def blend_faces(faces: List[Face]) -> Optional[Face]:
ValueError: If the embeddings have different shapes.
"""
embeddings = [face.embedding for face in faces]
embeddings: list[Any] = [face.embedding for face in faces]
if len(embeddings) > 0:
embedding_shape = embeddings[0].shape
@@ -592,7 +593,6 @@ def blend_faces(faces: List[Face]) -> Optional[Face]:
def swap_face(
reference_face: CV2ImgU8,
source_face: Face,
target_img: PILImage,
target_faces: List[Face],
@@ -604,7 +604,6 @@ def swap_face(
Swaps faces in the target image with the source face.
Args:
reference_face (CV2ImgU8): The reference face used for similarity comparison.
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.
@@ -614,7 +613,9 @@ def swap_face(
"""
return_result = ImageResult(target_img, {}, {})
target_img_cv2: CV2ImgU8 = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
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)
@@ -732,7 +733,6 @@ def process_image_unit(
save_img_debug(image, "Before swap")
result: ImageResult = swap_face(
reference_face=reference_face,
source_face=src_face,
target_img=current_image,
target_faces=target_faces,
scripts/faceswaplab_swapping/upcaled_inswapper_options.py
+4 -3
View File
@@ -1,20 +1,21 @@
from dataclasses import *
from typing import Optional
from client_api import api_utils
@dataclass
class InswappperOptions:
face_restorer_name: str = None
face_restorer_name: Optional[str] = None
restorer_visibility: float = 1
codeformer_weight: float = 1
upscaler_name: str = None
upscaler_name: Optional[str] = None
improved_mask: bool = False
color_corrections: bool = False
sharpen: bool = False
erosion_factor: float = 1.0
@staticmethod
def from_api_dto(dto: api_utils.InswappperOptions) -> "InswappperOptions":
def from_api_dto(dto: Optional[api_utils.InswappperOptions]) -> "InswappperOptions":
"""
Converts a InpaintingOptions object from an API DTO (Data Transfer Object).
scripts/faceswaplab_swapping/upscaled_inswapper.py
+15 -8
View File
@@ -1,10 +1,9 @@
from typing import Any, Tuple, Union
from typing import Any, Optional, Tuple, Union
import cv2
import numpy as np
from insightface.model_zoo.inswapper import INSwapper
from insightface.utils import face_align
from modules import processing, shared
from modules.shared import opts
from modules.upscaler import UpscalerData
from scripts.faceswaplab_postprocessing import upscaling
@@ -14,13 +13,14 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
from scripts.faceswaplab_swapping.facemask import generate_face_mask
from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2
from scripts.faceswaplab_utils.sd_utils import get_sd_option
from scripts.faceswaplab_utils.typing import CV2ImgU8, Face
from scripts.faceswaplab_utils.faceswaplab_logging import logger
def get_upscaler() -> UpscalerData:
def get_upscaler() -> Optional[UpscalerData]:
for upscaler in shared.sd_upscalers:
if upscaler.name == opts.data.get(
if upscaler.name == get_sd_option(
"faceswaplab_upscaled_swapper_upscaler", "LDSR"
):
return upscaler
@@ -130,8 +130,14 @@ class UpscaledINSwapper(INSwapper):
self.__dict__.update(inswapper.__dict__)
def upscale_and_restore(
self, img: CV2ImgU8, k: int = 2, inswapper_options: InswappperOptions = None
self,
img: CV2ImgU8,
k: int = 2,
inswapper_options: Optional[InswappperOptions] = None,
) -> CV2ImgU8:
if inswapper_options is None:
return img
pil_img = cv2_to_pil(img)
pp_options = PostProcessingOptions(
upscaler_name=inswapper_options.upscaler_name,
@@ -156,7 +162,7 @@ class UpscaledINSwapper(INSwapper):
target_face: Face,
source_face: Face,
paste_back: bool = True,
options: InswappperOptions = None,
options: Optional[InswappperOptions] = None,
) -> Union[CV2ImgU8, Tuple[CV2ImgU8, Any]]:
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
blob = cv2.dnn.blobFromImage(
@@ -166,9 +172,10 @@ class UpscaledINSwapper(INSwapper):
(self.input_mean, self.input_mean, self.input_mean),
swapRB=True,
)
latent = source_face.normed_embedding.reshape((1, -1))
latent = source_face.normed_embedding.reshape((1, -1)) # type: ignore
latent = np.dot(latent, self.emap)
latent /= np.linalg.norm(latent)
assert self.session is not None
pred = self.session.run(
self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent}
)[0]
@@ -274,7 +281,7 @@ class UpscaledINSwapper(INSwapper):
mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
mask_size = int(np.sqrt(mask_h * mask_w))
erosion_factor = options.erosion_factor
erosion_factor = options.erosion_factor if options else 1
k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))