IOPaint/lama_cleaner/parse_args.py

257 lines
8.5 KiB
Python
Raw Normal View History

2022-04-18 09:01:10 +02:00
import os
import imghdr
import argparse
2022-12-31 14:07:08 +01:00
from pathlib import Path
from loguru import logger
2022-04-18 09:01:10 +02:00
2023-03-19 15:40:23 +01:00
from lama_cleaner.const import *
2023-01-20 09:52:38 +01:00
from lama_cleaner.runtime import dump_environment_info
2022-04-18 09:01:10 +02:00
def parse_args():
2023-02-19 14:16:30 +01:00
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
2022-04-18 09:01:10 +02:00
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", default=8080, type=int)
2023-01-20 09:52:38 +01:00
2023-02-06 15:00:47 +01:00
parser.add_argument(
"--config-installer",
action="store_true",
help="Open config web page, mainly for windows installer",
)
parser.add_argument(
"--load-installer-config",
action="store_true",
help="Load all cmd args from installer config file",
)
parser.add_argument(
"--installer-config", default=None, help="Config file for windows installer"
)
2023-01-20 09:52:38 +01:00
2023-02-06 15:00:47 +01:00
parser.add_argument("--model", default=DEFAULT_MODEL, choices=AVAILABLE_MODELS)
2023-01-20 09:52:38 +01:00
parser.add_argument("--no-half", action="store_true", help=NO_HALF_HELP)
parser.add_argument("--cpu-offload", action="store_true", help=CPU_OFFLOAD_HELP)
parser.add_argument("--disable-nsfw", action="store_true", help=DISABLE_NSFW_HELP)
2023-02-06 15:00:47 +01:00
parser.add_argument(
"--sd-cpu-textencoder", action="store_true", help=SD_CPU_TEXTENCODER_HELP
)
parser.add_argument("--sd-controlnet", action="store_true", help=SD_CONTROLNET_HELP)
2023-05-02 08:09:40 +02:00
parser.add_argument(
"--sd-controlnet-method",
2023-05-13 07:45:27 +02:00
default=DEFAULT_CONTROLNET_METHOD,
2023-05-11 15:51:58 +02:00
choices=SD_CONTROLNET_CHOICES,
2023-05-02 08:09:40 +02:00
)
2023-03-29 16:05:34 +02:00
parser.add_argument("--sd-local-model-path", default=None, help=SD_LOCAL_MODEL_HELP)
2023-02-06 15:00:47 +01:00
parser.add_argument(
"--local-files-only", action="store_true", help=LOCAL_FILES_ONLY_HELP
)
parser.add_argument(
"--enable-xformers", action="store_true", help=ENABLE_XFORMERS_HELP
)
parser.add_argument(
"--device", default=DEFAULT_DEVICE, type=str, choices=AVAILABLE_DEVICES
)
2023-01-20 09:52:38 +01:00
parser.add_argument("--gui", action="store_true", help=GUI_HELP)
2023-02-06 15:00:47 +01:00
parser.add_argument(
"--no-gui-auto-close", action="store_true", help=NO_GUI_AUTO_CLOSE_HELP
)
2022-09-15 16:21:27 +02:00
parser.add_argument(
2023-01-20 09:52:38 +01:00
"--gui-size",
default=[1600, 1000],
nargs=2,
type=int,
help="Set window size for GUI",
2022-09-15 16:21:27 +02:00
)
2023-01-20 09:52:38 +01:00
parser.add_argument("--input", type=str, default=None, help=INPUT_HELP)
parser.add_argument("--output-dir", type=str, default=None, help=OUTPUT_DIR_HELP)
2023-02-06 15:00:47 +01:00
parser.add_argument(
"--model-dir", type=str, default=DEFAULT_MODEL_DIR, help=MODEL_DIR_HELP
)
parser.add_argument(
"--disable-model-switch",
action="store_true",
help="Disable model switch in frontend",
)
parser.add_argument(
2023-03-22 05:57:18 +01:00
"--quality",
default=95,
type=int,
help=QUALITY_HELP,
)
# Plugins
parser.add_argument(
"--enable-interactive-seg",
action="store_true",
2023-03-28 15:55:25 +02:00
help=INTERACTIVE_SEG_HELP,
2023-03-22 05:57:18 +01:00
)
2023-04-06 15:55:20 +02:00
parser.add_argument(
"--interactive-seg-model",
default="vit_l",
2023-04-16 04:35:51 +02:00
choices=AVAILABLE_INTERACTIVE_SEG_MODELS,
2023-05-02 08:09:40 +02:00
help=INTERACTIVE_SEG_MODEL_HELP,
2023-04-06 15:55:20 +02:00
)
parser.add_argument(
"--interactive-seg-device",
default="cpu",
2023-04-16 04:35:51 +02:00
choices=AVAILABLE_INTERACTIVE_SEG_DEVICES,
2023-04-06 15:55:20 +02:00
)
2023-03-22 05:57:18 +01:00
parser.add_argument(
"--enable-remove-bg",
action="store_true",
2023-03-28 15:55:25 +02:00
help=REMOVE_BG_HELP,
2023-03-22 05:57:18 +01:00
)
2023-05-09 13:07:12 +02:00
parser.add_argument(
"--enable-anime-seg",
action="store_true",
help=ANIMESEG_HELP,
)
2023-03-22 05:57:18 +01:00
parser.add_argument(
"--enable-realesrgan",
action="store_true",
2023-03-28 15:55:25 +02:00
help=REALESRGAN_HELP,
2023-03-22 05:57:18 +01:00
)
parser.add_argument(
2023-03-28 15:55:25 +02:00
"--realesrgan-device",
default="cpu",
type=str,
choices=REALESRGAN_AVAILABLE_DEVICES,
2023-03-25 03:15:44 +01:00
)
parser.add_argument(
"--realesrgan-model",
default=RealESRGANModelName.realesr_general_x4v3.value,
type=str,
choices=RealESRGANModelNameList,
)
2023-04-03 07:32:04 +02:00
parser.add_argument(
"--realesrgan-no-half",
action="store_true",
help="Disable half precision for RealESRGAN",
)
2023-03-28 15:55:25 +02:00
parser.add_argument("--enable-gfpgan", action="store_true", help=GFPGAN_HELP)
2023-03-26 07:39:09 +02:00
parser.add_argument(
2023-03-30 10:07:38 +02:00
"--gfpgan-device", default="cpu", type=str, choices=GFPGAN_AVAILABLE_DEVICES
)
parser.add_argument(
"--enable-restoreformer", action="store_true", help=RESTOREFORMER_HELP
)
parser.add_argument(
"--restoreformer-device",
default="cpu",
type=str,
choices=RESTOREFORMER_AVAILABLE_DEVICES,
2023-03-26 07:39:09 +02:00
)
2023-03-25 02:53:22 +01:00
parser.add_argument(
"--enable-gif",
action="store_true",
2023-03-28 15:55:25 +02:00
help=GIF_HELP,
2023-03-25 02:53:22 +01:00
)
2023-03-29 16:05:34 +02:00
parser.add_argument(
"--install-plugins-package",
action="store_true",
)
2023-03-22 05:57:18 +01:00
#########
2023-01-20 09:52:38 +01:00
# useless args
2023-02-19 14:16:30 +01:00
parser.add_argument("--debug", action="store_true", help=argparse.SUPPRESS)
parser.add_argument("--hf_access_token", default="", help=argparse.SUPPRESS)
2022-09-15 16:21:27 +02:00
parser.add_argument(
2023-02-19 14:16:30 +01:00
"--sd-disable-nsfw", action="store_true", help=argparse.SUPPRESS
2023-01-05 15:07:39 +01:00
)
2023-02-19 14:16:30 +01:00
parser.add_argument("--sd-run-local", action="store_true", help=argparse.SUPPRESS)
parser.add_argument(
2023-02-19 14:16:30 +01:00
"--sd-enable-xformers", action="store_true", help=argparse.SUPPRESS
)
2022-04-18 09:01:10 +02:00
args = parser.parse_args()
2023-01-14 14:36:32 +01:00
2023-01-20 09:52:38 +01:00
# collect system info to help debug
dump_environment_info()
2023-03-29 16:05:34 +02:00
if args.install_plugins_package:
from lama_cleaner.installer import install_plugins_package
install_plugins_package()
exit()
2023-01-20 09:52:38 +01:00
if args.config_installer:
if args.installer_config is None:
2023-02-06 15:00:47 +01:00
parser.error(
2023-05-02 08:09:40 +02:00
"args.config_installer==True, must set args.installer_config to store config file"
2023-02-06 15:00:47 +01:00
)
2023-01-20 09:52:38 +01:00
from lama_cleaner.web_config import main
2023-02-06 15:00:47 +01:00
2023-05-02 08:09:40 +02:00
logger.info("Launching installer web config page")
2023-01-20 09:52:38 +01:00
main(args.installer_config)
exit()
if args.load_installer_config:
if args.installer_config and not os.path.exists(args.installer_config):
parser.error(f"args.installer_config={args.installer_config} not exists")
logger.info(f"Loading installer config from {args.installer_config}")
_args = load_config(args.installer_config)
for k, v in vars(_args).items():
if k in vars(args):
setattr(args, k, v)
2023-01-14 15:00:34 +01:00
if args.device == "cuda":
2023-05-29 15:40:58 +02:00
import platform
if platform.system() == "Darwin":
logger.info("MacOS does not support cuda, use cpu instead")
setattr(args, "device", "cpu")
else:
import torch
if torch.cuda.is_available() is False:
parser.error(
"torch.cuda.is_available() is False, please use --device cpu or check your pytorch installation"
)
2023-02-06 15:00:47 +01:00
2023-03-29 16:05:34 +02:00
if args.sd_local_model_path and args.model == "sd1.5":
if not os.path.exists(args.sd_local_model_path):
parser.error(
f"invalid --sd-local-model-path: {args.sd_local_model_path} not exists"
)
if not os.path.isfile(args.sd_local_model_path):
parser.error(
f"invalid --sd-local-model-path: {args.sd_local_model_path} is a directory"
)
2023-03-22 05:57:18 +01:00
os.environ["U2NET_HOME"] = DEFAULT_MODEL_DIR
2023-01-20 09:52:38 +01:00
if args.model_dir and args.model_dir is not None:
2023-01-14 14:36:32 +01:00
if os.path.isfile(args.model_dir):
parser.error(f"invalid --model-dir: {args.model_dir} is a file")
if not os.path.exists(args.model_dir):
logger.info(f"Create model cache directory: {args.model_dir}")
Path(args.model_dir).mkdir(exist_ok=True, parents=True)
os.environ["XDG_CACHE_HOME"] = args.model_dir
2023-03-22 05:57:18 +01:00
os.environ["U2NET_HOME"] = args.model_dir
2023-01-14 14:36:32 +01:00
2023-01-20 09:52:38 +01:00
if args.input and args.input is not None:
2022-04-18 09:01:10 +02:00
if not os.path.exists(args.input):
parser.error(f"invalid --input: {args.input} not exists")
2022-12-31 14:07:08 +01:00
if os.path.isfile(args.input):
if imghdr.what(args.input) is None:
parser.error(f"invalid --input: {args.input} is not a valid image file")
else:
if args.output_dir is None:
2023-02-06 15:00:47 +01:00
parser.error(
f"invalid --input: {args.input} is a directory, --output-dir is required"
)
if args.output_dir is not None:
output_dir = Path(args.output_dir)
if not output_dir.exists():
logger.info(f"Creating output directory: {output_dir}")
output_dir.mkdir(parents=True)
else:
if not output_dir.is_dir():
parser.error(f"invalid --output-dir: {output_dir} is not a directory")
2022-04-18 09:01:10 +02:00
return args