From e6bd5af6a8e306a1cdef63402a77a980a04ad6e1 Mon Sep 17 00:00:00 2001 From: Grafting Rayman <156515434+GraftingRayman@users.noreply.github.com> Date: Fri, 17 Jan 2025 11:06:44 +0000 Subject: Add files via upload --- scripts/reactor_faceswap.py | 185 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 scripts/reactor_faceswap.py (limited to 'scripts/reactor_faceswap.py') diff --git a/scripts/reactor_faceswap.py b/scripts/reactor_faceswap.py new file mode 100644 index 0000000..7e6c03e --- /dev/null +++ b/scripts/reactor_faceswap.py @@ -0,0 +1,185 @@ +import os, glob + +from PIL import Image + +import modules.scripts as scripts +# from modules.upscaler import Upscaler, UpscalerData +from modules import scripts, scripts_postprocessing +from modules.processing import ( + StableDiffusionProcessing, + StableDiffusionProcessingImg2Img, +) +from modules.shared import state +from scripts.reactor_logger import logger +from scripts.reactor_swapper import ( + swap_face, + swap_face_many, + get_current_faces_model, + analyze_faces, + half_det_size, + providers +) +import folder_paths +import comfy.model_management as model_management + + +def get_models(): + models_path = os.path.join(folder_paths.models_dir,"insightface/*") + models = glob.glob(models_path) + models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")] + return models + + +class FaceSwapScript(scripts.Script): + + def process( + self, + p: StableDiffusionProcessing, + img, + enable, + source_faces_index, + faces_index, + model, + swap_in_source, + swap_in_generated, + gender_source, + gender_target, + face_model, + faces_order, + face_boost_enabled, + face_restore_model, + face_restore_visibility, + codeformer_weight, + interpolation, + ): + self.enable = enable + if self.enable: + + self.source = img + self.swap_in_generated = swap_in_generated + self.gender_source = gender_source + self.gender_target = gender_target + self.model = model + self.face_model = face_model + self.faces_order = faces_order + self.face_boost_enabled = face_boost_enabled + self.face_restore_model = face_restore_model + self.face_restore_visibility = face_restore_visibility + self.codeformer_weight = codeformer_weight + self.interpolation = interpolation + self.source_faces_index = [ + int(x) for x in source_faces_index.strip(",").split(",") if x.isnumeric() + ] + self.faces_index = [ + int(x) for x in faces_index.strip(",").split(",") if x.isnumeric() + ] + if len(self.source_faces_index) == 0: + self.source_faces_index = [0] + if len(self.faces_index) == 0: + self.faces_index = [0] + + if self.gender_source is None or self.gender_source == "no": + self.gender_source = 0 + elif self.gender_source == "female": + self.gender_source = 1 + elif self.gender_source == "male": + self.gender_source = 2 + + if self.gender_target is None or self.gender_target == "no": + self.gender_target = 0 + elif self.gender_target == "female": + self.gender_target = 1 + elif self.gender_target == "male": + self.gender_target = 2 + + # if self.source is not None: + if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source: + logger.status(f"Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) + + if len(p.init_images) == 1: + + result = swap_face( + self.source, + p.init_images[0], + source_faces_index=self.source_faces_index, + faces_index=self.faces_index, + model=self.model, + gender_source=self.gender_source, + gender_target=self.gender_target, + face_model=self.face_model, + faces_order=self.faces_order, + face_boost_enabled=self.face_boost_enabled, + face_restore_model=self.face_restore_model, + face_restore_visibility=self.face_restore_visibility, + codeformer_weight=self.codeformer_weight, + interpolation=self.interpolation, + ) + p.init_images[0] = result + + # for i in range(len(p.init_images)): + # if state.interrupted or model_management.processing_interrupted(): + # logger.status("Interrupted by User") + # break + # if len(p.init_images) > 1: + # logger.status(f"Swap in %s", i) + # result = swap_face( + # self.source, + # p.init_images[i], + # source_faces_index=self.source_faces_index, + # faces_index=self.faces_index, + # model=self.model, + # gender_source=self.gender_source, + # gender_target=self.gender_target, + # face_model=self.face_model, + # ) + # p.init_images[i] = result + + elif len(p.init_images) > 1: + result = swap_face_many( + self.source, + p.init_images, + source_faces_index=self.source_faces_index, + faces_index=self.faces_index, + model=self.model, + gender_source=self.gender_source, + gender_target=self.gender_target, + face_model=self.face_model, + faces_order=self.faces_order, + face_boost_enabled=self.face_boost_enabled, + face_restore_model=self.face_restore_model, + face_restore_visibility=self.face_restore_visibility, + codeformer_weight=self.codeformer_weight, + interpolation=self.interpolation, + ) + p.init_images = result + + logger.status("--Done!--") + # else: + # logger.error(f"Please provide a source face") + + def postprocess_batch(self, p, *args, **kwargs): + if self.enable: + images = kwargs["images"] + + def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args): + if self.enable and self.swap_in_generated: + if self.source is not None: + logger.status(f"Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) + image: Image.Image = script_pp.image + result = swap_face( + self.source, + image, + source_faces_index=self.source_faces_index, + faces_index=self.faces_index, + model=self.model, + upscale_options=self.upscale_options, + gender_source=self.gender_source, + gender_target=self.gender_target, + ) + try: + pp = scripts_postprocessing.PostprocessedImage(result) + pp.info = {} + p.extra_generation_params.update(pp.info) + script_pp.image = pp.image + except: + logger.error(f"Cannot create a result image") -- cgit v1.2.3