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,15 +109,19 @@ 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")
+        with gr.Tabs():
-            #     exit_btn.click(close_server)
+            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)
 
-        model = gr.Radio(AVAILABLE_MODELS, label="Model", value=init_config.model)
+                model = gr.Radio(
-        device = gr.Radio(AVAILABLE_DEVICES, label="Device", value=init_config.device)
+                    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})",
@@ -115,15 +136,52 @@ def main(config_file: str):
                         init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}"
                     )
 
-            model_dir = gr.Textbox(init_config.model_dir, label=f"{MODEL_DIR_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}"
+                        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}"
+                        init_config.output_dir,
+                        label=f"Output directory. {OUTPUT_DIR_HELP}",
                     )
 
-        with gr.Column():
+            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}"
                 )
@@ -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,
         )