diff --git a/scripts/faceswap.py b/scripts/faceswap.py index d6f44ee..b189aac 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -1,5 +1,6 @@ import gradio as gr import modules.scripts as scripts +from modules.upscaler import Upscaler, UpscalerData from modules import scripts, shared, images, scripts_postprocessing from modules.processing import ( StableDiffusionProcessing, @@ -11,7 +12,7 @@ import glob from modules.face_restoration import FaceRestoration from scripts.roop_logging import logger -from scripts.swapper import swap_face, ImageResult +from scripts.swapper import UpscaleOptions, swap_face, ImageResult from scripts.cimage import check_batch from scripts.roop_version import version_flag import os @@ -51,7 +52,15 @@ class FaceSwapScript(scripts.Script): face_restorer_visibility = gr.Slider( 0, 1, 1, step=0.1, label="Restore visibility" ) - + upscaler_name = gr.inputs.Dropdown( + choices=[upscaler.name for upscaler in shared.sd_upscalers], + label="Upscaler", + ) + upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Upscaler scale") + upscaler_visibility = gr.Slider( + 0, 1, 1, step=0.1, label="Upscaler visibility (if scale = 1)" + ) + models = get_models() if len(models) == 0: logger.warning( @@ -86,10 +95,20 @@ class FaceSwapScript(scripts.Script): model, face_restorer_name, face_restorer_visibility, + upscaler_name, + upscaler_scale, + upscaler_visibility, swap_in_source, swap_in_generated, ] + @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: @@ -97,6 +116,16 @@ class FaceSwapScript(scripts.Script): return face_restorer return None + @property + def upscale_options(self) -> UpscaleOptions: + return UpscaleOptions( + scale=self.upscaler_scale, + upscaler=self.upscaler, + face_restorer=self.face_restorer, + upscale_visibility=self.upscaler_visibility, + restorer_visibility=self.face_restorer_visibility, + ) + def process( self, p: StableDiffusionProcessing, @@ -106,13 +135,19 @@ class FaceSwapScript(scripts.Script): model, face_restorer_name, face_restorer_visibility, + upscaler_name, + upscaler_scale, + upscaler_visibility, swap_in_source, swap_in_generated, ): 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.enable = enable + self.upscaler_name = upscaler_name self.swap_in_generated = swap_in_generated self.model = model self.faces_index = { @@ -132,6 +167,7 @@ class FaceSwapScript(scripts.Script): p.init_images[i], faces_index=self.faces_index, model=self.model, + upscale_options=self.upscale_options, ) p.init_images[i] = result.image() else: @@ -151,6 +187,7 @@ class FaceSwapScript(scripts.Script): image, faces_index=self.faces_index, model=self.model, + upscale_options=self.upscale_options, ) pp = scripts_postprocessing.PostprocessedImage(result.image()) pp.info = {} diff --git a/scripts/swapper.py b/scripts/swapper.py index cb5cc18..aa5b633 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -14,11 +14,21 @@ import onnxruntime from scripts.cimage import convert_to_sd from modules.face_restoration import FaceRestoration, restore_faces +from modules.upscaler import Upscaler, UpscalerData from scripts.roop_logging import logger providers = onnxruntime.get_available_providers() +@dataclass +class UpscaleOptions: + scale: int = 1 + upscaler: UpscalerData = None + upscale_visibility: float = 0.5 + face_restorer: FaceRestoration = None + restorer_visibility: float = 0.5 + + def save_image(img: Image, filename: str): convert_to_sd(img).save(filename) @@ -67,6 +77,36 @@ def getFaceSwapModel(model_path: str): return FS_MODEL +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: + 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("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 + + def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)): face_analyser = copy.deepcopy(getAnalysisModel()) face_analyser.prepare(ctx_id=0, det_size=det_size) @@ -98,7 +138,9 @@ def swap_face( target_img: Image.Image, model: Union[str, None] = None, faces_index: Set[int] = {0}, + upscale_options: Union[UpscaleOptions, None] = None, ) -> ImageResult: + result_image = target_img fn = tempfile.NamedTemporaryFile(delete=False, suffix=".png") if model is not None: source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) @@ -117,8 +159,11 @@ def swap_face( logger.info(f"No target face found for {face_num}") result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) + if upscale_options is not None: + result_image = upscale_image(result_image, upscale_options) - save_image(result_image, fn.name) + # save_image(result_image, fn.name) else: logger.info("No source face found") + save_image(result_image, fn.name) return ImageResult(path=fn.name)