From 0885fb4dab6d23ccb3f1d3297e44fd701223d432 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 9 Sep 2022 14:37:03 -0400 Subject: [PATCH] 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)