Browse Source

Added memory clearance after deletion

Ritesh Gangnani 1 year ago
parent
commit
44db35fb1a
1 changed files with 5 additions and 1 deletions
  1. 5 1
      modules/sd_hijack.py

+ 5 - 1
modules/sd_hijack.py

@@ -1,3 +1,5 @@
+import gc 
+
 import torch
 import torch
 from torch.nn.functional import silu
 from torch.nn.functional import silu
 from types import MethodType
 from types import MethodType
@@ -190,7 +192,9 @@ class StableDiffusionModelHijack:
                 delattr(m.model.diffusion_model.output_blocks[0][1].transformer_blocks,i)
                 delattr(m.model.diffusion_model.output_blocks[0][1].transformer_blocks,i)
                 delattr(m.model.diffusion_model.output_blocks[1][1].transformer_blocks,i)
                 delattr(m.model.diffusion_model.output_blocks[1][1].transformer_blocks,i)
             delattr(m.model.diffusion_model.output_blocks[4][1].transformer_blocks,'1')
             delattr(m.model.diffusion_model.output_blocks[4][1].transformer_blocks,'1')
-            delattr(m.model.diffusion_model.output_blocks[5][1].transformer_blocks,'1') 
+            delattr(m.model.diffusion_model.output_blocks[5][1].transformer_blocks,'1')
+            torch.cuda.empty_cache()
+            gc.collect()
 
 
     def hijack(self, m):
     def hijack(self, m):
         conditioner = getattr(m, 'conditioner', None)
         conditioner = getattr(m, 'conditioner', None)