123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
import os
|
|
import imghdr
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from loguru import logger
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", default=8080, type=int)
|
|
parser.add_argument(
|
|
"--model",
|
|
default="lama",
|
|
choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2", "paint_by_example"],
|
|
)
|
|
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.")
|
|
parser.add_argument(
|
|
"--hf_access_token",
|
|
default="",
|
|
help="SD model no more need token: https://github.com/huggingface/diffusers/issues/1447",
|
|
)
|
|
parser.add_argument(
|
|
"--sd-disable-nsfw",
|
|
action="store_true",
|
|
help="Disable Stable Diffusion NSFW checker",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-nsfw",
|
|
action="store_true",
|
|
help="Disable Stable Diffusion/Paint By Example NSFW checker",
|
|
)
|
|
parser.add_argument(
|
|
"--sd-cpu-textencoder",
|
|
action="store_true",
|
|
help="Always run Stable Diffusion TextEncoder model on CPU",
|
|
)
|
|
parser.add_argument(
|
|
"--sd-run-local",
|
|
action="store_true",
|
|
help="SD model no more need token, use --local-files-only to set not connect to huggingface server",
|
|
)
|
|
parser.add_argument(
|
|
"--local-files-only",
|
|
action="store_true",
|
|
help="sd/paint_by_example model. Use local files only, not connect to huggingface server",
|
|
)
|
|
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"
|
|
)
|
|
parser.add_argument(
|
|
"--enable-xformers",
|
|
action="store_true",
|
|
help="sd/paint_by_example model. Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers"
|
|
)
|
|
parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu", "mps"])
|
|
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="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"
|
|
)
|
|
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"
|
|
)
|
|
parser.add_argument("--disable-model-switch", action="store_true", help="Disable model switch in frontend")
|
|
parser.add_argument("--debug", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
|
|
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")
|
|
|
|
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
|
|
|
|
if args.input is not None:
|
|
if not os.path.exists(args.input):
|
|
parser.error(f"invalid --input: {args.input} not exists")
|
|
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")
|
|
|
|
return args
|