IOPaint/lama_cleaner/parse_args.py

113 lines
4.4 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
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", default=8080, type=int)
2022-09-15 16:21:27 +02:00
parser.add_argument(
"--model",
default="lama",
2022-12-10 15:06:15 +01:00
choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2", "paint_by_example"],
2022-09-15 16:21:27 +02:00
)
2023-01-05 15:07:39 +01:00
parser.add_argument("--no-half", action="store_true", help="sd/paint_by_example model no half precision")
parser.add_argument("--cpu-offload", action="store_true",
help="sd/paint_by_example model, offloads all models to CPU, significantly reducing vRAM usage.")
2022-09-15 16:21:27 +02:00
parser.add_argument(
"--hf_access_token",
default="",
2023-01-08 13:59:55 +01:00
help="SD model no more need token: https://github.com/huggingface/diffusers/issues/1447",
2022-09-15 16:21:27 +02:00
)
2022-09-29 03:42:19 +02:00
parser.add_argument(
"--sd-disable-nsfw",
action="store_true",
2022-09-29 07:13:09 +02:00
help="Disable Stable Diffusion NSFW checker",
2022-09-29 06:20:55 +02:00
)
parser.add_argument(
"--sd-cpu-textencoder",
action="store_true",
help="Always run Stable Diffusion TextEncoder model on CPU",
2022-09-29 03:42:19 +02:00
)
2022-09-29 07:13:09 +02:00
parser.add_argument(
"--sd-run-local",
action="store_true",
2023-01-08 13:59:55 +01:00
help="SD model no more need token, use --local-files-only to set not connect to huggingface server",
2023-01-05 15:07:39 +01:00
)
parser.add_argument(
"--local-files-only",
action="store_true",
help="sd/paint_by_example model. Use local files only, not connect to huggingface server",
2022-09-29 07:13:09 +02:00
)
parser.add_argument(
"--sd-enable-xformers",
action="store_true",
help="Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers"
)
2022-12-04 13:49:52 +01:00
parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu", "mps"])
2022-04-18 09:01:10 +02:00
parser.add_argument("--gui", action="store_true", help="Launch as desktop app")
parser.add_argument(
"--gui-size",
default=[1600, 1000],
nargs=2,
type=int,
help="Set window size for GUI",
)
parser.add_argument(
2022-12-31 14:07:08 +01:00
"--input", type=str,
help="If input is image, it will be load by default. If input is directory, all images will be loaded to file manager"
)
parser.add_argument(
"--output-dir", type=str,
help="Only required when --input is directory. Output directory for all processed images"
2022-04-18 09:01:10 +02:00
)
2023-01-14 14:36:32 +01:00
parser.add_argument(
"--model-dir", type=str, default=None,
help="Model download directory (by setting XDG_CACHE_HOME environment variable), "
"by default model downloaded to ~/.cache"
)
2022-11-13 06:14:37 +01:00
parser.add_argument("--disable-model-switch", action="store_true", help="Disable model switch in frontend")
2022-04-18 09:01:10 +02:00
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()
2023-01-14 14:36:32 +01:00
2023-01-14 15:00:34 +01:00
if args.device == "cuda":
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-01-14 14:36:32 +01:00
if args.model_dir is not None:
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
2022-04-18 09:01:10 +02:00
if args.input is not None:
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:
parser.error(f"invalid --input: {args.input} is a directory, --output-dir is required")
else:
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