# Copyright 2023 Lincoln D. Stein and the InvokeAI Team

import argparse
import shlex
from argparse import ArgumentParser

# note that this includes both old sampler names and new scheduler names
# in order to be able to parse both 2.0 and 3.0-pre-nodes versions of invokeai.init
SAMPLER_CHOICES = [
    "ddim",
    "ddpm",
    "deis",
    "lms",
    "lms_k",
    "pndm",
    "heun",
    "heun_k",
    "euler",
    "euler_k",
    "euler_a",
    "kdpm_2",
    "kdpm_2_a",
    "dpmpp_2s",
    "dpmpp_2s_k",
    "dpmpp_2m",
    "dpmpp_2m_k",
    "dpmpp_2m_sde",
    "dpmpp_2m_sde_k",
    "dpmpp_sde",
    "dpmpp_sde_k",
    "unipc",
    "k_dpm_2_a",
    "k_dpm_2",
    "k_dpmpp_2_a",
    "k_dpmpp_2",
    "k_euler_a",
    "k_euler",
    "k_heun",
    "k_lms",
    "plms",
]

PRECISION_CHOICES = [
    "auto",
    "float32",
    "autocast",
    "float16",
]


class FileArgumentParser(ArgumentParser):
    """
    Supports reading defaults from an init file.
    """

    def convert_arg_line_to_args(self, arg_line):
        return shlex.split(arg_line, comments=True)


legacy_parser = FileArgumentParser(
    description="""
Generate images using Stable Diffusion.
    Use --web to launch the web interface.
    Use --from_file to load prompts from a file path or standard input ("-").
    Otherwise you will be dropped into an interactive command prompt (type -h for help.)
    Other command-line arguments are defaults that can usually be overridden
    prompt the command prompt.
    """,
    fromfile_prefix_chars="@",
)
general_group = legacy_parser.add_argument_group("General")
model_group = legacy_parser.add_argument_group("Model selection")
file_group = legacy_parser.add_argument_group("Input/output")
web_server_group = legacy_parser.add_argument_group("Web server")
render_group = legacy_parser.add_argument_group("Rendering")
postprocessing_group = legacy_parser.add_argument_group("Postprocessing")
deprecated_group = legacy_parser.add_argument_group("Deprecated options")

deprecated_group.add_argument("--laion400m")
deprecated_group.add_argument("--weights")  # deprecated
general_group.add_argument("--version", "-V", action="store_true", help="Print InvokeAI version number")
model_group.add_argument(
    "--root_dir",
    default=None,
    help='Path to directory containing "models", "outputs" and "configs". If not present will read from environment variable INVOKEAI_ROOT. Defaults to ~/invokeai.',
)
model_group.add_argument(
    "--config",
    "-c",
    "-config",
    dest="conf",
    default="./configs/models.yaml",
    help="Path to configuration file for alternate models.",
)
model_group.add_argument(
    "--model",
    help='Indicates which diffusion model to load (defaults to "default" stanza in configs/models.yaml)',
)
model_group.add_argument(
    "--weight_dirs",
    nargs="+",
    type=str,
    help="List of one or more directories that will be auto-scanned for new model weights to import",
)
model_group.add_argument(
    "--png_compression",
    "-z",
    type=int,
    default=6,
    choices=range(0, 9),
    dest="png_compression",
    help="level of PNG compression, from 0 (none) to 9 (maximum). Default is 6.",
)
model_group.add_argument(
    "-F",
    "--full_precision",
    dest="full_precision",
    action="store_true",
    help="Deprecated way to set --precision=float32",
)
model_group.add_argument(
    "--max_loaded_models",
    dest="max_loaded_models",
    type=int,
    default=2,
    help="Maximum number of models to keep in memory for fast switching, including the one in GPU",
)
model_group.add_argument(
    "--free_gpu_mem",
    dest="free_gpu_mem",
    action="store_true",
    help="Force free gpu memory before final decoding",
)
model_group.add_argument(
    "--sequential_guidance",
    dest="sequential_guidance",
    action="store_true",
    help="Calculate guidance in serial instead of in parallel, lowering memory requirement " "at the expense of speed",
)
model_group.add_argument(
    "--xformers",
    action=argparse.BooleanOptionalAction,
    default=True,
    help="Enable/disable xformers support (default enabled if installed)",
)
model_group.add_argument(
    "--always_use_cpu", dest="always_use_cpu", action="store_true", help="Force use of CPU even if GPU is available"
)
model_group.add_argument(
    "--precision",
    dest="precision",
    type=str,
    choices=PRECISION_CHOICES,
    metavar="PRECISION",
    help=f'Set model precision. Defaults to auto selected based on device. Options: {", ".join(PRECISION_CHOICES)}',
    default="auto",
)
model_group.add_argument(
    "--ckpt_convert",
    action=argparse.BooleanOptionalAction,
    dest="ckpt_convert",
    default=True,
    help="Deprecated option. Legacy ckpt files are now always converted to diffusers when loaded.",
)
model_group.add_argument(
    "--internet",
    action=argparse.BooleanOptionalAction,
    dest="internet_available",
    default=True,
    help="Indicate whether internet is available for just-in-time model downloading (default: probe automatically).",
)
model_group.add_argument(
    "--nsfw_checker",
    "--safety_checker",
    action=argparse.BooleanOptionalAction,
    dest="safety_checker",
    default=False,
    help="Check for and blur potentially NSFW images. Use --no-nsfw_checker to disable.",
)
model_group.add_argument(
    "--autoimport",
    default=None,
    type=str,
    help="Check the indicated directory for .ckpt/.safetensors weights files at startup and import directly",
)
model_group.add_argument(
    "--autoconvert",
    default=None,
    type=str,
    help="Check the indicated directory for .ckpt/.safetensors weights files at startup and import as optimized diffuser models",
)
model_group.add_argument(
    "--patchmatch",
    action=argparse.BooleanOptionalAction,
    default=True,
    help="Load the patchmatch extension for outpainting. Use --no-patchmatch to disable.",
)
file_group.add_argument(
    "--from_file",
    dest="infile",
    type=str,
    help="If specified, load prompts from this file",
)
file_group.add_argument(
    "--outdir",
    "-o",
    type=str,
    help="Directory to save generated images and a log of prompts and seeds. Default: ROOTDIR/outputs",
    default="outputs",
)
file_group.add_argument(
    "--prompt_as_dir",
    "-p",
    action="store_true",
    help="Place images in subdirectories named after the prompt.",
)
render_group.add_argument(
    "--fnformat",
    default="{prefix}.{seed}.png",
    type=str,
    help="Overwrite the filename format. You can use any argument as wildcard enclosed in curly braces. Default is {prefix}.{seed}.png",
)
render_group.add_argument("-s", "--steps", type=int, default=50, help="Number of steps")
render_group.add_argument(
    "-W",
    "--width",
    type=int,
    help="Image width, multiple of 64",
)
render_group.add_argument(
    "-H",
    "--height",
    type=int,
    help="Image height, multiple of 64",
)
render_group.add_argument(
    "-C",
    "--cfg_scale",
    default=7.5,
    type=float,
    help='Classifier free guidance (CFG) scale - higher numbers cause generator to "try" harder.',
)
render_group.add_argument(
    "--sampler",
    "-A",
    "-m",
    dest="sampler_name",
    type=str,
    choices=SAMPLER_CHOICES,
    metavar="SAMPLER_NAME",
    help=f'Set the default sampler. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
    default="k_lms",
)
render_group.add_argument(
    "--log_tokenization", "-t", action="store_true", help="shows how the prompt is split into tokens"
)
render_group.add_argument(
    "-f",
    "--strength",
    type=float,
    help="img2img strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely",
)
render_group.add_argument(
    "-T",
    "-fit",
    "--fit",
    action=argparse.BooleanOptionalAction,
    help="If specified, will resize the input image to fit within the dimensions of width x height (512x512 default)",
)

render_group.add_argument("--grid", "-g", action=argparse.BooleanOptionalAction, help="generate a grid")
render_group.add_argument(
    "--embedding_directory",
    "--embedding_path",
    dest="embedding_path",
    default="embeddings",
    type=str,
    help="Path to a directory containing .bin and/or .pt files, or a single .bin/.pt file. You may use subdirectories. (default is ROOTDIR/embeddings)",
)
render_group.add_argument(
    "--lora_directory",
    dest="lora_path",
    default="loras",
    type=str,
    help="Path to a directory containing LoRA files; subdirectories are not supported. (default is ROOTDIR/loras)",
)
render_group.add_argument(
    "--embeddings",
    action=argparse.BooleanOptionalAction,
    default=True,
    help="Enable embedding directory (default). Use --no-embeddings to disable.",
)
render_group.add_argument("--enable_image_debugging", action="store_true", help="Generates debugging image to display")
render_group.add_argument(
    "--karras_max",
    type=int,
    default=None,
    help="control the point at which the K* samplers will shift from using the Karras noise schedule (good for low step counts) to the LatentDiffusion noise schedule (good for high step counts). Set to 0 to use LatentDiffusion for all step values, and to a high value (e.g. 1000) to use Karras for all step values. [29].",
)
# Restoration related args
postprocessing_group.add_argument(
    "--no_restore",
    dest="restore",
    action="store_false",
    help="Disable face restoration with GFPGAN or codeformer",
)
postprocessing_group.add_argument(
    "--no_upscale",
    dest="esrgan",
    action="store_false",
    help="Disable upscaling with ESRGAN",
)
postprocessing_group.add_argument(
    "--esrgan_bg_tile",
    type=int,
    default=400,
    help="Tile size for background sampler, 0 for no tile during testing. Default: 400.",
)
postprocessing_group.add_argument(
    "--esrgan_denoise_str",
    type=float,
    default=0.75,
    help="esrgan denoise str. 0 is no denoise, 1 is max denoise.  Default: 0.75",
)
postprocessing_group.add_argument(
    "--gfpgan_model_path",
    type=str,
    default="./models/gfpgan/GFPGANv1.4.pth",
    help="Indicates the path to the GFPGAN model",
)
web_server_group.add_argument(
    "--web",
    dest="web",
    action="store_true",
    help="Start in web server mode.",
)
web_server_group.add_argument(
    "--web_develop",
    dest="web_develop",
    action="store_true",
    help="Start in web server development mode.",
)
web_server_group.add_argument(
    "--web_verbose",
    action="store_true",
    help="Enables verbose logging",
)
web_server_group.add_argument(
    "--cors",
    nargs="*",
    type=str,
    help="Additional allowed origins, comma-separated",
)
web_server_group.add_argument(
    "--host",
    type=str,
    default="127.0.0.1",
    help="Web server: Host or IP to listen on. Set to 0.0.0.0 to accept traffic from other devices on your network.",
)
web_server_group.add_argument("--port", type=int, default="9090", help="Web server: Port to listen on")
web_server_group.add_argument(
    "--certfile",
    type=str,
    default=None,
    help="Web server: Path to certificate file to use for SSL. Use together with --keyfile",
)
web_server_group.add_argument(
    "--keyfile",
    type=str,
    default=None,
    help="Web server: Path to private key file to use for SSL. Use together with --certfile",
)
web_server_group.add_argument(
    "--gui",
    dest="gui",
    action="store_true",
    help="Start InvokeAI GUI",
)