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]
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"

View File

@ -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",

View File

@ -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(