improve behavior and fix bugs in CLI history handling and completion

-if readline.set_auto_history() is not implemented, as in pyreadline3, will fall
back gracefully to automatic history saving. The only issue with this is that
-!history commands will be recorded in the history.
-!fetch on missing file no longer crashes script
-!history is now one of the autocomplete commands
-.dream_history now stored in output directory rather than ~user directory.

An important limitation of the last feature is that the history is
loaded and saved to the .dream_history file in the --outdir directory
specified at script launch time. It is not swapped around when the
--outdir is changed during the session.
This commit is contained in:
Lincoln Stein 2022-09-30 14:53:37 -04:00
parent fe28c5fbdc
commit 1a0cf1320b
3 changed files with 76 additions and 59 deletions

View File

@ -11,8 +11,7 @@ seeds:
import os
import re
import atexit
completer = None
from ldm.dream.args import Args
# ---------------readline utilities---------------------
try:
@ -21,10 +20,6 @@ try:
except:
readline_available = False
#to simulate what happens on windows systems, uncomment
# this line
#readline_available = False
IMG_EXTENSIONS = ('.png','.jpg','.jpeg')
COMMANDS = (
'--steps','-s',
@ -50,7 +45,7 @@ COMMANDS = (
'--skip_normalize','-x',
'--log_tokenization','-t',
'--hires_fix',
'!fix','!fetch',
'!fix','!fetch','!history',
)
IMG_PATH_COMMANDS = (
'--outdir[=\s]',
@ -72,6 +67,7 @@ class Completer:
self.matches = list()
self.default_dir = None
self.linebuffer = None
self.auto_history_active = True
return
def complete(self, text, state):
@ -110,7 +106,8 @@ class Completer:
'''
Pass thru to readline
'''
readline.add_history(line)
if not self.auto_history_active:
readline.add_history(line)
def remove_history_item(self,pos):
readline.remove_history_item(pos)
@ -248,29 +245,35 @@ class DummyCompleter(Completer):
def set_line(self,line):
print(f'# {line}')
if readline_available:
completer = Completer(COMMANDS)
readline.set_completer(
completer.complete
)
readline.set_auto_history(False)
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 bell-style visible')
readline.parse_and_bind('set show-all-if-ambiguous on')
histfile = os.path.join(os.path.expanduser('~'), '.dream_history')
try:
readline.read_history_file(histfile)
readline.set_history_length(1000)
except FileNotFoundError:
pass
atexit.register(readline.write_history_file, histfile)
def get_completer(opt:Args)->Completer:
if readline_available:
completer = Completer(COMMANDS)
else:
completer = DummyCompleter(COMMANDS)
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), '.dream_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

View File

@ -586,6 +586,9 @@ class Generate:
strength = opt.strength,
image_callback = callback,
)
elif tool is None:
print(f'* please provide at least one postprocessing option, such as -G or -U')
return None
else:
print(f'* postprocessing tool {tool} is not yet supported')
return None

View File

@ -9,7 +9,7 @@ import copy
import warnings
import time
sys.path.append('.') # corrects a weird problem on Macs
from ldm.dream.readline import completer
from ldm.dream.readline import get_completer
from ldm.dream.args import Args, metadata_dumps, metadata_from_png, dream_cmd_from_png
from ldm.dream.pngwriter import PngWriter
from ldm.dream.image_util import make_grid
@ -17,9 +17,6 @@ from ldm.dream.log import write_log
from omegaconf import OmegaConf
from backend.invoke_ai_web_server import InvokeAIWebServer
# The output counter labels each output and is keyed to the
# command-line history
output_cntr = completer.get_current_history_length()+1
def main():
"""Initialize command-line parsers and the diffusion model"""
@ -130,6 +127,12 @@ def main_loop(gen, opt, infile):
last_results = list()
model_config = OmegaConf.load(opt.conf)[opt.model]
# 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)
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')
@ -161,29 +164,34 @@ def main_loop(gen, opt, infile):
done = True
break
if command.startswith('!dream'): # in case a stored prompt still contains the !dream command
command = command.replace('!dream ','',1)
if command.startswith('!'):
subcommand = command[1:]
if command.startswith('!fix'):
command = command.replace('!fix ','',1)
operation = 'postprocess'
if subcommand.startswith('dream'): # in case a stored prompt still contains the !dream command
command = command.replace('!dream ','',1)
if command.startswith('!fetch'):
file_path = command.replace('!fetch ','',1)
retrieve_dream_command(opt,file_path)
continue
elif subcommand.startswith('fix'):
command = command.replace('!fix ','',1)
operation = 'postprocess'
if command == '!history':
completer.show_history()
continue
elif subcommand.startswith('fetch'):
file_path = command.replace('!fetch ','',1)
retrieve_dream_command(opt,file_path,completer)
continue
elif subcommand.startswith('history'):
completer.show_history()
continue
elif re.match('^(\d+)',subcommand):
command_no = re.match('^(\d+)',subcommand).groups()[0]
command = completer.get_line(int(command_no))
completer.set_line(command)
continue
else: # not a recognized subcommand, so give the --help text
command = '-h'
match = re.match('^!(\d+)',command)
if match:
command_no = match.groups()[0]
command = completer.get_line(int(command_no))
completer.set_line(command)
continue
if opt.parse_cmd(command) is None:
continue
@ -358,7 +366,6 @@ def main_loop(gen, opt, infile):
print('Outputs:')
log_path = os.path.join(current_outdir, 'dream_log')
global output_cntr
output_cntr = write_log(results, log_path ,('txt', 'md'), output_cntr)
print()
if operation == 'postprocess':
@ -500,7 +507,7 @@ def split_variations(variations_string) -> list:
else:
return parts
def retrieve_dream_command(opt,file_path):
def retrieve_dream_command(opt,file_path,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,
@ -512,7 +519,11 @@ def retrieve_dream_command(opt,file_path):
path = os.path.join(opt.outdir,basename)
else:
path = file_path
cmd = dream_cmd_from_png(path)
try:
cmd = dream_cmd_from_png(path)
except FileNotFoundError:
print(f'** {path}: file not found')
return
completer.set_line(cmd)
if __name__ == '__main__':