diff --git a/invokeai/app/services/model_manager_initializer.py b/invokeai/app/services/model_manager_initializer.py index 8fd152a00a..3d30577c55 100644 --- a/invokeai/app/services/model_manager_initializer.py +++ b/invokeai/app/services/model_manager_initializer.py @@ -2,7 +2,6 @@ import os import sys import torch from argparse import Namespace -from invokeai.backend import Args from omegaconf import OmegaConf from pathlib import Path from typing import types @@ -13,7 +12,7 @@ from ...backend import ModelManager from ...backend.util import choose_precision, choose_torch_device # TODO: Replace with an abstract class base ModelManagerBase -def get_model_manager(config: Args, logger: types.ModuleType) -> ModelManager: +def get_model_manager(config: InvokeAISettings, logger: types.ModuleType) -> ModelManager: model_config = config.model_conf_path if not model_config.exists(): report_model_error( @@ -44,7 +43,7 @@ def get_model_manager(config: Args, logger: types.ModuleType) -> ModelManager: else choose_precision(device) model_manager = ModelManager( - OmegaConf.load(config.conf), + OmegaConf.load(config.model_conf_path), precision=precision, device_type=device, max_loaded_models=config.max_loaded_models, diff --git a/invokeai/backend/__init__.py b/invokeai/backend/__init__.py index dd126a322d..0e1b6d3a0d 100644 --- a/invokeai/backend/__init__.py +++ b/invokeai/backend/__init__.py @@ -1,7 +1,6 @@ """ Initialization file for invokeai.backend """ -from .generate import Generate from .generator import ( InvokeAIGeneratorBasicParams, InvokeAIGenerator, @@ -12,5 +11,3 @@ from .generator import ( ) from .model_management import ModelManager, SDModelComponent from .safety_checker import SafetyChecker -from .args import Args -from .globals import Globals diff --git a/invokeai/backend/args.py b/invokeai/backend/args.py deleted file mode 100644 index eb8b396ee0..0000000000 --- a/invokeai/backend/args.py +++ /dev/null @@ -1,1387 +0,0 @@ -"""Helper class for dealing with image generation arguments. - -The Args class parses both the command line (shell) arguments, as well as the -command string passed at the invoke> prompt. It serves as the definitive repository -of all the arguments used by Generate and their default values, and implements the -preliminary metadata standards discussed here: - -https://github.com/lstein/stable-diffusion/issues/266 - -To use: - opt = Args() - - # Read in the command line options: - # this returns a namespace object like the underlying argparse library) - # You do not have to use the return value, but you can check it against None - # to detect illegal arguments on the command line. - args = opt.parse_args() - if not args: - print('oops') - sys.exit(-1) - - # read in a command passed to the invoke> prompt: - opts = opt.parse_cmd('do androids dream of electric sheep? -H256 -W1024 -n4') - - # The Args object acts like a namespace object - print(opt.model) - -You can set attributes in the usual way, use vars(), etc.: - - opt.model = 'something-else' - do_something(**vars(a)) - -It is helpful in saving metadata: - - # To get a json representation of all the values, allowing - # you to override any values dynamically - j = opt.json(seed=42) - - # To get the prompt string with the switches, allowing you - # to override any values dynamically - j = opt.dream_prompt_str(seed=42) - -If you want to access the namespace objects from the shell args or the -parsed command directly, you may use the values returned from the -original calls to parse_args() and parse_cmd(), or get them later -using the _arg_switches and _cmd_switches attributes. This can be -useful if both the args and the command contain the same attribute and -you wish to apply logic as to which one to use. For example: - - a = Args() - args = a.parse_args() - opts = a.parse_cmd(string) - do_grid = args.grid or opts.grid - -To add new attributes, edit the _create_arg_parser() and -_create_dream_cmd_parser() methods. - -**Generating and retrieving sd-metadata** - -To generate a dict representing RFC266 metadata: - - metadata = metadata_dumps(opt,) - -This will generate an RFC266 dictionary that can then be turned into a JSON -and written to the PNG file. The optional seeds, weights, model_hash and -postprocesser arguments are not available to the opt object and so must be -provided externally. See how invoke.py does it. - -Note that this function was originally called format_metadata() and a wrapper -is provided that issues a deprecation notice. - -To retrieve a (series of) opt objects corresponding to the metadata, do this: - - opt_list = metadata_loads(metadata) - -The metadata should be pulled out of the PNG image. pngwriter has a method -retrieve_metadata that will do this, or you can do it in one swell foop -with metadata_from_png(): - - opt_list = metadata_from_png('/path/to/image_file.png') -""" - -import argparse -import base64 -import copy -import functools -import hashlib -import json -import os -import pydoc -import re -import shlex -import sys -from argparse import Namespace -from pathlib import Path -from typing import List - -import invokeai.version -import invokeai.backend.util.logging as logger -from invokeai.backend.image_util import retrieve_metadata - -from .globals import Globals -from .prompting import split_weighted_subprompts - -APP_ID = invokeai.version.__app_id__ -APP_NAME = invokeai.version.__app_name__ -APP_VERSION = invokeai.version.__version__ - -SAMPLER_CHOICES = [ - "ddim", - "k_dpm_2_a", - "k_dpm_2", - "k_dpmpp_2_a", - "k_dpmpp_2", - "k_euler_a", - "k_euler", - "k_heun", - "k_lms", - "plms", - # diffusers: - "pndm", -] - -PRECISION_CHOICES = [ - "auto", - "float32", - "autocast", - "float16", -] - - -class ArgFormatter(argparse.RawTextHelpFormatter): - # use defined argument order to display usage - def _format_usage(self, usage, actions, groups, prefix): - if prefix is None: - prefix = "usage: " - - # if usage is specified, use that - if usage is not None: - usage = usage % dict(prog=self._prog) - - # if no optionals or positionals are available, usage is just prog - elif usage is None and not actions: - usage = "invoke>" - elif usage is None: - prog = "invoke>" - # build full usage string - action_usage = self._format_actions_usage(actions, groups) # NEW - usage = " ".join([s for s in [prog, action_usage] if s]) - # omit the long line wrapping code - # prefix with 'usage:' - return "%s%s\n\n" % (prefix, usage) - - -class PagingArgumentParser(argparse.ArgumentParser): - """ - A custom ArgumentParser that uses pydoc to page its output. - It also supports reading defaults from an init file. - """ - - def print_help(self, file=None): - text = self.format_help() - pydoc.pager(text) - - def convert_arg_line_to_args(self, arg_line): - return shlex.split(arg_line, comments=True) - - -class Args(object): - def __init__(self, arg_parser=None, cmd_parser=None): - """ - Initialize new Args class. It takes two optional arguments, an argparse - parser for switches given on the shell command line, and an argparse - parser for switches given on the invoke> CLI line. If one or both are - missing, it creates appropriate parsers internally. - """ - self._arg_parser = arg_parser or self._create_arg_parser() - self._cmd_parser = cmd_parser or self._create_dream_cmd_parser() - self._arg_switches = self.parse_cmd("") # fill in defaults - self._cmd_switches = self.parse_cmd("") # fill in defaults - - def parse_args(self, args: List[str] = None): - """Parse the shell switches and store.""" - sysargs = args if args is not None else sys.argv[1:] - try: - # pre-parse before we do any initialization to get root directory - # and intercept --version request - switches = self._arg_parser.parse_args(sysargs) - if switches.version: - print(f"{APP_NAME} {APP_VERSION}") - sys.exit(0) - - logger.info("Initializing, be patient...") - Globals.root = Path(os.path.abspath(switches.root_dir or Globals.root)) - Globals.try_patchmatch = switches.patchmatch - - # now use root directory to find the init file - initfile = os.path.expanduser(os.path.join(Globals.root, Globals.initfile)) - legacyinit = os.path.expanduser("~/.invokeai") - if os.path.exists(initfile): - logger.info( - f"Initialization file {initfile} found. Loading...", - ) - sysargs.insert(0, f"@{initfile}") - elif os.path.exists(legacyinit): - logger.warning( - f"Old initialization file found at {legacyinit}. This location is deprecated. Please move it to {Globals.root}/invokeai.init." - ) - sysargs.insert(0, f"@{legacyinit}") - Globals.log_tokenization = self._arg_parser.parse_args( - sysargs - ).log_tokenization - - self._arg_switches = self._arg_parser.parse_args(sysargs) - return self._arg_switches - except Exception as e: - logger.error(f"An exception has occurred: {e}") - return None - - def parse_cmd(self, cmd_string): - """Parse a invoke>-style command string""" - # handle the case in which the first token is a switch - if cmd_string.startswith("-"): - prompt = "" - switches = cmd_string - # handle the case in which the prompt is enclosed by quotes - elif cmd_string.startswith('"'): - a = shlex.split(cmd_string, comments=True) - prompt = a[0] - switches = shlex.join(a[1:]) - else: - # no initial quote, so get everything up to the first thing - # that looks like a switch - if cmd_string.startswith("-"): - prompt = "" - switches = cmd_string - else: - match = re.match("^(.+?)\s(--?[a-zA-Z].+)", cmd_string) - if match: - prompt, switches = match.groups() - else: - prompt = cmd_string - switches = "" - try: - self._cmd_switches = self._cmd_parser.parse_args( - shlex.split(switches, comments=True) - ) - if not getattr(self._cmd_switches, "prompt"): - setattr(self._cmd_switches, "prompt", prompt) - return self._cmd_switches - except: - return None - - def json(self, **kwargs): - return json.dumps(self.to_dict(**kwargs)) - - def to_dict(self, **kwargs): - a = vars(self) - a.update(kwargs) - return a - - # Isn't there a more automated way of doing this? - # Ideally we get the switch strings out of the argparse objects, - # but I don't see a documented API for this. - def dream_prompt_str(self, **kwargs): - """Normalized dream_prompt.""" - a = vars(self) - a.update(kwargs) - switches = list() - prompt = a["prompt"] - prompt.replace('"', '\\"') - switches.append(prompt) - switches.append(f'-s {a["steps"]}') - switches.append(f'-S {a["seed"]}') - switches.append(f'-W {a["width"]}') - switches.append(f'-H {a["height"]}') - switches.append(f'-C {a["cfg_scale"]}') - if a["karras_max"] is not None: - switches.append(f'--karras_max {a["karras_max"]}') - if a["perlin"] > 0: - switches.append(f'--perlin {a["perlin"]}') - if a["threshold"] > 0: - switches.append(f'--threshold {a["threshold"]}') - if a["grid"]: - switches.append("--grid") - if a["seamless"]: - switches.append("--seamless") - if a["hires_fix"]: - switches.append("--hires_fix") - if a["h_symmetry_time_pct"]: - switches.append(f'--h_symmetry_time_pct {a["h_symmetry_time_pct"]}') - if a["v_symmetry_time_pct"]: - switches.append(f'--v_symmetry_time_pct {a["v_symmetry_time_pct"]}') - - # img2img generations have parameters relevant only to them and have special handling - if a["init_img"] and len(a["init_img"]) > 0: - switches.append(f'-I {a["init_img"]}') - switches.append(f'-A {a["sampler_name"]}') - if a["fit"]: - switches.append("--fit") - if a["init_mask"] and len(a["init_mask"]) > 0: - switches.append(f'-M {a["init_mask"]}') - if a["init_color"] and len(a["init_color"]) > 0: - switches.append(f'--init_color {a["init_color"]}') - if a["strength"] and a["strength"] > 0: - switches.append(f'-f {a["strength"]}') - if a["inpaint_replace"]: - switches.append("--inpaint_replace") - if a["text_mask"]: - switches.append(f'-tm {" ".join([str(u) for u in a["text_mask"]])}') - else: - switches.append(f'-A {a["sampler_name"]}') - - # facetool-specific parameters, only print if running facetool - if a["facetool_strength"]: - switches.append(f'-G {a["facetool_strength"]}') - switches.append(f'-ft {a["facetool"]}') - if a["facetool"] == "codeformer": - switches.append(f'-cf {a["codeformer_fidelity"]}') - - if a["outcrop"]: - switches.append(f'-c {" ".join([str(u) for u in a["outcrop"]])}') - - # esrgan-specific parameters - if a["upscale"]: - switches.append(f'-U {" ".join([str(u) for u in a["upscale"]])}') - - # embiggen parameters - if a["embiggen"]: - switches.append(f'--embiggen {" ".join([str(u) for u in a["embiggen"]])}') - if a["embiggen_tiles"]: - switches.append( - f'--embiggen_tiles {" ".join([str(u) for u in a["embiggen_tiles"]])}' - ) - if a["embiggen_strength"]: - switches.append(f'--embiggen_strength {a["embiggen_strength"]}') - - # outpainting parameters - if a["out_direction"]: - switches.append(f'-D {" ".join([str(u) for u in a["out_direction"]])}') - - # LS: slight semantic drift which needs addressing in the future: - # 1. Variations come out of the stored metadata as a packed string with the keyword "variations" - # 2. However, they come out of the CLI (and probably web) with the keyword "with_variations" and - # in broken-out form. Variation (1) should be changed to comply with (2) - if a["with_variations"] and len(a["with_variations"]) > 0: - formatted_variations = ",".join( - f"{seed}:{weight}" for seed, weight in (a["with_variations"]) - ) - switches.append(f"-V {formatted_variations}") - if "variations" in a and len(a["variations"]) > 0: - switches.append(f'-V {a["variations"]}') - return " ".join(switches) - - def __getattribute__(self, name): - """ - Returns union of command-line arguments and dream_prompt arguments, - with the latter superseding the former. - """ - cmd_switches = None - arg_switches = None - try: - cmd_switches = object.__getattribute__(self, "_cmd_switches") - arg_switches = object.__getattribute__(self, "_arg_switches") - except AttributeError: - pass - - if cmd_switches and arg_switches and name == "__dict__": - return self._merge_dict( - arg_switches.__dict__, - cmd_switches.__dict__, - ) - try: - return object.__getattribute__(self, name) - except AttributeError: - pass - - if not hasattr(cmd_switches, name) and not hasattr(arg_switches, name): - raise AttributeError - - value_arg, value_cmd = (None, None) - try: - value_cmd = getattr(cmd_switches, name) - except AttributeError: - pass - try: - value_arg = getattr(arg_switches, name) - except AttributeError: - pass - - # here is where we can pick and choose which to use - # default behavior is to choose the dream_command value over - # the arg value. For example, the --grid and --individual options are a little - # funny because of their push/pull relationship. This is how to handle it. - if name == "grid": - if cmd_switches.individual: - return False - else: - return value_cmd or value_arg - return value_cmd if value_cmd is not None else value_arg - - def __setattr__(self, name, value): - if name.startswith("_"): - object.__setattr__(self, name, value) - else: - self._cmd_switches.__dict__[name] = value - - def _merge_dict(self, dict1, dict2): - new_dict = {} - for k in set(list(dict1.keys()) + list(dict2.keys())): - value1 = dict1.get(k, None) - value2 = dict2.get(k, None) - new_dict[k] = value2 if value2 is not None else value1 - return new_dict - - def _create_init_file(self, initfile: str): - with open(initfile, mode="w", encoding="utf-8") as f: - f.write( - """# InvokeAI initialization file -# Put frequently-used startup commands here, one or more per line -# Examples: -# --web --host=0.0.0.0 -# --steps 20 -# -Ak_euler_a -C10.0 -""" - ) - - def _create_arg_parser(self): - """ - This defines all the arguments used on the command line when you launch - the CLI or web backend. - """ - parser = PagingArgumentParser( - 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 = parser.add_argument_group("General") - model_group = parser.add_argument_group("Model selection") - file_group = parser.add_argument_group("Input/output") - web_server_group = parser.add_argument_group("Web server") - render_group = parser.add_argument_group("Rendering") - postprocessing_group = parser.add_argument_group("Postprocessing") - deprecated_group = parser.add_argument_group("Deprecated options") - - deprecated_group.add_argument("--laion400m") - deprecated_group.add_argument("--weights") # deprecated - deprecated_group.add_argument( - "--ckpt_convert", - action=argparse.BooleanOptionalAction, - dest="ckpt_convert", - default=True, - help="Load legacy ckpt files as diffusers (deprecated; always true now).", - ) - - 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, 10), - 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( - "--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="(DEPRECATED - NONFUNCTIONAL). 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( - "--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", - ) - return parser - - # This creates the parser that processes commands on the invoke> command line - def _create_dream_cmd_parser(self): - parser = PagingArgumentParser( - formatter_class=ArgFormatter, - description=""" - *Image generation* - invoke> a fantastic alien landscape -W576 -H512 -s60 -n4 - - *postprocessing* - !fix applies upscaling/facefixing to a previously-generated image. - invoke> !fix 0000045.4829112.png -G1 -U4 -ft codeformer - - *embeddings* - invoke> !triggers -- return all trigger phrases contained in loaded embedding files - - *History manipulation* - !fetch retrieves the command used to generate an earlier image. Provide - a directory wildcard and the name of a file to write and all the commands - used to generate the images in the directory will be written to that file. - invoke> !fetch 0000015.8929913.png - invoke> a fantastic alien landscape -W 576 -H 512 -s 60 -A plms -C 7.5 - invoke> !fetch /path/to/images/*.png prompts.txt - - !replay /path/to/prompts.txt - Replays all the prompts contained in the file prompts.txt. - - !history lists all the commands issued during the current session. - - !NN retrieves the NNth command from the history - - *Model manipulation* - !models -- list models in configs/models.yaml - !switch -- switch to model named - !import_model /path/to/weights/file.ckpt -- adds a .ckpt model to your config - !import_model /path/to/weights/ -- interactively import models from a directory - !import_model http://path_to_model.ckpt -- downloads and adds a .ckpt model to your config - !import_model hakurei/waifu-diffusion -- downloads and adds a diffusers model to your config - !optimize_model -- converts a .ckpt model to a diffusers model - !convert_model /path/to/weights/file.ckpt -- converts a .ckpt file path to a diffusers model - !edit_model -- edit a model's description - !del_model -- delete a model - """, - ) - render_group = parser.add_argument_group("General rendering") - img2img_group = parser.add_argument_group("Image-to-image and inpainting") - inpainting_group = parser.add_argument_group("Inpainting") - outpainting_group = parser.add_argument_group("Outpainting and outcropping") - variation_group = parser.add_argument_group("Creating and combining variations") - postprocessing_group = parser.add_argument_group("Post-processing") - special_effects_group = parser.add_argument_group("Special effects") - deprecated_group = parser.add_argument_group("Deprecated options") - render_group.add_argument( - "--prompt", - default="", - help="prompt string", - ) - render_group.add_argument("-s", "--steps", type=int, help="Number of steps") - render_group.add_argument( - "-S", - "--seed", - type=int, - default=None, - help="Image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc", - ) - render_group.add_argument( - "-n", - "--iterations", - type=int, - default=1, - help="Number of samplings to perform (slower, but will provide seeds for individual images)", - ) - 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", - type=float, - help='Classifier free guidance (CFG) scale - higher numbers cause generator to "try" harder.', - ) - render_group.add_argument( - "--threshold", - default=0.0, - type=float, - help='Latent threshold for classifier free guidance (CFG) - prevent generator from "trying" too hard. Use positive values, 0 disables.', - ) - render_group.add_argument( - "--perlin", - default=0.0, - type=float, - help="Perlin noise scale (0.0 - 1.0) - add perlin noise to the initialization instead of the usual gaussian noise.", - ) - render_group.add_argument( - "--h_symmetry_time_pct", - default=None, - type=float, - help="Horizontal symmetry point (0.0 - 1.0) - apply horizontal symmetry at this point in image generation.", - ) - render_group.add_argument( - "--v_symmetry_time_pct", - default=None, - type=float, - help="Vertical symmetry point (0.0 - 1.0) - apply vertical symmetry at this point in image generation.", - ) - 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( - "--grid", - "-g", - action=argparse.BooleanOptionalAction, - help="generate a grid", - ) - render_group.add_argument( - "-i", - "--individual", - action="store_true", - help="override command-line --grid setting and generate individual images", - ) - render_group.add_argument( - "-x", - "--skip_normalize", - action="store_true", - help="Skip subprompt weight normalization", - ) - render_group.add_argument( - "-A", - "-m", - "--sampler", - dest="sampler_name", - type=str, - choices=SAMPLER_CHOICES, - metavar="SAMPLER_NAME", - help=f'Switch to a different sampler. Supported samplers: {", ".join(SAMPLER_CHOICES)}', - ) - render_group.add_argument( - "-t", - "--log_tokenization", - action="store_true", - help="shows how the prompt is split into tokens", - ) - render_group.add_argument( - "--outdir", - "-o", - type=str, - help="Directory to save generated images and a log of prompts and seeds", - ) - render_group.add_argument( - "--hires_fix", - action="store_true", - dest="hires_fix", - help="Create hires image using img2img to prevent duplicated objects", - ) - render_group.add_argument( - "--save_intermediates", - type=int, - default=0, - dest="save_intermediates", - help='Save every nth intermediate image into an "intermediates" directory within the output directory', - ) - render_group.add_argument( - "--png_compression", - "-z", - type=int, - choices=range(0, 10), - dest="png_compression", - help="level of PNG compression, from 0 (none) to 9 (maximum). [6]", - ) - 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].", - ) - img2img_group.add_argument( - "-I", - "--init_img", - type=str, - help="Path to input image for img2img mode (supersedes width and height)", - ) - img2img_group.add_argument( - "-tm", - "--text_mask", - nargs="+", - type=str, - help='Use the clipseg classifier to generate the mask area for inpainting. Provide a description of the area to mask ("a mug"), optionally followed by the confidence level threshold (0-1.0; defaults to 0.5).', - default=None, - ) - img2img_group.add_argument( - "--init_color", - type=str, - help="Path to reference image for color correction (used for repeated img2img and inpainting)", - ) - img2img_group.add_argument( - "-T", - "-fit", - "--fit", - action="store_true", - help="If specified, will resize the input image to fit within the dimensions of width x height (512x512 default)", - ) - img2img_group.add_argument( - "-f", - "--strength", - type=float, - help="img2img strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely", - ) - inpainting_group.add_argument( - "-M", - "--init_mask", - type=str, - help="Path to input mask for inpainting mode (supersedes width and height)", - ) - inpainting_group.add_argument( - "--invert_mask", - action="store_true", - help="Invert the mask", - ) - inpainting_group.add_argument( - "-r", - "--inpaint_replace", - type=float, - default=0.0, - help="when inpainting, adjust how aggressively to replace the part of the picture under the mask, from 0.0 (a gentle merge) to 1.0 (replace entirely)", - ) - outpainting_group.add_argument( - "-c", - "--outcrop", - nargs="+", - type=str, - metavar=("direction", "pixels"), - help="Outcrop the image with one or more direction/pixel pairs: e.g. -c top 64 bottom 128 left 64 right 64", - ) - outpainting_group.add_argument( - "--force_outpaint", - action="store_true", - default=False, - help="Force outpainting if you have no inpainting mask to pass", - ) - outpainting_group.add_argument( - "--seam_size", - type=int, - default=0, - help="When outpainting, size of the mask around the seam between original and outpainted image", - ) - outpainting_group.add_argument( - "--seam_blur", - type=int, - default=0, - help="When outpainting, the amount to blur the seam inwards", - ) - outpainting_group.add_argument( - "--seam_strength", - type=float, - default=0.7, - help="When outpainting, the img2img strength to use when filling the seam. Values around 0.7 work well", - ) - outpainting_group.add_argument( - "--seam_steps", - type=int, - default=10, - help="When outpainting, the number of steps to use to fill the seam. Low values (~10) work well", - ) - outpainting_group.add_argument( - "--tile_size", - type=int, - default=32, - help="When outpainting, the tile size to use for filling outpaint areas", - ) - postprocessing_group.add_argument( - "--new_prompt", - type=str, - help="Change the text prompt applied during postprocessing (default, use original generation prompt)", - ) - postprocessing_group.add_argument( - "-ft", - "--facetool", - type=str, - default="gfpgan", - help="Select the face restoration AI to use: gfpgan, codeformer", - ) - postprocessing_group.add_argument( - "-G", - "--facetool_strength", - "--gfpgan_strength", - type=float, - help="The strength at which to apply the face restoration to the result.", - default=0.0, - ) - postprocessing_group.add_argument( - "-cf", - "--codeformer_fidelity", - type=float, - help="Used along with CodeFormer. Takes values between 0 and 1. 0 produces high quality but low accuracy. 1 produces high accuracy but low quality.", - default=0.75, - ) - postprocessing_group.add_argument( - "-U", - "--upscale", - nargs="+", - type=float, - help="Scale factor (1, 2, 3, 4, etc..) for upscaling final output followed by upscaling strength (0-1.0). If strength not specified, defaults to 0.75", - default=None, - ) - postprocessing_group.add_argument( - "--save_original", - "-save_orig", - action="store_true", - help="Save original. Use it when upscaling to save both versions.", - ) - postprocessing_group.add_argument( - "--embiggen", - "-embiggen", - nargs="+", - type=float, - help="Arbitrary upscaling using img2img. Provide scale factor (0.75), optionally followed by strength (0.75) and tile overlap proportion (0.25).", - default=None, - ) - postprocessing_group.add_argument( - "--embiggen_tiles", - "-embiggen_tiles", - nargs="+", - type=int, - help="For embiggen, provide list of tiles to process and replace onto the image e.g. `1 3 5`.", - default=None, - ) - postprocessing_group.add_argument( - "--embiggen_strength", - "-embiggen_strength", - type=float, - help="The strength of the embiggen img2img step, defaults to 0.4", - default=None, - ) - special_effects_group.add_argument( - "--seamless", - action="store_true", - help="Change the model to seamless tiling (circular) mode", - ) - special_effects_group.add_argument( - "--seamless_axes", - default=["x", "y"], - type=list[str], - help="Specify which axes to use circular convolution on.", - ) - variation_group.add_argument( - "-v", - "--variation_amount", - default=0.0, - type=float, - help="If > 0, generates variations on the initial seed instead of random seeds per iteration. Must be between 0 and 1. Higher values will be more different.", - ) - variation_group.add_argument( - "-V", - "--with_variations", - default=None, - type=str, - help="list of variations to apply, in the format `seed:weight,seed:weight,...", - ) - render_group.add_argument( - "--use_mps_noise", - action="store_true", - dest="use_mps_noise", - help="Simulate noise on M1 systems to get the same results", - ) - deprecated_group.add_argument( - "-D", - "--out_direction", - nargs="+", - type=str, - metavar=("direction", "pixels"), - help="Older outcropping system. Direction to extend the given image (left|right|top|bottom). If a distance pixel value is not specified it defaults to half the image size", - ) - return parser - - -def format_metadata(**kwargs): - logger.warning("format_metadata() is deprecated. Please use metadata_dumps()") - return metadata_dumps(kwargs) - - -def metadata_dumps(opt, seeds=[], model_hash=None, postprocessing=None): - """ - Given an Args object, returns a dict containing the keys and - structure of the proposed stable diffusion metadata standard - https://github.com/lstein/stable-diffusion/discussions/392 - This is intended to be turned into JSON and stored in the - "sd - """ - - # top-level metadata minus `image` or `images` - metadata = { - "model": "stable diffusion", - "model_id": opt.model, - "model_hash": model_hash, - "app_id": APP_ID, - "app_version": APP_VERSION, - } - - # # add some RFC266 fields that are generated internally, and not as - # # user args - image_dict = opt.to_dict(postprocessing=postprocessing) - - # remove any image keys not mentioned in RFC #266 - rfc266_img_fields = [ - "type", - "postprocessing", - "sampler", - "prompt", - "seed", - "variations", - "steps", - "cfg_scale", - "threshold", - "perlin", - "step_number", - "width", - "height", - "extra", - "strength", - "seamless" "init_img", - "init_mask", - "facetool", - "facetool_strength", - "upscale", - "h_symmetry_time_pct", - "v_symmetry_time_pct", - ] - rfc_dict = {} - - for item in image_dict.items(): - key, value = item - if key in rfc266_img_fields: - rfc_dict[key] = value - - # semantic drift - rfc_dict["sampler"] = image_dict.get("sampler_name", None) - - # display weighted subprompts (liable to change) - if opt.prompt: - subprompts = split_weighted_subprompts(opt.prompt) - subprompts = [{"prompt": x[0], "weight": x[1]} for x in subprompts] - rfc_dict["prompt"] = subprompts - - # 'variations' should always exist and be an array, empty or consisting of {'seed': seed, 'weight': weight} pairs - rfc_dict["variations"] = ( - [{"seed": x[0], "weight": x[1]} for x in opt.with_variations] - if opt.with_variations - else [] - ) - - # if variations are present then we need to replace 'seed' with 'orig_seed' - if hasattr(opt, "first_seed"): - rfc_dict["seed"] = opt.first_seed - - if opt.init_img: - rfc_dict["type"] = "img2img" - rfc_dict["strength_steps"] = rfc_dict.pop("strength") - rfc_dict["orig_hash"] = calculate_init_img_hash(opt.init_img) - rfc_dict["inpaint_replace"] = opt.inpaint_replace - else: - rfc_dict["type"] = "txt2img" - rfc_dict.pop("strength") - - if len(seeds) == 0 and opt.seed: - seeds = [opt.seed] - - if opt.grid: - images = [] - for seed in seeds: - rfc_dict["seed"] = seed - images.append(copy.copy(rfc_dict)) - metadata["images"] = images - else: - # there should only ever be a single seed if we did not generate a grid - assert len(seeds) == 1, "Expected a single seed" - rfc_dict["seed"] = seeds[0] - metadata["image"] = rfc_dict - - return metadata - - -@functools.lru_cache(maxsize=50) -def args_from_png(png_file_path) -> list[Args]: - """ - Given the path to a PNG file created by invoke.py, - retrieves a list of Args objects containing the image - data. - """ - try: - meta = retrieve_metadata(png_file_path) - except AttributeError: - return [legacy_metadata_load({}, png_file_path)] - - try: - return metadata_loads(meta) - except: - return [legacy_metadata_load(meta, png_file_path)] - - -@functools.lru_cache(maxsize=50) -def metadata_from_png(png_file_path) -> Args: - """ - Given the path to a PNG file created by dream.py, retrieves - an Args object containing the image metadata. Note that this - returns a single Args object, not multiple. - """ - args_list = args_from_png(png_file_path) - return args_list[0] if len(args_list) > 0 else Args() # empty args - - -def dream_cmd_from_png(png_file_path): - opt = metadata_from_png(png_file_path) - return opt.dream_prompt_str() - - -def metadata_loads(metadata) -> list: - """ - Takes the dictionary corresponding to RFC266 (https://github.com/lstein/stable-diffusion/issues/266) - and returns a series of opt objects for each of the images described in the dictionary. Note that this - returns a list, and not a single object. See metadata_from_png() for a more convenient function for - files that contain a single image. - """ - results = [] - try: - if "images" in metadata["sd-metadata"]: - images = metadata["sd-metadata"]["images"] - else: - images = [metadata["sd-metadata"]["image"]] - for image in images: - # repack the prompt and variations - if "prompt" in image: - image["prompt"] = repack_prompt(image["prompt"]) - if "variations" in image: - image["variations"] = ",".join( - [ - ":".join([str(x["seed"]), str(x["weight"])]) - for x in image["variations"] - ] - ) - # fix a bit of semantic drift here - image["sampler_name"] = image.pop("sampler") - opt = Args() - opt._cmd_switches = Namespace(**image) - results.append(opt) - except Exception: - import sys - import traceback - - logger.error("Could not read metadata") - print(traceback.format_exc(), file=sys.stderr) - return results - - -def repack_prompt(prompt_list: list) -> str: - # in the common case of no weighting syntax, just return the prompt as is - if len(prompt_list) > 1: - return ",".join( - [":".join([x["prompt"], str(x["weight"])]) for x in prompt_list] - ) - else: - return prompt_list[0]["prompt"] - - -# image can either be a file path on disk or a base64-encoded -# representation of the file's contents -def calculate_init_img_hash(image_string): - prefix = "data:image/png;base64," - hash = None - if image_string.startswith(prefix): - imagebase64 = image_string[len(prefix) :] - imagedata = base64.b64decode(imagebase64) - with open("outputs/test.png", "wb") as file: - file.write(imagedata) - sha = hashlib.sha256() - sha.update(imagedata) - hash = sha.hexdigest() - else: - hash = sha256(image_string) - return hash - - -# Bah. This should be moved somewhere else... -def sha256(path): - sha = hashlib.sha256() - with open(path, "rb") as f: - while True: - data = f.read(65536) - if not data: - break - sha.update(data) - return sha.hexdigest() - - -def legacy_metadata_load(meta, pathname) -> Args: - opt = Args() - if "Dream" in meta and len(meta["Dream"]) > 0: - dream_prompt = meta["Dream"] - opt.parse_cmd(dream_prompt) - else: # if nothing else, we can get the seed - match = re.search("\d+\.(\d+)", pathname) - if match: - seed = match.groups()[0] - opt.seed = seed - else: - opt.prompt = "" - opt.seed = 0 - return opt diff --git a/invokeai/backend/generate.py b/invokeai/backend/generate.py deleted file mode 100644 index 4f3df60f1c..0000000000 --- a/invokeai/backend/generate.py +++ /dev/null @@ -1,1250 +0,0 @@ -# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein) -# Derived from source code carrying the following copyrights -# Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich -# Copyright (c) 2022 Robin Rombach and Patrick Esser and contributors - -import gc -import importlib -import logging -import os -import random -import re -import sys -import time -import traceback -from typing import List - -import cv2 -import diffusers -import numpy as np -import skimage -import torch -import transformers -from PIL import Image, ImageOps -from accelerate.utils import set_seed -from diffusers.pipeline_utils import DiffusionPipeline -from diffusers.utils.import_utils import is_xformers_available -from omegaconf import OmegaConf -from pathlib import Path - -import invokeai.backend.util.logging as logger -from .args import metadata_from_png -from .generator import infill_methods -from .globals import Globals, global_cache_dir -from .image_util import InitImageResizer, PngWriter, Txt2Mask, configure_model_padding -from .model_management import ModelManager -from .safety_checker import SafetyChecker -from .prompting import get_uc_and_c_and_ec -from .prompting.conditioning import log_tokenization -from .stable_diffusion import HuggingFaceConceptsLibrary -from .util import choose_precision, choose_torch_device - -def fix_func(orig): - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - - def new_func(*args, **kw): - device = kw.get("device", "mps") - kw["device"] = "cpu" - return orig(*args, **kw).to(device) - - return new_func - return orig - - -torch.rand = fix_func(torch.rand) -torch.rand_like = fix_func(torch.rand_like) -torch.randn = fix_func(torch.randn) -torch.randn_like = fix_func(torch.randn_like) -torch.randint = fix_func(torch.randint) -torch.randint_like = fix_func(torch.randint_like) -torch.bernoulli = fix_func(torch.bernoulli) -torch.multinomial = fix_func(torch.multinomial) - -# this is fallback model in case no default is defined -FALLBACK_MODEL_NAME = "stable-diffusion-1.5" - -"""Simplified text to image API for stable diffusion/latent diffusion - -Example Usage: - -from ldm.generate import Generate - -# Create an object with default values -gr = Generate('stable-diffusion-1.4') - -# do the slow model initialization -gr.load_model() - -# Do the fast inference & image generation. Any options passed here -# override the default values assigned during class initialization -# Will call load_model() if the model was not previously loaded and so -# may be slow at first. -# The method returns a list of images. Each row of the list is a sub-list of [filename,seed] -results = gr.prompt2png(prompt = "an astronaut riding a horse", - outdir = "./outputs/samples", - iterations = 3) - -for row in results: - print(f'filename={row[0]}') - print(f'seed ={row[1]}') - -# Same thing, but using an initial image. -results = gr.prompt2png(prompt = "an astronaut riding a horse", - outdir = "./outputs/, - iterations = 3, - init_img = "./sketches/horse+rider.png") - -for row in results: - print(f'filename={row[0]}') - print(f'seed ={row[1]}') - -# Same thing, but we return a series of Image objects, which lets you manipulate them, -# combine them, and save them under arbitrary names - -results = gr.prompt2image(prompt = "an astronaut riding a horse" - outdir = "./outputs/") -for row in results: - im = row[0] - seed = row[1] - im.save(f'./outputs/samples/an_astronaut_riding_a_horse-{seed}.png') - im.thumbnail(100,100).save('./outputs/samples/astronaut_thumb.jpg') - -Note that the old txt2img() and img2img() calls are deprecated but will -still work. - -The full list of arguments to Generate() are: -gr = Generate( - # these values are set once and shouldn't be changed - conf:str = path to configuration file ('configs/models.yaml') - model:str = symbolic name of the model in the configuration file - precision:float = float precision to be used - safety_checker:bool = activate safety checker [False] - - # this value is sticky and maintained between generation calls - sampler_name:str = ['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_dpmpp_2', 'k_dpmpp_2_a', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms'] // k_lms - - # these are deprecated - use conf and model instead - weights = path to model weights ('models/ldm/stable-diffusion-v1/model.ckpt') - config = path to model configuration ('configs/stable-diffusion/v1-inference.yaml') - ) - -""" - - -class Generate: - """Generate class - Stores default values for multiple configuration items - """ - - def __init__( - self, - model=None, - conf="configs/models.yaml", - embedding_path=None, - sampler_name="k_lms", - ddim_eta=0.0, # deterministic - full_precision=False, - precision="auto", - outdir="outputs/img-samples", - gfpgan=None, - codeformer=None, - esrgan=None, - free_gpu_mem: bool = False, - safety_checker: bool = False, - max_loaded_models: int = 2, - # these are deprecated; if present they override values in the conf file - weights=None, - config=None, - ): - mconfig = OmegaConf.load(conf) - self.height = None - self.width = None - self.model_manager = None - self.iterations = 1 - self.steps = 50 - self.cfg_scale = 7.5 - self.sampler_name = sampler_name - self.ddim_eta = ddim_eta # same seed always produces same image - self.precision = precision - self.strength = 0.75 - self.seamless = False - self.seamless_axes = {"x", "y"} - self.hires_fix = False - self.embedding_path = embedding_path - self.model = None # empty for now - self.model_hash = None - self.sampler = None - self.device = None - self.max_memory_allocated = 0 - self.memory_allocated = 0 - self.session_peakmem = 0 - self.base_generator = None - self.seed = None - self.outdir = outdir - self.gfpgan = gfpgan - self.codeformer = codeformer - self.esrgan = esrgan - self.free_gpu_mem = free_gpu_mem - self.max_loaded_models = (max_loaded_models,) - self.size_matters = True # used to warn once about large image sizes and VRAM - self.txt2mask = None - self.safety_checker = None - self.karras_max = None - self.infill_method = None - - # Note that in previous versions, there was an option to pass the - # device to Generate(). However the device was then ignored, so - # it wasn't actually doing anything. This logic could be reinstated. - self.device = torch.device(choose_torch_device()) - logger.info(f"Using device_type {self.device.type}") - if full_precision: - if self.precision != "auto": - raise ValueError("Remove --full_precision / -F if using --precision") - logger.warning("Please remove deprecated --full_precision / -F") - logger.warning("If auto config does not work you can use --precision=float32") - self.precision = "float32" - if self.precision == "auto": - self.precision = choose_precision(self.device) - Globals.full_precision = self.precision == "float32" - - if is_xformers_available(): - if torch.cuda.is_available() and not Globals.disable_xformers: - logger.info("xformers memory-efficient attention is available and enabled") - else: - logger.info( - "xformers memory-efficient attention is available but disabled" - ) - else: - logger.info("xformers not installed") - - # model caching system for fast switching - self.model_manager = ModelManager( - mconfig, - self.device, - self.precision, - max_loaded_models=max_loaded_models, - sequential_offload=self.free_gpu_mem, - embedding_path=Path(self.embedding_path), - ) - # don't accept invalid models - fallback = self.model_manager.default_model() or FALLBACK_MODEL_NAME - model = model or fallback - if not self.model_manager.valid_model(model): - logger.warning( - f'"{model}" is not a known model name; falling back to {fallback}.' - ) - model = None - self.model_name = model or fallback - - # for VRAM usage statistics - self.session_peakmem = ( - torch.cuda.max_memory_allocated(self.device) if self._has_cuda else None - ) - transformers.logging.set_verbosity_error() - - # gets rid of annoying messages about random seed - logging.getLogger("pytorch_lightning").setLevel(logging.ERROR) - - # load safety checker if requested - if safety_checker: - logger.info("Initializing NSFW checker") - self.safety_checker = SafetyChecker(self.device) - else: - logger.info("NSFW checker is disabled") - - def prompt2png(self, prompt, outdir, **kwargs): - """ - Takes a prompt and an output directory, writes out the requested number - of PNG files, and returns an array of [[filename,seed],[filename,seed]...] - Optional named arguments are the same as those passed to Generate and prompt2image() - """ - results = self.prompt2image(prompt, **kwargs) - pngwriter = PngWriter(outdir) - prefix = pngwriter.unique_prefix() - outputs = [] - for image, seed in results: - name = f"{prefix}.{seed}.png" - path = pngwriter.save_image_and_prompt_to_png( - image, dream_prompt=f"{prompt} -S{seed}", name=name - ) - outputs.append([path, seed]) - return outputs - - def txt2img(self, prompt, **kwargs): - outdir = kwargs.pop("outdir", self.outdir) - return self.prompt2png(prompt, outdir, **kwargs) - - def img2img(self, prompt, **kwargs): - outdir = kwargs.pop("outdir", self.outdir) - assert ( - "init_img" in kwargs - ), "call to img2img() must include the init_img argument" - return self.prompt2png(prompt, outdir, **kwargs) - - def prompt2image( - self, - # these are common - prompt, - iterations=None, - steps=None, - seed=None, - cfg_scale=None, - ddim_eta=None, - skip_normalize=False, - image_callback=None, - step_callback=None, - width=None, - height=None, - sampler_name=None, - seamless=False, - seamless_axes={"x", "y"}, - log_tokenization=False, - with_variations=None, - variation_amount=0.0, - threshold=0.0, - perlin=0.0, - h_symmetry_time_pct=None, - v_symmetry_time_pct=None, - karras_max=None, - outdir=None, - # these are specific to img2img and inpaint - init_img=None, - init_mask=None, - text_mask=None, - invert_mask=False, - fit=False, - strength=None, - init_color=None, - # these are specific to embiggen (which also relies on img2img args) - embiggen=None, - embiggen_tiles=None, - embiggen_strength=None, - # these are specific to GFPGAN/ESRGAN - gfpgan_strength=0, - facetool=None, - facetool_strength=0, - codeformer_fidelity=None, - save_original=False, - upscale=None, - upscale_denoise_str=0.75, - # this is specific to inpainting and causes more extreme inpainting - inpaint_replace=0.0, - # This controls the size at which inpaint occurs (scaled up for inpaint, then back down for the result) - inpaint_width=None, - inpaint_height=None, - # This will help match inpainted areas to the original image more smoothly - mask_blur_radius: int = 8, - # Set this True to handle KeyboardInterrupt internally - catch_interrupts=False, - hires_fix=False, - use_mps_noise=False, - # Seam settings for outpainting - seam_size: int = 0, - seam_blur: int = 0, - seam_strength: float = 0.7, - seam_steps: int = 10, - tile_size: int = 32, - infill_method=None, - force_outpaint: bool = False, - enable_image_debugging=False, - **args, - ): # eat up additional cruft - self.clear_cuda_stats() - """ - ldm.generate.prompt2image() is the common entry point for txt2img() and img2img() - It takes the following arguments: - prompt // prompt string (no default) - iterations // iterations (1); image count=iterations - steps // refinement steps per iteration - seed // seed for random number generator - width // width of image, in multiples of 64 (512) - height // height of image, in multiples of 64 (512) - cfg_scale // how strongly the prompt influences the image (7.5) (must be >1) - seamless // whether the generated image should tile - hires_fix // whether the Hires Fix should be applied during generation - init_img // path to an initial image - init_mask // path to a mask for the initial image - text_mask // a text string that will be used to guide clipseg generation of the init_mask - invert_mask // boolean, if true invert the mask - strength // strength for noising/unnoising init_img. 0.0 preserves image exactly, 1.0 replaces it completely - facetool_strength // strength for GFPGAN/CodeFormer. 0.0 preserves image exactly, 1.0 replaces it completely - ddim_eta // image randomness (eta=0.0 means the same seed always produces the same image) - step_callback // a function or method that will be called each step - image_callback // a function or method that will be called each time an image is generated - with_variations // a weighted list [(seed_1, weight_1), (seed_2, weight_2), ...] of variations which should be applied before doing any generation - variation_amount // optional 0-1 value to slerp from -S noise to random noise (allows variations on an image) - threshold // optional value >=0 to add thresholding to latent values for k-diffusion samplers (0 disables) - perlin // optional 0-1 value to add a percentage of perlin noise to the initial noise - h_symmetry_time_pct // optional 0-1 value that indicates the time at which horizontal symmetry is applied - v_symmetry_time_pct // optional 0-1 value that indicates the time at which vertical symmetry is applied - embiggen // scale factor relative to the size of the --init_img (-I), followed by ESRGAN upscaling strength (0-1.0), followed by minimum amount of overlap between tiles as a decimal ratio (0 - 1.0) or number of pixels - embiggen_tiles // list of tiles by number in order to process and replace onto the image e.g. `0 2 4` - embiggen_strength // strength for embiggen. 0.0 preserves image exactly, 1.0 replaces it completely - - To use the step callback, define a function that receives two arguments: - - Image GPU data - - The step number - - To use the image callback, define a function of method that receives two arguments, an Image object - and the seed. You can then do whatever you like with the image, including converting it to - different formats and manipulating it. For example: - - def process_image(image,seed): - image.save(f{'images/seed.png'}) - - The code used to save images to a directory can be found in ldm/invoke/pngwriter.py. - It contains code to create the requested output directory, select a unique informative - name for each image, and write the prompt into the PNG metadata. - """ - # TODO: convert this into a getattr() loop - steps = steps or self.steps - width = width or self.width - height = height or self.height - seamless = seamless or self.seamless - seamless_axes = seamless_axes or self.seamless_axes - hires_fix = hires_fix or self.hires_fix - cfg_scale = cfg_scale or self.cfg_scale - ddim_eta = ddim_eta or self.ddim_eta - iterations = iterations or self.iterations - strength = strength or self.strength - outdir = outdir or self.outdir - self.seed = seed - self.log_tokenization = log_tokenization - self.step_callback = step_callback - self.karras_max = karras_max - self.infill_method = ( - infill_method or infill_methods()[0], - ) # The infill method to use - with_variations = [] if with_variations is None else with_variations - - # will instantiate the model or return it from cache - model = self.set_model(self.model_name) - - # self.width and self.height are set by set_model() - # to the width and height of the image training set - width = width or self.width - height = height or self.height - - if isinstance(model, DiffusionPipeline): - configure_model_padding(model.unet, seamless, seamless_axes) - configure_model_padding(model.vae, seamless, seamless_axes) - else: - configure_model_padding(model, seamless, seamless_axes) - - assert cfg_scale > 1.0, "CFG_Scale (-C) must be >1.0" - assert threshold >= 0.0, "--threshold must be >=0.0" - assert ( - 0.0 < strength <= 1.0 - ), "img2img and inpaint strength can only work with 0.0 < strength < 1.0" - assert ( - 0.0 <= variation_amount <= 1.0 - ), "-v --variation_amount must be in [0.0, 1.0]" - assert 0.0 <= perlin <= 1.0, "--perlin must be in [0.0, 1.0]" - assert (embiggen == None and embiggen_tiles == None) or ( - (embiggen != None or embiggen_tiles != None) and init_img != None - ), "Embiggen requires an init/input image to be specified" - - if len(with_variations) > 0 or variation_amount > 1.0: - assert seed is not None, "seed must be specified when using with_variations" - if variation_amount == 0.0: - assert ( - iterations == 1 - ), "when using --with_variations, multiple iterations are only possible when using --variation_amount" - assert all( - 0 <= weight <= 1 for _, weight in with_variations - ), f"variation weights must be in [0.0, 1.0]: got {[weight for _, weight in with_variations]}" - - width, height, _ = self._resolution_check(width, height, log=True) - assert ( - inpaint_replace >= 0.0 and inpaint_replace <= 1.0 - ), "inpaint_replace must be between 0.0 and 1.0" - - if sampler_name and (sampler_name != self.sampler_name): - self.sampler_name = sampler_name - self._set_scheduler() - - # apply the concepts library to the prompt - prompt = self.huggingface_concepts_library.replace_concepts_with_triggers( - prompt, - lambda concepts: self.load_huggingface_concepts(concepts), - self.model.textual_inversion_manager.get_all_trigger_strings(), - ) - - tic = time.time() - if self._has_cuda(): - torch.cuda.reset_peak_memory_stats() - - results = list() - - try: - uc, c, extra_conditioning_info = get_uc_and_c_and_ec( - prompt, - model=self.model, - skip_normalize_legacy_blend=skip_normalize, - log_tokens=self.log_tokenization, - ) - - init_image, mask_image = self._make_images( - init_img, - init_mask, - width, - height, - fit=fit, - text_mask=text_mask, - invert_mask=invert_mask, - force_outpaint=force_outpaint, - ) - - # TODO: Hacky selection of operation to perform. Needs to be refactored. - generator = self.select_generator( - init_image, mask_image, embiggen, hires_fix, force_outpaint - ) - - generator.set_variation(self.seed, variation_amount, with_variations) - generator.use_mps_noise = use_mps_noise - - results = generator.generate( - prompt, - iterations=iterations, - seed=self.seed, - sampler=self.sampler, - steps=steps, - cfg_scale=cfg_scale, - conditioning=(uc, c, extra_conditioning_info), - ddim_eta=ddim_eta, - image_callback=image_callback, # called after the final image is generated - step_callback=step_callback, # called after each intermediate image is generated - width=width, - height=height, - init_img=init_img, # embiggen needs to manipulate from the unmodified init_img - init_image=init_image, # notice that init_image is different from init_img - mask_image=mask_image, - strength=strength, - threshold=threshold, - perlin=perlin, - h_symmetry_time_pct=h_symmetry_time_pct, - v_symmetry_time_pct=v_symmetry_time_pct, - embiggen=embiggen, - embiggen_tiles=embiggen_tiles, - embiggen_strength=embiggen_strength, - inpaint_replace=inpaint_replace, - mask_blur_radius=mask_blur_radius, - safety_checker=self.safety_checker, - seam_size=seam_size, - seam_blur=seam_blur, - seam_strength=seam_strength, - seam_steps=seam_steps, - tile_size=tile_size, - infill_method=infill_method, - force_outpaint=force_outpaint, - inpaint_height=inpaint_height, - inpaint_width=inpaint_width, - enable_image_debugging=enable_image_debugging, - free_gpu_mem=self.free_gpu_mem, - clear_cuda_cache=self.clear_cuda_cache, - ) - - if init_color: - self.correct_colors( - image_list=results, - reference_image_path=init_color, - image_callback=image_callback, - ) - - if upscale is not None or facetool_strength > 0: - self.upscale_and_reconstruct( - results, - upscale=upscale, - upscale_denoise_str=upscale_denoise_str, - facetool=facetool, - strength=facetool_strength, - codeformer_fidelity=codeformer_fidelity, - save_original=save_original, - image_callback=image_callback, - ) - - except KeyboardInterrupt: - # Clear the CUDA cache on an exception - self.clear_cuda_cache() - - if catch_interrupts: - logger.warning("Interrupted** Partial results will be returned.") - else: - raise KeyboardInterrupt - except RuntimeError: - # Clear the CUDA cache on an exception - self.clear_cuda_cache() - - print(traceback.format_exc(), file=sys.stderr) - logger.info("Could not generate image.") - - toc = time.time() - logger.info("Usage stats:") - logger.info(f"{len(results)} image(s) generated in "+"%4.2fs" % (toc - tic)) - self.print_cuda_stats() - return results - - def gather_cuda_stats(self): - if self._has_cuda(): - self.max_memory_allocated = max( - self.max_memory_allocated, torch.cuda.max_memory_allocated(self.device) - ) - self.memory_allocated = max( - self.memory_allocated, torch.cuda.memory_allocated(self.device) - ) - self.session_peakmem = max( - self.session_peakmem, torch.cuda.max_memory_allocated(self.device) - ) - - def clear_cuda_cache(self): - if self._has_cuda(): - self.gather_cuda_stats() - # Run garbage collection prior to emptying the CUDA cache - gc.collect() - torch.cuda.empty_cache() - - def clear_cuda_stats(self): - self.max_memory_allocated = 0 - self.memory_allocated = 0 - - def print_cuda_stats(self): - if self._has_cuda(): - self.gather_cuda_stats() - logger.info( - "Max VRAM used for this generation: "+ - "%4.2fG. " % (self.max_memory_allocated / 1e9)+ - "Current VRAM utilization: "+ - "%4.2fG" % (self.memory_allocated / 1e9) - ) - - logger.info( - "Max VRAM used since script start: " + - "%4.2fG" % (self.session_peakmem / 1e9) - ) - - # this needs to be generalized to all sorts of postprocessors, which should be wrapped - # in a nice harmonized call signature. For now we have a bunch of if/elses! - def apply_postprocessor( - self, - image_path, - tool="gfpgan", # one of 'upscale', 'gfpgan', 'codeformer', 'outpaint', or 'embiggen' - facetool_strength=0.0, - codeformer_fidelity=0.75, - upscale=None, - upscale_denoise_str=0.75, - out_direction=None, - outcrop=[], - save_original=True, # to get new name - callback=None, - opt=None, - ): - # retrieve the seed from the image; - seed = None - prompt = None - - args = metadata_from_png(image_path) - seed = opt.seed or args.seed - if seed is None or seed < 0: - seed = random.randrange(0, np.iinfo(np.uint32).max) - - prompt = opt.prompt or args.prompt or "" - logger.info(f'using seed {seed} and prompt "{prompt}" for {image_path}') - - # try to reuse the same filename prefix as the original file. - # we take everything up to the first period - prefix = None - m = re.match(r"^([^.]+)\.", os.path.basename(image_path)) - if m: - prefix = m.groups()[0] - - # face fixers and esrgan take an Image, but embiggen takes a path - image = Image.open(image_path) - - # used by multiple postfixers - # todo: cross-attention control - uc, c, extra_conditioning_info = get_uc_and_c_and_ec( - prompt, - model=self.model, - skip_normalize_legacy_blend=opt.skip_normalize, - log_tokens=log_tokenization, - ) - - if tool in ("gfpgan", "codeformer", "upscale"): - if tool == "gfpgan": - facetool = "gfpgan" - elif tool == "codeformer": - facetool = "codeformer" - elif tool == "upscale": - facetool = "gfpgan" # but won't be run - facetool_strength = 0 - return self.upscale_and_reconstruct( - [[image, seed]], - facetool=facetool, - strength=facetool_strength, - codeformer_fidelity=codeformer_fidelity, - save_original=save_original, - upscale=upscale, - upscale_denoise_str=upscale_denoise_str, - image_callback=callback, - prefix=prefix, - ) - - elif tool == "outcrop": - from .restoration.outcrop import Outcrop - - extend_instructions = {} - for direction, pixels in _pairwise(opt.outcrop): - try: - extend_instructions[direction] = int(pixels) - except ValueError: - logger.warning( - 'invalid extension instruction. Use ..., as in "top 64 left 128 right 64 bottom 64"' - ) - - opt.seed = seed - opt.prompt = prompt - - if len(extend_instructions) > 0: - restorer = Outcrop( - image, - self, - ) - return restorer.process( - extend_instructions, - opt=opt, - orig_opt=args, - image_callback=callback, - prefix=prefix, - ) - - elif tool == "embiggen": - # fetch the metadata from the image - generator = self.select_generator(embiggen=True) - opt.strength = opt.embiggen_strength or 0.40 - logger.info( - f"Setting img2img strength to {opt.strength} for happy embiggening" - ) - generator.generate( - prompt, - sampler=self.sampler, - steps=opt.steps, - cfg_scale=opt.cfg_scale, - ddim_eta=self.ddim_eta, - conditioning=(uc, c, extra_conditioning_info), - init_img=image_path, # not the Image! (sigh) - init_image=image, # embiggen wants both! (sigh) - strength=opt.strength, - width=opt.width, - height=opt.height, - embiggen=opt.embiggen, - embiggen_tiles=opt.embiggen_tiles, - embiggen_strength=opt.embiggen_strength, - image_callback=callback, - clear_cuda_cache=self.clear_cuda_cache, - ) - elif tool == "outpaint": - from .restoration.outpaint import Outpaint - - restorer = Outpaint(image, self) - return restorer.process(opt, args, image_callback=callback, prefix=prefix) - - elif tool is None: - logger.warning( - "please provide at least one postprocessing option, such as -G or -U" - ) - return None - else: - logger.warning(f"postprocessing tool {tool} is not yet supported") - return None - - def select_generator( - self, - init_image: Image.Image = None, - mask_image: Image.Image = None, - embiggen: bool = False, - hires_fix: bool = False, - force_outpaint: bool = False, - ): - if hires_fix: - return self._make_txt2img2img() - - if embiggen is not None: - return self._make_embiggen() - - if ((init_image is not None) and (mask_image is not None)) or force_outpaint: - return self._make_inpaint() - - if init_image is not None: - return self._make_img2img() - - return self._make_txt2img() - - def _make_images( - self, - img, - mask, - width, - height, - fit=False, - text_mask=None, - invert_mask=False, - force_outpaint=False, - ): - init_image = None - init_mask = None - if not img: - return None, None - - image = self._load_img(img) - - if image.width < self.width and image.height < self.height: - logger.warning( - f"img2img and inpainting may produce unexpected results with initial images smaller than {self.width}x{self.height} in both dimensions" - ) - - # if image has a transparent area and no mask was provided, then try to generate mask - if self._has_transparency(image): - self._transparency_check_and_warning(image, mask, force_outpaint) - init_mask = self._create_init_mask(image, width, height, fit=fit) - - if (image.width * image.height) > ( - self.width * self.height - ) and self.size_matters: - logger.info( - "This input is larger than your defaults. If you run out of memory, please use a smaller image." - ) - self.size_matters = False - - init_image = self._create_init_image(image, width, height, fit=fit) - - if mask: - mask_image = self._load_img(mask) - init_mask = self._create_init_mask(mask_image, width, height, fit=fit) - - elif text_mask: - init_mask = self._txt2mask(image, text_mask, width, height, fit=fit) - - if init_mask and invert_mask: - init_mask = ImageOps.invert(init_mask) - - return init_image, init_mask - - def _make_base(self): - return self._load_generator("", "Generator") - - def _make_txt2img(self): - return self._load_generator(".txt2img", "Txt2Img") - - def _make_img2img(self): - return self._load_generator(".img2img", "Img2Img") - - def _make_embiggen(self): - return self._load_generator(".embiggen", "Embiggen") - - def _make_txt2img2img(self): - return self._load_generator(".txt2img2img", "Txt2Img2Img") - - def _make_inpaint(self): - return self._load_generator(".inpaint", "Inpaint") - - def _load_generator(self, module, class_name): - mn = f"invokeai.backend.generator{module}" - cn = class_name - module = importlib.import_module(mn) - constructor = getattr(module, cn) - return constructor(self.model, self.precision) - - def load_model(self): - """ - preload model identified in self.model_name - """ - return self.set_model(self.model_name) - - def set_model(self, model_name): - """ - Given the name of a model defined in models.yaml, will load and initialize it - and return the model object. Previously-used models will be cached. - - If the passed model_name is invalid, raises a KeyError. - If the model fails to load for some reason, will attempt to load the previously- - loaded model (if any). If that fallback fails, will raise an AssertionError - """ - if self.model_name == model_name and self.model is not None: - return self.model - - previous_model_name = self.model_name - - # the model cache does the loading and offloading - cache = self.model_manager - if not cache.valid_model(model_name): - raise KeyError( - f'** "{model_name}" is not a known model name. Cannot change.' - ) - - cache.print_vram_usage() - - # have to get rid of all references to model in order - # to free it from GPU memory - self.model = None - self.sampler = None - self.generators = {} - gc.collect() - try: - model_data = cache.get_model(model_name) - except Exception as e: - logger.warning(f"model {model_name} could not be loaded: {str(e)}") - print(traceback.format_exc(), file=sys.stderr) - if previous_model_name is None: - raise e - logger.warning("trying to reload previous model") - model_data = cache.get_model(previous_model_name) # load previous - if model_data is None: - raise e - model_name = previous_model_name - - self.model = model_data["model"] - self.width = model_data["width"] - self.height = model_data["height"] - self.model_hash = model_data["hash"] - - # uncache generators so they pick up new models - self.generators = {} - - set_seed(random.randrange(0, np.iinfo(np.uint32).max)) - self.model_name = model_name - self._set_scheduler() # requires self.model_name to be set first - return self.model - - def load_huggingface_concepts(self, concepts: list[str]): - self.model.textual_inversion_manager.load_huggingface_concepts(concepts) - - @property - def huggingface_concepts_library(self) -> HuggingFaceConceptsLibrary: - return self.model.textual_inversion_manager.hf_concepts_library - - @property - def embedding_trigger_strings(self) -> List[str]: - return self.model.textual_inversion_manager.get_all_trigger_strings() - - def correct_colors(self, image_list, reference_image_path, image_callback=None): - reference_image = Image.open(reference_image_path) - correction_target = cv2.cvtColor(np.asarray(reference_image), cv2.COLOR_RGB2LAB) - for r in image_list: - image, seed = r - image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2LAB) - image = skimage.exposure.match_histograms( - image, correction_target, channel_axis=2 - ) - image = Image.fromarray( - cv2.cvtColor(image, cv2.COLOR_LAB2RGB).astype("uint8") - ) - if image_callback is not None: - image_callback(image, seed) - else: - r[0] = image - - def upscale_and_reconstruct( - self, - image_list, - facetool="gfpgan", - upscale=None, - upscale_denoise_str=0.75, - strength=0.0, - codeformer_fidelity=0.75, - save_original=False, - image_callback=None, - prefix=None, - ): - results = [] - for r in image_list: - image, seed, _ = r - try: - if strength > 0: - if self.gfpgan is not None or self.codeformer is not None: - if facetool == "gfpgan": - if self.gfpgan is None: - logger.info( - "GFPGAN not found. Face restoration is disabled." - ) - else: - image = self.gfpgan.process(image, strength, seed) - if facetool == "codeformer": - if self.codeformer is None: - logger.info( - "CodeFormer not found. Face restoration is disabled." - ) - else: - cf_device = ( - "cpu" if str(self.device) == "mps" else self.device - ) - image = self.codeformer.process( - image=image, - strength=strength, - device=cf_device, - seed=seed, - fidelity=codeformer_fidelity, - ) - else: - logger.info("Face Restoration is disabled.") - if upscale is not None: - if self.esrgan is not None: - if len(upscale) < 2: - upscale.append(0.75) - image = self.esrgan.process( - image, - upscale[1], - seed, - int(upscale[0]), - denoise_str=upscale_denoise_str, - ) - else: - logger.info("ESRGAN is disabled. Image not upscaled.") - except Exception as e: - logger.info( - f"Error running RealESRGAN or GFPGAN. Your image was not upscaled.\n{e}" - ) - - if image_callback is not None: - image_callback(image, seed, upscaled=True, use_prefix=prefix) - else: - r[0] = image - - results.append([image, seed]) - - return results - - def apply_textmask( - self, image_path: str, prompt: str, callback, threshold: float = 0.5 - ): - assert os.path.exists( - image_path - ), f'** "{image_path}" not found. Please enter the name of an existing image file to mask **' - basename, _ = os.path.splitext(os.path.basename(image_path)) - if self.txt2mask is None: - self.txt2mask = Txt2Mask(device=self.device, refined=True) - segmented = self.txt2mask.segment(image_path, prompt) - trans = segmented.to_transparent() - inverse = segmented.to_transparent(invert=True) - mask = segmented.to_mask(threshold) - - path_filter = re.compile(r'[<>:"/\\|?*]') - safe_prompt = path_filter.sub("_", prompt)[:50].rstrip(" .") - - callback(trans, f"{safe_prompt}.deselected", use_prefix=basename) - callback(inverse, f"{safe_prompt}.selected", use_prefix=basename) - callback(mask, f"{safe_prompt}.masked", use_prefix=basename) - - # to help WebGUI - front end to generator util function - def sample_to_image(self, samples): - return self._make_base().sample_to_image(samples) - - def sample_to_lowres_estimated_image(self, samples): - return self._make_base().sample_to_lowres_estimated_image(samples) - - def is_legacy_model(self, model_name) -> bool: - return self.model_manager.is_legacy(model_name) - - def _set_scheduler(self): - default = self.model.scheduler - - # See https://github.com/huggingface/diffusers/issues/277#issuecomment-1371428672 - scheduler_map = dict( - ddim=diffusers.DDIMScheduler, - dpmpp_2=diffusers.DPMSolverMultistepScheduler, - k_dpm_2=diffusers.KDPM2DiscreteScheduler, - k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler, - # DPMSolverMultistepScheduler is technically not `k_` anything, as it is neither - # the k-diffusers implementation nor included in EDM (Karras 2022), but we can - # provide an alias for compatibility. - k_dpmpp_2=diffusers.DPMSolverMultistepScheduler, - k_euler=diffusers.EulerDiscreteScheduler, - k_euler_a=diffusers.EulerAncestralDiscreteScheduler, - k_heun=diffusers.HeunDiscreteScheduler, - k_lms=diffusers.LMSDiscreteScheduler, - plms=diffusers.PNDMScheduler, - ) - - if self.sampler_name in scheduler_map: - sampler_class = scheduler_map[self.sampler_name] - msg = ( - f"Setting Sampler to {self.sampler_name} ({sampler_class.__name__})" - ) - self.sampler = sampler_class.from_config(self.model.scheduler.config) - else: - msg = ( - f" Unsupported Sampler: {self.sampler_name} "+ - f"Defaulting to {default}" - ) - self.sampler = default - - logger.info(msg) - - if not hasattr(self.sampler, "uses_inpainting_model"): - # FIXME: terrible kludge! - self.sampler.uses_inpainting_model = lambda: False - - def _load_img(self, img) -> Image: - if isinstance(img, Image.Image): - image = img - logger.info(f"using provided input image of size {image.width}x{image.height}") - elif isinstance(img, str): - assert os.path.exists(img), f"{img}: File not found" - - image = Image.open(img) - logger.info( - f"loaded input image of size {image.width}x{image.height} from {img}" - ) - else: - image = Image.open(img) - logger.info(f"loaded input image of size {image.width}x{image.height}") - image = ImageOps.exif_transpose(image) - return image - - def _create_init_image(self, image: Image.Image, width, height, fit=True): - if image.mode != "RGBA": - image = image.convert("RGBA") - image = ( - self._fit_image(image, (width, height)) - if fit - else self._squeeze_image(image) - ) - return image - - def _create_init_mask(self, image, width, height, fit=True): - # convert into a black/white mask - image = self._image_to_mask(image) - image = image.convert("RGB") - image = ( - self._fit_image(image, (width, height)) - if fit - else self._squeeze_image(image) - ) - return image - - # The mask is expected to have the region to be inpainted - # with alpha transparency. It converts it into a black/white - # image with the transparent part black. - def _image_to_mask(self, mask_image: Image.Image, invert=False) -> Image: - # Obtain the mask from the transparency channel - if mask_image.mode == "L": - mask = mask_image - elif mask_image.mode in ("RGB", "P"): - mask = mask_image.convert("L") - else: - # Obtain the mask from the transparency channel - mask = Image.new(mode="L", size=mask_image.size, color=255) - mask.putdata(mask_image.getdata(band=3)) - if invert: - mask = ImageOps.invert(mask) - return mask - - def _txt2mask( - self, image: Image, text_mask: list, width, height, fit=True - ) -> Image: - prompt = text_mask[0] - confidence_level = text_mask[1] if len(text_mask) > 1 else 0.5 - if self.txt2mask is None: - self.txt2mask = Txt2Mask(device=self.device) - - segmented = self.txt2mask.segment(image, prompt) - mask = segmented.to_mask(float(confidence_level)) - mask = mask.convert("RGB") - mask = ( - self._fit_image(mask, (width, height)) if fit else self._squeeze_image(mask) - ) - return mask - - def _has_transparency(self, image): - if image.info.get("transparency", None) is not None: - return True - if image.mode == "P": - transparent = image.info.get("transparency", -1) - for _, index in image.getcolors(): - if index == transparent: - return True - elif image.mode == "RGBA": - extrema = image.getextrema() - if extrema[3][0] < 255: - return True - return False - - def _check_for_erasure(self, image: Image.Image) -> bool: - if image.mode not in ("RGBA", "RGB"): - return False - width, height = image.size - pixdata = image.load() - colored = 0 - for y in range(height): - for x in range(width): - if pixdata[x, y][3] == 0: - r, g, b, _ = pixdata[x, y] - if (r, g, b) != (0, 0, 0) and (r, g, b) != (255, 255, 255): - colored += 1 - return colored == 0 - - def _transparency_check_and_warning(self, image, mask, force_outpaint=False): - if not mask: - logger.info( - "Initial image has transparent areas. Will inpaint in these regions." - ) - if (not force_outpaint) and self._check_for_erasure(image): - logger.info( - "Colors underneath the transparent region seem to have been erased.\n" + - "Inpainting will be suboptimal. Please preserve the colors when making\n" + - "a transparency mask, or provide mask explicitly using --init_mask (-M)." - ) - - def _squeeze_image(self, image): - x, y, resize_needed = self._resolution_check(image.width, image.height) - if resize_needed: - return InitImageResizer(image).resize(x, y) - return image - - def _fit_image(self, image, max_dimensions): - w, h = max_dimensions - logger.info(f"image will be resized to fit inside a box {w}x{h} in size.") - # note that InitImageResizer does the multiple of 64 truncation internally - image = InitImageResizer(image).resize(width=w, height=h) - logger.info( - f"after adjusting image dimensions to be multiples of 64, init image is {image.width}x{image.height}" - ) - return image - - def _resolution_check(self, width, height, log=False): - resize_needed = False - w, h = map( - lambda x: x - x % 64, (width, height) - ) # resize to integer multiple of 64 - if h != height or w != width: - if log: - logger.info( - f"Provided width and height must be multiples of 64. Auto-resizing to {w}x{h}" - ) - height = h - width = w - resize_needed = True - return width, height, resize_needed - - def _has_cuda(self): - return self.device.type == "cuda" - - def write_intermediate_images(self, modulus, path): - counter = -1 - if not os.path.exists(path): - os.makedirs(path) - - def callback(img): - nonlocal counter - counter += 1 - if counter % modulus != 0: - return - image = self.sample_to_image(img) - image.save(os.path.join(path, f"{counter:03}.png"), "PNG") - - return callback - - -def _pairwise(iterable): - "s -> (s0, s1), (s2, s3), (s4, s5), ..." - a = iter(iterable) - return zip(a, a) diff --git a/invokeai/backend/globals.py b/invokeai/backend/globals.py deleted file mode 100644 index a3636ad5f8..0000000000 --- a/invokeai/backend/globals.py +++ /dev/null @@ -1,135 +0,0 @@ -""" -invokeai.backend.globals defines a small number of global variables that would -otherwise have to be passed through long and complex call chains. - -It defines a Namespace object named "Globals" that contains -the attributes: - - - root - the root directory under which "models" and "outputs" can be found - - initfile - path to the initialization file - - try_patchmatch - option to globally disable loading of 'patchmatch' module - - always_use_cpu - force use of CPU even if GPU is available -""" - -import os -import os.path as osp -from argparse import Namespace -from pathlib import Path -from typing import Union -from pydantic import BaseSettings - -Globals = Namespace() - -# Where to look for the initialization file and other key components -Globals.initfile = "invokeai.init" -Globals.models_file = "models.yaml" -Globals.models_dir = "models" -Globals.config_dir = "configs" -Globals.autoscan_dir = "weights" -Globals.converted_ckpts_dir = "converted_ckpts" - -# Set the default root directory. This can be overwritten by explicitly -# passing the `--root ` argument on the command line. -# logic is: -# 1) use INVOKEAI_ROOT environment variable (no check for this being a valid directory) -# 2) use VIRTUAL_ENV environment variable, with a check for initfile being there -# 3) use ~/invokeai - -if os.environ.get("INVOKEAI_ROOT"): - Globals.root = osp.abspath(os.environ.get("INVOKEAI_ROOT")) -elif ( - os.environ.get("VIRTUAL_ENV") - and Path(os.environ.get("VIRTUAL_ENV"), "..", Globals.initfile).exists() -): - Globals.root = osp.abspath(osp.join(os.environ.get("VIRTUAL_ENV"), "..")) -else: - Globals.root = osp.abspath(osp.expanduser("~/invokeai")) - -# Try loading patchmatch -Globals.try_patchmatch = True - -# Use CPU even if GPU is available (main use case is for debugging MPS issues) -Globals.always_use_cpu = False - -# Whether the internet is reachable for dynamic downloads -# The CLI will test connectivity at startup time. -Globals.internet_available = True - -# Whether to disable xformers -Globals.disable_xformers = False - -# Low-memory tradeoff for guidance calculations. -Globals.sequential_guidance = False - -# whether we are forcing full precision -Globals.full_precision = False - -# whether we should convert ckpt files into diffusers models on the fly -Globals.ckpt_convert = True - -# logging tokenization everywhere -Globals.log_tokenization = False - - -def global_config_file() -> Path: - return Path(Globals.root, Globals.config_dir, Globals.models_file) - - -def global_config_dir() -> Path: - return Path(Globals.root, Globals.config_dir) - - -def global_models_dir() -> Path: - return Path(Globals.root, Globals.models_dir) - - -def global_autoscan_dir() -> Path: - return Path(Globals.root, Globals.autoscan_dir) - - -def global_converted_ckpts_dir() -> Path: - return Path(global_models_dir(), Globals.converted_ckpts_dir) - - -def global_set_root(root_dir: Union[str, Path]): - Globals.root = root_dir - - -def global_cache_dir(subdir: Union[str, Path] = "") -> Path: - """ - Returns Path to the model cache directory. If a subdirectory - is provided, it will be appended to the end of the path, allowing - for Hugging Face-style conventions. Currently, Hugging Face has - moved all models into the "hub" subfolder, so for any pretrained - HF model, use: - global_cache_dir('hub') - - The legacy location for transformers used to be global_cache_dir('transformers') - and global_cache_dir('diffusers') for diffusers. - """ - home: str = os.getenv("HF_HOME") - - if home is None: - home = os.getenv("XDG_CACHE_HOME") - - if home is not None: - # Set `home` to $XDG_CACHE_HOME/huggingface, which is the default location mentioned in Hugging Face Hub Client Library. - # See: https://huggingface.co/docs/huggingface_hub/main/en/package_reference/environment_variables#xdgcachehome - home += os.sep + "huggingface" - - if home is not None: - return Path(home, subdir) - else: - return Path(Globals.root, "models", subdir) - -def copy_conf_to_globals(conf: Union[dict,BaseSettings]): - ''' - Given a dict or dict-like object, copy its keys and - values into the Globals Namespace. This is a transitional - workaround until we remove Globals entirely. - ''' - if isinstance(conf,BaseSettings): - conf = conf.dict() - for key in conf.keys(): - if key is not None: - setattr(Globals,key,conf[key]) diff --git a/invokeai/backend/image_util/patchmatch.py b/invokeai/backend/image_util/patchmatch.py index 5b5dd75f68..07d07b6bea 100644 --- a/invokeai/backend/image_util/patchmatch.py +++ b/invokeai/backend/image_util/patchmatch.py @@ -6,7 +6,9 @@ be suppressed or deferred """ import numpy as np import invokeai.backend.util.logging as logger -from invokeai.backend.globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig + +config = InvokeAIAppConfig() class PatchMatch: """ @@ -23,7 +25,7 @@ class PatchMatch: def _load_patch_match(self): if self.tried_load: return - if Globals.try_patchmatch: + if config.try_patchmatch: from patchmatch import patch_match as pm if pm.patchmatch_available: diff --git a/invokeai/backend/image_util/txt2mask.py b/invokeai/backend/image_util/txt2mask.py index 248f19d81d..0c193aa554 100644 --- a/invokeai/backend/image_util/txt2mask.py +++ b/invokeai/backend/image_util/txt2mask.py @@ -33,11 +33,11 @@ from PIL import Image, ImageOps from transformers import AutoProcessor, CLIPSegForImageSegmentation import invokeai.backend.util.logging as logger -from invokeai.backend.globals import global_cache_dir +from invokeai.app.services.config import InvokeAIAppConfig CLIPSEG_MODEL = "CIDAS/clipseg-rd64-refined" CLIPSEG_SIZE = 352 - +config = InvokeAIAppConfig() class SegmentedGrayscale(object): def __init__(self, image: Image, heatmap: torch.Tensor): @@ -88,10 +88,10 @@ class Txt2Mask(object): # BUG: we are not doing anything with the device option at this time self.device = device self.processor = AutoProcessor.from_pretrained( - CLIPSEG_MODEL, cache_dir=global_cache_dir("hub") + CLIPSEG_MODEL, cache_dir=config.cache_dir ) self.model = CLIPSegForImageSegmentation.from_pretrained( - CLIPSEG_MODEL, cache_dir=global_cache_dir("hub") + CLIPSEG_MODEL, cache_dir=config.cache_dir ) @torch.no_grad() diff --git a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py index 8aec5a01d9..42c3b4671d 100644 --- a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py @@ -26,7 +26,7 @@ import torch from safetensors.torch import load_file import invokeai.backend.util.logging as logger -from invokeai.backend.globals import global_cache_dir, global_config_dir +from invokeai.app.services.config import InvokeAIAppConfig from .model_manager import ModelManager, SDLegacyType @@ -73,6 +73,7 @@ from transformers import ( from ..stable_diffusion import StableDiffusionGeneratorPipeline +config = InvokeAIAppConfig() def shave_segments(path, n_shave_prefix_segments=1): """ @@ -842,7 +843,7 @@ def convert_ldm_bert_checkpoint(checkpoint, config): def convert_ldm_clip_checkpoint(checkpoint): text_model = CLIPTextModel.from_pretrained( - "openai/clip-vit-large-patch14", cache_dir=global_cache_dir("hub") + "openai/clip-vit-large-patch14", cache_dir=config.cache_dir ) keys = list(checkpoint.keys()) @@ -897,7 +898,7 @@ textenc_pattern = re.compile("|".join(protected.keys())) def convert_paint_by_example_checkpoint(checkpoint): - cache_dir = global_cache_dir("hub") + cache_dir = config.cache_dir config = CLIPVisionConfig.from_pretrained( "openai/clip-vit-large-patch14", cache_dir=cache_dir ) @@ -969,7 +970,7 @@ def convert_paint_by_example_checkpoint(checkpoint): def convert_open_clip_checkpoint(checkpoint): - cache_dir = global_cache_dir("hub") + cache_dir = config.cache_dir text_model = CLIPTextModel.from_pretrained( "stabilityai/stable-diffusion-2", subfolder="text_encoder", cache_dir=cache_dir ) @@ -1105,7 +1106,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt( else: checkpoint = load_file(checkpoint_path) - cache_dir = global_cache_dir("hub") + cache_dir = config.cache_dir pipeline_class = ( StableDiffusionGeneratorPipeline if return_generator_pipeline @@ -1129,25 +1130,23 @@ def load_pipeline_from_original_stable_diffusion_ckpt( if model_type == SDLegacyType.V2_v: original_config_file = ( - global_config_dir() / "stable-diffusion" / "v2-inference-v.yaml" + config.legacy_conf_path / "v2-inference-v.yaml" ) if global_step == 110000: # v2.1 needs to upcast attention upcast_attention = True elif model_type == SDLegacyType.V2_e: original_config_file = ( - global_config_dir() / "stable-diffusion" / "v2-inference.yaml" + config.legacy_conf_path / "v2-inference.yaml" ) elif model_type == SDLegacyType.V1_INPAINT: original_config_file = ( - global_config_dir() - / "stable-diffusion" - / "v1-inpainting-inference.yaml" + config.legacy_conf_path / "v1-inpainting-inference.yaml" ) elif model_type == SDLegacyType.V1: original_config_file = ( - global_config_dir() / "stable-diffusion" / "v1-inference.yaml" + config.legacy_conf_path / "v1-inference.yaml" ) else: @@ -1297,7 +1296,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt( ) safety_checker = StableDiffusionSafetyChecker.from_pretrained( "CompVis/stable-diffusion-safety-checker", - cache_dir=global_cache_dir("hub"), + cache_dir=config.cache_dir, ) feature_extractor = AutoFeatureExtractor.from_pretrained( "CompVis/stable-diffusion-safety-checker", cache_dir=cache_dir diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index a0a899a319..f78da5a232 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -36,8 +36,6 @@ from omegaconf import OmegaConf from omegaconf.dictconfig import DictConfig from picklescan.scanner import scan_file_path -from invokeai.backend.globals import Globals, global_cache_dir - from transformers import ( CLIPTextModel, CLIPTokenizer, @@ -49,9 +47,9 @@ from diffusers.pipelines.stable_diffusion.safety_checker import ( from ..stable_diffusion import ( StableDiffusionGeneratorPipeline, ) +from invokeai.app.services.config import InvokeAIAppConfig from ..util import CUDA_DEVICE, ask_user, download_with_resume - class SDLegacyType(Enum): V1 = auto() V1_INPAINT = auto() @@ -70,6 +68,7 @@ class SDModelComponent(Enum): feature_extractor="feature_extractor" DEFAULT_MAX_MODELS = 2 +config = InvokeAIAppConfig() class ModelManager(object): """ @@ -292,7 +291,7 @@ class ModelManager(object): """ # if we are converting legacy files automatically, then # there are no legacy ckpts! - if Globals.ckpt_convert: + if config.ckpt_convert: return False info = self.model_info(model_name) if "weights" in info and info["weights"].endswith((".ckpt", ".safetensors")): @@ -502,13 +501,13 @@ class ModelManager(object): # TODO: scan weights maybe? pipeline_args: dict[str, Any] = dict( - safety_checker=None, local_files_only=not Globals.internet_available + safety_checker=None, local_files_only=not config.internet_available ) if "vae" in mconfig and mconfig["vae"] is not None: if vae := self._load_vae(mconfig["vae"]): pipeline_args.update(vae=vae) if not isinstance(name_or_path, Path): - pipeline_args.update(cache_dir=global_cache_dir("hub")) + pipeline_args.update(cache_dir=config.cache_dir) if using_fp16: pipeline_args.update(torch_dtype=torch.float16) fp_args_list = [{"revision": "fp16"}, {}] @@ -561,9 +560,9 @@ class ModelManager(object): height = mconfig.height if not os.path.isabs(config): - config = os.path.join(Globals.root, config) + config = os.path.join(config.root, config) if not os.path.isabs(weights): - weights = os.path.normpath(os.path.join(Globals.root, weights)) + weights = os.path.normpath(os.path.join(config.root, weights)) # Convert to diffusers and return a diffusers pipeline self.logger.info(f"Converting legacy checkpoint {model_name} into a diffusers model...") @@ -581,7 +580,7 @@ class ModelManager(object): vae_path = ( vae if os.path.isabs(vae) - else os.path.normpath(os.path.join(Globals.root, vae)) + else os.path.normpath(os.path.join(config.root, vae)) ) if self._has_cuda(): torch.cuda.empty_cache() @@ -616,7 +615,7 @@ class ModelManager(object): if "path" in mconfig and mconfig["path"] is not None: path = Path(mconfig["path"]) if not path.is_absolute(): - path = Path(Globals.root, path).resolve() + path = Path(config.root, path).resolve() return path elif "repo_id" in mconfig: return mconfig["repo_id"] @@ -864,25 +863,16 @@ class ModelManager(object): model_type = self.probe_model_type(checkpoint) if model_type == SDLegacyType.V1: self.logger.debug("SD-v1 model detected") - model_config_file = Path( - Globals.root, "configs/stable-diffusion/v1-inference.yaml" - ) + model_config_file = config.legacy_conf_path / "v1-inference.yaml" elif model_type == SDLegacyType.V1_INPAINT: self.logger.debug("SD-v1 inpainting model detected") - model_config_file = Path( - Globals.root, - "configs/stable-diffusion/v1-inpainting-inference.yaml", - ) + model_config_file = config.legacy_conf_path / "v1-inpainting-inference.yaml", elif model_type == SDLegacyType.V2_v: self.logger.debug("SD-v2-v model detected") - model_config_file = Path( - Globals.root, "configs/stable-diffusion/v2-inference-v.yaml" - ) + model_config_file = config.legacy_conf_path / "v2-inference-v.yaml" elif model_type == SDLegacyType.V2_e: self.logger.debug("SD-v2-e model detected") - model_config_file = Path( - Globals.root, "configs/stable-diffusion/v2-inference.yaml" - ) + model_config_file = config.legacy_conf_path / "v2-inference.yaml" elif model_type == SDLegacyType.V2: self.logger.warning( f"{thing} is a V2 checkpoint file, but its parameterization cannot be determined. Please provide configuration file path." @@ -909,9 +899,7 @@ class ModelManager(object): self.logger.debug(f"Using VAE file {vae_path.name}") vae = None if vae_path else dict(repo_id="stabilityai/sd-vae-ft-mse") - diffuser_path = Path( - Globals.root, "models", Globals.converted_ckpts_dir, model_path.stem - ) + diffuser_path = config.root / "models/converted_ckpts" / model_path.stem model_name = self.convert_and_import( model_path, diffusers_path=diffuser_path, @@ -1044,9 +1032,7 @@ class ModelManager(object): """ yaml_str = OmegaConf.to_yaml(self.config) if not os.path.isabs(config_file_path): - config_file_path = os.path.normpath( - os.path.join(Globals.root, config_file_path) - ) + config_file_path = config.model_conf_path tmpfile = os.path.join(os.path.dirname(config_file_path), "new_config.tmp") with open(tmpfile, "w", encoding="utf-8") as outfile: outfile.write(self.preamble()) @@ -1078,7 +1064,7 @@ class ModelManager(object): """ # Three transformer models to check: bert, clip and safety checker, and # the diffusers as well - models_dir = Path(Globals.root, "models") + models_dir = config.root / "models" legacy_locations = [ Path( models_dir, @@ -1090,8 +1076,8 @@ class ModelManager(object): "openai/clip-vit-large-patch14/models--openai--clip-vit-large-patch14", ), ] - legacy_locations.extend(list(global_cache_dir("diffusers").glob("*"))) - + legacy_cache_dir = config.cache_dir / "../diffusers" + legacy_locations.extend(list(legacy_cache_dir.glob("*"))) legacy_layout = False for model in legacy_locations: legacy_layout = legacy_layout or model.exists() @@ -1113,7 +1099,7 @@ class ModelManager(object): # transformer files get moved into the hub directory if cls._is_huggingface_hub_directory_present(): - hub = global_cache_dir("hub") + hub = config.cache_dir else: hub = models_dir / "hub" @@ -1152,12 +1138,12 @@ class ModelManager(object): if str(source).startswith(("http:", "https:", "ftp:")): dest_directory = Path(dest_directory) if not dest_directory.is_absolute(): - dest_directory = Globals.root / dest_directory + dest_directory = config.root / dest_directory dest_directory.mkdir(parents=True, exist_ok=True) resolved_path = download_with_resume(str(source), dest_directory) else: if not os.path.isabs(source): - source = os.path.join(Globals.root, source) + source = config.root / source resolved_path = Path(source) return resolved_path @@ -1208,7 +1194,7 @@ class ModelManager(object): path = name_or_path else: owner, repo = name_or_path.split("/") - path = Path(global_cache_dir("hub") / f"models--{owner}--{repo}") + path = Path(config.cache_dir / f"models--{owner}--{repo}") if not path.exists(): return None hashpath = path / "checksum.sha256" @@ -1269,8 +1255,8 @@ class ModelManager(object): using_fp16 = self.precision == "float16" vae_args.update( - cache_dir=global_cache_dir("hub"), - local_files_only=not Globals.internet_available, + cache_dir=config.cache_dir, + local_files_only=not config.internet_available, ) self.logger.debug(f"Loading diffusers VAE from {name_or_path}") @@ -1308,7 +1294,7 @@ class ModelManager(object): @classmethod def _delete_model_from_cache(cls,repo_id): - cache_info = scan_cache_dir(global_cache_dir("hub")) + cache_info = scan_cache_dir(config.cache_dir) # I'm sure there is a way to do this with comprehensions # but the code quickly became incomprehensible! @@ -1327,7 +1313,7 @@ class ModelManager(object): def _abs_path(path: str | Path) -> Path: if path is None or Path(path).is_absolute(): return path - return Path(Globals.root, path).resolve() + return Path(config.root, path).resolve() @staticmethod def _is_huggingface_hub_directory_present() -> bool: diff --git a/invokeai/backend/prompting/conditioning.py b/invokeai/backend/prompting/conditioning.py index d9130ace04..549308bf0a 100644 --- a/invokeai/backend/prompting/conditioning.py +++ b/invokeai/backend/prompting/conditioning.py @@ -19,11 +19,12 @@ from compel.prompt_parser import ( ) import invokeai.backend.util.logging as logger -from invokeai.backend.globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig from ..stable_diffusion import InvokeAIDiffuserComponent from ..util import torch_dtype +config = InvokeAIAppConfig() def get_uc_and_c_and_ec( prompt_string, model, log_tokens=False, skip_normalize_legacy_blend=False @@ -61,7 +62,7 @@ def get_uc_and_c_and_ec( negative_prompt_string ) - if log_tokens or getattr(Globals, "log_tokenization", False): + if log_tokens or config.log_tokenization: log_tokenization(positive_prompt, negative_prompt, tokenizer=tokenizer) c, options = compel.build_conditioning_tensor_for_prompt_object(positive_prompt) diff --git a/invokeai/backend/restoration/codeformer.py b/invokeai/backend/restoration/codeformer.py index 5b578af082..72ed54048e 100644 --- a/invokeai/backend/restoration/codeformer.py +++ b/invokeai/backend/restoration/codeformer.py @@ -6,7 +6,8 @@ import numpy as np import torch import invokeai.backend.util.logging as logger -from ..globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig +config = InvokeAIAppConfig() pretrained_model_url = ( "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth" @@ -18,7 +19,7 @@ class CodeFormerRestoration: self, codeformer_dir="models/codeformer", codeformer_model_path="codeformer.pth" ) -> None: if not os.path.isabs(codeformer_dir): - codeformer_dir = os.path.join(Globals.root, codeformer_dir) + codeformer_dir = os.path.join(config.root, codeformer_dir) self.model_path = os.path.join(codeformer_dir, codeformer_model_path) self.codeformer_model_exists = os.path.isfile(self.model_path) @@ -72,7 +73,7 @@ class CodeFormerRestoration: use_parse=True, device=device, model_rootpath=os.path.join( - Globals.root, "models", "gfpgan", "weights" + config.root, "models", "gfpgan", "weights" ), ) face_helper.clean_all() diff --git a/invokeai/backend/restoration/gfpgan.py b/invokeai/backend/restoration/gfpgan.py index b5c0278362..f049369116 100644 --- a/invokeai/backend/restoration/gfpgan.py +++ b/invokeai/backend/restoration/gfpgan.py @@ -7,13 +7,14 @@ import torch from PIL import Image import invokeai.backend.util.logging as logger -from invokeai.backend.globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig +config = InvokeAIAppConfig() class GFPGAN: def __init__(self, gfpgan_model_path="models/gfpgan/GFPGANv1.4.pth") -> None: if not os.path.isabs(gfpgan_model_path): gfpgan_model_path = os.path.abspath( - os.path.join(Globals.root, gfpgan_model_path) + os.path.join(config.root, gfpgan_model_path) ) self.model_path = gfpgan_model_path self.gfpgan_model_exists = os.path.isfile(self.model_path) @@ -33,7 +34,7 @@ class GFPGAN: warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=UserWarning) cwd = os.getcwd() - os.chdir(os.path.join(Globals.root, "models")) + os.chdir(os.path.join(config.root, "models")) try: from gfpgan import GFPGANer diff --git a/invokeai/backend/safety_checker.py b/invokeai/backend/safety_checker.py index 3003981888..00c41e10ef 100644 --- a/invokeai/backend/safety_checker.py +++ b/invokeai/backend/safety_checker.py @@ -15,9 +15,11 @@ from transformers import AutoFeatureExtractor import invokeai.assets.web as web_assets import invokeai.backend.util.logging as logger -from .globals import global_cache_dir +from invokeai.app.services.config import InvokeAIAppConfig from .util import CPU_DEVICE +config = InvokeAIAppConfig() + class SafetyChecker(object): CAUTION_IMG = "caution.png" @@ -29,7 +31,7 @@ class SafetyChecker(object): try: safety_model_id = "CompVis/stable-diffusion-safety-checker" - safety_model_path = global_cache_dir("hub") + safety_model_path = config.cache_dir self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( safety_model_id, local_files_only=True, diff --git a/invokeai/backend/stable_diffusion/concepts_lib.py b/invokeai/backend/stable_diffusion/concepts_lib.py index ebbcc9c3e9..75621193bc 100644 --- a/invokeai/backend/stable_diffusion/concepts_lib.py +++ b/invokeai/backend/stable_diffusion/concepts_lib.py @@ -18,15 +18,15 @@ from huggingface_hub import ( ) import invokeai.backend.util.logging as logger -from invokeai.backend.globals import Globals - +from invokeai.app.services.config import InvokeAIAppConfig +config = InvokeAIAppConfig() class HuggingFaceConceptsLibrary(object): def __init__(self, root=None): """ Initialize the Concepts object. May optionally pass a root directory. """ - self.root = root or Globals.root + self.root = root or config.root self.hf_api = HfApi() self.local_concepts = dict() self.concept_list = None @@ -58,7 +58,7 @@ class HuggingFaceConceptsLibrary(object): self.concept_list.extend(list(local_concepts_to_add)) return self.concept_list return self.concept_list - elif Globals.internet_available is True: + elif config.internet_available is True: try: models = self.hf_api.list_models( filter=ModelFilter(model_name="sd-concepts-library/") diff --git a/invokeai/backend/stable_diffusion/diffusers_pipeline.py b/invokeai/backend/stable_diffusion/diffusers_pipeline.py index 94ec9da7e8..758fd10a74 100644 --- a/invokeai/backend/stable_diffusion/diffusers_pipeline.py +++ b/invokeai/backend/stable_diffusion/diffusers_pipeline.py @@ -33,8 +33,7 @@ from torchvision.transforms.functional import resize as tv_resize from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer from typing_extensions import ParamSpec -from invokeai.backend.globals import Globals - +from invokeai.app.services.config import InvokeAIAppConfig from ..util import CPU_DEVICE, normalize_device from .diffusion import ( AttentionMapSaver, @@ -44,6 +43,7 @@ from .diffusion import ( from .offloading import FullyLoadedModelGroup, LazilyLoadedModelGroup, ModelGroup from .textual_inversion_manager import TextualInversionManager +config = InvokeAIAppConfig() @dataclass class PipelineIntermediateState: @@ -351,7 +351,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): if ( torch.cuda.is_available() and is_xformers_available() - and not Globals.disable_xformers + and not config.disable_xformers ): self.enable_xformers_memory_efficient_attention() else: diff --git a/invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py b/invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py index b0c85e9fd3..174d306445 100644 --- a/invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py +++ b/invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py @@ -9,7 +9,7 @@ from diffusers.models.attention_processor import AttentionProcessor from typing_extensions import TypeAlias import invokeai.backend.util.logging as logger -from invokeai.backend.globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig from .cross_attention_control import ( Arguments, @@ -31,6 +31,7 @@ ModelForwardCallback: TypeAlias = Union[ Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], ] +config = InvokeAIAppConfig() @dataclass(frozen=True) class PostprocessingSettings: @@ -77,7 +78,7 @@ class InvokeAIDiffuserComponent: self.is_running_diffusers = is_running_diffusers self.model_forward_callback = model_forward_callback self.cross_attention_control_context = None - self.sequential_guidance = Globals.sequential_guidance + self.sequential_guidance = config.sequential_guidance @contextmanager def custom_attention_context( diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index c70a43ff09..d3da346069 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -4,17 +4,16 @@ from contextlib import nullcontext import torch from torch import autocast - -from invokeai.backend.globals import Globals +from invokeai.app.services.config import InvokeAIAppConfig CPU_DEVICE = torch.device("cpu") CUDA_DEVICE = torch.device("cuda") MPS_DEVICE = torch.device("mps") - +config = InvokeAIAppConfig() def choose_torch_device() -> torch.device: """Convenience routine for guessing which GPU device to run model on""" - if Globals.always_use_cpu: + if config.always_use_cpu: return CPU_DEVICE if torch.cuda.is_available(): return torch.device("cuda") @@ -33,7 +32,7 @@ def choose_precision(device: torch.device) -> str: def torch_dtype(device: torch.device) -> torch.dtype: - if Globals.full_precision: + if config.full_precision: return torch.float32 if choose_precision(device) == "float16": return torch.float16 diff --git a/invokeai/frontend/CLI/CLI.py b/invokeai/frontend/CLI/CLI.py deleted file mode 100644 index aa0c4bea5f..0000000000 --- a/invokeai/frontend/CLI/CLI.py +++ /dev/null @@ -1,1291 +0,0 @@ -import os -import re -import shlex -import sys -import traceback -from argparse import Namespace -from pathlib import Path -from typing import Union - -import click -from compel import PromptParser - -if sys.platform == "darwin": - os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" - -import pyparsing # type: ignore - -import invokeai.version as invokeai -import invokeai.backend.util.logging as logger - -from ...backend import Generate, ModelManager -from ...backend.args import Args, dream_cmd_from_png, metadata_dumps, metadata_from_png -from ...backend.globals import Globals, global_config_dir -from ...backend.image_util import ( - PngWriter, - make_grid, - retrieve_metadata, - write_metadata, -) -from ...backend.stable_diffusion import PipelineIntermediateState -from ...backend.util import url_attachment_name, write_log -from .readline import Completer, get_completer - -# global used in multiple functions (fix) -infile = None - - -def main(): - """Initialize command-line parsers and the diffusion model""" - global infile - - opt = Args() - args = opt.parse_args() - if not args: - sys.exit(-1) - - if args.laion400m: - print( - "--laion400m flag has been deprecated. Please use --model laion400m instead." - ) - sys.exit(-1) - if args.weights: - print( - "--weights argument has been deprecated. Please edit ./configs/models.yaml, and select the weights using --model instead." - ) - sys.exit(-1) - if args.max_loaded_models is not None: - if args.max_loaded_models <= 0: - print("--max_loaded_models must be >= 1; using 1") - args.max_loaded_models = 1 - - # alert - setting a few globals here - Globals.try_patchmatch = args.patchmatch - Globals.always_use_cpu = args.always_use_cpu - Globals.internet_available = args.internet_available and check_internet() - Globals.disable_xformers = not args.xformers - Globals.sequential_guidance = args.sequential_guidance - Globals.ckpt_convert = True # always true now - - # run any post-install patches needed - run_patches() - - logger.info(f"Internet connectivity is {Globals.internet_available}") - - if not args.conf: - config_file = os.path.join(Globals.root, "configs", "models.yaml") - if not os.path.exists(config_file): - report_model_error( - opt, FileNotFoundError(f"The file {config_file} could not be found.") - ) - - logger.info(f"{invokeai.__app_name__}, version {invokeai.__version__}") - logger.info(f'InvokeAI runtime directory is "{Globals.root}"') - - # loading here to avoid long delays on startup - # these two lines prevent a horrible warning message from appearing - # when the frozen CLIP tokenizer is imported - import transformers # type: ignore - - transformers.logging.set_verbosity_error() - import diffusers - - diffusers.logging.set_verbosity_error() - - # Loading Face Restoration and ESRGAN Modules - gfpgan, codeformer, esrgan = load_face_restoration(opt) - - # normalize the config directory relative to root - if not os.path.isabs(opt.conf): - opt.conf = os.path.normpath(os.path.join(Globals.root, opt.conf)) - - if opt.embeddings: - if not os.path.isabs(opt.embedding_path): - embedding_path = os.path.normpath( - os.path.join(Globals.root, opt.embedding_path) - ) - else: - embedding_path = opt.embedding_path - else: - embedding_path = None - - # migrate legacy models - ModelManager.migrate_models() - - # load the infile as a list of lines - if opt.infile: - try: - if os.path.isfile(opt.infile): - infile = open(opt.infile, "r", encoding="utf-8") - elif opt.infile == "-": # stdin - infile = sys.stdin - else: - raise FileNotFoundError(f"{opt.infile} not found.") - except (FileNotFoundError, IOError) as e: - logger.critical('Aborted',exc_info=True) - sys.exit(-1) - - # creating a Generate object: - try: - gen = Generate( - conf=opt.conf, - model=opt.model, - sampler_name=opt.sampler_name, - embedding_path=embedding_path, - full_precision=opt.full_precision, - precision=opt.precision, - gfpgan=gfpgan, - codeformer=codeformer, - esrgan=esrgan, - free_gpu_mem=opt.free_gpu_mem, - safety_checker=opt.safety_checker, - max_loaded_models=opt.max_loaded_models, - ) - except (FileNotFoundError, TypeError, AssertionError) as e: - report_model_error(opt, e) - except (IOError, KeyError): - logger.critical("Aborted",exc_info=True) - sys.exit(-1) - - if opt.seamless: - logger.info("Changed to seamless tiling mode") - - # preload the model - try: - gen.load_model() - except KeyError: - pass - except Exception as e: - report_model_error(opt, e) - - # try to autoconvert new models - if path := opt.autoconvert: - gen.model_manager.heuristic_import( - str(path), commit_to_conf=opt.conf - ) - - # web server loops forever - if opt.web or opt.gui: - invoke_ai_web_server_loop(gen, gfpgan, codeformer, esrgan) - sys.exit(0) - - if not infile: - print( - "\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)" - ) - - try: - main_loop(gen, opt) - except KeyboardInterrupt: - print( - f'\nGoodbye!\nYou can start InvokeAI again by running the "invoke.bat" (or "invoke.sh") script from {Globals.root}' - ) - except Exception: - logger.error("An error occurred",exc_info=True) - -# TODO: main_loop() has gotten busy. Needs to be refactored. -def main_loop(gen, opt): - """prompt/read/execute loop""" - global infile - done = False - doneAfterInFile = infile is not None - path_filter = re.compile(r'[<>:"/\\|?*]') - last_results = list() - - # The readline completer reads history from the .dream_history file located in the - # output directory specified at the time of script launch. We do not currently support - # changing the history file midstream when the output directory is changed. - completer = get_completer(opt, models=gen.model_manager.list_models()) - set_default_output_dir(opt, completer) - if gen.model: - add_embedding_terms(gen, completer) - output_cntr = completer.get_current_history_length() + 1 - - # os.pathconf is not available on Windows - if hasattr(os, "pathconf"): - path_max = os.pathconf(opt.outdir, "PC_PATH_MAX") - name_max = os.pathconf(opt.outdir, "PC_NAME_MAX") - else: - path_max = 260 - name_max = 255 - - while not done: - operation = "generate" - - try: - command = get_next_command(infile, gen.model_name) - except EOFError: - done = infile is None or doneAfterInFile - infile = None - continue - - # skip empty lines - if not command.strip(): - continue - - if command.startswith(("#", "//")): - continue - - if len(command.strip()) == 1 and command.startswith("q"): - done = True - break - - if not command.startswith("!history"): - completer.add_history(command) - - if command.startswith("!"): - command, operation = do_command(command, gen, opt, completer) - - if operation is None: - continue - - if opt.parse_cmd(command) is None: - continue - - if opt.init_img: - try: - if not opt.prompt: - oldargs = metadata_from_png(opt.init_img) - opt.prompt = oldargs.prompt - logger.info(f'Retrieved old prompt "{opt.prompt}" from {opt.init_img}') - except (OSError, AttributeError, KeyError): - pass - - if len(opt.prompt) == 0: - opt.prompt = "" - - # width and height are set by model if not specified - if not opt.width: - opt.width = gen.width - if not opt.height: - opt.height = gen.height - - # retrieve previous value of init image if requested - if opt.init_img is not None and re.match("^-\\d+$", opt.init_img): - try: - opt.init_img = last_results[int(opt.init_img)][0] - logger.info(f"Reusing previous image {opt.init_img}") - except IndexError: - logger.info(f"No previous initial image at position {opt.init_img} found") - opt.init_img = None - continue - - # the outdir can change with each command, so we adjust it here - set_default_output_dir(opt, completer) - - # try to relativize pathnames - for attr in ("init_img", "init_mask", "init_color"): - if getattr(opt, attr) and not os.path.exists(getattr(opt, attr)): - basename = getattr(opt, attr) - path = os.path.join(opt.outdir, basename) - setattr(opt, attr, path) - - # retrieve previous value of seed if requested - # Exception: for postprocess operations negative seed values - # mean "discard the original seed and generate a new one" - # (this is a non-obvious hack and needs to be reworked) - if opt.seed is not None and opt.seed < 0 and operation != "postprocess": - try: - opt.seed = last_results[opt.seed][1] - logger.info(f"Reusing previous seed {opt.seed}") - except IndexError: - logger.info(f"No previous seed at position {opt.seed} found") - opt.seed = None - continue - - if opt.strength is None: - opt.strength = 0.75 if opt.out_direction is None else 0.83 - - if opt.with_variations is not None: - opt.with_variations = split_variations(opt.with_variations) - - if opt.prompt_as_dir and operation == "generate": - # sanitize the prompt to a valid folder name - subdir = path_filter.sub("_", opt.prompt)[:name_max].rstrip(" .") - - # truncate path to maximum allowed length - # 39 is the length of '######.##########.##########-##.png', plus two separators and a NUL - subdir = subdir[: (path_max - 39 - len(os.path.abspath(opt.outdir)))] - current_outdir = os.path.join(opt.outdir, subdir) - - logger.info('Writing files to directory: "' + current_outdir + '"') - - # make sure the output directory exists - if not os.path.exists(current_outdir): - os.makedirs(current_outdir) - else: - if not os.path.exists(opt.outdir): - os.makedirs(opt.outdir) - current_outdir = opt.outdir - - # Here is where the images are actually generated! - last_results = [] - try: - file_writer = PngWriter(current_outdir) - results = [] # list of filename, prompt pairs - grid_images = dict() # seed -> Image, only used if `opt.grid` - prior_variations = opt.with_variations or [] - prefix = file_writer.unique_prefix() - step_callback = ( - make_step_callback(gen, opt, prefix) - if opt.save_intermediates > 0 - else None - ) - - def image_writer( - image, - seed, - upscaled=False, - first_seed=None, - use_prefix=None, - prompt_in=None, - attention_maps_image=None, - ): - # note the seed is the seed of the current image - # the first_seed is the original seed that noise is added to - # when the -v switch is used to generate variations - nonlocal prior_variations - nonlocal prefix - - path = None - if opt.grid: - grid_images[seed] = image - - elif operation == "mask": - filename = f"{prefix}.{use_prefix}.{seed}.png" - tm = opt.text_mask[0] - th = opt.text_mask[1] if len(opt.text_mask) > 1 else 0.5 - formatted_dream_prompt = ( - f"!mask {opt.input_file_path} -tm {tm} {th}" - ) - path = file_writer.save_image_and_prompt_to_png( - image=image, - dream_prompt=formatted_dream_prompt, - metadata={}, - name=filename, - compress_level=opt.png_compression, - ) - results.append([path, formatted_dream_prompt]) - - else: - if use_prefix is not None: - prefix = use_prefix - postprocessed = upscaled if upscaled else operation == "postprocess" - opt.prompt = ( - gen.huggingface_concepts_library.replace_triggers_with_concepts( - opt.prompt or prompt_in - ) - ) # to avoid the problem of non-unique concept triggers - filename, formatted_dream_prompt = prepare_image_metadata( - opt, - prefix, - seed, - operation, - prior_variations, - postprocessed, - first_seed, - ) - path = file_writer.save_image_and_prompt_to_png( - image=image, - dream_prompt=formatted_dream_prompt, - metadata=metadata_dumps( - opt, - seeds=[ - seed - if opt.variation_amount == 0 - and len(prior_variations) == 0 - else first_seed - ], - model_hash=gen.model_hash, - ), - name=filename, - compress_level=opt.png_compression, - ) - - # update rfc metadata - if operation == "postprocess": - tool = re.match( - "postprocess:(\w+)", opt.last_operation - ).groups()[0] - add_postprocessing_to_metadata( - opt, - opt.input_file_path, - filename, - tool, - formatted_dream_prompt, - ) - - if (not postprocessed) or opt.save_original: - # only append to results if we didn't overwrite an earlier output - results.append([path, formatted_dream_prompt]) - - # so that the seed autocompletes (on linux|mac when -S or --seed specified - if completer and operation == "generate": - completer.add_seed(seed) - completer.add_seed(first_seed) - last_results.append([path, seed]) - - if operation == "generate": - catch_ctrl_c = ( - infile is None - ) # if running interactively, we catch keyboard interrupts - opt.last_operation = "generate" - try: - gen.prompt2image( - image_callback=image_writer, - step_callback=step_callback, - catch_interrupts=catch_ctrl_c, - **vars(opt), - ) - except (PromptParser.ParsingException, pyparsing.ParseException): - logger.error("An error occurred while processing your prompt",exc_info=True) - elif operation == "postprocess": - logger.info(f"fixing {opt.prompt}") - opt.last_operation = do_postprocess(gen, opt, image_writer) - - elif operation == "mask": - logger.info(f"generating masks from {opt.prompt}") - do_textmask(gen, opt, image_writer) - - if opt.grid and len(grid_images) > 0: - grid_img = make_grid(list(grid_images.values())) - grid_seeds = list(grid_images.keys()) - first_seed = last_results[0][1] - filename = f"{prefix}.{first_seed}.png" - formatted_dream_prompt = opt.dream_prompt_str( - seed=first_seed, grid=True, iterations=len(grid_images) - ) - formatted_dream_prompt += f" # {grid_seeds}" - metadata = metadata_dumps( - opt, seeds=grid_seeds, model_hash=gen.model_hash - ) - path = file_writer.save_image_and_prompt_to_png( - image=grid_img, - dream_prompt=formatted_dream_prompt, - metadata=metadata, - name=filename, - ) - results = [[path, formatted_dream_prompt]] - - except AssertionError: - logger.error(e) - continue - - except OSError as e: - logger.error(e) - continue - - print("Outputs:") - log_path = os.path.join(current_outdir, "invoke_log") - output_cntr = write_log(results, log_path, ("txt", "md"), output_cntr) - print() - - print( - f'\nGoodbye!\nYou can start InvokeAI again by running the "invoke.bat" (or "invoke.sh") script from {Globals.root}' - ) - - -# TO DO: remove repetitive code and the awkward command.replace() trope -# Just do a simple parse of the command! -def do_command(command: str, gen, opt: Args, completer) -> tuple: - global infile - operation = "generate" # default operation, alternative is 'postprocess' - command = command.replace("\\", "/") # windows - - if command.startswith( - "!dream" - ): # in case a stored prompt still contains the !dream command - command = command.replace("!dream ", "", 1) - - elif command.startswith("!fix"): - command = command.replace("!fix ", "", 1) - operation = "postprocess" - - elif command.startswith("!mask"): - command = command.replace("!mask ", "", 1) - operation = "mask" - - elif command.startswith("!switch"): - model_name = command.replace("!switch ", "", 1) - try: - gen.set_model(model_name) - add_embedding_terms(gen, completer) - except KeyError as e: - logger.error(e) - except Exception as e: - report_model_error(opt, e) - completer.add_history(command) - operation = None - - elif command.startswith("!models"): - gen.model_manager.print_models() - completer.add_history(command) - operation = None - - elif command.startswith("!import"): - path = shlex.split(command) - if len(path) < 2: - logger.warning( - "please provide (1) a URL to a .ckpt file to import; (2) a local path to a .ckpt file; or (3) a diffusers repository id in the form stabilityai/stable-diffusion-2-1" - ) - else: - try: - import_model(path[1], gen, opt, completer) - completer.add_history(command) - except KeyboardInterrupt: - print("\n") - operation = None - - elif command.startswith(("!convert", "!optimize")): - path = shlex.split(command) - if len(path) < 2: - logger.warning("please provide the path to a .ckpt or .safetensors model") - else: - try: - convert_model(path[1], gen, opt, completer) - completer.add_history(command) - except KeyboardInterrupt: - print("\n") - operation = None - - elif command.startswith("!edit"): - path = shlex.split(command) - if len(path) < 2: - logger.warning("please provide the name of a model") - else: - edit_model(path[1], gen, opt, completer) - completer.add_history(command) - operation = None - - elif command.startswith("!del"): - path = shlex.split(command) - if len(path) < 2: - logger.warning("please provide the name of a model") - else: - del_config(path[1], gen, opt, completer) - completer.add_history(command) - operation = None - - elif command.startswith("!fetch"): - file_path = command.replace("!fetch", "", 1).strip() - retrieve_dream_command(opt, file_path, completer) - completer.add_history(command) - operation = None - - elif command.startswith("!replay"): - file_path = command.replace("!replay", "", 1).strip() - file_path = os.path.join(opt.outdir, file_path) - if infile is None and os.path.isfile(file_path): - infile = open(file_path, "r", encoding="utf-8") - completer.add_history(command) - operation = None - - elif command.startswith("!trigger"): - print("Embedding trigger strings: ", ", ".join(gen.embedding_trigger_strings)) - operation = None - - elif command.startswith("!history"): - completer.show_history() - operation = None - - elif command.startswith("!search"): - search_str = command.replace("!search", "", 1).strip() - completer.show_history(search_str) - operation = None - - elif command.startswith("!clear"): - completer.clear_history() - operation = None - - elif re.match("^!(\d+)", command): - command_no = re.match("^!(\d+)", command).groups()[0] - command = completer.get_line(int(command_no)) - completer.set_line(command) - operation = None - - else: # not a recognized command, so give the --help text - command = "-h" - return command, operation - - -def set_default_output_dir(opt: Args, completer: Completer): - """ - If opt.outdir is relative, we add the root directory to it - normalize the outdir relative to root and make sure it exists. - """ - if not os.path.isabs(opt.outdir): - opt.outdir = os.path.normpath(os.path.join(Globals.root, opt.outdir)) - if not os.path.exists(opt.outdir): - os.makedirs(opt.outdir) - completer.set_default_dir(opt.outdir) - - -def import_model(model_path: str, gen, opt, completer): - """ - model_path can be (1) a URL to a .ckpt file; (2) a local .ckpt file path; - (3) a huggingface repository id; or (4) a local directory containing a - diffusers model. - """ - default_name = Path(model_path).stem - model_name = None - model_desc = None - - if ( - Path(model_path).is_dir() - and not (Path(model_path) / "model_index.json").exists() - ): - pass - else: - if model_path.startswith(("http:", "https:")): - try: - default_name = url_attachment_name(model_path) - default_name = Path(default_name).stem - except Exception: - logger.warning(f"A problem occurred while assigning the name of the downloaded model",exc_info=True) - model_name, model_desc = _get_model_name_and_desc( - gen.model_manager, - completer, - model_name=default_name, - ) - imported_name = gen.model_manager.heuristic_import( - model_path, - model_name=model_name, - description=model_desc, - ) - - if not imported_name: - if config_file := _pick_configuration_file(completer): - imported_name = gen.model_manager.heuristic_import( - model_path, - model_name=model_name, - description=model_desc, - model_config_file=config_file, - ) - if not imported_name: - logger.error("Aborting import.") - return - - if not _verify_load(imported_name, gen): - logger.error("model failed to load. Discarding configuration entry") - gen.model_manager.del_model(imported_name) - return - if click.confirm("Make this the default model?", default=False): - gen.model_manager.set_default_model(imported_name) - - gen.model_manager.commit(opt.conf) - completer.update_models(gen.model_manager.list_models()) - logger.info(f"{imported_name} successfully installed") - -def _pick_configuration_file(completer)->Path: - print( -""" -Please select the type of this model: -[1] A Stable Diffusion v1.x ckpt/safetensors model -[2] A Stable Diffusion v1.x inpainting ckpt/safetensors model -[3] A Stable Diffusion v2.x base model (512 pixels) -[4] A Stable Diffusion v2.x v-predictive model (768 pixels) -[5] Other (you will be prompted to enter the config file path) -[Q] I have no idea! Skip the import. -""") - choices = [ - global_config_dir() / 'stable-diffusion' / x - for x in [ - 'v1-inference.yaml', - 'v1-inpainting-inference.yaml', - 'v2-inference.yaml', - 'v2-inference-v.yaml', - ] - ] - - ok = False - while not ok: - try: - choice = input('select 0-5, Q > ').strip() - if choice.startswith(('q','Q')): - return - if choice == '5': - completer.complete_extensions(('.yaml')) - choice = Path(input('Select config file for this model> ').strip()).absolute() - completer.complete_extensions(None) - ok = choice.exists() - else: - choice = choices[int(choice)-1] - ok = True - except (ValueError, IndexError): - print(f'{choice} is not a valid choice') - except EOFError: - return - return choice - -def _verify_load(model_name: str, gen) -> bool: - logger.info("Verifying that new model loads...") - current_model = gen.model_name - try: - if not gen.set_model(model_name): - return - except Exception as e: - logger.warning(f"model failed to load: {str(e)}") - logger.warning( - "** note that importing 2.X checkpoints is not supported. Please use !convert_model instead." - ) - return False - if click.confirm("Keep model loaded?", default=True): - gen.set_model(model_name) - else: - logger.info("Restoring previous model") - gen.set_model(current_model) - return True - - -def _get_model_name_and_desc( - model_manager, completer, model_name: str = "", model_description: str = "" -): - model_name = _get_model_name(model_manager.list_models(), completer, model_name) - model_description = model_description or f"Imported model {model_name}" - completer.set_line(model_description) - model_description = ( - input(f"Description for this model [{model_description}]: ").strip() - or model_description - ) - return model_name, model_description - -def convert_model(model_name_or_path: Union[Path, str], gen, opt, completer): - model_name_or_path = model_name_or_path.replace("\\", "/") # windows - manager = gen.model_manager - ckpt_path = None - original_config_file = None - if model_name_or_path == gen.model_name: - logger.warning("Can't convert the active model. !switch to another model first. **") - return - elif model_info := manager.model_info(model_name_or_path): - if "weights" in model_info: - ckpt_path = Path(model_info["weights"]) - original_config_file = Path(model_info["config"]) - model_name = model_name_or_path - model_description = model_info["description"] - vae_path = model_info.get("vae") - else: - logger.warning(f"{model_name_or_path} is not a legacy .ckpt weights file") - return - model_name = manager.convert_and_import( - ckpt_path, - diffusers_path=Path( - Globals.root, "models", Globals.converted_ckpts_dir, model_name_or_path - ), - model_name=model_name, - model_description=model_description, - original_config_file=original_config_file, - vae_path=vae_path, - ) - else: - try: - import_model(model_name_or_path, gen, opt, completer) - except KeyboardInterrupt: - return - - manager.commit(opt.conf) - if click.confirm(f"Delete the original .ckpt file at {ckpt_path}?", default=False): - ckpt_path.unlink(missing_ok=True) - logger.warning(f"{ckpt_path} deleted") - - -def del_config(model_name: str, gen, opt, completer): - current_model = gen.model_name - if model_name == current_model: - logger.warning("Can't delete active model. !switch to another model first. **") - return - if model_name not in gen.model_manager.config: - logger.warning(f"Unknown model {model_name}") - return - - if not click.confirm( - f"Remove {model_name} from the list of models known to InvokeAI?", default=True - ): - return - - delete_completely = click.confirm( - "Completely remove the model file or directory from disk?", default=False - ) - gen.model_manager.del_model(model_name, delete_files=delete_completely) - gen.model_manager.commit(opt.conf) - logger.warning(f"{model_name} deleted") - completer.update_models(gen.model_manager.list_models()) - - -def edit_model(model_name: str, gen, opt, completer): - manager = gen.model_manager - if not (info := manager.model_info(model_name)): - logger.warning(f"** Unknown model {model_name}") - return - print() - logger.info(f"Editing model {model_name} from configuration file {opt.conf}") - new_name = _get_model_name(manager.list_models(), completer, model_name) - - for attribute in info.keys(): - if type(info[attribute]) != str: - continue - if attribute == "format": - continue - completer.set_line(info[attribute]) - info[attribute] = input(f"{attribute}: ") or info[attribute] - - if info["format"] == "diffusers": - vae = info.get("vae", dict(repo_id=None, path=None, subfolder=None)) - completer.set_line(vae.get("repo_id") or "stabilityai/sd-vae-ft-mse") - vae["repo_id"] = input("External VAE repo_id: ").strip() or None - if not vae["repo_id"]: - completer.set_line(vae.get("path") or "") - vae["path"] = ( - input("Path to a local diffusers VAE model (usually none): ").strip() - or None - ) - completer.set_line(vae.get("subfolder") or "") - vae["subfolder"] = ( - input("Name of subfolder containing the VAE model (usually none): ").strip() - or None - ) - info["vae"] = vae - - if new_name != model_name: - manager.del_model(model_name) - - # this does the update - manager.add_model(new_name, info, True) - - if click.confirm("Make this the default model?", default=False): - manager.set_default_model(new_name) - manager.commit(opt.conf) - completer.update_models(manager.list_models()) - logger.info("Model successfully updated") - - -def _get_model_name(existing_names, completer, default_name: str = "") -> str: - done = False - completer.set_line(default_name) - while not done: - model_name = input(f"Short name for this model [{default_name}]: ").strip() - if len(model_name) == 0: - model_name = default_name - if not re.match("^[\w._+:/-]+$", model_name): - logger.warning( - 'model name must contain only words, digits and the characters "._+:/-" **' - ) - elif model_name != default_name and model_name in existing_names: - logger.warning(f"the name {model_name} is already in use. Pick another.") - else: - done = True - return model_name - - -def do_textmask(gen, opt, callback): - image_path = opt.prompt - if not os.path.exists(image_path): - image_path = os.path.join(opt.outdir, image_path) - assert os.path.exists( - image_path - ), '** "{opt.prompt}" not found. Please enter the name of an existing image file to mask **' - assert ( - opt.text_mask is not None and len(opt.text_mask) >= 1 - ), "** Please provide a text mask with -tm **" - opt.input_file_path = image_path - tm = opt.text_mask[0] - threshold = float(opt.text_mask[1]) if len(opt.text_mask) > 1 else 0.5 - gen.apply_textmask( - image_path=image_path, - prompt=tm, - threshold=threshold, - callback=callback, - ) - - -def do_postprocess(gen, opt, callback): - file_path = opt.prompt # treat the prompt as the file pathname - if opt.new_prompt is not None: - opt.prompt = opt.new_prompt - else: - opt.prompt = None - - if os.path.dirname(file_path) == "": # basename given - file_path = os.path.join(opt.outdir, file_path) - - opt.input_file_path = file_path - - tool = None - if opt.facetool_strength > 0: - tool = opt.facetool - elif opt.embiggen: - tool = "embiggen" - elif opt.upscale: - tool = "upscale" - elif opt.out_direction: - tool = "outpaint" - elif opt.outcrop: - tool = "outcrop" - opt.save_original = True # do not overwrite old image! - opt.last_operation = f"postprocess:{tool}" - try: - gen.apply_postprocessor( - image_path=file_path, - tool=tool, - facetool_strength=opt.facetool_strength, - codeformer_fidelity=opt.codeformer_fidelity, - save_original=opt.save_original, - upscale=opt.upscale, - upscale_denoise_str=opt.esrgan_denoise_str, - out_direction=opt.out_direction, - outcrop=opt.outcrop, - callback=callback, - opt=opt, - ) - except OSError: - logger.error(f"{file_path}: file could not be read",exc_info=True) - return - except (KeyError, AttributeError): - logger.error(f"an error occurred while applying the {tool} postprocessor",exc_info=True) - return - return opt.last_operation - - -def add_postprocessing_to_metadata(opt, original_file, new_file, tool, command): - original_file = ( - original_file - if os.path.exists(original_file) - else os.path.join(opt.outdir, original_file) - ) - new_file = ( - new_file if os.path.exists(new_file) else os.path.join(opt.outdir, new_file) - ) - try: - meta = retrieve_metadata(original_file)["sd-metadata"] - except AttributeError: - try: - meta = retrieve_metadata(new_file)["sd-metadata"] - except AttributeError: - meta = {} - - if "image" not in meta: - meta = metadata_dumps(opt, seeds=[opt.seed])["image"] - meta["image"] = {} - img_data = meta.get("image") - pp = img_data.get("postprocessing", []) or [] - pp.append( - { - "tool": tool, - "dream_command": command, - } - ) - meta["image"]["postprocessing"] = pp - write_metadata(new_file, meta) - - -def prepare_image_metadata( - opt, - prefix, - seed, - operation="generate", - prior_variations=[], - postprocessed=False, - first_seed=None, -): - if postprocessed and opt.save_original: - filename = choose_postprocess_name(opt, prefix, seed) - else: - wildcards = dict(opt.__dict__) - wildcards["prefix"] = prefix - wildcards["seed"] = seed - try: - filename = opt.fnformat.format(**wildcards) - except KeyError as e: - logger.error( - f"The filename format contains an unknown key '{e.args[0]}'. Will use {{prefix}}.{{seed}}.png' instead" - ) - filename = f"{prefix}.{seed}.png" - except IndexError: - logger.error( - "The filename format is broken or complete. Will use '{prefix}.{seed}.png' instead" - ) - filename = f"{prefix}.{seed}.png" - - if opt.variation_amount > 0: - first_seed = first_seed or seed - this_variation = [[seed, opt.variation_amount]] - opt.with_variations = prior_variations + this_variation - formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed) - elif len(prior_variations) > 0: - formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed) - elif operation == "postprocess": - formatted_dream_prompt = "!fix " + opt.dream_prompt_str( - seed=seed, prompt=opt.input_file_path - ) - else: - formatted_dream_prompt = opt.dream_prompt_str(seed=seed) - return filename, formatted_dream_prompt - - -def choose_postprocess_name(opt, prefix, seed) -> str: - match = re.search("postprocess:(\w+)", opt.last_operation) - if match: - modifier = match.group( - 1 - ) # will look like "gfpgan", "upscale", "outpaint" or "embiggen" - else: - modifier = "postprocessed" - - counter = 0 - filename = None - available = False - while not available: - if counter == 0: - filename = f"{prefix}.{seed}.{modifier}.png" - else: - filename = f"{prefix}.{seed}.{modifier}-{counter:02d}.png" - available = not os.path.exists(os.path.join(opt.outdir, filename)) - counter += 1 - return filename - - -def get_next_command(infile=None, model_name="no model") -> str: # command string - if infile is None: - command = input(f"({model_name}) invoke> ").strip() - else: - command = infile.readline() - if not command: - raise EOFError - else: - command = command.strip() - if len(command) > 0: - print(f"#{command}") - return command - - -def invoke_ai_web_server_loop(gen: Generate, gfpgan, codeformer, esrgan): - print("\n* --web was specified, starting web server...") - from invokeai.backend.web import InvokeAIWebServer - - # Change working directory to the stable-diffusion directory - os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) - - invoke_ai_web_server = InvokeAIWebServer( - generate=gen, gfpgan=gfpgan, codeformer=codeformer, esrgan=esrgan - ) - - try: - invoke_ai_web_server.run() - except KeyboardInterrupt: - pass - - -def add_embedding_terms(gen, completer): - """ - Called after setting the model, updates the autocompleter with - any terms loaded by the embedding manager. - """ - trigger_strings = gen.model.textual_inversion_manager.get_all_trigger_strings() - completer.add_embedding_terms(trigger_strings) - - -def split_variations(variations_string) -> list: - # shotgun parsing, woo - parts = [] - broken = False # python doesn't have labeled loops... - for part in variations_string.split(","): - seed_and_weight = part.split(":") - if len(seed_and_weight) != 2: - logger.warning(f'Could not parse with_variation part "{part}"') - broken = True - break - try: - seed = int(seed_and_weight[0]) - weight = float(seed_and_weight[1]) - except ValueError: - logger.warning(f'Could not parse with_variation part "{part}"') - broken = True - break - parts.append([seed, weight]) - if broken: - return None - elif len(parts) == 0: - return None - else: - return parts - - -def load_face_restoration(opt): - try: - gfpgan, codeformer, esrgan = None, None, None - if opt.restore or opt.esrgan: - from invokeai.backend.restoration import Restoration - - restoration = Restoration() - if opt.restore: - gfpgan, codeformer = restoration.load_face_restore_models( - opt.gfpgan_model_path - ) - else: - logger.info("Face restoration disabled") - if opt.esrgan: - esrgan = restoration.load_esrgan(opt.esrgan_bg_tile) - else: - logger.info("Upscaling disabled") - else: - logger.info("Face restoration and upscaling disabled") - except (ModuleNotFoundError, ImportError): - print(traceback.format_exc(), file=sys.stderr) - logger.info("You may need to install the ESRGAN and/or GFPGAN modules") - return gfpgan, codeformer, esrgan - - -def make_step_callback(gen, opt, prefix): - destination = os.path.join(opt.outdir, "intermediates", prefix) - os.makedirs(destination, exist_ok=True) - logger.info(f"Intermediate images will be written into {destination}") - - def callback(state: PipelineIntermediateState): - latents = state.latents - step = state.step - if step % opt.save_intermediates == 0 or step == opt.steps - 1: - filename = os.path.join(destination, f"{step:04}.png") - image = gen.sample_to_lowres_estimated_image(latents) - image = image.resize((image.size[0] * 8, image.size[1] * 8)) - image.save(filename, "PNG") - - return callback - - -def retrieve_dream_command(opt, command, completer): - """ - Given a full or partial path to a previously-generated image file, - will retrieve and format the dream command used to generate the image, - and pop it into the readline buffer (linux, Mac), or print out a comment - for cut-and-paste (windows) - - Given a wildcard path to a folder with image png files, - will retrieve and format the dream command used to generate the images, - and save them to a file commands.txt for further processing - """ - if len(command) == 0: - return - - tokens = command.split() - dir, basename = os.path.split(tokens[0]) - if len(dir) == 0: - path = os.path.join(opt.outdir, basename) - else: - path = tokens[0] - - if len(tokens) > 1: - return write_commands(opt, path, tokens[1]) - - cmd = "" - try: - cmd = dream_cmd_from_png(path) - except OSError: - logger.error(f"{tokens[0]}: file could not be read") - except (KeyError, AttributeError, IndexError): - logger.error(f"{tokens[0]}: file has no metadata") - except: - logger.error(f"{tokens[0]}: file could not be processed") - if len(cmd) > 0: - completer.set_line(cmd) - -def write_commands(opt, file_path: str, outfilepath: str): - dir, basename = os.path.split(file_path) - try: - paths = sorted(list(Path(dir).glob(basename))) - except ValueError: - logger.error(f'"{basename}": unacceptable pattern') - return - - commands = [] - cmd = None - for path in paths: - try: - cmd = dream_cmd_from_png(path) - except (KeyError, AttributeError, IndexError): - logger.error(f"{path}: file has no metadata") - except: - logger.error(f"{path}: file could not be processed") - if cmd: - commands.append(f"# {path}") - commands.append(cmd) - if len(commands) > 0: - dir, basename = os.path.split(outfilepath) - if len(dir) == 0: - outfilepath = os.path.join(opt.outdir, basename) - with open(outfilepath, "w", encoding="utf-8") as f: - f.write("\n".join(commands)) - logger.info(f"File {outfilepath} with commands created") - - -def report_model_error(opt: Namespace, e: Exception): - logger.warning(f'An error occurred while attempting to initialize the model: "{str(e)}"') - logger.warning( - "This can be caused by a missing or corrupted models file, and can sometimes be fixed by (re)installing the models." - ) - yes_to_all = os.environ.get("INVOKE_MODEL_RECONFIGURE") - if yes_to_all: - logger.warning( - "Reconfiguration is being forced by environment variable INVOKE_MODEL_RECONFIGURE" - ) - else: - if not click.confirm( - "Do you want to run invokeai-configure script to select and/or reinstall models?", - default=False, - ): - return - - logger.info("invokeai-configure is launching....\n") - - # Match arguments that were set on the CLI - # only the arguments accepted by the configuration script are parsed - root_dir = ["--root", opt.root_dir] if opt.root_dir is not None else [] - config = ["--config", opt.conf] if opt.conf is not None else [] - previous_args = sys.argv - sys.argv = ["invokeai-configure"] - sys.argv.extend(root_dir) - sys.argv.extend(config) - if yes_to_all is not None: - for arg in yes_to_all.split(): - sys.argv.append(arg) - - from ..install import invokeai_configure - - invokeai_configure() - logger.warning("InvokeAI will now restart") - sys.argv = previous_args - main() # would rather do a os.exec(), but doesn't exist? - sys.exit(0) - - -def check_internet() -> bool: - """ - Return true if the internet is reachable. - It does this by pinging huggingface.co. - """ - import urllib.request - - host = "http://huggingface.co" - try: - urllib.request.urlopen(host, timeout=1) - return True - except: - return False - -# This routine performs any patch-ups needed after installation -def run_patches(): - # install ckpt configuration files that may have been added to the - # distro after original root directory configuration - import invokeai.configs as conf - from shutil import copyfile - - root_configs = Path(global_config_dir(), 'stable-diffusion') - repo_configs = Path(conf.__path__[0], 'stable-diffusion') - if not root_configs.exists(): - os.makedirs(root_configs, exist_ok=True) - for src in repo_configs.iterdir(): - dest = root_configs / src.name - if not dest.exists(): - copyfile(src, dest) - -if __name__ == "__main__": - main() diff --git a/invokeai/frontend/CLI/readline.py b/invokeai/frontend/CLI/readline.py deleted file mode 100644 index 5a877ae810..0000000000 --- a/invokeai/frontend/CLI/readline.py +++ /dev/null @@ -1,497 +0,0 @@ -""" -Readline helper functions for invoke.py. -You may import the global singleton `completer` to get access to the -completer object itself. This is useful when you want to autocomplete -seeds: - - from invokeai.frontend.CLI.readline import completer - completer.add_seed(18247566) - completer.add_seed(9281839) -""" -import atexit -import os -import re - -from ...backend.args import Args -from ...backend.globals import Globals -from ...backend.stable_diffusion import HuggingFaceConceptsLibrary - -# ---------------readline utilities--------------------- -try: - import readline - - readline_available = True -except (ImportError, ModuleNotFoundError) as e: - print(f"** An error occurred when loading the readline module: {str(e)}") - readline_available = False - -IMG_EXTENSIONS = (".png", ".jpg", ".jpeg", ".PNG", ".JPG", ".JPEG", ".gif", ".GIF") -WEIGHT_EXTENSIONS = (".ckpt", ".vae", ".safetensors") -TEXT_EXTENSIONS = (".txt", ".TXT") -CONFIG_EXTENSIONS = (".yaml", ".yml") -COMMANDS = ( - "--steps", - "-s", - "--seed", - "-S", - "--iterations", - "-n", - "--width", - "-W", - "--height", - "-H", - "--cfg_scale", - "-C", - "--threshold", - "--perlin", - "--grid", - "-g", - "--individual", - "-i", - "--save_intermediates", - "--init_img", - "-I", - "--init_mask", - "-M", - "--init_color", - "--strength", - "-f", - "--variants", - "-v", - "--outdir", - "-o", - "--sampler", - "-A", - "-m", - "--embedding_path", - "--device", - "--grid", - "-g", - "--facetool", - "-ft", - "--facetool_strength", - "-G", - "--codeformer_fidelity", - "-cf", - "--upscale", - "-U", - "-save_orig", - "--save_original", - "--log_tokenization", - "-t", - "--hires_fix", - "--inpaint_replace", - "-r", - "--png_compression", - "-z", - "--text_mask", - "-tm", - "--h_symmetry_time_pct", - "--v_symmetry_time_pct", - "!fix", - "!fetch", - "!replay", - "!history", - "!search", - "!clear", - "!models", - "!switch", - "!import_model", - "!optimize_model", - "!convert_model", - "!edit_model", - "!del_model", - "!mask", - "!triggers", -) -MODEL_COMMANDS = ( - "!switch", - "!edit_model", - "!del_model", -) -CKPT_MODEL_COMMANDS = ("!optimize_model",) -WEIGHT_COMMANDS = ( - "!import_model", - "!convert_model", -) -IMG_PATH_COMMANDS = ("--outdir[=\s]",) -TEXT_PATH_COMMANDS = ("!replay",) -IMG_FILE_COMMANDS = ( - "!fix", - "!fetch", - "!mask", - "--init_img[=\s]", - "-I", - "--init_mask[=\s]", - "-M", - "--init_color[=\s]", - "--embedding_path[=\s]", -) - -path_regexp = "(" + "|".join(IMG_PATH_COMMANDS + IMG_FILE_COMMANDS) + ")\s*\S*$" -weight_regexp = "(" + "|".join(WEIGHT_COMMANDS) + ")\s*\S*$" -text_regexp = "(" + "|".join(TEXT_PATH_COMMANDS) + ")\s*\S*$" - - -class Completer(object): - def __init__(self, options, models={}): - self.options = sorted(options) - self.models = models - self.seeds = set() - self.matches = list() - self.default_dir = None - self.linebuffer = None - self.auto_history_active = True - self.extensions = None - self.concepts = None - self.embedding_terms = set() - return - - def complete(self, text, state): - """ - Completes invoke command line. - BUG: it doesn't correctly complete files that have spaces in the name. - """ - buffer = readline.get_line_buffer() - - if state == 0: - # extensions defined, so go directly into path completion mode - if self.extensions is not None: - self.matches = self._path_completions(text, state, self.extensions) - - # looking for an image file - elif re.search(path_regexp, buffer): - do_shortcut = re.search("^" + "|".join(IMG_FILE_COMMANDS), buffer) - self.matches = self._path_completions( - text, state, IMG_EXTENSIONS, shortcut_ok=do_shortcut - ) - - # looking for a seed - elif re.search("(-S\s*|--seed[=\s])\d*$", buffer): - self.matches = self._seed_completions(text, state) - - # looking for an embedding concept - elif re.search("<[\w-]*$", buffer): - self.matches = self._concept_completions(text, state) - - # looking for a model - elif re.match("^" + "|".join(MODEL_COMMANDS), buffer): - self.matches = self._model_completions(text, state) - - # looking for a ckpt model - elif re.match("^" + "|".join(CKPT_MODEL_COMMANDS), buffer): - self.matches = self._model_completions(text, state, ckpt_only=True) - - elif re.search(weight_regexp, buffer): - self.matches = self._path_completions( - text, - state, - WEIGHT_EXTENSIONS, - default_dir=Globals.root, - ) - - elif re.search(text_regexp, buffer): - self.matches = self._path_completions(text, state, TEXT_EXTENSIONS) - - # This is the first time for this text, so build a match list. - elif text: - self.matches = [s for s in self.options if s and s.startswith(text)] - else: - self.matches = self.options[:] - - # Return the state'th item from the match list, - # if we have that many. - try: - response = self.matches[state] - except IndexError: - response = None - return response - - def complete_extensions(self, extensions: list): - """ - If called with a list of extensions, will force completer - to do file path completions. - """ - self.extensions = extensions - - def add_history(self, line): - """ - Pass thru to readline - """ - if not self.auto_history_active: - readline.add_history(line) - - def clear_history(self): - """ - Pass clear_history() thru to readline - """ - readline.clear_history() - - def search_history(self, match: str): - """ - Like show_history() but only shows items that - contain the match string. - """ - self.show_history(match) - - def remove_history_item(self, pos): - readline.remove_history_item(pos) - - def add_seed(self, seed): - """ - Add a seed to the autocomplete list for display when -S is autocompleted. - """ - if seed is not None: - self.seeds.add(str(seed)) - - def set_default_dir(self, path): - self.default_dir = path - - def set_options(self, options): - self.options = options - - def get_line(self, index): - try: - line = self.get_history_item(index) - except IndexError: - return None - return line - - def get_current_history_length(self): - return readline.get_current_history_length() - - def get_history_item(self, index): - return readline.get_history_item(index) - - def show_history(self, match=None): - """ - Print the session history using the pydoc pager - """ - import pydoc - - lines = list() - h_len = self.get_current_history_length() - if h_len < 1: - print("") - return - - for i in range(0, h_len): - line = self.get_history_item(i + 1) - if match and match not in line: - continue - lines.append(f"[{i+1}] {line}") - pydoc.pager("\n".join(lines)) - - def set_line(self, line) -> None: - """ - Set the default string displayed in the next line of input. - """ - self.linebuffer = line - readline.redisplay() - - def update_models(self, models: dict) -> None: - """ - update our list of models - """ - self.models = models - - def _seed_completions(self, text, state): - m = re.search("(-S\s?|--seed[=\s]?)(\d*)", text) - if m: - switch = m.groups()[0] - partial = m.groups()[1] - else: - switch = "" - partial = text - - matches = list() - for s in self.seeds: - if s.startswith(partial): - matches.append(switch + s) - matches.sort() - return matches - - def add_embedding_terms(self, terms: list[str]): - self.embedding_terms = set(terms) - if self.concepts: - self.embedding_terms.update(set(self.concepts.list_concepts())) - - def _concept_completions(self, text, state): - if self.concepts is None: - # cache Concepts() instance so we can check for updates in concepts_list during runtime. - self.concepts = HuggingFaceConceptsLibrary() - self.embedding_terms.update(set(self.concepts.list_concepts())) - else: - self.embedding_terms.update(set(self.concepts.list_concepts())) - - partial = text[1:] # this removes the leading '<' - if len(partial) == 0: - return list(self.embedding_terms) # whole dump - think if user wants this! - - matches = list() - for concept in self.embedding_terms: - if concept.startswith(partial): - matches.append(f"<{concept}>") - matches.sort() - return matches - - def _model_completions(self, text, state, ckpt_only=False): - m = re.search("(!switch\s+)(\w*)", text) - if m: - switch = m.groups()[0] - partial = m.groups()[1] - else: - switch = "" - partial = text - matches = list() - for s in self.models: - format = self.models[s]["format"] - if format == "vae": - continue - if ckpt_only and format != "ckpt": - continue - if s.startswith(partial): - matches.append(switch + s) - matches.sort() - return matches - - def _pre_input_hook(self): - if self.linebuffer: - readline.insert_text(self.linebuffer) - readline.redisplay() - self.linebuffer = None - - def _path_completions( - self, text, state, extensions, shortcut_ok=True, default_dir: str = "" - ): - # separate the switch from the partial path - match = re.search("^(-\w|--\w+=?)(.*)", text) - if match is None: - switch = None - partial_path = text - else: - switch, partial_path = match.groups() - - partial_path = partial_path.lstrip() - - matches = list() - path = os.path.expanduser(partial_path) - - if os.path.isdir(path): - dir = path - elif os.path.dirname(path) != "": - dir = os.path.dirname(path) - else: - dir = default_dir if os.path.exists(default_dir) else "" - path = os.path.join(dir, path) - - dir_list = os.listdir(dir or ".") - if shortcut_ok and os.path.exists(self.default_dir) and dir == "": - dir_list += os.listdir(self.default_dir) - - for node in dir_list: - if node.startswith(".") and len(node) > 1: - continue - full_path = os.path.join(dir, node) - - if not (node.endswith(extensions) or os.path.isdir(full_path)): - continue - - if path and not full_path.startswith(path): - continue - - if switch is None: - match_path = os.path.join(dir, node) - matches.append( - match_path + "/" if os.path.isdir(full_path) else match_path - ) - elif os.path.isdir(full_path): - matches.append( - switch + os.path.join(os.path.dirname(full_path), node) + "/" - ) - elif node.endswith(extensions): - matches.append(switch + os.path.join(os.path.dirname(full_path), node)) - - return matches - - -class DummyCompleter(Completer): - def __init__(self, options): - super().__init__(options) - self.history = list() - - def add_history(self, line): - self.history.append(line) - - def clear_history(self): - self.history = list() - - def get_current_history_length(self): - return len(self.history) - - def get_history_item(self, index): - return self.history[index - 1] - - def remove_history_item(self, index): - return self.history.pop(index - 1) - - def set_line(self, line): - print(f"# {line}") - - -def generic_completer(commands: list) -> Completer: - if readline_available: - completer = Completer(commands, []) - readline.set_completer(completer.complete) - readline.set_pre_input_hook(completer._pre_input_hook) - readline.set_completer_delims(" ") - readline.parse_and_bind("tab: complete") - readline.parse_and_bind("set print-completions-horizontally off") - readline.parse_and_bind("set page-completions on") - readline.parse_and_bind("set skip-completed-text on") - readline.parse_and_bind("set show-all-if-ambiguous on") - else: - completer = DummyCompleter(commands) - return completer - - -def get_completer(opt: Args, models=[]) -> Completer: - if readline_available: - completer = Completer(COMMANDS, models) - - readline.set_completer(completer.complete) - # pyreadline3 does not have a set_auto_history() method - try: - readline.set_auto_history(False) - completer.auto_history_active = False - except: - completer.auto_history_active = True - readline.set_pre_input_hook(completer._pre_input_hook) - readline.set_completer_delims(" ") - readline.parse_and_bind("tab: complete") - readline.parse_and_bind("set print-completions-horizontally off") - readline.parse_and_bind("set page-completions on") - readline.parse_and_bind("set skip-completed-text on") - readline.parse_and_bind("set show-all-if-ambiguous on") - - outdir = os.path.expanduser(opt.outdir) - if os.path.isabs(outdir): - histfile = os.path.join(outdir, ".invoke_history") - else: - histfile = os.path.join(Globals.root, outdir, ".invoke_history") - try: - readline.read_history_file(histfile) - readline.set_history_length(1000) - except FileNotFoundError: - pass - except OSError: # file likely corrupted - newname = f"{histfile}.old" - print( - f"## Your history file {histfile} couldn't be loaded and may be corrupted. Renaming it to {newname}" - ) - os.replace(histfile, newname) - atexit.register(readline.write_history_file, histfile) - - else: - completer = DummyCompleter(COMMANDS) - return completer diff --git a/invokeai/frontend/CLI/sd_metadata.py b/invokeai/frontend/CLI/sd_metadata.py deleted file mode 100644 index c26907a18e..0000000000 --- a/invokeai/frontend/CLI/sd_metadata.py +++ /dev/null @@ -1,30 +0,0 @@ -''' -This is a modularized version of the sd-metadata.py script, -which retrieves and prints the metadata from a series of generated png files. -''' -import sys -import json -from invokeai.backend.image_util import retrieve_metadata - - -def print_metadata(): - if len(sys.argv) < 2: - print("Usage: file2prompt.py ...") - print("This script opens up the indicated invoke.py-generated PNG file(s) and prints out their metadata.") - exit(-1) - - filenames = sys.argv[1:] - for f in filenames: - try: - metadata = retrieve_metadata(f) - print(f'{f}:\n',json.dumps(metadata['sd-metadata'], indent=4)) - except FileNotFoundError: - sys.stderr.write(f'{f} not found\n') - continue - except PermissionError: - sys.stderr.write(f'{f} could not be opened due to inadequate permissions\n') - continue - -if __name__== '__main__': - print_metadata() - diff --git a/pyproject.toml b/pyproject.toml index 0ab1e0ede6..df1476186f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,6 @@ dependencies = [ "textual_inversion.py" = "invokeai.frontend.training:invokeai_textual_inversion" # modern entrypoints -"invokeai" = "invokeai.frontend.CLI:invokeai_command_line_interface" "invokeai-configure" = "invokeai.frontend.install:invokeai_configure" "invokeai-merge" = "invokeai.frontend.merge:invokeai_merge_diffusers" "invokeai-ti" = "invokeai.frontend.training:invokeai_textual_inversion"