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