From 1a0cf1320bde16d9be6f7a7a6aad4fd677992404 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 30 Sep 2022 14:53:37 -0400 Subject: [PATCH] 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. --- ldm/dream/readline.py | 69 ++++++++++++++++++++++--------------------- ldm/generate.py | 3 ++ scripts/dream.py | 63 +++++++++++++++++++++++---------------- 3 files changed, 76 insertions(+), 59 deletions(-) diff --git a/ldm/dream/readline.py b/ldm/dream/readline.py index 55a6026d15..014ffeaaa3 100644 --- a/ldm/dream/readline.py +++ b/ldm/dream/readline.py @@ -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 diff --git a/ldm/generate.py b/ldm/generate.py index 9a958ee65f..c0936fdccf 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -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 diff --git a/scripts/dream.py b/scripts/dream.py index 0ffef91e07..7a9b26a995 100644 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -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__':