huge changes, inpainting in faces unit, change faces processing, change api, refactor, requires further testing

scripts/faceswaplab_postprocessing/i2i_pp.py

@@ -1,83 +0,0 @@
-from scripts.faceswaplab_utils.faceswaplab_logging import logger
-from PIL import Image
-from modules import shared
-from scripts.faceswaplab_utils import imgutils
-from modules import shared, processing
-from modules.processing import StableDiffusionProcessingImg2Img
-from scripts.faceswaplab_postprocessing.postprocessing_options import (
-    PostProcessingOptions,
-)
-from modules import sd_models
-
-from scripts.faceswaplab_swapping import swapper
-
-
-def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image:
-    if pp.inpainting_denoising_strengh == 0:
-        logger.info("Discard inpainting denoising strength is 0")
-        return img
-
-    try:
-        logger.info(
-            f"""Inpainting face
-Sampler : {pp.inpainting_sampler}
-inpainting_denoising_strength : {pp.inpainting_denoising_strengh}
-inpainting_steps : {pp.inpainting_steps}
-"""
-        )
-        if not isinstance(pp.inpainting_sampler, str):
-            pp.inpainting_sampler = "Euler"
-
-        logger.info("send faces to image to image")
-        img = img.copy()
-        faces = swapper.get_faces(imgutils.pil_to_cv2(img))
-        if faces:
-            for face in faces:
-                bbox = face.bbox.astype(int)
-                mask = imgutils.create_mask(img, bbox)
-                prompt = pp.inpainting_prompt.replace(
-                    "[gender]", "man" if face["gender"] == 1 else "woman"
-                )
-                negative_prompt = pp.inpainting_negative_prompt.replace(
-                    "[gender]", "man" if face["gender"] == 1 else "woman"
-                )
-                logger.info("Denoising prompt : %s", prompt)
-                logger.info("Denoising strenght : %s", pp.inpainting_denoising_strengh)
-
-                i2i_kwargs = {
-                    "sampler_name": pp.inpainting_sampler,
-                    "do_not_save_samples": True,
-                    "steps": pp.inpainting_steps,
-                    "width": img.width,
-                    "inpainting_fill": 1,
-                    "inpaint_full_res": True,
-                    "height": img.height,
-                    "mask": mask,
-                    "prompt": prompt,
-                    "negative_prompt": negative_prompt,
-                    "denoising_strength": pp.inpainting_denoising_strengh,
-                }
-                current_model_checkpoint = shared.opts.sd_model_checkpoint
-                if pp.inpainting_model and pp.inpainting_model != "Current":
-                    # Change checkpoint
-                    shared.opts.sd_model_checkpoint = pp.inpainting_model
-                    sd_models.select_checkpoint
-                    sd_models.load_model()
-                i2i_p = StableDiffusionProcessingImg2Img([img], **i2i_kwargs)
-                i2i_processed = processing.process_images(i2i_p)
-                if pp.inpainting_model and pp.inpainting_model != "Current":
-                    # Restore checkpoint
-                    shared.opts.sd_model_checkpoint = current_model_checkpoint
-                    sd_models.select_checkpoint
-                    sd_models.load_model()
-
-                images = i2i_processed.images
-                if len(images) > 0:
-                    img = images[0]
-        return img
-    except Exception as e:
-        logger.error("Failed to apply img2img to face : %s", e)
-        import traceback
-
-        traceback.print_exc()
-        raise e

scripts/faceswaplab_postprocessing/postprocessing.py

@@ -4,8 +4,9 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
     InpaintingWhen,
 )
-from scripts.faceswaplab_postprocessing.i2i_pp import img2img_diffusion
+from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
 from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_face
+import traceback
 
 
 def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
@@ -19,7 +20,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING
         ):
             logger.debug("Inpaint before upscale")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                img=result_image, options=pp_options.inpainting_options
+            )
         result_image = upscale_img(result_image, pp_options)
 
         if (
@@ -27,7 +30,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE
         ):
             logger.debug("Inpaint before restore")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                result_image, pp_options.inpainting_options
+            )
 
         result_image = restore_face(result_image, pp_options)
 
@@ -36,9 +41,11 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.AFTER_ALL
         ):
             logger.debug("Inpaint after all")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                result_image, pp_options.inpainting_options
+            )
 
     except Exception as e:
-        logger.error("Failed to upscale %s", e)
-
+        logger.error("Failed to post-process %s", e)
+        traceback.print_exc()
     return result_image

scripts/faceswaplab_postprocessing/postprocessing_options.py

@@ -3,6 +3,8 @@ from modules.upscaler import UpscalerData
 from dataclasses import dataclass
 from modules import shared
 from enum import Enum
+from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
+from client_api import api_utils
 
 
 class InpaintingWhen(Enum):
@@ -22,13 +24,10 @@ class PostProcessingOptions:
     scale: float = 1
     upscale_visibility: float = 0.5
 
-    inpainting_denoising_strengh: float = 0
-    inpainting_prompt: str = ""
-    inpainting_negative_prompt: str = ""
-    inpainting_steps: int = 20
-    inpainting_sampler: str = "Euler"
     inpainting_when: InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING
-    inpainting_model: str = "Current"
+
+    # (Don't use optional for this or gradio parsing will fail) :
+    inpainting_options: InpaintingOptions = None
 
     @property
     def upscaler(self) -> UpscalerData:
@@ -43,3 +42,28 @@ class PostProcessingOptions:
             if face_restorer.name() == self.face_restorer_name:
                 return face_restorer
         return None
+
+    @staticmethod
+    def from_api_dto(
+        options: api_utils.PostProcessingOptions,
+    ) -> "PostProcessingOptions":
+        """
+        Converts a PostProcessingOptions object from an API DTO (Data Transfer Object).
+
+        :param options: An object of api_utils.PostProcessingOptions representing the
+                        post-processing options as received from the API.
+        :return: A PostProcessingOptions instance containing the translated values
+                from the API DTO.
+        """
+        return PostProcessingOptions(
+            face_restorer_name=options.face_restorer_name,
+            restorer_visibility=options.restorer_visibility,
+            codeformer_weight=options.codeformer_weight,
+            upscaler_name=options.upscaler_name,
+            scale=options.scale,
+            upscale_visibility=options.upscaler_visibility,
+            inpainting_when=InpaintingWhen(options.inpainting_when.value),
+            inpainting_options=InpaintingOptions.from_api_dto(
+                options.inpainting_options
+            ),
+        )

scripts/faceswaplab_postprocessing/upscaling.py

@@ -5,11 +5,12 @@ from scripts.faceswaplab_utils.faceswaplab_logging import logger
 from PIL import Image
 import numpy as np
 from modules import codeformer_model
+from scripts.faceswaplab_utils.typing import *
 
 
-def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
+def upscale_img(image: PILImage, pp_options: PostProcessingOptions) -> PILImage:
     if pp_options.upscaler is not None and pp_options.upscaler.name != "None":
-        original_image = image.copy()
+        original_image: PILImage = image.copy()
         logger.info(
             "Upscale with %s scale = %s",
             pp_options.upscaler.name,
@@ -18,7 +19,12 @@ def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image.
         result_image = pp_options.upscaler.scaler.upscale(
             image, pp_options.scale, pp_options.upscaler.data_path
         )
-        if pp_options.scale == 1:
+
+        # FIXME : Could be better (managing images whose dimensions are not multiples of 16)
+        if pp_options.scale == 1 and original_image.size == result_image.size:
+            logger.debug(
+                "Sizes orig=%s, result=%s", original_image.size, result_image.size
+            )
             result_image = Image.blend(
                 original_image, result_image, pp_options.upscale_visibility
             )