diff --git a/extensions-builtin/LDSR/ldsr_model_arch.py b/extensions-builtin/LDSR/ldsr_model_arch.py index 5299f2c29..4153139ba 100644 --- a/extensions-builtin/LDSR/ldsr_model_arch.py +++ b/extensions-builtin/LDSR/ldsr_model_arch.py @@ -111,10 +111,7 @@ class LDSR: eta = 1.0 gc.collect() - if torch.cuda.is_available: - torch.cuda.empty_cache() - if devices.backend == 'ipex': - torch.xpu.empty_cache() + devices.torch_gc() im_og = image width_og, height_og = im_og.size @@ -150,8 +147,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 239307f89..009b3537a 100644 --- a/extensions-builtin/ScuNET/scripts/scunet_model.py +++ b/extensions-builtin/ScuNET/scripts/scunet_model.py @@ -88,7 +88,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler): def do_upscale(self, img: PIL.Image.Image, selected_file): - torch.cuda.empty_cache() + devices.torch_gc() model = self.load_model(selected_file) if model is None: @@ -112,7 +112,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 f8cd77e69..7d786de23 100644 --- a/extensions-builtin/SwinIR/scripts/swinir_model.py +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -40,13 +40,7 @@ class UpscalerSwinIR(Upscaler): return img model = model.to(device_swinir, dtype=devices.dtype) img = upscale(img, model) - try: - if devices.backend == 'ipex': - torch.xpu.empty_cache() - else: - 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 92763e078..e6e75b219 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -101,10 +101,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 - if devices.backend == 'ipex': - torch.xpu.empty_cache() - else: - torch.cuda.empty_cache() + devices.torch_gc() except Exception as error: print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr) restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1)) diff --git a/modules/devices.py b/modules/devices.py index b67bc8d80..90cd05f6f 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -193,7 +193,7 @@ else: if backend == 'ipex': #Fix broken functions with ipex from modules.sd_hijack_utils import CondFunc - torch.cuda.empty_cache = torch.xpu.empty_cache + torch.cuda.empty_cache = torch_gc #Functions with dtype errors: CondFunc('torch.nn.modules.GroupNorm.forward',