diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index 8aad3557af4..e592e34f835 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -545,7 +545,9 @@ def _load_model_from_config(self, config, ckpt): return model def _load_img(self, path): - image = Image.open(path).convert('RGB') + with Image.open(path) as img: + image = img.convert("RGB") + w, h = image.size print(f'loaded input image of size ({w}, {h}) from {path}') w, h = map( diff --git a/scripts/dream.py b/scripts/dream.py index f6feb10adc8..32e31dd7a2c 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -64,13 +64,16 @@ def main(): # gets rid of annoying messages about random seed logging.getLogger('pytorch_lightning').setLevel(logging.ERROR) + # load the infile as a list of lines infile = None - try: - if opt.infile is not None: - infile = open(opt.infile, 'r') - except FileNotFoundError as e: - print(e) - exit(-1) + if opt.infile: + if os.path.isfile(opt.infile): + with open(opt.infile, "r") as file: + infile = file.read() + infile = infile.split("\n") + else: + print(f"WARNING: '{opt.infile}' not found. Aborting.") + sys.exit(-1) # exit does not work on every os, sys.exit does afaik # preload the model t2i.load_model() @@ -115,29 +118,29 @@ def main(): ) log_path = os.path.join(opt.outdir, 'dream_log.txt') - with open(log_path, 'a') as log: - cmd_parser = create_cmd_parser() - main_loop(t2i, opt.outdir, cmd_parser, log, infile) - log.close() - if infile: - infile.close() + cmd_parser = create_cmd_parser() + main_loop(t2i, opt.outdir, cmd_parser, log_path, infile) -def main_loop(t2i, outdir, parser, log, infile): +def main_loop(t2i, outdir, parser, log_path, infile): """prompt/read/execute loop""" done = False last_seeds = [] while not done: - try: - command = infile.readline() if infile else input('dream> ') - except EOFError: - done = True - break + if not infile: + command = input("dream> ") + else: + try: + # get the next line of the infile + command = infile.pop(0) + except IndexError: + done = True + break - if infile and len(command) == 0: - done = True - break + # skip empty lines + if not command.strip(): + continue if command.startswith(('#', '//')): continue @@ -152,9 +155,6 @@ def main_loop(t2i, outdir, parser, log, infile): print(str(e)) continue - if len(elements) == 0: - continue - if elements[0] == 'q': done = True break @@ -239,7 +239,7 @@ def main_loop(t2i, outdir, parser, log, infile): continue print('Outputs:') - write_log_message(t2i, normalized_prompt, results, log) + write_log_message(t2i, normalized_prompt, results, log_path) print('goodbye!') @@ -309,19 +309,14 @@ def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400): # return variants -def write_log_message(t2i, prompt, results, logfile): +### the t2i variable doesn't seem to be necessary here. maybe remove it? +def write_log_message(t2i, prompt, results, log_path): """logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata""" - last_seed = None - img_num = 1 - seenit = {} - - for r in results: - seed = r[1] - log_message = f'{r[0]}: {prompt} -S{seed}' + log_lines = [f"{r[0]}: {prompt} -S{r[1]}\n" for r in results] + print(*log_lines, sep="") - print(log_message) - logfile.write(log_message + '\n') - logfile.flush() + with open(log_path, "a") as file: + file.writelines(log_lines) def create_argv_parser():