From e734859128c302d709cf7e75cffdff0623f90655 Mon Sep 17 00:00:00 2001 From: Qing Date: Tue, 28 Mar 2023 21:55:25 +0800 Subject: [PATCH] update web_config --- lama_cleaner/const.py | 20 ++++ lama_cleaner/parse_args.py | 28 +++-- lama_cleaner/plugins/realesrgan.py | 10 +- lama_cleaner/web_config.py | 160 ++++++++++++++++++++--------- 4 files changed, 151 insertions(+), 67 deletions(-) diff --git a/lama_cleaner/const.py b/lama_cleaner/const.py index 1aeea29..01df33f 100644 --- a/lama_cleaner/const.py +++ b/lama_cleaner/const.py @@ -1,4 +1,5 @@ import os +from enum import Enum MPS_SUPPORT_MODELS = [ "instruct_pix2pix", @@ -86,3 +87,22 @@ Prevent backend auto close after the GUI window closed. QUALITY_HELP = """ Quality of image encoding, 0-100. Default is 95, higher quality will generate larger file size. """ + + +class RealESRGANModelName(str, Enum): + realesr_general_x4v3 = "realesr-general-x4v3" + RealESRGAN_x4plus = "RealESRGAN_x4plus" + RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B" + + +RealESRGANModelNameList = [e.value for e in RealESRGANModelName] + +INTERACTIVE_SEG_HELP = "Enable interactive segmentation. Always run on CPU" +REMOVE_BG_HELP = "Enable remove background. Always run on CPU" +REALESRGAN_HELP = "Enable realesrgan super resolution" +REALESRGAN_AVAILABLE_DEVICES = ["cpu", "cuda", "mps"] +GFPGAN_HELP = ( + "Enable GFPGAN face restore. To enhance background, use with --enable-realesrgan" +) +GFPGAN_AVAILABLE_DEVICES = ["cpu", "cuda", "mps"] +GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image" diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index c710b51..6d3de56 100644 --- a/lama_cleaner/parse_args.py +++ b/lama_cleaner/parse_args.py @@ -6,7 +6,6 @@ from pathlib import Path from loguru import logger from lama_cleaner.const import * -from lama_cleaner.plugins.realesrgan import RealESRGANModelName, RealESRGANModelNameList from lama_cleaner.runtime import dump_environment_info @@ -80,20 +79,23 @@ def parse_args(): parser.add_argument( "--enable-interactive-seg", action="store_true", - help="Enable interactive segmentation. Always run on CPU", + help=INTERACTIVE_SEG_HELP, ) parser.add_argument( "--enable-remove-bg", action="store_true", - help="Enable remove background. Always run on CPU", + help=REMOVE_BG_HELP, ) parser.add_argument( "--enable-realesrgan", action="store_true", - help="Enable realesrgan super resolution", + help=REALESRGAN_HELP, ) parser.add_argument( - "--realesrgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"] + "--realesrgan-device", + default="cpu", + type=str, + choices=REALESRGAN_AVAILABLE_DEVICES, ) parser.add_argument( "--realesrgan-model", @@ -101,18 +103,14 @@ def parse_args(): type=str, choices=RealESRGANModelNameList, ) - parser.add_argument( - "--enable-gfpgan", - action="store_true", - help="Enable GFPGAN face restore", - ) + parser.add_argument("--enable-gfpgan", action="store_true", help=GFPGAN_HELP) parser.add_argument( "--gfpgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"] ) parser.add_argument( "--enable-gif", action="store_true", - help="Enable GIF plugin", + help=GIF_HELP, ) ######### @@ -200,4 +198,12 @@ def parse_args(): if not output_dir.is_dir(): parser.error(f"invalid --output-dir: {output_dir} is not a directory") + if args.enable_gfpgan: + if args.enable_realesrgan: + logger.info("Use realesrgan as GFPGAN background upscaler") + else: + logger.info( + f"GFPGAN no background upscaler, use --enable-realesrgan to enable it" + ) + return args diff --git a/lama_cleaner/plugins/realesrgan.py b/lama_cleaner/plugins/realesrgan.py index a408b51..bba2d07 100644 --- a/lama_cleaner/plugins/realesrgan.py +++ b/lama_cleaner/plugins/realesrgan.py @@ -3,19 +3,11 @@ from enum import Enum import cv2 from loguru import logger +from lama_cleaner.const import RealESRGANModelName from lama_cleaner.helper import download_model from lama_cleaner.plugins.base_plugin import BasePlugin -class RealESRGANModelName(str, Enum): - realesr_general_x4v3 = "realesr-general-x4v3" - RealESRGAN_x4plus = "RealESRGAN_x4plus" - RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B" - - -RealESRGANModelNameList = [e.value for e in RealESRGANModelName] - - class RealESRGANUpscaler(BasePlugin): name = "RealESRGAN" diff --git a/lama_cleaner/web_config.py b/lama_cleaner/web_config.py index 98e33cd..a544ce8 100644 --- a/lama_cleaner/web_config.py +++ b/lama_cleaner/web_config.py @@ -28,6 +28,15 @@ class Config(BaseModel): model_dir: str = DEFAULT_MODEL_DIR input: str = None output_dir: str = None + # plugins + enable_interactive_seg: bool = False + enable_remove_bg: bool = False + enable_realesrgan: bool = False + realesrgan_device: str = "cpu" + realesrgan_model: str = RealESRGANModelName.realesr_general_x4v3.value + enable_gfpgan: bool = False + gfpgan_device: str = "cpu" + enable_gif: bool = False def load_config(installer_config: str): @@ -56,6 +65,14 @@ def save_config( input, output_dir, quality, + enable_interactive_seg, + enable_remove_bg, + enable_realesrgan, + realesrgan_device, + realesrgan_model, + enable_gfpgan, + gfpgan_device, + enable_gif, ): config = Config(**locals()) print(config) @@ -92,57 +109,98 @@ def main(config_file: str): with gr.Column(scale=1): save_btn = gr.Button(value="Save configurations") message = gr.HTML() - # with gr.Column(scale=0, min_width=100): - # exit_btn = gr.Button(value="Close") - # exit_btn.click(close_server) - with gr.Row(): - host = gr.Textbox(init_config.host, label="Host") - port = gr.Number(init_config.port, label="Port", precision=0) - model = gr.Radio(AVAILABLE_MODELS, label="Model", value=init_config.model) - device = gr.Radio(AVAILABLE_DEVICES, label="Device", value=init_config.device) - quality = gr.Slider( - value=95, - label=f"Image Quality ({QUALITY_HELP})", - minimum=75, - maximum=100, - step=1, - ) + with gr.Tabs(): + with gr.Tab("Common"): + with gr.Row(): + host = gr.Textbox(init_config.host, label="Host") + port = gr.Number(init_config.port, label="Port", precision=0) - with gr.Column(): - gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}") - no_gui_auto_close = gr.Checkbox( - init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}" - ) + model = gr.Radio( + AVAILABLE_MODELS, label="Model", value=init_config.model + ) + device = gr.Radio( + AVAILABLE_DEVICES, label="Device", value=init_config.device + ) + quality = gr.Slider( + value=95, + label=f"Image Quality ({QUALITY_HELP})", + minimum=75, + maximum=100, + step=1, + ) - model_dir = gr.Textbox(init_config.model_dir, label=f"{MODEL_DIR_HELP}") - input = gr.Textbox( - init_config.input, label=f"Input file or directory. {INPUT_HELP}" - ) - output_dir = gr.Textbox( - init_config.output_dir, label=f"Output directory. {OUTPUT_DIR_HELP}" - ) + with gr.Column(): + gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}") + no_gui_auto_close = gr.Checkbox( + init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}" + ) - with gr.Column(): - sd_controlnet = gr.Checkbox( - init_config.sd_controlnet, label=f"{SD_CONTROLNET_HELP}" - ) - no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}") - cpu_offload = gr.Checkbox( - init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}" - ) - disable_nsfw = gr.Checkbox( - init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}" - ) - sd_cpu_textencoder = gr.Checkbox( - init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}" - ) - enable_xformers = gr.Checkbox( - init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}" - ) - local_files_only = gr.Checkbox( - init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}" - ) + with gr.Column(): + model_dir = gr.Textbox( + init_config.model_dir, label=f"{MODEL_DIR_HELP}" + ) + input = gr.Textbox( + init_config.input, + label=f"Input file or directory. {INPUT_HELP}", + ) + output_dir = gr.Textbox( + init_config.output_dir, + label=f"Output directory. {OUTPUT_DIR_HELP}", + ) + + with gr.Tab("Plugins"): + enable_interactive_seg = gr.Checkbox( + init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP + ) + enable_remove_bg = gr.Checkbox( + init_config.enable_remove_bg, label=REMOVE_BG_HELP + ) + with gr.Row(): + enable_realesrgan = gr.Checkbox( + init_config.enable_realesrgan, label=REALESRGAN_HELP + ) + realesrgan_device = gr.Radio( + REALESRGAN_AVAILABLE_DEVICES, + label="RealESRGAN Device", + value=init_config.realesrgan_device, + ) + realesrgan_model = gr.Radio( + RealESRGANModelNameList, + label="RealESRGAN model", + value=init_config.realesrgan_model, + ) + with gr.Row(): + enable_gfpgan = gr.Checkbox( + init_config.enable_gfpgan, label=GFPGAN_HELP + ) + gfpgan_device = gr.Radio( + GFPGAN_AVAILABLE_DEVICES, + label="GFPGAN Device", + value=init_config.gfpgan_device, + ) + enable_gif = gr.Checkbox(init_config.enable_gif, label=GIF_HELP) + + with gr.Tab("Diffusion Model"): + sd_controlnet = gr.Checkbox( + init_config.sd_controlnet, label=f"{SD_CONTROLNET_HELP}" + ) + no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}") + cpu_offload = gr.Checkbox( + init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}" + ) + disable_nsfw = gr.Checkbox( + init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}" + ) + sd_cpu_textencoder = gr.Checkbox( + init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}" + ) + enable_xformers = gr.Checkbox( + init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}" + ) + local_files_only = gr.Checkbox( + init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}" + ) save_btn.click( save_config, @@ -164,6 +222,14 @@ def main(config_file: str): input, output_dir, quality, + enable_interactive_seg, + enable_remove_bg, + enable_realesrgan, + realesrgan_device, + realesrgan_model, + enable_gfpgan, + gfpgan_device, + enable_gif, ], message, )