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 @@
- 
+ 
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 @@
- 
+ 
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