Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

avoid crash due to dangling batch_size reference

Browse Source
This commit is contained in:
Lincoln Stein 2022-08-31 00:56:12 -04:00
parent bdb0651eb2
commit 3be1cee17c
1 changed files with 2 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
ldm/dream/pngwriter.py
View File

@ -17,9 +17,8 @@ from PIL import Image, PngImagePlugin
class PngWriter: class PngWriter:
def __init__(self, outdir, prompt=None, batch_size=1): def __init__(self, outdir, prompt=None):
self.outdir = outdir self.outdir = outdir
self.batch_size = batch_size
self.prompt = prompt self.prompt = prompt
self.filepath = None self.filepath = None
self.files_written = [] self.files_written = []
@ -50,10 +49,7 @@ class PngWriter:
) )
basecount = int(filename.split('.', 1)[0]) basecount = int(filename.split('.', 1)[0])
basecount += 1 basecount += 1
if self.batch_size > 1: filename = f'{basecount:06}.{seed}.png'
filename = f'{basecount:06}.{seed}.01.png'
else:
filename = f'{basecount:06}.{seed}.png'
return os.path.join(self.outdir, filename) return os.path.join(self.outdir, filename)
else: else:
@ -69,11 +65,6 @@ class PngWriter:
series += 1 series += 1
filename = f'{basecount:06}.{seed}.png' filename = f'{basecount:06}.{seed}.png'
path = os.path.join(self.outdir, filename) path = os.path.join(self.outdir, filename)
if self.batch_size > 1 or os.path.exists(path):
if upscaled:
break
filename = f'{basecount:06}.{seed}.{series:02}.png'
path = os.path.join(self.outdir, filename)
finished = not os.path.exists(path) finished = not os.path.exists(path)
return os.path.join(self.outdir, filename) return os.path.join(self.outdir, filename)