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

Author: rmagur1203

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,

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)
 

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: