update web_config

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"

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(
+    parser.add_argument("--enable-gfpgan", action="store_true", help=GFPGAN_HELP)
-        "--enable-gfpgan",
-        action="store_true",
-        help="Enable GFPGAN face restore",
-    )
     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

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"
 

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)
+        with gr.Tabs():
-        device = gr.Radio(AVAILABLE_DEVICES, label="Device", value=init_config.device)
+            with gr.Tab("Common"):
-        quality = gr.Slider(
+                with gr.Row():
-            value=95,
+                    host = gr.Textbox(init_config.host, label="Host")
-            label=f"Image Quality ({QUALITY_HELP})",
+                    port = gr.Number(init_config.port, label="Port", precision=0)
-            minimum=75,
-            maximum=100,
-            step=1,
-        )
 
-        with gr.Column():
+                model = gr.Radio(
-            gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}")
+                    AVAILABLE_MODELS, label="Model", value=init_config.model
-            no_gui_auto_close = gr.Checkbox(
+                )
-                init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}"
+                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}")
+                with gr.Column():
-            input = gr.Textbox(
+                    gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}")
-                init_config.input, label=f"Input file or directory. {INPUT_HELP}"
+                    no_gui_auto_close = gr.Checkbox(
-            )
+                        init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}"
-            output_dir = gr.Textbox(
+                    )
-                init_config.output_dir, label=f"Output directory. {OUTPUT_DIR_HELP}"
-            )
 
-        with gr.Column():
+                with gr.Column():
-            sd_controlnet = gr.Checkbox(
+                    model_dir = gr.Textbox(
-                init_config.sd_controlnet, label=f"{SD_CONTROLNET_HELP}"
+                        init_config.model_dir, label=f"{MODEL_DIR_HELP}"
-            )
+                    )
-            no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}")
+                    input = gr.Textbox(
-            cpu_offload = gr.Checkbox(
+                        init_config.input,
-                init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}"
+                        label=f"Input file or directory. {INPUT_HELP}",
-            )
+                    )
-            disable_nsfw = gr.Checkbox(
+                    output_dir = gr.Textbox(
-                init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}"
+                        init_config.output_dir,
-            )
+                        label=f"Output directory. {OUTPUT_DIR_HELP}",
-            sd_cpu_textencoder = gr.Checkbox(
+                    )
-                init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}"
+
-            )
+            with gr.Tab("Plugins"):
-            enable_xformers = gr.Checkbox(
+                enable_interactive_seg = gr.Checkbox(
-                init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}"
+                    init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP
-            )
+                )
-            local_files_only = gr.Checkbox(
+                enable_remove_bg = gr.Checkbox(
-                init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}"
+                    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,
         )