diff options
| -rw-r--r-- | extensions-builtin/LDSR/ldsr_model_arch.py | 8 | ||||
| -rw-r--r-- | extensions-builtin/ScuNET/scripts/scunet_model.py | 4 | ||||
| -rw-r--r-- | extensions-builtin/SwinIR/scripts/swinir_model.py | 5 | ||||
| -rw-r--r-- | modules/codeformer_model.py | 2 | ||||
| -rw-r--r-- | modules/devices.py | 3 | ||||
| -rw-r--r-- | modules/sd_models.py | 1 |
6 files changed, 10 insertions, 13 deletions
diff --git a/extensions-builtin/LDSR/ldsr_model_arch.py b/extensions-builtin/LDSR/ldsr_model_arch.py index 7f450086..7cac36ce 100644 --- a/extensions-builtin/LDSR/ldsr_model_arch.py +++ b/extensions-builtin/LDSR/ldsr_model_arch.py @@ -12,7 +12,7 @@ import safetensors.torch from ldm.models.diffusion.ddim import DDIMSampler 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 @@ -112,8 +112,7 @@ class LDSR: gc.collect() - if torch.cuda.is_available: - torch.cuda.empty_cache() + devices.torch_gc() im_og = image width_og, height_og = im_og.size @@ -150,8 +149,7 @@ class LDSR: del model gc.collect() - if torch.cuda.is_available: - torch.cuda.empty_cache() + devices.torch_gc() return a diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/extensions-builtin/ScuNET/scripts/scunet_model.py index ffef26b2..167d2f64 100644 --- a/extensions-builtin/ScuNET/scripts/scunet_model.py +++ b/extensions-builtin/ScuNET/scripts/scunet_model.py @@ -85,7 +85,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler): def do_upscale(self, img: PIL.Image.Image, selected_file): - torch.cuda.empty_cache() + devices.torch_gc() try: 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 np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy() del torch_img, torch_output - torch.cuda.empty_cache() + devices.torch_gc() output = np_output.transpose((1, 2, 0)) # CHW to HWC output = output[:, :, ::-1] # BGR to RGB diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py index c6bc53a8..c2c2a43c 100644 --- a/extensions-builtin/SwinIR/scripts/swinir_model.py +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -42,10 +42,7 @@ class UpscalerSwinIR(Upscaler): return img model = model.to(device_swinir, dtype=devices.dtype) img = upscale(img, model) - try: - torch.cuda.empty_cache() - except Exception: - pass + devices.torch_gc() return img def load_model(self, path, scale=4): diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index f293acf5..da42b5e9 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -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))
diff --git a/modules/devices.py b/modules/devices.py index 620ed1a6..c5ad950f 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -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(): diff --git a/modules/sd_models.py b/modules/sd_models.py index f65f4e36..653c4cc0 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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()}.")
|