From 29b348ece1ccc767e08d393e84125c54d2500cb1 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 23 Mar 2023 22:03:02 -0400 Subject: [PATCH] fix corrupted outputs/.next_prefix file - Since 2.3.2 invokeai stores the next PNG file's numeric prefix in a file named `.next_prefix` in the outputs directory. This avoids the overhead of doing a directory listing to find out what file number comes next. - The code uses advisory locking to prevent corruption of this file in the event that multiple invokeai's try to access it simultaneously, but some users have experienced corruption of the file nevertheless. - This PR addresses the problem by detecting a potentially corrupted `.next_prefix` file and falling back to the directory listing method. A fixed version of the file is then written out. - Closes #3001 --- ldm/invoke/pngwriter.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/ldm/invoke/pngwriter.py b/ldm/invoke/pngwriter.py index da5af82aa8..cdd9107f20 100644 --- a/ldm/invoke/pngwriter.py +++ b/ldm/invoke/pngwriter.py @@ -30,14 +30,17 @@ class PngWriter: prefix = self._unused_prefix() else: with open(next_prefix_file,'r') as file: - prefix=int(file.readline() or int(self._unused_prefix())-1) - prefix+=1 + prefix = 0 + try: + prefix=int(file.readline()) + except (TypeError, ValueError): + prefix=self._unused_prefix() with open(next_prefix_file,'w') as file: - file.write(str(prefix)) + file.write(str(prefix+1)) return f'{prefix:06}' # gives the next unique prefix in outdir - def _unused_prefix(self): + def _unused_prefix(self)->int: # sort reverse alphabetically until we find max+1 dirlist = sorted(os.listdir(self.outdir), reverse=True) # find the first filename that matches our pattern or return 000000.0.png @@ -45,8 +48,7 @@ class PngWriter: (f for f in dirlist if re.match('^(\d+)\..*\.png', f)), '0000000.0.png', ) - basecount = int(existing_name.split('.', 1)[0]) + 1 - return f'{basecount:06}' + return int(existing_name.split('.', 1)[0]) + 1 # saves image named _image_ to outdir/name, writing metadata from prompt # returns full path of output