Reduce more memories on free_gpu_mem option (#1915)

* Enhance free_gpu_mem option Unload cond_stage_model on free_gpu_mem option is setted * Enhance free_gpu_mem option Unload cond_stage_model on free_gpu_mem option is setted
invoke-ai · Dec 11, 2022 · bd0c0d7 · bd0c0d7
1 parent f745f78
commit bd0c0d7
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/ldm/generate.py b/ldm/generate.py
@@ -453,6 +453,11 @@ def process_image(image,seed):
         init_image = None
         mask_image = None
 
+
+        if self.free_gpu_mem and self.model.cond_stage_model.device != self.model.device:
+            self.model.cond_stage_model.device = self.model.device
+            self.model.cond_stage_model.to(self.model.device)
+
         try:
             uc, c, extra_conditioning_info = get_uc_and_c_and_ec(
                 prompt, model =self.model,

diff --git a/ldm/invoke/generator/txt2img.py b/ldm/invoke/generator/txt2img.py
@@ -6,6 +6,7 @@
 import numpy as  np
 from ldm.invoke.generator.base import Generator
 from ldm.models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
+import gc
 
 
 class Txt2Img(Generator):
@@ -55,7 +56,11 @@ def make_image(x_T):
             )
 
             if self.free_gpu_mem:
-                self.model.model.to("cpu")
+                self.model.model.to('cpu')
+                self.model.cond_stage_model.device = 'cpu'
+                self.model.cond_stage_model.to('cpu')
+                gc.collect()
+                torch.cuda.empty_cache()
 
             return self.sample_to_image(samples)
 

diff --git a/ldm/invoke/generator/txt2img2img.py b/ldm/invoke/generator/txt2img2img.py
@@ -100,7 +100,11 @@ def make_image(x_T):
             )
 
             if self.free_gpu_mem:
-                self.model.model.to("cpu")
+                self.model.model.to('cpu')
+                self.model.cond_stage_model.device = 'cpu'
+                self.model.cond_stage_model.to('cpu')
+                gc.collect()
+                torch.cuda.empty_cache()
 
             return self.sample_to_image(samples)
 
@@ -142,7 +146,7 @@ def inpaint_make_image(x_T):
                 **kwargs
                 )
             return result[0][0]
-            
+
         if sampler.uses_inpainting_model():
             return inpaint_make_image
         else: