From 0baea4521005088bda6b8d320b53610864825a9b Mon Sep 17 00:00:00 2001 From: Qing Date: Sun, 16 Apr 2023 10:35:51 +0800 Subject: [PATCH] fix parse_args --- lama_cleaner/const.py | 1 + lama_cleaner/parse_args.py | 5 +++-- lama_cleaner/web_config.py | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/lama_cleaner/const.py b/lama_cleaner/const.py index b831ae4..5caf7b0 100644 --- a/lama_cleaner/const.py +++ b/lama_cleaner/const.py @@ -106,6 +106,7 @@ class RealESRGANModelName(str, Enum): RealESRGANModelNameList = [e.value for e in RealESRGANModelName] INTERACTIVE_SEG_HELP = "Enable interactive segmentation using Segment Anything." +INTERACTIVE_SEG_MODEL_HELP = "Model size: vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed." AVAILABLE_INTERACTIVE_SEG_MODELS = ["vit_b", "vit_l", "vit_h"] AVAILABLE_INTERACTIVE_SEG_DEVICES = ["cuda", "cpu", "mps"] REMOVE_BG_HELP = "Enable remove background. Always run on CPU" diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index 9b72727..5af1cb2 100644 --- a/lama_cleaner/parse_args.py +++ b/lama_cleaner/parse_args.py @@ -85,12 +85,13 @@ def parse_args(): parser.add_argument( "--interactive-seg-model", default="vit_l", - help=AVAILABLE_INTERACTIVE_SEG_MODELS, + choices=AVAILABLE_INTERACTIVE_SEG_MODELS, + help=INTERACTIVE_SEG_MODEL_HELP ) parser.add_argument( "--interactive-seg-device", default="cpu", - help=AVAILABLE_INTERACTIVE_SEG_DEVICES, + choices=AVAILABLE_INTERACTIVE_SEG_DEVICES, ) parser.add_argument( "--enable-remove-bg", diff --git a/lama_cleaner/web_config.py b/lama_cleaner/web_config.py index e2f85e9..39cba72 100644 --- a/lama_cleaner/web_config.py +++ b/lama_cleaner/web_config.py @@ -123,7 +123,7 @@ def main(config_file: str): ) interactive_seg_model = gr.Radio( AVAILABLE_INTERACTIVE_SEG_MODELS, - label="Segment Anything models. Model size: vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed.", + label=f"Segment Anything models. {INTERACTIVE_SEG_MODEL_HELP}", value=init_config.interactive_seg_model, ) interactive_seg_device = gr.Radio(