From 3be1cee17cd0801d3d7e18b8547400dad2729f94 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Wed, 31 Aug 2022 00:56:12 -0400
Subject: [PATCH] avoid crash due to dangling batch_size reference

---
 ldm/dream/pngwriter.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/ldm/dream/pngwriter.py b/ldm/dream/pngwriter.py
index a75395d917..bf2488bcbd 100644
--- a/ldm/dream/pngwriter.py
+++ b/ldm/dream/pngwriter.py
@@ -17,9 +17,8 @@ from PIL import Image, PngImagePlugin
 
 
 class PngWriter:
-    def __init__(self, outdir, prompt=None, batch_size=1):
+    def __init__(self, outdir, prompt=None):
         self.outdir = outdir
-        self.batch_size = batch_size
         self.prompt = prompt
         self.filepath = None
         self.files_written = []
@@ -50,10 +49,7 @@ class PngWriter:
             )
             basecount = int(filename.split('.', 1)[0])
             basecount += 1
-            if self.batch_size > 1:
-                filename = f'{basecount:06}.{seed}.01.png'
-            else:
-                filename = f'{basecount:06}.{seed}.png'
+            filename = f'{basecount:06}.{seed}.png'
             return os.path.join(self.outdir, filename)
 
         else:
@@ -69,11 +65,6 @@ class PngWriter:
                 series += 1
                 filename = f'{basecount:06}.{seed}.png'
                 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)
             return os.path.join(self.outdir, filename)