implement locking when acquiring next output file prefix

This commit is contained in:
Lincoln Stein 2023-03-04 09:13:17 -05:00
parent 3c64fad379
commit 117f70e1ec
2 changed files with 21 additions and 4 deletions

File diff suppressed because one or more lines are too long

View File

@ -9,6 +9,8 @@ Exports function retrieve_metadata(path)
import os
import re
import json
from pathlib import Path
from filelock import FileLock
from PIL import PngImagePlugin, Image
# -------------------image generation utils-----
@ -19,8 +21,23 @@ class PngWriter:
self.outdir = outdir
os.makedirs(outdir, exist_ok=True)
def unique_prefix(self)->str:
next_prefix_file = Path(self.outdir,'.next_prefix')
next_prefix_lock = Path(self.outdir,'.next_prefix.lock')
prefix = 0
with FileLock(next_prefix_lock):
if not next_prefix_file.exists():
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
with open(next_prefix_file,'w') as file:
file.write(str(prefix))
return f'{prefix:06}'
# gives the next unique prefix in outdir
def unique_prefix(self):
def _unused_prefix(self):
# 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
@ -94,11 +111,11 @@ class PromptFormatter:
# to do: put model name into the t2i object
# switches.append(f'--model{t2i.model_name}')
if opt.seamless or t2i.seamless:
switches.append(f'--seamless')
switches.append('--seamless')
if opt.init_img:
switches.append(f'-I{opt.init_img}')
if opt.fit:
switches.append(f'--fit')
switches.append('--fit')
if opt.strength and opt.init_img is not None:
switches.append(f'-f{opt.strength or t2i.strength}')
if opt.gfpgan_strength: