UPDATE: Postprocess logic improvement

+ Version UP

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)

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}")

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":
+    if upscale_options.do_restore_first:
-        original_image = result_image.copy()
+        if upscale_options.face_restorer is not None:
-        logger.info(
+            original_image = result_image.copy()
-            "Upscale with %s scale = %s",
+            logger.info("1. Restore face with %s", upscale_options.face_restorer.name())
-            upscale_options.upscaler.name,
+            numpy_image = np.array(result_image)
-            upscale_options.scale,
+            numpy_image = upscale_options.face_restorer.restore(numpy_image)
-        )
+            restored_image = Image.fromarray(numpy_image)
-        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
+                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