lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-02-20 14:53:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

revise space

Browse Source
This commit is contained in:
layerdiffusion
2024-08-30 16:48:11 -07:00
parent ec7917bd16
commit 8e9b124456
1 changed files with 1 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
spaces.py
View File

@@ -114,24 +114,17 @@ class GPUObject:
def __enter__(self):
self.original_init = torch.nn.Module.__init__
self.original_to = torch.nn.Module.to
def patched_init(module, *args, **kwargs):
self.module_list.append(module)
return self.original_init(module, *args, **kwargs)
def patched_to(module, *args, **kwargs):
self.module_list.append(module)
return self.original_to(module, *args, **kwargs)
torch.nn.Module.__init__ = patched_init
torch.nn.Module.to = patched_to
return self
def __exit__(self, exc_type, exc_val, exc_tb):
torch.nn.Module.__init__ = self.original_init
torch.nn.Module.to = self.original_to
self.module_list = set(self.module_list)
self.module_list = list(set(self.module_list))
self.to(device=torch.device('cpu'))
memory_management.soft_empty_cache()
return