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(