Refactor file handling

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(

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():