From bee2c4ee2d27a1424ad26861237d86c37d528412 Mon Sep 17 00:00:00 2001 From: Gourieff <777@lovemet.ru> Date: Fri, 24 Nov 2023 13:09:58 +0700 Subject: [PATCH] UPDATE: Safetensors Face Models --- README.md | 2 +- README_RU.md | 2 +- example/api_example.py | 3 + example/api_external.curl | 5 +- example/api_external.json | 5 +- scripts/reactor_api.py | 8 +- scripts/reactor_faceswap.py | 407 ++++++++++++++++++++++++++++-------- scripts/reactor_globals.py | 17 ++ scripts/reactor_helpers.py | 46 +++- scripts/reactor_swapper.py | 106 +++++++--- scripts/reactor_version.py | 2 +- 11 files changed, 479 insertions(+), 124 deletions(-) diff --git a/README.md b/README.md index 70ae680..8de80b1 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ logo - ![Version](https://img.shields.io/badge/version-0.5.1_beta1-green?style=for-the-badge&labelColor=darkgreen) + ![Version](https://img.shields.io/badge/version-0.5.1_beta2-green?style=for-the-badge&labelColor=darkgreen) Buy Me a Coffee at ko-fi.com diff --git a/README_RU.md b/README_RU.md index 4a9e121..20dcb41 100644 --- a/README_RU.md +++ b/README_RU.md @@ -2,7 +2,7 @@ logo - ![Version](https://img.shields.io/badge/версия-0.5.1_beta1-green?style=for-the-badge&labelColor=darkgreen) + ![Version](https://img.shields.io/badge/версия-0.5.1_beta2-green?style=for-the-badge&labelColor=darkgreen) Buy Me a Coffee at ko-fi.com diff --git a/example/api_example.py b/example/api_example.py index 72ec048..f3485c1 100644 --- a/example/api_example.py +++ b/example/api_example.py @@ -43,6 +43,9 @@ args=[ False, #18 Source Image Hash Check, True - by default False, #19 Target Image Hash Check, False - by default "CUDA", #20 CPU or CUDA (if you have it), CPU - by default + True, #21 Face Mask Correction + 1, #22 Select Source, 0 - Image, 1 - Face Model + "elena.safetensors", #23 Filename of the face model (from "models/reactor/faces"), e.g. elena.safetensors ] # The args for ReActor can be found by diff --git a/example/api_external.curl b/example/api_external.curl index 4fd0b1b..1c235e7 100644 --- a/example/api_external.curl +++ b/example/api_external.curl @@ -19,5 +19,8 @@ curl -X POST \ "gender_target": 0, "save_to_file": 1, "result_file_path": "", - "device": "CUDA" + "device": "CUDA", + "mask_face": 1, + "select_source": 1, + "face_model": "elena.safetensors" }' diff --git a/example/api_external.json b/example/api_external.json index 67b4d9a..1766427 100644 --- a/example/api_external.json +++ b/example/api_external.json @@ -15,5 +15,8 @@ "gender_target": 0, "save_to_file": 1, "result_file_path": "", - "device": "CUDA" + "device": "CUDA", + "mask_face": 1, + "select_source": 1, + "face_model": "elena.safetensors" } \ No newline at end of file diff --git a/scripts/reactor_api.py b/scripts/reactor_api.py index f0cc133..b5df7fa 100644 --- a/scripts/reactor_api.py +++ b/scripts/reactor_api.py @@ -71,7 +71,10 @@ def reactor_api(_: gr.Blocks, app: FastAPI): gender_target: int = Body(0,title="Gender Detection (Target) (0 - No, 1 - Female Only, 2 - Male Only)"), save_to_file: int = Body(0,title="Save Result to file, 0 - No, 1 - Yes"), result_file_path: str = Body("",title="(if 'save_to_file = 1') Result file path"), - device: str = Body("CPU",title="CPU or CUDA (if you have it)") + device: str = Body("CPU",title="CPU or CUDA (if you have it)"), + mask_face: int = Body(0,title="Face Mask Correction, 1 - True, 0 - False"), + select_source: int = Body(0,title="Select Source, 0 - Image, 1 - Face Model"), + face_model: str = Body("None",title="Filename of the face model (from 'models/reactor/faces'), e.g. elena.safetensors") ): s_image = api.decode_base64_to_image(source_image) t_image = api.decode_base64_to_image(target_image) @@ -80,11 +83,12 @@ def reactor_api(_: gr.Blocks, app: FastAPI): gender_s = gender_source gender_t = gender_target restore_first_bool = True if restore_first == 1 else False + mask_face = True if mask_face == 1 else False up_options = EnhancementOptions(do_restore_first=restore_first_bool, scale=scale, upscaler=get_upscaler(upscaler), upscale_visibility=upscale_visibility,face_restorer=get_face_restorer(face_restorer),restorer_visibility=restorer_visibility,codeformer_weight=codeformer_weight) use_model = get_full_model(model) if use_model is None: Exception("Model not found") - result = swap_face(s_image, t_image, use_model, sf_index, f_index, up_options, gender_s, gender_t, True, True, device) + result = swap_face(s_image, t_image, use_model, sf_index, f_index, up_options, gender_s, gender_t, True, True, device, mask_face, select_source, face_model) if save_to_file == 1: if result_file_path == "": result_file_path = default_file_path() diff --git a/scripts/reactor_faceswap.py b/scripts/reactor_faceswap.py index d2fbb9e..623e275 100644 --- a/scripts/reactor_faceswap.py +++ b/scripts/reactor_faceswap.py @@ -28,10 +28,16 @@ except: model_path = os.path.abspath("models") from scripts.reactor_logger import logger -from scripts.reactor_swapper import EnhancementOptions, swap_face, check_process_halt, reset_messaged +from scripts.reactor_swapper import ( + EnhancementOptions, + swap_face, + check_process_halt, + reset_messaged, + build_face_model +) from scripts.reactor_version import version_flag, app_title from scripts.console_log_patch import apply_logging_patch -from scripts.reactor_helpers import make_grid, get_image_path, set_Device +from scripts.reactor_helpers import make_grid, get_image_path, set_Device, get_model_names, get_facemodels from scripts.reactor_globals import DEVICE, DEVICE_LIST @@ -61,12 +67,69 @@ class FaceSwapScript(scripts.Script): def ui(self, is_img2img): with gr.Accordion(f"{app_title}", open=False): + + def update_fm_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=get_model_names(get_facemodels) + ) + def update_upscalers_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=[upscaler.name for upscaler in shared.sd_upscalers] + ) + def update_models_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=get_models() + ) + + # TAB MAIN with gr.Tab("Main"): with gr.Column(): - img = gr.Image(type="pil") + img = gr.Image( + type="pil", + label="Source Image", + ) + # face_model = gr.File( + # file_types=[".safetensors"], + # label="Face Model", + # show_label=True, + # ) enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}") - save_original = gr.Checkbox(False, label="Save Original", info="Save the original image(s) made before swapping; If you use \"img2img\" - this option will affect with \"Swap in generated\" only") - mask_face = gr.Checkbox(False, label="Face Mask Correction", info="Apply this option if you see some pixelation around face contours") + gr.Markdown("
") + with gr.Row(): + select_source = gr.Radio( + ["Image","Face Model"], + value="Image", + label="Select Source", + type="index", + scale=1, + ) + face_models = get_model_names(get_facemodels) + face_model = gr.Dropdown( + choices=face_models, + label="Choose Face Model", + value="None", + scale=2, + ) + fm_update = gr.Button( + value="🔄", + variant="tool", + ) + fm_update.click( + update_fm_list, + inputs=[face_model], + outputs=[face_model], + ) + setattr(face_model, "do_not_save_to_config", True) + save_original = gr.Checkbox( + False, + label="Save Original", + info="Save the original image(s) made before swapping; If you use \"img2img\" - this option will affect with \"Swap in generated\" only" + ) + mask_face = gr.Checkbox( + False, + label="Face Mask Correction", + info="Apply this option if you see some pixelation around face contours" + ) gr.Markdown("
") gr.Markdown("Source Image (above):") with gr.Row(): @@ -120,18 +183,30 @@ class FaceSwapScript(scripts.Script): True, label="Swap in generated image", visible=is_img2img, - ) + ) + + # TAB UPSCALE with gr.Tab("Upscale"): restore_first = gr.Checkbox( True, label="1. Restore Face -> 2. Upscale (-Uncheck- if you want vice versa)", info="Postprocessing Order" ) - upscaler_name = gr.Dropdown( - choices=[upscaler.name for upscaler in shared.sd_upscalers], - label="Upscaler", - value="None", - info="Won't scale if you choose -Swap in Source- via img2img, only 1x-postprocessing will affect (texturing, denoising, restyling etc.)" + with gr.Row(): + upscaler_name = gr.Dropdown( + choices=[upscaler.name for upscaler in shared.sd_upscalers], + label="Upscaler", + value="None", + info="Won't scale if you choose -Swap in Source- via img2img, only 1x-postprocessing will affect (texturing, denoising, restyling etc.)" + ) + upscalers_update = gr.Button( + value="🔄", + variant="tool", + ) + upscalers_update.click( + update_upscalers_list, + inputs=[upscaler_name], + outputs=[upscaler_name], ) gr.Markdown("
") with gr.Row(): @@ -139,6 +214,30 @@ class FaceSwapScript(scripts.Script): upscaler_visibility = gr.Slider( 0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)" ) + + # TAB TOOLS + with gr.Tab("Tools 🆕"): + with gr.Tab("Face Models"): + gr.Markdown("Load an image containing one person, name it and click 'Build and Save'") + img_fm = gr.Image( + type="pil", + label="Load Image to build Face Model", + ) + with gr.Row(equal_height=True): + fm_name = gr.Textbox( + value="", + placeholder="Please type any name (e.g. Elena)", + label="Face Model Name", + ) + save_fm_btn = gr.Button("Build and Save") + save_fm = gr.Markdown("You can find saved models in 'models/reactor/faces'") + save_fm_btn.click( + build_face_model, + inputs=[img_fm, fm_name], + outputs=[save_fm], + ) + + # TAB SETTINGS with gr.Tab("Settings"): models = get_models() with gr.Row(visible=EP_is_visible): @@ -161,21 +260,30 @@ class FaceSwapScript(scripts.Script): with gr.Row(): if len(models) == 0: logger.warning( - "You should at least have one model in models directory, please read the doc here : https://github.com/Gourieff/sd-webui-reactor/" + "You should at least have one model in models directory, please read the doc here: https://github.com/Gourieff/sd-webui-reactor/" ) model = gr.Dropdown( choices=models, - label="Model not found, please download one and reload WebUI", + label="Model not found, please download one and refresh the list" ) else: model = gr.Dropdown( choices=models, label="Model", value=models[0] ) + models_update = gr.Button( + value="🔄", + variant="tool", + ) + models_update.click( + update_models_list, + inputs=[model], + outputs=[model], + ) console_logging_level = gr.Radio( ["No log", "Minimum", "Default"], value="Minimum", label="Console Log Level", - type="index", + type="index" ) gr.Markdown("
") with gr.Row(): @@ -189,6 +297,8 @@ class FaceSwapScript(scripts.Script): label="Target Image Hash Check", info="Affects if you use Extras tab or img2img with only 'Swap in source image' on." ) + + gr.Markdown("by Eugene Gourieff") return [ img, @@ -213,6 +323,8 @@ class FaceSwapScript(scripts.Script): target_hash_check, device, mask_face, + select_source, + face_model, ] @@ -267,10 +379,14 @@ class FaceSwapScript(scripts.Script): target_hash_check, device, mask_face, + select_source, + face_model, ): self.enable = enable if self.enable: + logger.debug("*** Start process") + reset_messaged() if check_process_halt(): return @@ -295,6 +411,8 @@ class FaceSwapScript(scripts.Script): self.target_hash_check = target_hash_check self.device = device self.mask_face = mask_face + self.select_source = select_source + self.face_model = face_model if self.gender_source is None or self.gender_source == "No": self.gender_source = 0 if self.gender_target is None or self.gender_target == "No": @@ -318,9 +436,11 @@ class FaceSwapScript(scripts.Script): if self.mask_face is None: self.mask_face = False + logger.debug("*** Set Device") set_Device(self.device) - if self.source is not None: + if (self.source is not None and self.select_source == 0) or ((self.face_model is not None and self.face_model != "None") and self.select_source == 1): + logger.debug("*** Log patch") apply_logging_patch(console_logging_level) if isinstance(p, StableDiffusionProcessingImg2Img) and self.swap_in_source: logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) @@ -341,6 +461,8 @@ class FaceSwapScript(scripts.Script): target_hash_check=self.target_hash_check, device=self.device, mask_face=self.mask_face, + select_source=self.select_source, + face_model = self.face_model, ) p.init_images[i] = result # result_path = get_image_path(p.init_images[i], p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], "txt", p=p, suffix="-swapped") @@ -353,10 +475,13 @@ class FaceSwapScript(scripts.Script): else: logger.error("Please provide a source face") + return def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args): if self.enable: + logger.debug("*** Check postprocess") + reset_messaged() if check_process_halt(): return @@ -373,42 +498,44 @@ class FaceSwapScript(scripts.Script): if self.swap_in_generated: logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) - if self.source is not None: - for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)): - if check_process_halt(): - postprocess_run = False - break - if len(orig_images) > 1: - logger.status("Swap in %s", i) - result, output, swapped = swap_face( - self.source, - img, - source_faces_index=self.source_faces_index, - faces_index=self.faces_index, - model=self.model, - enhancement_options=self.enhancement_options, - gender_source=self.gender_source, - gender_target=self.gender_target, - source_hash_check=self.source_hash_check, - target_hash_check=self.target_hash_check, - device=self.device, - mask_face=self.mask_face, - ) - if result is not None and swapped > 0: - result_images.append(result) - suffix = "-swapped" - try: - img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix) - except: - logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)") - elif result is None: - logger.error("Cannot create a result image") - - # if len(output) != 0: - # split_fullfn = os.path.splitext(img_path[0]) - # fullfn = split_fullfn[0] + ".txt" - # with open(fullfn, 'w', encoding="utf8") as f: - # f.writelines(output) + # if self.source is not None: + for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)): + if check_process_halt(): + postprocess_run = False + break + if len(orig_images) > 1: + logger.status("Swap in %s", i) + result, output, swapped = swap_face( + self.source, + img, + source_faces_index=self.source_faces_index, + faces_index=self.faces_index, + model=self.model, + enhancement_options=self.enhancement_options, + gender_source=self.gender_source, + gender_target=self.gender_target, + source_hash_check=self.source_hash_check, + target_hash_check=self.target_hash_check, + device=self.device, + mask_face=self.mask_face, + select_source=self.select_source, + face_model = self.face_model, + ) + if result is not None and swapped > 0: + result_images.append(result) + suffix = "-swapped" + try: + img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix) + except: + logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)") + elif result is None: + logger.error("Cannot create a result image") + + # if len(output) != 0: + # split_fullfn = os.path.splitext(img_path[0]) + # fullfn = split_fullfn[0] + ".txt" + # with open(fullfn, 'w', encoding="utf8") as f: + # f.writelines(output) if shared.opts.return_grid and len(result_images) > 2 and postprocess_run: grid = make_grid(result_images) @@ -423,11 +550,14 @@ class FaceSwapScript(scripts.Script): def postprocess_batch(self, p, *args, **kwargs): if self.enable and not self.save_original: + logger.debug("*** Check postprocess_batch") images = kwargs["images"] def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args): if self.enable and self.swap_in_generated and not self.save_original: + logger.debug("*** Check postprocess_image") + current_job_number = shared.state.job_no + 1 job_count = shared.state.job_count if current_job_number == job_count: @@ -435,36 +565,38 @@ class FaceSwapScript(scripts.Script): if check_process_halt(): return - if self.source is not None: - logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) - image: Image.Image = script_pp.image - result, output, swapped = swap_face( - self.source, - image, - source_faces_index=self.source_faces_index, - faces_index=self.faces_index, - model=self.model, - enhancement_options=self.enhancement_options, - gender_source=self.gender_source, - gender_target=self.gender_target, - source_hash_check=self.source_hash_check, - target_hash_check=self.target_hash_check, - device=self.device, - mask_face=self.mask_face, - ) - try: - pp = scripts_postprocessing.PostprocessedImage(result) - pp.info = {} - p.extra_generation_params.update(pp.info) - script_pp.image = pp.image + # if (self.source is not None and self.select_source == 0) or ((self.face_model is not None and self.face_model != "None") and self.select_source == 1): + logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) + image: Image.Image = script_pp.image + result, output, swapped = swap_face( + self.source, + image, + source_faces_index=self.source_faces_index, + faces_index=self.faces_index, + model=self.model, + enhancement_options=self.enhancement_options, + gender_source=self.gender_source, + gender_target=self.gender_target, + source_hash_check=self.source_hash_check, + target_hash_check=self.target_hash_check, + device=self.device, + mask_face=self.mask_face, + select_source=self.select_source, + face_model = self.face_model, + ) + try: + pp = scripts_postprocessing.PostprocessedImage(result) + pp.info = {} + p.extra_generation_params.update(pp.info) + script_pp.image = pp.image - # if len(output) != 0: - # result_path = get_image_path(script_pp.image, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "txt", p=p, suffix="-swapped") - # if len(output) != 0: - # with open(result_path, 'w', encoding="utf8") as f: - # f.writelines(output) - except: - logger.error("Cannot create a result image") + # if len(output) != 0: + # result_path = get_image_path(script_pp.image, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "txt", p=p, suffix="-swapped") + # if len(output) != 0: + # with open(result_path, 'w', encoding="utf8") as f: + # f.writelines(output) + except: + logger.error("Cannot create a result image") class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): @@ -473,11 +605,56 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): def ui(self): with gr.Accordion(f"{app_title}", open=False): + + def update_fm_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=get_model_names(get_facemodels) + ) + def update_upscalers_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=[upscaler.name for upscaler in shared.sd_upscalers] + ) + def update_models_list(selected: str): + return gr.Dropdown.update( + value=selected, choices=get_models() + ) + + # TAB MAIN with gr.Tab("Main"): with gr.Column(): img = gr.Image(type="pil") enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}") - mask_face = gr.Checkbox(False, label="Face Mask Correction", info="Apply this option if you see some pixelation around face contours") + # gr.Markdown("
") + with gr.Row(): + select_source = gr.Radio( + ["Image","Face Model"], + value="Image", + label="Select Source", + type="index", + scale=1, + ) + face_models = get_model_names(get_facemodels) + face_model = gr.Dropdown( + choices=face_models, + label="Choose Face Model", + value="None", + scale=2, + ) + fm_update = gr.Button( + value="🔄", + variant="tool", + ) + fm_update.click( + update_fm_list, + inputs=[face_model], + outputs=[face_model], + ) + setattr(face_model, "do_not_save_to_config", True) + mask_face = gr.Checkbox( + False, + label="Face Mask Correction", + info="Apply this option if you see some pixelation around face contours" + ) gr.Markdown("Source Image (above):") with gr.Row(): source_faces_index = gr.Textbox( @@ -519,23 +696,58 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): 0, 1, 0.5, step=0.1, label="CodeFormer Weight", info="0 = maximum effect, 1 = minimum effect" ) + # TAB UPSCALE with gr.Tab("Upscale"): restore_first = gr.Checkbox( True, label="1. Restore Face -> 2. Upscale (-Uncheck- if you want vice versa)", info="Postprocessing Order" ) - upscaler_name = gr.Dropdown( - choices=[upscaler.name for upscaler in shared.sd_upscalers], - label="Upscaler", - value="None", - info="Won't scale if you choose -Swap in Source- via img2img, only 1x-postprocessing will affect (texturing, denoising, restyling etc.)" + with gr.Row(): + upscaler_name = gr.Dropdown( + choices=[upscaler.name for upscaler in shared.sd_upscalers], + label="Upscaler", + value="None", + info="Won't scale if you choose -Swap in Source- via img2img, only 1x-postprocessing will affect (texturing, denoising, restyling etc.)" + ) + upscalers_update = gr.Button( + value="🔄", + variant="tool", + ) + upscalers_update.click( + update_upscalers_list, + inputs=[upscaler_name], + outputs=[upscaler_name], ) with gr.Row(): upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Scale by") upscaler_visibility = gr.Slider( 0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)" ) + + # TAB TOOLS + with gr.Tab("Tools 🆕"): + with gr.Tab("Face Models"): + gr.Markdown("Load an image containing one person, name it and click 'Build and Save'") + img_fm = gr.Image( + type="pil", + label="Load Image to build Face Model", + ) + with gr.Row(equal_height=True): + fm_name = gr.Textbox( + value="", + placeholder="Please type any name (e.g. Elena)", + label="Face Model Name", + ) + save_fm_btn = gr.Button("Build and Save") + save_fm = gr.Markdown("You can find saved models in 'models/reactor/faces'") + save_fm_btn.click( + build_face_model, + inputs=[img_fm, fm_name], + outputs=[save_fm], + ) + + # TAB SETTINGS with gr.Tab("Settings"): models = get_models() with gr.Row(visible=EP_is_visible): @@ -558,22 +770,33 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): with gr.Row(): if len(models) == 0: logger.warning( - "You should at least have one model in models directory, please read the doc here : https://github.com/Gourieff/sd-webui-reactor/" + "You should at least have one model in models directory, please read the doc here: https://github.com/Gourieff/sd-webui-reactor/" ) model = gr.Dropdown( choices=models, - label="Model not found, please download one and reload WebUI", + label="Model not found, please download one and refresh the list", ) else: model = gr.Dropdown( choices=models, label="Model", value=models[0] ) + models_update = gr.Button( + value="🔄", + variant="tool", + ) + models_update.click( + update_models_list, + inputs=[model], + outputs=[model], + ) console_logging_level = gr.Radio( ["No log", "Minimum", "Default"], value="Minimum", label="Console Log Level", type="index", ) + + gr.Markdown("by Eugene Gourieff") args = { 'img': img, @@ -593,6 +816,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): 'codeformer_weight': codeformer_weight, 'device': device, 'mask_face': mask_face, + 'select_source': select_source, + 'face_model': face_model, } return args @@ -643,6 +868,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): self.codeformer_weight = args['codeformer_weight'] self.device = args['device'] self.mask_face = args['mask_face'] + self.select_source = args['select_source'] + self.face_model = args['face_model'] if self.gender_source is None or self.gender_source == "No": self.gender_source = 0 if self.gender_target is None or self.gender_target == "No": @@ -667,7 +894,7 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): set_Device(self.device) - if self.source is not None: + if (self.source is not None and self.select_source == 0) or ((self.face_model is not None and self.face_model != "None") and self.select_source == 1): apply_logging_patch(self.console_logging_level) logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) image: Image.Image = pp.image @@ -684,6 +911,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): target_hash_check=True, device=self.device, mask_face=self.mask_face, + select_source=self.select_source, + face_model=self.face_model, ) try: pp.info["ReActor"] = True diff --git a/scripts/reactor_globals.py b/scripts/reactor_globals.py index 96e6d7d..7aea387 100644 --- a/scripts/reactor_globals.py +++ b/scripts/reactor_globals.py @@ -1,10 +1,27 @@ import os from pathlib import Path +try: + from modules.paths_internal import models_path +except: + try: + from modules.paths import models_path + except: + models_path = os.path.abspath("models") + IS_RUN: bool = False BASE_PATH = os.path.join(Path(__file__).parents[1]) DEVICE_LIST: list = ["CPU", "CUDA"] +MODELS_PATH = models_path +REACTOR_MODELS_PATH = os.path.join(models_path, "reactor") +FACE_MODELS_PATH = os.path.join(REACTOR_MODELS_PATH, "faces") + +if not os.path.exists(REACTOR_MODELS_PATH): + os.makedirs(REACTOR_MODELS_PATH) + if not os.path.exists(FACE_MODELS_PATH): + os.makedirs(FACE_MODELS_PATH) + def updateDevice(): try: LAST_DEVICE_PATH = os.path.join(BASE_PATH, "last_device.txt") diff --git a/scripts/reactor_helpers.py b/scripts/reactor_helpers.py index 21cc90f..cfce05f 100644 --- a/scripts/reactor_helpers.py +++ b/scripts/reactor_helpers.py @@ -1,14 +1,17 @@ -import os +import os, glob from collections import Counter from PIL import Image from math import isqrt, ceil from typing import List import logging import hashlib +import torch +from safetensors.torch import save_file, safe_open +from insightface.app.common import Face from modules.images import FilenameGenerator, get_next_sequence_number from modules import shared, script_callbacks -from scripts.reactor_globals import DEVICE, BASE_PATH +from scripts.reactor_globals import DEVICE, BASE_PATH, FACE_MODELS_PATH def set_Device(value): global DEVICE @@ -133,3 +136,42 @@ def addLoggingLevel(levelName, levelNum, methodName=None): def get_image_md5hash(image: Image.Image): md5hash = hashlib.md5(image.tobytes()) return md5hash.hexdigest() + +def save_face_model(face: Face, filename: str) -> None: + try: + tensors = { + "bbox": torch.tensor(face["bbox"]), + "kps": torch.tensor(face["kps"]), + "det_score": torch.tensor(face["det_score"]), + "landmark_3d_68": torch.tensor(face["landmark_3d_68"]), + "pose": torch.tensor(face["pose"]), + "landmark_2d_106": torch.tensor(face["landmark_2d_106"]), + "embedding": torch.tensor(face["embedding"]), + "gender": torch.tensor(face["gender"]), + "age": torch.tensor(face["age"]), + } + save_file(tensors, filename) + # print(f"Face model has been saved to '{filename}'") + except Exception as e: + print(f"Error: {e}") + +def load_face_model(filename: str): + face = {} + model_path = os.path.join(FACE_MODELS_PATH, filename) + with safe_open(model_path, framework="pt") as f: + for k in f.keys(): + face[k] = f.get_tensor(k).numpy() + return Face(face) + +def get_facemodels(): + models_path = os.path.join(FACE_MODELS_PATH, "*") + models = glob.glob(models_path) + models = [x for x in models if x.endswith(".safetensors")] + return models + +def get_model_names(get_models): + models = get_models() + names = ["None"] + for x in models: + names.append(os.path.basename(x)) + return names diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py index c244724..65d2434 100644 --- a/scripts/reactor_swapper.py +++ b/scripts/reactor_swapper.py @@ -8,8 +8,12 @@ import numpy as np from PIL import Image import insightface +from insightface.app.common import Face + +from scripts.reactor_globals import FACE_MODELS_PATH +from scripts.reactor_helpers import get_image_md5hash, get_Device, save_face_model, load_face_model +from scripts.console_log_patch import apply_logging_patch -from scripts.reactor_helpers import get_image_md5hash, get_Device from modules.face_restoration import FaceRestoration try: # A1111 from modules import codeformer_model @@ -26,7 +30,7 @@ except: try: from modules.paths import models_path except: - model_path = os.path.abspath("models") + models_path = os.path.abspath("models") import warnings @@ -78,10 +82,11 @@ def check_process_halt(msgforced: bool = False): FS_MODEL = None +ANALYSIS_MODEL = None MASK_MODEL = None + CURRENT_FS_MODEL_PATH = None CURRENT_MASK_MODEL_PATH = None -ANALYSIS_MODEL = None SOURCE_FACES = None SOURCE_IMAGE_HASH = None @@ -108,8 +113,6 @@ def getFaceSwapModel(model_path: str): return FS_MODEL - - def restore_face(image: Image, enhancement_options: EnhancementOptions): result_image = image @@ -173,6 +176,7 @@ def enhance_image(image: Image, enhancement_options: EnhancementOptions): result_image = restore_face(result_image, enhancement_options) return result_image + def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementOptions,target_img_orig:Image.Image,entire_mask_image:Image.Image)->Image.Image: result_image = image @@ -309,6 +313,8 @@ def swap_face( target_hash_check: bool = False, device: str = "CPU", mask_face: bool = False, + select_source: int = 0, + face_model: str = "None", ): global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, PROVIDERS result_image = target_img @@ -333,40 +339,56 @@ def swap_face( source_img = Image.open(io.BytesIO(img_bytes)) - source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) + target_img_orig = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) entire_mask_image = np.zeros_like(np.array(target_img)) + output: List = [] output_info: str = "" swapped = 0 - if source_hash_check: + if select_source == 0 and 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_hash_check: - 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_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) + 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: + 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 + + else: logger.status("Analyzing Source Image...") source_faces = analyze_faces(source_img) - SOURCE_FACES = source_faces - elif source_image_same: - logger.status("Using Ready Source Face(s) Model...") - source_faces = SOURCE_FACES - + + elif select_source == 1 and (face_model is not None and face_model != "None"): + source_face_model = [load_face_model(face_model)] + if source_face_model is not None: + source_faces_index = [0] + source_faces = source_face_model + logger.status("Using Loaded Source Face Model...") + else: + logger.error(f"Cannot load Face Model File: {face_model}.safetensors") else: - logger.status("Analyzing Source Image...") - source_faces = analyze_faces(source_img) + logger.error("Cannot detect any Source") if source_faces is not None: @@ -390,7 +412,7 @@ def swap_face( target_faces = analyze_faces(target_img) TARGET_FACES = target_faces elif target_image_same: - logger.status("Using Ready Target Face(s) Model...") + logger.status("Using Hashed Target Face(s) Model...") target_faces = TARGET_FACES else: @@ -398,7 +420,13 @@ def swap_face( target_faces = analyze_faces(target_img) logger.status("Detecting Source Face, Index = %s", source_faces_index[0]) - source_face, wrong_gender, source_age, source_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source) + if select_source == 0 and source_img is not None: + source_face, wrong_gender, source_age, source_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source) + else: + source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]] + wrong_gender = 0 + source_age = source_face["age"] + source_gender = "Female" if source_face["gender"] == 0 else "Male" if source_age != "None" or source_gender != "None": logger.status("Detected: -%s- y.o. %s", source_age, source_gender) @@ -491,3 +519,29 @@ def swap_face( logger.status("No source face(s) found") return result_image, output, swapped + + +def build_face_model(image: Image.Image, name: str): + if image is None: + error_msg = "Please load an Image" + logger.error(error_msg) + return error_msg + if name is None: + error_msg = "Please filled out the 'Face Model Name' field" + logger.error(error_msg) + return error_msg + apply_logging_patch(1) + image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + logger.status("Building Face Model...") + face_model = analyze_faces(image)[0] + if face_model is not None: + face_model_path = os.path.join(FACE_MODELS_PATH, name + ".safetensors") + save_face_model(face_model,face_model_path) + logger.status("--Done!--") + done_msg = f"Face model has been saved to '{face_model_path}'" + logger.status(done_msg) + return done_msg + else: + no_face_msg = "No face found, please try another image" + logger.error(no_face_msg) + return no_face_msg diff --git a/scripts/reactor_version.py b/scripts/reactor_version.py index 3230c48..1c1c0c8 100644 --- a/scripts/reactor_version.py +++ b/scripts/reactor_version.py @@ -1,5 +1,5 @@ app_title = "ReActor" -version_flag = "v0.5.1-b1" +version_flag = "v0.5.1-b2" from scripts.reactor_logger import logger, get_Run, set_Run