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

@@ -11,8 +11,7 @@ seeds:
 import os
 import re
 import atexit
-
+from ldm.dream.args import Args
-completer = None
 
 # ---------------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',
@@ -49,7 +44,7 @@ COMMANDS = (
     '-save_orig','--save_original',
     '--skip_normalize','-x',
     '--log_tokenization','-t',
-    '!fix','!fetch',
+    '!fix','!fetch','!history',
     )
 IMG_PATH_COMMANDS = (
     '--outdir[=\s]',
@@ -71,6 +66,7 @@ class Completer:
         self.matches     = list()
         self.default_dir = None
         self.linebuffer  = None
+        self.auto_history_active = True
         return
 
     def complete(self, text, state):
@@ -109,7 +105,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)
@@ -247,29 +244,35 @@ class DummyCompleter(Completer):
     def set_line(self,line):
         print(f'# {line}')
 
-if readline_available:
+def get_completer(opt:Args)->Completer:
-    completer = Completer(COMMANDS)
+    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)
 
-else:
+        readline.set_completer(
-    completer = DummyCompleter(COMMANDS)
+            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

ldm/generate.py

@@ -583,6 +583,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

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
+        if command.startswith('!'):
-            command = command.replace('!dream ','',1)
+            subcommand = command[1:]
 
-        if command.startswith('!fix'):
+            if subcommand.startswith('dream'):   # in case a stored prompt still contains the !dream command
-            command = command.replace('!fix ','',1)
+                command = command.replace('!dream ','',1)
-            operation = 'postprocess'
 
-        if command.startswith('!fetch'):
+            elif subcommand.startswith('fix'):
-            file_path = command.replace('!fetch ','',1)
+                command = command.replace('!fix ','',1)
-            retrieve_dream_command(opt,file_path)
+                operation = 'postprocess'
-            continue
 
-        if command == '!history':
+            elif subcommand.startswith('fetch'):
-            completer.show_history()
+                file_path = command.replace('!fetch ','',1)
-            continue
+                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__':