IOPaint/lama_cleaner/parse_args.py

import os
import imghdr
import argparse
from pathlib import Path

from loguru import logger

from lama_cleaner.const import (
    AVAILABLE_MODELS,
    NO_HALF_HELP,
    CPU_OFFLOAD_HELP,
    DISABLE_NSFW_HELP,
    SD_CPU_TEXTENCODER_HELP,
    LOCAL_FILES_ONLY_HELP,
    AVAILABLE_DEVICES,
    ENABLE_XFORMERS_HELP,
    MODEL_DIR_HELP,
    OUTPUT_DIR_HELP,
    INPUT_HELP,
    GUI_HELP,
    DEFAULT_DEVICE,
    NO_GUI_AUTO_CLOSE_HELP,
    DEFAULT_MODEL_DIR,
    DEFAULT_MODEL,
    MPS_SUPPORT_MODELS,
)
from lama_cleaner.runtime import dump_environment_info


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(
        "--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"
    )

    parser.add_argument("--model", default=DEFAULT_MODEL, choices=AVAILABLE_MODELS)
    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)
    parser.add_argument(
        "--sd-cpu-textencoder", action="store_true", help=SD_CPU_TEXTENCODER_HELP
    )
    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
    )
    parser.add_argument("--gui", action="store_true", help=GUI_HELP)
    parser.add_argument(
        "--no-gui-auto-close", action="store_true", help=NO_GUI_AUTO_CLOSE_HELP
    )
    parser.add_argument(
        "--gui-size",
        default=[1600, 1000],
        nargs=2,
        type=int,
        help="Set window size for GUI",
    )
    parser.add_argument("--input", type=str, default=None, help=INPUT_HELP)
    parser.add_argument("--output-dir", type=str, default=None, help=OUTPUT_DIR_HELP)
    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("--debug", action="store_true")

    # useless args
    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(
        "--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(
        "--sd-enable-xformers",
        action="store_true",
        help="Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers",
    )

    args = parser.parse_args()

    # collect system info to help debug
    dump_environment_info()

    if args.config_installer:
        if args.installer_config is None:
            parser.error(
                f"args.config_installer==True, must set args.installer_config to store config file"
            )
        from lama_cleaner.web_config import main

        logger.info(f"Launching installer web config page")
        main(args.installer_config)
        exit()

    if args.load_installer_config:
        from lama_cleaner.web_config import load_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)

    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.device == "mps":
        if args.model not in MPS_SUPPORT_MODELS:
            parser.error(
                f"mps only support: {MPS_SUPPORT_MODELS}, but got {args.model}"
            )

    if args.model_dir and 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 and 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