fix parse_args

This commit is contained in:
Qing 2023-04-16 10:35:51 +08:00
parent 0d89c37ef1
commit 0baea45210
3 changed files with 5 additions and 3 deletions

View File

@ -106,6 +106,7 @@ class RealESRGANModelName(str, Enum):
RealESRGANModelNameList = [e.value for e in RealESRGANModelName] RealESRGANModelNameList = [e.value for e in RealESRGANModelName]
INTERACTIVE_SEG_HELP = "Enable interactive segmentation using Segment Anything." 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_MODELS = ["vit_b", "vit_l", "vit_h"]
AVAILABLE_INTERACTIVE_SEG_DEVICES = ["cuda", "cpu", "mps"] AVAILABLE_INTERACTIVE_SEG_DEVICES = ["cuda", "cpu", "mps"]
REMOVE_BG_HELP = "Enable remove background. Always run on CPU" REMOVE_BG_HELP = "Enable remove background. Always run on CPU"

View File

@ -85,12 +85,13 @@ def parse_args():
parser.add_argument( parser.add_argument(
"--interactive-seg-model", "--interactive-seg-model",
default="vit_l", default="vit_l",
help=AVAILABLE_INTERACTIVE_SEG_MODELS, choices=AVAILABLE_INTERACTIVE_SEG_MODELS,
help=INTERACTIVE_SEG_MODEL_HELP
) )
parser.add_argument( parser.add_argument(
"--interactive-seg-device", "--interactive-seg-device",
default="cpu", default="cpu",
help=AVAILABLE_INTERACTIVE_SEG_DEVICES, choices=AVAILABLE_INTERACTIVE_SEG_DEVICES,
) )
parser.add_argument( parser.add_argument(
"--enable-remove-bg", "--enable-remove-bg",

View File

@ -123,7 +123,7 @@ def main(config_file: str):
) )
interactive_seg_model = gr.Radio( interactive_seg_model = gr.Radio(
AVAILABLE_INTERACTIVE_SEG_MODELS, 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, value=init_config.interactive_seg_model,
) )
interactive_seg_device = gr.Radio( interactive_seg_device = gr.Radio(