From cb71d05d98075d7bdbc9924f0c1e53377c675314 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 13 Jun 2023 20:08:42 +0300 Subject: [PATCH] Add torch.xpu.empty_cache() to LDSR and SwinIR --- extensions-builtin/LDSR/ldsr_model_arch.py | 2 ++ extensions-builtin/SwinIR/scripts/swinir_model.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/extensions-builtin/LDSR/ldsr_model_arch.py b/extensions-builtin/LDSR/ldsr_model_arch.py index 17a7a4a97..41d97d071 100644 --- a/extensions-builtin/LDSR/ldsr_model_arch.py +++ b/extensions-builtin/LDSR/ldsr_model_arch.py @@ -113,6 +113,8 @@ class LDSR: gc.collect() if torch.cuda.is_available: torch.cuda.empty_cache() + if shared.cmd_opts.use_ipex: + torch.xpu.empty_cache() im_og = image width_og, height_og = im_og.size diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py index 085c2ecc3..c3c78ac82 100644 --- a/extensions-builtin/SwinIR/scripts/swinir_model.py +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -7,7 +7,7 @@ from tqdm.rich import tqdm from swinir_model_arch import SwinIR as net from swinir_model_arch_v2 import Swin2SR as net2 from modules import modelloader, devices, script_callbacks, shared -from modules.shared import opts, state +from modules.shared import cmd_opts, opts, state from modules.upscaler import Upscaler, UpscalerData @@ -44,6 +44,8 @@ class UpscalerSwinIR(Upscaler): torch.cuda.empty_cache() except Exception: pass + if cmd_opts.use_ipex: + torch.xpu.empty_cache() return img def load_model(self, path, scale=4):