From 0885fb4dab6d23ccb3f1d3297e44fd701223d432 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Fri, 9 Sep 2022 14:37:03 -0400
Subject: [PATCH 1/2] added ansi escapes to highlight key parts of CLI session

---
 scripts/dream.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/scripts/dream.py b/scripts/dream.py
index 56c7ed769b1..42327832450 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -15,6 +15,11 @@
 from ldm.dream.image_util import make_grid
 from omegaconf import OmegaConf
 
+# Placeholder to be replaced with proper class that tracks the
+# outputs and associates with the prompt that generated them.
+# Just want to get the formatting look right for now.
+output_cntr = 0
+
 def main():
     """Initialize command-line parsers and the diffusion model"""
     arg_parser = create_argv_parser()
@@ -292,16 +297,18 @@ def image_writer(image, seed, upscaled=False):
             print(e)
             continue
 
-        print('Outputs:')
+        print('\033[1mOutputs:\033[0m')
         log_path = os.path.join(current_outdir, 'dream_log.txt')
         write_log_message(results, log_path)
 
-    print('goodbye!')
+    print('goodbye!\033[0m')
 
 
 def get_next_command(infile=None) -> str: #command string
     if infile is None:
-        command = input('dream> ')
+        print('\033[1m') # add some boldface
+        command = input('dream> ')   
+        print('\033[0m',end='')
     else:
         command = infile.readline()
         if not command:
@@ -339,8 +346,11 @@ def dream_server_loop(t2i, host, port, outdir):
 
 def write_log_message(results, log_path):
     """logs the name of the output image, prompt, and prompt args to the terminal and log file"""
+    global output_cntr
     log_lines = [f'{path}: {prompt}\n' for path, prompt in results]
-    print(*log_lines, sep='')
+    for l in log_lines:
+        output_cntr += 1
+        print(f'\033[1m[{output_cntr}]\033[0m {l}',end='')
 
     with open(log_path, 'a', encoding='utf-8') as file:
         file.writelines(log_lines)

From 0e296acc649678879f92559d14b8c4569ccf2636 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Fri, 9 Sep 2022 17:26:22 -0400
Subject: [PATCH 2/2] adjust exception handling so that ^C will abort when
 reading prompts from a file

---
 ldm/generate.py  | 6 +++++-
 scripts/dream.py | 3 ++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/ldm/generate.py b/ldm/generate.py
index 3a81087d5f5..27f89bb4d6e 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -117,6 +117,7 @@ def __init__(
             seamless              = False,
             embedding_path        = None,
             device_type           = 'cuda',
+            ignore_ctrl_c         = False,
     ):
         self.iterations               = iterations
         self.width                    = width
@@ -134,6 +135,7 @@ def __init__(
         self.seamless                 = seamless
         self.embedding_path           = embedding_path
         self.device_type              = device_type
+        self.ignore_ctrl_c            = ignore_ctrl_c    # note, this logic probably doesn't belong here...
         self.model                    = None     # empty for now
         self.sampler                  = None
         self.device                   = None
@@ -210,7 +212,7 @@ def prompt2image(
             **args,
     ):   # eat up additional cruft
         """
-        ldm.prompt2image() is the common entry point for txt2img() and img2img()
+        ldm.generate.prompt2image() is the common entry point for txt2img() and img2img()
         It takes the following arguments:
            prompt                          // prompt string (no default)
            iterations                      // iterations (1); image count=iterations
@@ -341,6 +343,8 @@ def process_image(image,seed):
 
         except KeyboardInterrupt:
             print('*interrupted*')
+            if not self.ignore_ctrl_c:
+                raise KeyboardInterrupt
             print(
                 '>> Partial results will be returned; if --grid was requested, nothing will be returned.'
             )
diff --git a/scripts/dream.py b/scripts/dream.py
index 42327832450..0b1e80f66d8 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -68,7 +68,8 @@ def main():
         # this is solely for recreating the prompt
         seamless       = opt.seamless,
         embedding_path = opt.embedding_path,
-        device_type    = opt.device
+        device_type    = opt.device,
+        ignore_ctrl_c  = opt.infile is None,
     )
 
     # make sure the output directory exists