IOPaint/lama_cleaner/parse_args.py

63 lines
2.1 KiB
Python
Raw Normal View History

2022-04-18 09:01:10 +02:00
import os
import imghdr
import argparse
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-09-25 15:27:12 +02:00
choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.4", "cv2"],
2022-09-15 16:21:27 +02:00
)
parser.add_argument(
"--hf_access_token",
default="",
2022-09-29 06:20:55 +02:00
help="Huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens",
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",
help="After first time Stable Diffusion model downloaded, you can add this arg and remove --hf_access_token",
)
2022-04-18 09:01:10 +02:00
parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu"])
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(
"--input", type=str, help="Path to image you want to load by default"
)
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()
if args.input is not None:
if not os.path.exists(args.input):
parser.error(f"invalid --input: {args.input} not exists")
if imghdr.what(args.input) is None:
parser.error(f"invalid --input: {args.input} is not a valid image file")
2022-09-29 07:13:09 +02:00
if args.model.startswith("sd") and not args.sd_run_local:
2022-09-15 16:21:27 +02:00
if not args.hf_access_token.startswith("hf_"):
parser.error(
f"sd(stable-diffusion) model requires huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens"
)
2022-04-18 09:01:10 +02:00
return args