diff --git a/scripts/cimage.py b/scripts/cimage.py index 3892c51..35f09a7 100644 --- a/scripts/cimage.py +++ b/scripts/cimage.py @@ -1,57 +1,9 @@ -from typing import List, Union, Dict, Set, Tuple +import tempfile +from ifnude import detect -from diffusers.pipelines.stable_diffusion.safety_checker import ( - StableDiffusionSafetyChecker, -) -from transformers import AutoFeatureExtractor -import torch -from PIL import Image, ImageFilter -import numpy as np - -safety_model_id: str = "CompVis/stable-diffusion-safety-checker" -safety_feature_extractor: AutoFeatureExtractor = None -safety_checker: StableDiffusionSafetyChecker = None - - -def numpy_to_pil(images: np.ndarray) -> List[Image.Image]: - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - -def check_image(x_image: np.ndarray) -> Tuple[np.ndarray, List[bool]]: - global safety_feature_extractor, safety_checker - - if safety_feature_extractor is None: - safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id) - safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id) - - safety_checker_input = safety_feature_extractor( - images=numpy_to_pil(x_image), return_tensors="pt" - ) - x_checked_image, hs = safety_checker( - images=x_image, clip_input=safety_checker_input.pixel_values - ) - - return x_checked_image, hs - - -def check_batch(x: torch.Tensor) -> torch.Tensor: - x_samples_ddim_numpy = x.cpu().permute(0, 2, 3, 1).numpy() - x_checked_image, _ = check_image(x_samples_ddim_numpy) - x = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2) - return x - - -def convert_to_sd(img: Image) -> Image: - _, hs = check_image(np.array(img)) - if any(hs): - img = ( - img.resize((int(img.width * 0.1), int(img.height * 0.1))) - .resize(img.size, Image.BOX) - .filter(ImageFilter.BLUR) - ) - return img +def convert_to_sd(img): + shapes = [] + chunks = detect(img) + for chunk in chunks: + shapes.append(chunk["score"] > 0.7) + return [any(shapes), tempfile.NamedTemporaryFile(delete=False, suffix=".png")] diff --git a/scripts/faceswap.py b/scripts/faceswap.py index e79eb9c..6d5bf22 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -13,7 +13,6 @@ from modules.face_restoration import FaceRestoration from scripts.roop_logging import logger from scripts.swapper import UpscaleOptions, swap_face, ImageResult -from scripts.cimage import check_batch from scripts.roop_version import version_flag import os @@ -179,8 +178,7 @@ class FaceSwapScript(scripts.Script): def postprocess_batch(self, p, *args, **kwargs): if self.enable: - images = kwargs["images"] - images[:] = check_batch(images)[:] + return images def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args): if self.enable and self.swap_in_generated: diff --git a/scripts/swapper.py b/scripts/swapper.py index 4a2f4b5..2a3e79e 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -28,40 +28,7 @@ class UpscaleOptions: face_restorer: FaceRestoration = None restorer_visibility: float = 0.5 - -def save_image(img: Image, filename: str): - convert_to_sd(img).save(filename) - - -def cosine_distance(vector1: np.ndarray, vector2: np.ndarray) -> float: - vec1 = vector1.flatten() - vec2 = vector2.flatten() - - dot_product = np.dot(vec1, vec2) - norm1 = np.linalg.norm(vec1) - norm2 = np.linalg.norm(vec2) - - cosine_distance = 1 - (dot_product / (norm1 * norm2)) - return cosine_distance - - -def cosine_similarity(test_vec: np.ndarray, source_vecs: List[np.ndarray]) -> float: - cos_dist = sum(cosine_distance(test_vec, source_vec) for source_vec in source_vecs) - average_cos_dist = cos_dist / len(source_vecs) - return average_cos_dist - - -ANALYSIS_MODEL = None - - -def getAnalysisModel(): - global ANALYSIS_MODEL - if ANALYSIS_MODEL is None: - ANALYSIS_MODEL = insightface.app.FaceAnalysis( - name="buffalo_l", providers=providers - ) - return ANALYSIS_MODEL - +ANALYSIS_MODEL = insightface.app.FaceAnalysis(name="buffalo_l", providers=providers) FS_MODEL = None CURRENT_FS_MODEL_PATH = None @@ -108,7 +75,7 @@ def upscale_image(image: Image, upscale_options: UpscaleOptions): def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)): - face_analyser = copy.deepcopy(getAnalysisModel()) + face_analyser = copy.deepcopy(ANALYSIS_MODEL) face_analyser.prepare(ctx_id=0, det_size=det_size) face = face_analyser.get(img_data) @@ -140,8 +107,10 @@ def swap_face( faces_index: Set[int] = {0}, upscale_options: Union[UpscaleOptions, None] = None, ) -> ImageResult: - fn = tempfile.NamedTemporaryFile(delete=False, suffix=".png") - if model is not None: + result_image = target_img + converted = convert_to_sd(target_img) + scale, fn = converted[0], converted[1] + if model is not None and not scale: source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) source_face = get_face_single(source_img, face_index=0) @@ -160,8 +129,7 @@ def swap_face( result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) if upscale_options is not None: result_image = upscale_image(result_image, upscale_options) - - save_image(result_image, fn.name) else: logger.info("No source face found") + result_image.save(fn.name) return ImageResult(path=fn.name)