mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
116415b3fc
- Script will now offer the user the ability to create a minimal models.yaml and then gracefully exit. - Closes #1420
383 lines
11 KiB
Python
383 lines
11 KiB
Python
"""
|
|
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 ldm.invoke.readline import completer
|
|
completer.add_seed(18247566)
|
|
completer.add_seed(9281839)
|
|
"""
|
|
import os
|
|
import re
|
|
import atexit
|
|
from ldm.invoke.args import Args
|
|
|
|
# ---------------readline utilities---------------------
|
|
try:
|
|
import readline
|
|
readline_available = True
|
|
except (ImportError,ModuleNotFoundError):
|
|
readline_available = False
|
|
|
|
IMG_EXTENSIONS = ('.png','.jpg','.jpeg','.PNG','.JPG','.JPEG','.gif','.GIF')
|
|
WEIGHT_EXTENSIONS = ('.ckpt','.bae')
|
|
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',
|
|
'--skip_normalize','-x',
|
|
'--log_tokenization','-t',
|
|
'--hires_fix',
|
|
'--inpaint_replace','-r',
|
|
'--png_compression','-z',
|
|
'--text_mask','-tm',
|
|
'!fix','!fetch','!replay','!history','!search','!clear',
|
|
'!models','!switch','!import_model','!edit_model','!del_model',
|
|
'!mask',
|
|
)
|
|
MODEL_COMMANDS = (
|
|
'!switch',
|
|
'!edit_model',
|
|
'!del_model',
|
|
)
|
|
WEIGHT_COMMANDS = (
|
|
'!import_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 = sorted(models)
|
|
self.seeds = set()
|
|
self.matches = list()
|
|
self.default_dir = None
|
|
self.linebuffer = None
|
|
self.auto_history_active = True
|
|
self.extensions = None
|
|
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 a model
|
|
elif re.match('^'+'|'.join(MODEL_COMMANDS),buffer):
|
|
self.matches= self._model_completions(text, state)
|
|
|
|
elif re.search(weight_regexp,buffer):
|
|
self.matches = self._path_completions(text, state, WEIGHT_EXTENSIONS)
|
|
|
|
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 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 add_model(self,model_name:str)->None:
|
|
'''
|
|
add a model name to the completion list
|
|
'''
|
|
self.models.append(model_name)
|
|
|
|
def del_model(self,model_name:str)->None:
|
|
'''
|
|
removes a model name from the completion list
|
|
'''
|
|
self.models.remove(model_name)
|
|
|
|
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 _model_completions(self, text, state):
|
|
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:
|
|
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):
|
|
# 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 = ''
|
|
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 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 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')
|
|
|
|
histfile = os.path.join(os.path.expanduser(opt.outdir), '.invoke_history')
|
|
try:
|
|
readline.read_history_file(histfile)
|
|
readline.set_history_length(1000)
|
|
except FileNotFoundError:
|
|
pass
|
|
atexit.register(readline.write_history_file, histfile)
|
|
|
|
else:
|
|
completer = DummyCompleter(COMMANDS)
|
|
return completer
|