From 3ab9d0288347ec3f0ef6759954de7c8a258c89fa Mon Sep 17 00:00:00 2001
From: Jonathan <34005131+JPPhoto@users.noreply.github.com>
Date: Tue, 21 Feb 2023 11:12:24 -0600
Subject: [PATCH] Fixed embiggening crash due to clear_cuda_cache not being
 passed on and bad cuda stats initialization. (#2756)

---
 ldm/generate.py                  | 5 ++++-
 ldm/invoke/generator/base.py     | 2 +-
 ldm/invoke/generator/embiggen.py | 2 ++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/ldm/generate.py b/ldm/generate.py
index 76629ff2f9..aab6ca33e2 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -178,7 +178,9 @@ class Generate:
         self.model_hash = None
         self.sampler = None
         self.device = None
-        self.session_peakmem = None
+        self.max_memory_allocated = 0
+        self.memory_allocated = 0
+        self.session_peakmem = 0
         self.base_generator = None
         self.seed = None
         self.outdir = outdir
@@ -781,6 +783,7 @@ class Generate:
                 embiggen_tiles=opt.embiggen_tiles,
                 embiggen_strength=opt.embiggen_strength,
                 image_callback=callback,
+                clear_cuda_cache=self.clear_cuda_cache,
             )
         elif tool == "outpaint":
             from ldm.invoke.restoration.outpaint import Outpaint
diff --git a/ldm/invoke/generator/base.py b/ldm/invoke/generator/base.py
index d89fb48aff..00e395bf83 100644
--- a/ldm/invoke/generator/base.py
+++ b/ldm/invoke/generator/base.py
@@ -126,7 +126,7 @@ class Generator:
                 seed = self.new_seed()
 
                 # Free up memory from the last generation.
-                clear_cuda_cache = kwargs['clear_cuda_cache'] or None
+                clear_cuda_cache = kwargs['clear_cuda_cache'] if 'clear_cuda_cache' in kwargs else None
                 if clear_cuda_cache is not None:
                     clear_cuda_cache()
 
diff --git a/ldm/invoke/generator/embiggen.py b/ldm/invoke/generator/embiggen.py
index e7505aab2f..0a06f90b03 100644
--- a/ldm/invoke/generator/embiggen.py
+++ b/ldm/invoke/generator/embiggen.py
@@ -346,6 +346,7 @@ class Embiggen(Generator):
                 newinitimage = torch.from_numpy(newinitimage)
                 newinitimage = 2.0 * newinitimage - 1.0
                 newinitimage = newinitimage.to(self.model.device)
+                clear_cuda_cache = kwargs['clear_cuda_cache'] if 'clear_cuda_cache' in kwargs else None
 
                 tile_results = gen_img2img.generate(
                     prompt,
@@ -363,6 +364,7 @@ class Embiggen(Generator):
                     init_image     = newinitimage,    # notice that init_image is different from init_img
                     mask_image     = None,
                     strength       = strength,
+                    clear_cuda_cache     = clear_cuda_cache
                 )
 
                 emb_tile_store.append(tile_results[0][0])