stable-diffusion-webui/extensions/sd-webui-reactor/scripts/reactor_faceswap.py

import os, glob
import gradio as gr
from PIL import Image

from typing import List

import modules.scripts as scripts
from modules.upscaler import Upscaler, UpscalerData
from modules import scripts, shared, images, scripts_postprocessing, ui_components
from modules.processing import (
    Processed,
    StableDiffusionProcessing,
    StableDiffusionProcessingImg2Img,
)
from modules.face_restoration import FaceRestoration
from modules.images import save_image

from reactor_ui import (
    ui_main,
    ui_upscale,
    ui_tools,
    ui_settings,
    ui_detection,
)
from scripts.reactor_logger import logger
from scripts.reactor_swapper import (
    EnhancementOptions,
    DetectionOptions,
    swap_face,
    check_process_halt,
    reset_messaged,
)
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,
    set_Device,
    get_SDNEXT,
)
from scripts.reactor_globals import SWAPPER_MODELS_PATH #, DEVICE, DEVICE_LIST


class FaceSwapScript(scripts.Script):
    def title(self):
        return f"{app_title}"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        with ui_components.InputAccordion(False, label=f"{app_title}") as enable:
        # with gr.Accordion(f"{app_title}", open=False):

            # def on_files_upload_uncheck_so(selected: bool):
            #     global SAVE_ORIGINAL
            #     SAVE_ORIGINAL = selected
            #     return gr.Checkbox.update(value=False,visible=False)
            # def on_files_clear():
            #     clear_faces_list()
            #     return gr.Checkbox.update(value=SAVE_ORIGINAL,visible=True)

            # SD.Next fix
            if get_SDNEXT():
                enable = gr.Checkbox(False, label="Enable")
            
            # enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
            gr.Markdown(f"<sup>The Fast and Simple FaceSwap Extension - {version_flag}</sup>")

            # TAB MAIN
            msgs: dict = {
                "extra_multiple_source": "",
            }
            img, imgs, select_source, face_model, source_folder, save_original, mask_face, source_faces_index, gender_source, faces_index, gender_target, face_restorer_name, face_restorer_visibility, codeformer_weight, swap_in_source, swap_in_generated, random_image = ui_main.show(is_img2img=is_img2img, **msgs)
            
            # TAB DETECTION
            det_thresh, det_maxnum = ui_detection.show()
            
            # TAB UPSCALE
            restore_first, upscaler_name, upscaler_scale, upscaler_visibility, upscale_force = ui_upscale.show()

            # TAB TOOLS
            ui_tools.show()
            
            # TAB SETTINGS
            model, device, console_logging_level, source_hash_check, target_hash_check = ui_settings.show()
            
            gr.Markdown("<span style='display:block;text-align:right;padding:3px;font-size:0.666em;margin-bottom:-12px;'>by <a style='font-weight:normal' href='https://github.com/Gourieff' target='_blank'>Eugene Gourieff</a></span>")

        return [
            img,
            enable,
            source_faces_index,
            faces_index,
            model,
            face_restorer_name,
            face_restorer_visibility,
            restore_first,
            upscaler_name,
            upscaler_scale,
            upscaler_visibility,
            swap_in_source,
            swap_in_generated,
            console_logging_level,
            gender_source,
            gender_target,
            save_original,
            codeformer_weight,
            source_hash_check,
            target_hash_check,
            device,
            mask_face,
            select_source,
            face_model,
            source_folder,
            imgs,
            random_image,
            upscale_force,
            det_thresh,
            det_maxnum
        ]


    @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

    @property
    def enhancement_options(self) -> EnhancementOptions:
        return EnhancementOptions(
            do_restore_first=self.restore_first,
            scale=self.upscaler_scale,
            upscaler=self.upscaler,
            face_restorer=self.face_restorer,
            upscale_visibility=self.upscaler_visibility,
            restorer_visibility=self.face_restorer_visibility,
            codeformer_weight=self.codeformer_weight,
            upscale_force=self.upscale_force
        )
    
    @property
    def detection_options(self) -> DetectionOptions:
        return DetectionOptions(
            det_thresh=self.det_thresh,
            det_maxnum=self.det_maxnum
        )

    def process(
        self,
        p: StableDiffusionProcessing,
        img,
        enable,
        source_faces_index,
        faces_index,
        model,
        face_restorer_name,
        face_restorer_visibility,
        restore_first,
        upscaler_name,
        upscaler_scale,
        upscaler_visibility,
        swap_in_source,
        swap_in_generated,
        console_logging_level,
        gender_source,
        gender_target,
        save_original,
        codeformer_weight,
        source_hash_check,
        target_hash_check,
        device,
        mask_face,
        select_source,
        face_model,
        source_folder,
        imgs,
        random_image,
        upscale_force,
        det_thresh,
        det_maxnum
    ):
        self.enable = enable
        if self.enable:

            logger.debug("*** Start process")

            reset_messaged()
            if check_process_halt():
                return
            
            global SWAPPER_MODELS_PATH
            self.source = img
            self.face_restorer_name = face_restorer_name
            self.upscaler_scale = upscaler_scale
            self.upscaler_visibility = upscaler_visibility
            self.face_restorer_visibility = face_restorer_visibility
            self.restore_first = restore_first
            self.upscaler_name = upscaler_name  
            self.swap_in_source = swap_in_source
            self.swap_in_generated = swap_in_generated
            self.model = os.path.join(SWAPPER_MODELS_PATH,model)
            self.console_logging_level = console_logging_level
            self.gender_source = gender_source
            self.gender_target = gender_target
            self.save_original = save_original
            self.codeformer_weight = codeformer_weight
            self.source_hash_check = source_hash_check
            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
            self.source_folder = source_folder
            self.source_imgs = imgs
            self.random_image = random_image
            self.upscale_force = upscale_force
            self.det_thresh=det_thresh
            self.det_maxnum=det_maxnum
            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":
                self.gender_target = 0
            self.source_faces_index = [
                int(x) for x in source_faces_index.strip(",").split(",") if x.isnumeric()
            ]
            self.faces_index = [
                int(x) for x in faces_index.strip(",").split(",") if x.isnumeric()
            ]
            if len(self.source_faces_index) == 0:
                self.source_faces_index = [0]
            if len(self.faces_index) == 0:
                self.faces_index = [0]
            if self.save_original is None:
                self.save_original = False
            if self.source_hash_check is None:
                self.source_hash_check = True
            if self.target_hash_check is None:
                self.target_hash_check = False
            if self.mask_face is None:
                self.mask_face = False
            if self.random_image is None:
                self.random_image = False
            if self.upscale_force is None:
                self.upscale_force = False
            
            if shared.state.job_count > 0:
                # logger.debug(f"Job count: {shared.state.job_count}")
                self.face_restorer_visibility = shared.opts.data['restorer_visibility'] if 'restorer_visibility' in shared.opts.data.keys() else face_restorer_visibility
                self.codeformer_weight = shared.opts.data['codeformer_weight'] if 'codeformer_weight' in shared.opts.data.keys() else codeformer_weight
                self.mask_face = shared.opts.data['mask_face'] if 'mask_face' in shared.opts.data.keys() else mask_face
                self.face_model = shared.opts.data['face_model'] if 'face_model' in shared.opts.data.keys() else face_model

            logger.debug("*** Set Device")
            set_Device(self.device)

            if (self.save_original is None or not self.save_original) and (self.select_source == 2 or self.source_imgs is not None):
                p.do_not_save_samples = True
            
            if ((self.source is not None or self.source_imgs 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) or ((self.source_folder is not None and self.source_folder != "") and self.select_source == 2):
                logger.debug("*** Log patch")
                apply_logging_patch(console_logging_level)
                
                if isinstance(p, StableDiffusionProcessingImg2Img) and self.swap_in_source:

                    logger.debug("*** Check process")

                    logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)

                    for i in range(len(p.init_images)):
                        if len(p.init_images) > 1:
                            logger.status("Swap in %s", i)
                        result, output, swapped = swap_face(
                            self.source,
                            p.init_images[i],
                            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,
                            source_folder = None,
                            source_imgs = None,
                            random_image = False,
                            detection_options=self.detection_options,
                        )
                        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")
                        # if len(output) != 0:
                        #     with open(result_path, 'w', encoding="utf8") as f:
                        #         f.writelines(output)

                        if shared.state.interrupted or shared.state.skipped:
                            return
            
            else:
                logger.error("Please provide a source face")
                return

    def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args):
        if self.enable:

            logger.debug("*** Check postprocess - before IF")

            reset_messaged()
            if check_process_halt():
                return

            if self.save_original or ((self.select_source == 2 and self.source_folder is not None and self.source_folder != "") or (self.select_source == 0 and self.source_imgs is not None and self.source is None)):

                logger.debug("*** Check postprocess - after IF")

                postprocess_run: bool = True

                orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
                orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]

                result_images: List = processed.images
                # result_info: List = processed.infotexts

                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:
                        # self.source_folder = None
                        self.source_imgs = 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,
                            source_folder = self.source_folder,
                            source_imgs = self.source_imgs,
                            random_image = self.random_image,
                            detection_options=self.detection_options,
                        )

                        if self.select_source == 2 or (self.select_source == 0 and self.source_imgs is not None and self.source is None):
                            if len(result) > 0 and swapped > 0:
                                # result_images.extend(result)
                                if self.save_original:
                                    result_images.extend(result)
                                else:
                                    result_images = result
                                suffix = "-swapped"
                                for i,x in enumerate(result):
                                    try:
                                        img_path = save_image(result[i], 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 len(result) == 0:
                                logger.error("Cannot create a result image")

                        else:
                            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)
                    result_images.insert(0, grid)
                    try:
                        save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=info, short_filename=not shared.opts.grid_extended_filename, p=p, grid=True)
                    except:
                        logger.error("Cannot save a grid - please, check SD WebUI Settings (Saving and Paths)")
                
                processed.images = result_images
                # processed.infotexts = result_info
            
            elif self.select_source == 0 and self.source is not None and self.source_imgs is not None:

                logger.debug("*** Check postprocess - after ELIF")

                if self.result is not None:
                    orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
                    processed.images = [self.result]
                    try:
                        img_path = save_image(self.result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png", info=orig_infotexts[0], p=p, suffix="")
                    except:
                        logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
                else:
                    logger.error("Cannot create a result image")

    
    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 and ((self.select_source == 0 and self.source is not None) or self.select_source == 1):

            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:
                reset_messaged()
            if check_process_halt():
                return
            
            # 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,
                source_folder = None,
                source_imgs = None,
                random_image = False,
                detection_options=self.detection_options,
            )
            self.result = result
            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")


class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
    name = 'ReActor'
    order = 20000

    def ui(self):
        with ui_components.InputAccordion(False, label=f"{app_title}") as enable:
        # with gr.Accordion(f"{app_title}", open=False):
            
            # SD.Next fix
            if get_SDNEXT():
                enable = gr.Checkbox(False, label="Enable")

            # enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
            gr.Markdown(f"<span style='display:block;font-size:0.75em;margin-bottom:-24px;'>The Fast and Simple FaceSwap Extension - {version_flag}</span>")

            # TAB MAIN
            msgs: dict = {
                "extra_multiple_source": " | Сomparison grid as a result",
            }
            img, imgs, select_source, face_model, source_folder, save_original, mask_face, source_faces_index, gender_source, faces_index, gender_target, face_restorer_name, face_restorer_visibility, codeformer_weight, swap_in_source, swap_in_generated, random_image = ui_main.show(is_img2img=False, show_br=False, **msgs)
            
            # TAB DETECTION
            det_thresh, det_maxnum = ui_detection.show()
            
            # TAB UPSCALE
            restore_first, upscaler_name, upscaler_scale, upscaler_visibility, upscale_force = ui_upscale.show(show_br=False)
                        
            # TAB TOOLS
            ui_tools.show()
                        
            # TAB SETTINGS
            model, device, console_logging_level, source_hash_check, target_hash_check = ui_settings.show(hash_check_block=False)
                        
            gr.Markdown("<span style='display:block;text-align:right;padding-right:3px;font-size:0.666em;margin: -9px 0'>by <a style='font-weight:normal' href='https://github.com/Gourieff' target='_blank'>Eugene Gourieff</a></span>")

        args = {
            'img': img,
            'enable': enable,
            'source_faces_index': source_faces_index,
            'faces_index': faces_index,
            'model': model,
            'face_restorer_name': face_restorer_name,
            'face_restorer_visibility': face_restorer_visibility,
            'restore_first': restore_first,
            'upscaler_name': upscaler_name,
            'upscaler_scale': upscaler_scale,
            'upscaler_visibility': upscaler_visibility,
            'console_logging_level': console_logging_level,
            'gender_source': gender_source,
            'gender_target': gender_target,
            'codeformer_weight': codeformer_weight,
            'device': device,
            'mask_face': mask_face,
            'select_source': select_source,
            'face_model': face_model,
            'source_folder': source_folder,
            'imgs': imgs,
            'random_image': random_image,
            'upscale_force': upscale_force,
            'det_thresh': det_thresh,
            'det_maxnum': det_maxnum,
        }
        return args

    @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

    @property
    def enhancement_options(self) -> EnhancementOptions:
        return EnhancementOptions(
            do_restore_first=self.restore_first,
            scale=self.upscaler_scale,
            upscaler=self.upscaler,
            face_restorer=self.face_restorer,
            upscale_visibility=self.upscaler_visibility,
            restorer_visibility=self.face_restorer_visibility,
            codeformer_weight=self.codeformer_weight,
            upscale_force=self.upscale_force,
        )
    
    @property
    def detection_options(self) -> DetectionOptions:
        return DetectionOptions(
            det_thresh=self.det_thresh,
            det_maxnum=self.det_maxnum
        )

    def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
        if args['enable']:
            reset_messaged()
            if check_process_halt():
                return

            global SWAPPER_MODELS_PATH
            self.source = args['img']
            self.face_restorer_name = args['face_restorer_name']
            self.upscaler_scale = args['upscaler_scale']
            self.upscaler_visibility = args['upscaler_visibility']
            self.face_restorer_visibility = args['face_restorer_visibility']
            self.restore_first = args['restore_first']
            self.upscaler_name = args['upscaler_name']
            self.model = os.path.join(SWAPPER_MODELS_PATH, args['model'])
            self.console_logging_level = args['console_logging_level']
            self.gender_source = args['gender_source']
            self.gender_target = args['gender_target']
            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']
            self.source_folder = args['source_folder']
            self.source_imgs = args['imgs']
            self.random_image = args['random_image']
            self.upscale_force = args['upscale_force']
            self.det_thresh = args['det_thresh']
            self.det_maxnum = args['det_maxnum']
            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":
                self.gender_target = 0
            self.source_faces_index = [
                int(x) for x in args['source_faces_index'].strip(",").split(",") if x.isnumeric()
            ]
            self.faces_index = [
                int(x) for x in args['faces_index'].strip(",").split(",") if x.isnumeric()
            ]
            if len(self.source_faces_index) == 0:
                self.source_faces_index = [0]
            if len(self.faces_index) == 0:
                self.faces_index = [0]
            if self.mask_face is None:
                self.mask_face = False
            if self.random_image is None:
                self.random_image = False
            if self.upscale_force is None:
                self.upscale_force = False

            current_job_number = shared.state.job_no + 1
            job_count = shared.state.job_count
            if current_job_number == job_count:
                reset_messaged()

            set_Device(self.device)

            logger.debug("We're here: process() 1")
            
            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) or ((self.source_folder is not None and self.source_folder != "") and self.select_source == 2) or ((self.source_imgs is not None and self.source is None) and self.select_source == 0):

                logger.debug("We're here: process() 2")

                if self.source is not None and self.select_source == 0:
                    self.source_imgs = None

                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)
                # if self.select_source != 2:
                image: Image.Image = pp.image

                # Extract alpha channel
                logger.debug(f"image = {image}")
                if image.mode == 'RGBA':
                    _, _, _, alpha = image.split()
                else:
                    alpha = None
                logger.debug(f"alpha = {alpha}")

                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=True,
                    target_hash_check=True,
                    device=self.device,
                    mask_face=self.mask_face,
                    select_source=self.select_source,
                    face_model=self.face_model,
                    source_folder=self.source_folder,
                    source_imgs=self.source_imgs,
                    random_image=self.random_image,
                    detection_options=self.detection_options,
                )
                if self.select_source == 2 or (self.select_source == 0 and self.source_imgs is not None and self.source is None):
                    if len(result) > 0 and swapped > 0:
                        image = result[0]
                        if len(result) > 1:
                            grid = make_grid(result)
                            result.insert(0, grid)
                            image = grid
                        pp.info["ReActor"] = True
                        pp.image = image
                        logger.status("---Done!---")
                    else:
                        logger.error("Cannot create a result image")
                else:
                    try:
                        pp.info["ReActor"] = True

                        if alpha is not None:
                            logger.debug(f"result = {result}")
                            result = result.convert("RGBA")
                            result.putalpha(alpha)
                            logger.debug(f"result_alpha = {result}")

                        pp.image = result
                        logger.status("---Done!---")
                    except Exception:
                        logger.error("Cannot create a result image")
            else:
                logger.error("Please provide a source face")