StableDiffusion/WebUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added torch.mps.empty_cache() to torch_gc()

Browse Source 
changed a bunch of places that use torch.cuda.empty_cache() to use torch_gc() instead
This commit is contained in:
AUTOMATIC1111 2023-07-08 17:13:18 +03:00
parent e161b5a025
commit da8916f926
6 changed files with 10 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
extensions-builtin/LDSR/ldsr_model_arch.py
View File

@ -12,7 +12,7 @@ import safetensors.torch
from ldm.models.diffusion.ddim import DDIMSampler from ldm.models.diffusion.ddim import DDIMSampler
from ldm.util import instantiate_from_config, ismap from ldm.util import instantiate_from_config, ismap
from modules import shared, sd_hijack from modules import shared, sd_hijack, devices
cached_ldsr_model: torch.nn.Module = None cached_ldsr_model: torch.nn.Module = None
@ -112,8 +112,7 @@ class LDSR:
gc.collect() gc.collect()
if torch.cuda.is_available: devices.torch_gc()
torch.cuda.empty_cache()
im_og = image im_og = image
width_og, height_og = im_og.size width_og, height_og = im_og.size
@ -150,8 +149,7 @@ class LDSR:
del model del model
gc.collect() gc.collect()
if torch.cuda.is_available: devices.torch_gc()
torch.cuda.empty_cache()
return a return a

4
extensions-builtin/ScuNET/scripts/scunet_model.py
View File

@ -85,7 +85,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
def do_upscale(self, img: PIL.Image.Image, selected_file): def do_upscale(self, img: PIL.Image.Image, selected_file):
torch.cuda.empty_cache() devices.torch_gc()
try: try:
model = self.load_model(selected_file) model = self.load_model(selected_file)
@ -110,7 +110,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any
np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy() np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy()
del torch_img, torch_output del torch_img, torch_output
torch.cuda.empty_cache() devices.torch_gc()
output = np_output.transpose((1, 2, 0)) # CHW to HWC output = np_output.transpose((1, 2, 0)) # CHW to HWC
output = output[:, :, ::-1] # BGR to RGB output = output[:, :, ::-1] # BGR to RGB

5
extensions-builtin/SwinIR/scripts/swinir_model.py
View File

@ -42,10 +42,7 @@ class UpscalerSwinIR(Upscaler):
return img return img
model = model.to(device_swinir, dtype=devices.dtype) model = model.to(device_swinir, dtype=devices.dtype)
img = upscale(img, model) img = upscale(img, model)
try: devices.torch_gc()
torch.cuda.empty_cache()
except Exception:
pass
return img return img
def load_model(self, path, scale=4): def load_model(self, path, scale=4):

2
modules/codeformer_model.py
View File

@ -99,7 +99,7 @@ def setup_model(dirname):
output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0] output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1)) restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output del output
torch.cuda.empty_cache() devices.torch_gc()
except Exception: except Exception:
errors.report('Failed inference for CodeFormer', exc_info=True) errors.report('Failed inference for CodeFormer', exc_info=True)
restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1)) restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

3
modules/devices.py
View File

@ -49,10 +49,13 @@ def get_device_for(task):
def torch_gc(): def torch_gc():
if torch.cuda.is_available(): if torch.cuda.is_available():
with torch.cuda.device(get_cuda_device_string()): with torch.cuda.device(get_cuda_device_string()):
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
elif has_mps() and hasattr(torch.mps, 'empty_cache'):
torch.mps.empty_cache()
def enable_tf32(): def enable_tf32():

1
modules/sd_models.py
View File

@ -590,7 +590,6 @@ def unload_model_weights(sd_model=None, info=None):
sd_model = None sd_model = None
gc.collect() gc.collect()
devices.torch_gc() devices.torch_gc()
torch.cuda.empty_cache()
print(f"Unloaded weights {timer.summary()}.") print(f"Unloaded weights {timer.summary()}.")