first commit

scripts/faceswaplab_postprocessing/i2i_pp.py

@@ -0,0 +1,76 @@
+from modules.face_restoration import FaceRestoration
+from modules.upscaler import UpscalerData
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from PIL import Image
+import numpy as np
+from modules import shared
+from scripts.faceswaplab_utils import imgutils
+from modules import shared, processing, codeformer_model
+from modules.processing import (StableDiffusionProcessingImg2Img)
+from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
+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  :
+        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) :
+            pass
+
+        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

@@ -0,0 +1,27 @@
+from modules.face_restoration import FaceRestoration
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from PIL import Image
+from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
+from scripts.faceswaplab_postprocessing.i2i_pp import img2img_diffusion
+from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_face
+
+
+def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
+    result_image = image
+    try :
+        if pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value :
+            result_image = img2img_diffusion(image, pp_options)
+        result_image = upscale_img(result_image, pp_options)
+
+        if pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value :
+            result_image = img2img_diffusion(image,pp_options)
+
+        result_image = restore_face(result_image, pp_options)
+              
+        if pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value :
+            result_image = img2img_diffusion(image,pp_options)
+
+    except Exception as e:
+        logger.error("Failed to upscale %s", e)
+
+    return result_image

scripts/faceswaplab_postprocessing/postprocessing_options.py

@@ -0,0 +1,43 @@
+from modules.face_restoration import FaceRestoration
+from modules.upscaler import UpscalerData
+from dataclasses import dataclass
+from modules import shared
+from enum import Enum
+
+class InpaintingWhen(Enum):
+    NEVER = "Never"
+    BEFORE_UPSCALING = "Before Upscaling/all"
+    BEFORE_RESTORE_FACE = "After Upscaling/Before Restore Face"
+    AFTER_ALL = "After All"
+
+@dataclass
+class PostProcessingOptions:
+    face_restorer_name: str = ""
+    restorer_visibility: float = 0.5
+    codeformer_weight: float = 1
+
+    upscaler_name: str = ""
+    scale: int = 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"
+    
+    @property
+    def upscaler(self) -> UpscalerData:
+        for upscaler in shared.sd_upscalers:
+            if upscaler.name == self.upscaler_name:
+                return upscaler
+        return None
+
+    @property
+    def face_restorer(self) -> FaceRestoration:
+        for face_restorer in shared.face_restorers:
+            if face_restorer.name() == self.face_restorer_name:
+                return face_restorer
+        return None

scripts/faceswaplab_postprocessing/upscaling.py

@@ -0,0 +1,42 @@
+
+from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from PIL import Image
+import numpy as np
+from modules import shared, processing, codeformer_model
+
+def upscale_img(image : Image.Image, pp_options :PostProcessingOptions) -> Image.Image :
+    if pp_options.upscaler is not None and pp_options.upscaler.name != "None":
+        original_image = image.copy()
+        logger.info(
+            "Upscale with %s scale = %s",
+            pp_options.upscaler.name,
+            pp_options.scale,
+        )
+        result_image = pp_options.upscaler.scaler.upscale(
+            image, pp_options.scale, pp_options.upscaler.data_path
+        )
+        if pp_options.scale == 1:
+            result_image = Image.blend(
+                original_image, result_image, pp_options.upscale_visibility
+            )
+        return result_image
+    return image
+
+def restore_face(image : Image.Image, pp_options : PostProcessingOptions) -> Image.Image :
+    
+    if pp_options.face_restorer is not None:
+        original_image = image.copy()
+        logger.info("Restore face with %s", pp_options.face_restorer.name())
+        numpy_image = np.array(image)
+        if pp_options.face_restorer_name == "CodeFormer" :
+            numpy_image = codeformer_model.codeformer.restore(numpy_image, w=pp_options.codeformer_weight)
+        else :
+            numpy_image = pp_options.face_restorer.restore(numpy_image)
+
+        restored_image = Image.fromarray(numpy_image)
+        result_image = Image.blend(
+            original_image, restored_image, pp_options.restorer_visibility
+        )
+        return result_image
+    return image