From 2154dd2349df970db79a5e73f2c8d11213acc66e Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Thu, 6 Oct 2022 10:54:05 -0400
Subject: [PATCH] prevent crashes due to uninitialized free_gpu_mem

---
 ldm/generate.py  | 4 +++-
 scripts/dream.py | 5 ++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/ldm/generate.py b/ldm/generate.py
index ea5ac88245..fc40fa6152 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -174,7 +174,8 @@ class Generate:
             config                = None,
             gfpgan=None,
             codeformer=None,
-            esrgan=None
+            esrgan=None,
+            free_gpu_mem=False,
     ):
         models              = OmegaConf.load(conf)
         mconfig             = models[model]
@@ -201,6 +202,7 @@ class Generate:
         self.gfpgan = gfpgan
         self.codeformer = codeformer
         self.esrgan = esrgan
+        self.free_gpu_mem = free_gpu_mem
 
         # Note that in previous versions, there was an option to pass the
         # device to Generate(). However the device was then ignored, so
diff --git a/scripts/dream.py b/scripts/dream.py
index cd0d7c33b4..b1882db827 100644
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -75,7 +75,8 @@ def main():
             precision      = opt.precision,
             gfpgan=gfpgan,
             codeformer=codeformer,
-            esrgan=esrgan
+            esrgan=esrgan,
+            free_gpu_mem=opt.free_gpu_mem,
             )
     except (FileNotFoundError, IOError, KeyError) as e:
         print(f'{e}. Aborting.')
@@ -104,8 +105,6 @@ def main():
 
     # preload the model
     gen.load_model()
-    #set additional option
-    gen.free_gpu_mem = opt.free_gpu_mem
 
     # web server loops forever
     if opt.web or opt.gui: