fixed filenames

Author: lstein

README.md

@@ -86,6 +86,21 @@ completely). The default is 0.75, and ranges from 0.25-0.75 give interesting res
 
 ## Changes
 
+* v1.05 (22 August 2022 - after the drop)
+   * Filenames now use the following formats:
+       000010.95183149.png      -- Two files produced by the same command (e.g. -n2),
+       000010.26742632.png      -- distinguished by a different seed.
+
+       000011.455191342.01.png  -- Two files produced by the same command using
+       000011.455191342.02.png  -- a batch size>1 (e.g. -b2). They have the same seed.
+
+       000011.4160627868.grid#1-4.png  -- a grid of four images (-g); the whole grid can
+                                          be regenerated with the indicated key
+
+    * It should no longer be possible for one image to overwrite another
+    * You can use the "cd" and "pwd" commands at the dream> prompt to set and retrieve
+      the path of the output directory.
+     
 * v1.04 (22 August 2022 - after the drop)
    * Updated README to reflect installation of the released weights.
    * Suppressed very noisy and inconsequential warning when loading the frozen CLIP

ldm/simplet2i.py

@@ -60,6 +60,7 @@
 from contextlib import contextmanager, nullcontext
 import time
 import math
+import re
 
 from ldm.util import instantiate_from_config
 from ldm.models.diffusion.ddim     import DDIMSampler
@@ -171,7 +172,6 @@ def txt2img(self,prompt,outdir=None,batch_size=None,iterations=None,
 
         # make directories and establish names for the output files
         os.makedirs(outdir, exist_ok=True)
-        base_count = len(os.listdir(outdir))-1
 
         start_code = None
         if self.fixed_code:
@@ -185,7 +185,7 @@ def txt2img(self,prompt,outdir=None,batch_size=None,iterations=None,
         sampler         = self.sampler
         images = list()
         seeds  = list()
-
+        filename = None
         tic    = time.time()
 
         with torch.no_grad():
@@ -218,10 +218,11 @@ def txt2img(self,prompt,outdir=None,batch_size=None,iterations=None,
                             if not grid:
                                 for x_sample in x_samples_ddim:
                                     x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                                    filename = os.path.join(outdir, f"{base_count:05}.png")
+                                    filename = self._unique_filename(outdir,previousname=filename,
+                                                                     seed=seed,isbatch=(batch_size>1))
+                                    assert not os.path.exists(filename)
                                     Image.fromarray(x_sample.astype(np.uint8)).save(filename)
                                     images.append([filename,seed])
-                                    base_count += 1
                             else:
                                 all_samples.append(x_samples_ddim)
                                 seeds.append(seed)
@@ -283,7 +284,6 @@ def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=Non
 
         # make directories and establish names for the output files
         os.makedirs(outdir, exist_ok=True)
-        base_count = len(os.listdir(outdir))-1
 
         assert os.path.isfile(init_img)
         init_image = self._load_img(init_img).to(self.device)
@@ -304,7 +304,8 @@ def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=Non
 
         images = list()
         seeds  = list()
-
+        filename = None
+        
         tic    = time.time()
 
         with torch.no_grad():
@@ -333,10 +334,10 @@ def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=Non
                             if not grid:
                                 for x_sample in x_samples:
                                     x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                                    filename = os.path.join(outdir, f"{base_count:05}.png")
+                                    filename = self._unique_filename(outdir,filename,seed=seed,isbatch=(batch_size>1))
+                                    assert not os.path.exists(filename)
                                     Image.fromarray(x_sample.astype(np.uint8)).save(filename)
                                     images.append([filename,seed])
-                                    base_count += 1
                             else:
                                 all_samples.append(x_samples)
                                 seeds.append(seed)
@@ -357,7 +358,6 @@ def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=Non
 
     def _make_grid(self,samples,seeds,batch_size,iterations,outdir):
         images = list()
-        base_count = len(os.listdir(outdir))-1
         n_rows = batch_size if batch_size>1 else int(math.sqrt(batch_size * iterations))
         # save as grid
         grid = torch.stack(samples, 0)
@@ -366,7 +366,7 @@ def _make_grid(self,samples,seeds,batch_size,iterations,outdir):
 
         # to image
         grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
-        filename = os.path.join(outdir, f"{base_count:05}.png")
+        filename = self._unique_filename(outdir,seed=seeds[0],grid_count=batch_size*iterations)
         Image.fromarray(grid.astype(np.uint8)).save(filename)
         for s in seeds:
             images.append([filename,s])
@@ -430,3 +430,40 @@ def _load_img(self,path):
         image = image[None].transpose(0, 3, 1, 2)
         image = torch.from_numpy(image)
         return 2.*image - 1.
+
+    def _unique_filename(self,outdir,previousname=None,seed=0,isbatch=False,grid_count=None):
+        revision = 1
+
+        if previousname is None:
+            # count up until we find an unfilled slot
+            dir_list  = [a.split('.',1)[0] for a in os.listdir(outdir)]
+            uniques   = dict.fromkeys(dir_list,True)
+            basecount = 1
+            while f'{basecount:06}' in uniques:
+                basecount += 1
+            if grid_count is not None:
+                grid_label = f'grid#1-{grid_count}'
+                filename = f'{basecount:06}.{seed}.{grid_label}.png'
+            elif isbatch:
+                filename = f'{basecount:06}.{seed}.01.png'
+            else:
+                filename = f'{basecount:06}.{seed}.png'
+            
+            return os.path.join(outdir,filename)
+
+        else:
+            previousname = os.path.basename(previousname)
+            x = re.match('^(\d+)\..*\.png',previousname)
+            if not x:
+                return self._unique_filename(outdir,previousname,seed)
+
+            basecount = int(x.groups()[0])
+            series = 0 
+            finished = False
+            while not finished:
+                series += 1
+                filename = f'{basecount:06}.{seed}.png'
+                if isbatch or os.path.exists(os.path.join(outdir,filename)):
+                    filename = f'{basecount:06}.{seed}.{series:02}.png'
+                finished = not os.path.exists(os.path.join(outdir,filename))
+            return os.path.join(outdir,filename)

scripts/dream.py

@@ -72,7 +72,7 @@ def main():
         t2i.load_model()
     print("\n* Initialization done! Awaiting your command (-h for help, 'q' to quit, 'cd' to change output dir, 'pwd' to print output dir)...")
 
-    log_path   = os.path.join(opt.outdir,"dream_log.txt")
+    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,cmd_parser,log)
@@ -103,9 +103,12 @@ def main_loop(t2i,parser,log):
             done = True
             break
 
-        if elements[0]=='cd' and len(elements)>1 and os.path.exists(elements[1]):
-            print(f"setting image output directory to {elements[1]}")
-            t2i.outdir=elements[1]
+        if elements[0]=='cd' and len(elements)>1:
+            if os.path.exists(elements[1]):
+                print(f"setting image output directory to {elements[1]}")
+                t2i.outdir=elements[1]
+            else:
+                print(f"directory {elements[1]} does not exist")
             continue
 
         if elements[0]=='pwd':
@@ -166,7 +169,13 @@ def write_log_message(t2i,opt,results,logfile):
     img_num    = 1
     batch_size = opt.batch_size or t2i.batch_size
     seenit     = {}
-    
+
+    seeds = [a[1] for a in results]
+    if batch_size > 1:
+        seeds = f"(seeds for each batch row: {seeds})"
+    else:
+        seeds = f"(seeds for individual images: {seeds})"
+
     for r in results:
         seed = r[1]
         log_message = (f'{r[0]}: {prompt_str} -S{seed}')
@@ -185,7 +194,10 @@ def write_log_message(t2i,opt,results,logfile):
         if r[0] not in seenit:
             seenit[r[0]] = True
             try:
-                _write_prompt_to_png(r[0],f'{prompt_str} -S{seed}')
+                if opt.grid:
+                    _write_prompt_to_png(r[0],f'{prompt_str} -g -S{seed} {seeds}')
+                else:
+                    _write_prompt_to_png(r[0],f'{prompt_str} -S{seed}')
             except FileNotFoundError:
                 print(f"Could not open file '{r[0]}' for reading")