From 58f00fe7d5cae0ca66167fe6850a531f2e58a5db Mon Sep 17 00:00:00 2001 From: Gourieff <777@lovemet.ru> Date: Fri, 30 Jun 2023 14:53:36 +0700 Subject: [PATCH] UPDATE: Postprocess logic improvement + Version UP --- scripts/faceswap.py | 17 ++++++++-- scripts/roop_version.py | 4 +-- scripts/swapper.py | 69 ++++++++++++++++++++++++++++------------- 3 files changed, 63 insertions(+), 27 deletions(-) diff --git a/scripts/faceswap.py b/scripts/faceswap.py index a973914..5cf57b7 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -26,13 +26,13 @@ def get_models(): class FaceSwapScript(scripts.Script): def title(self): - return f"nsfw-roop" + return f"NSFW-Roop" def show(self, is_img2img): return scripts.AlwaysVisible def ui(self, is_img2img): - with gr.Accordion(f"nsfw-roop {version_flag}", open=False): + with gr.Accordion(f"NSFW-Roop {version_flag}", open=False): with gr.Column(): img = gr.inputs.Image(type="pil") enable = gr.Checkbox(False, placeholder="enable", label="Enable") @@ -56,6 +56,12 @@ class FaceSwapScript(scripts.Script): face_restorer_visibility = gr.Slider( 0, 1, 1, step=0.1, label="Restore visibility" ) + restore_first = gr.Checkbox( + True, + placeholder="Restore face, than Upscale", + label="1. Restore face -> 2. Upscale (-Uncheck- if you want vice versa)", + visible=is_img2img, + ) upscaler_name = gr.inputs.Dropdown( choices=[upscaler.name for upscaler in shared.sd_upscalers], label="Upscaler", @@ -100,6 +106,7 @@ class FaceSwapScript(scripts.Script): model, face_restorer_name, face_restorer_visibility, + restore_first, upscaler_name, upscaler_scale, upscaler_visibility, @@ -107,6 +114,7 @@ class FaceSwapScript(scripts.Script): swap_in_generated, ] + @property def upscaler(self) -> UpscalerData: for upscaler in shared.sd_upscalers: @@ -124,6 +132,7 @@ class FaceSwapScript(scripts.Script): @property def upscale_options(self) -> UpscaleOptions: return UpscaleOptions( + do_restore_first = self.restore_first, scale=self.upscaler_scale, upscaler=self.upscaler, face_restorer=self.face_restorer, @@ -141,6 +150,7 @@ class FaceSwapScript(scripts.Script): model, face_restorer_name, face_restorer_visibility, + restore_first, upscaler_name, upscaler_scale, upscaler_visibility, @@ -153,6 +163,7 @@ class FaceSwapScript(scripts.Script): self.upscaler_visibility = upscaler_visibility self.face_restorer_visibility = face_restorer_visibility self.enable = enable + self.restore_first = restore_first self.upscaler_name = upscaler_name self.swap_in_generated = swap_in_generated self.model = model @@ -169,7 +180,7 @@ class FaceSwapScript(scripts.Script): if self.enable: if self.source is not None: if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source: - logger.info(f"nsfw-roop enabled, face index %s", self.faces_index) + logger.info(f"NSFW-Roop is enabled, face index %s", self.faces_index) for i in range(len(p.init_images)): logger.info(f"Swap in source %s", i) diff --git a/scripts/roop_version.py b/scripts/roop_version.py index 0ac5860..bfba7c4 100644 --- a/scripts/roop_version.py +++ b/scripts/roop_version.py @@ -1,5 +1,5 @@ -version_flag = "v0.0.4" +version_flag = "v0.1.0" from scripts.roop_logging import logger -logger.info(f"nsfw-roop {version_flag}") +logger.info(f"NSFW-Roop {version_flag}") diff --git a/scripts/swapper.py b/scripts/swapper.py index 559d36d..a996caf 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -21,6 +21,7 @@ providers = onnxruntime.get_available_providers() @dataclass class UpscaleOptions: + do_restore_first: bool = True scale: int = 1 upscaler: UpscalerData = None upscale_visibility: float = 0.5 @@ -62,30 +63,54 @@ def getFaceSwapModel(model_path: str): def upscale_image(image: Image, upscale_options: UpscaleOptions): result_image = image - if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None": - original_image = result_image.copy() - logger.info( - "Upscale with %s scale = %s", - upscale_options.upscaler.name, - upscale_options.scale, - ) - result_image = upscale_options.upscaler.scaler.upscale( - image, upscale_options.scale, upscale_options.upscaler.data_path - ) - if upscale_options.scale == 1: + if upscale_options.do_restore_first: + if upscale_options.face_restorer is not None: + original_image = result_image.copy() + logger.info("1. Restore face with %s", upscale_options.face_restorer.name()) + numpy_image = np.array(result_image) + numpy_image = upscale_options.face_restorer.restore(numpy_image) + restored_image = Image.fromarray(numpy_image) result_image = Image.blend( - original_image, result_image, upscale_options.upscale_visibility + original_image, restored_image, upscale_options.restorer_visibility + ) + if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None": + original_image = result_image.copy() + logger.info( + "2. Upscale with %s scale = %s", + upscale_options.upscaler.name, + upscale_options.scale, + ) + result_image = upscale_options.upscaler.scaler.upscale( + original_image, upscale_options.scale, upscale_options.upscaler.data_path + ) + if upscale_options.scale == 1: + result_image = Image.blend( + original_image, result_image, upscale_options.upscale_visibility + ) + else: + if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None": + original_image = result_image.copy() + logger.info( + "1. Upscale with %s scale = %s", + upscale_options.upscaler.name, + upscale_options.scale, + ) + result_image = upscale_options.upscaler.scaler.upscale( + image, upscale_options.scale, upscale_options.upscaler.data_path + ) + if upscale_options.scale == 1: + result_image = Image.blend( + original_image, result_image, upscale_options.upscale_visibility + ) + if upscale_options.face_restorer is not None: + original_image = result_image.copy() + logger.info("2. Restore face with %s", upscale_options.face_restorer.name()) + numpy_image = np.array(result_image) + numpy_image = upscale_options.face_restorer.restore(numpy_image) + restored_image = Image.fromarray(numpy_image) + result_image = Image.blend( + original_image, restored_image, upscale_options.restorer_visibility ) - - if upscale_options.face_restorer is not None: - original_image = result_image.copy() - logger.info("Restore face with %s", upscale_options.face_restorer.name()) - numpy_image = np.array(result_image) - numpy_image = upscale_options.face_restorer.restore(numpy_image) - restored_image = Image.fromarray(numpy_image) - result_image = Image.blend( - original_image, restored_image, upscale_options.restorer_visibility - ) return result_image