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

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

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]
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output
torch.cuda.empty_cache()
devices.torch_gc()
except Exception:
errors.report('Failed inference for CodeFormer', exc_info=True)
restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

View File

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

View File

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