From 6142cf25ccf4dcee4d91fec57166502a9853de1f Mon Sep 17 00:00:00 2001
From: Peter Baylies <pbaylies@gmail.com>
Date: Wed, 21 Sep 2022 09:12:32 -0400
Subject: [PATCH]  * Bug fixes

---
 .gitignore                    | 2 +-
 ldm/dream/restoration/base.py | 2 +-
 ldm/dream/server.py           | 2 +-
 scripts/dream.py              | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4c0ec37174..da74df3a1a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,7 +1,7 @@
 # ignore default image save location and model symbolic link
 outputs/
 models/ldm/stable-diffusion-v1/model.ckpt
-ldm/restoration/codeformer/weights
+ldm/dream/restoration/codeformer/weights
 
 # ignore the Anaconda/Miniconda installer used while building Docker image
 anaconda.sh
diff --git a/ldm/dream/restoration/base.py b/ldm/dream/restoration/base.py
index 539301d802..9037bc40cb 100644
--- a/ldm/dream/restoration/base.py
+++ b/ldm/dream/restoration/base.py
@@ -27,7 +27,7 @@ class Restoration():
         return CodeFormerRestoration()
 
     # Upscale Models
-    def load_ersgan(self):
+    def load_esrgan(self):
         from ldm.dream.restoration.realesrgan import ESRGAN
         esrgan = ESRGAN(self.esrgan_bg_tile)
         print('>> ESRGAN Initialized')
diff --git a/ldm/dream/server.py b/ldm/dream/server.py
index cf10ebf024..46002a9b27 100644
--- a/ldm/dream/server.py
+++ b/ldm/dream/server.py
@@ -163,7 +163,7 @@ class DreamServer(BaseHTTPRequestHandler):
         # is complete. The upscaling replaces the original file, so the second
         # entry should not be inserted into the image list.
         # LS: This repeats code in dream.py
-        def image_done(image, seed, upscaled=False):
+        def image_done(image, seed, upscaled=False, first_seed=None):
             name = f'{prefix}.{seed}.png'
             iter_opt  = copy.copy(opt)
             if opt.variation_amount > 0:
diff --git a/scripts/dream.py b/scripts/dream.py
index dc1757dce2..9efa1829c6 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -47,14 +47,14 @@ def main():
     # Loading Face Restoration and ESRGAN Modules
     try:
         gfpgan, codeformer, esrgan = None, None, None
-        from ldm.dream.restoration import Restoration
+        from ldm.dream.restoration.base import Restoration
         restoration = Restoration(opt.gfpgan_dir, opt.gfpgan_model_path, opt.esrgan_bg_tile)
         if opt.restore:
             gfpgan, codeformer = restoration.load_face_restore_models()
         else:
             print('>> Face restoration disabled')
         if opt.esrgan:
-            esrgan = restoration.load_ersgan()
+            esrgan = restoration.load_esrgan()
         else:
             print('>> Upscaling disabled')
     except (ModuleNotFoundError, ImportError):