Commit 44db35fb authored by Ritesh Gangnani's avatar Ritesh Gangnani

Added memory clearance after deletion

parent ff1609f9
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: ...@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment