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)