From d44608f5f544cb5a99060b3777603ca965a7e5e9 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Thu, 25 Jan 2024 13:20:43 -0800
Subject: [PATCH] i

---
 extensions-builtin/LDSR/scripts/ldsr_model.py     | 3 ++-
 extensions-builtin/SwinIR/scripts/swinir_model.py | 4 +++-
 modules/dat_model.py                              | 3 ++-
 modules/esrgan_model.py                           | 3 ++-
 modules/hat_model.py                              | 3 ++-
 modules/realesrgan_model.py                       | 4 +++-
 6 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/extensions-builtin/LDSR/scripts/ldsr_model.py b/extensions-builtin/LDSR/scripts/ldsr_model.py
index bd78dece..79b01606 100644
--- a/extensions-builtin/LDSR/scripts/ldsr_model.py
+++ b/extensions-builtin/LDSR/scripts/ldsr_model.py
@@ -1,7 +1,7 @@
 import os
 
 from modules.modelloader import load_file_from_url
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 from ldsr_model_arch import LDSR
 from modules import shared, script_callbacks, errors
 import sd_hijack_autoencoder  # noqa: F401
@@ -49,6 +49,7 @@ class UpscalerLDSR(Upscaler):
         return LDSR(model, yaml)
 
     def do_upscale(self, img, path):
+        prepare_free_memory()
         try:
             ldsr = self.load_model(path)
         except Exception:
diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py
index 16bf9b79..d96e27c8 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/extensions-builtin/SwinIR/scripts/swinir_model.py
@@ -5,7 +5,7 @@ import torch
 from PIL import Image
 
 from modules import devices, modelloader, script_callbacks, shared, upscaler_utils
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 
 SWINIR_MODEL_URL = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth"
 
@@ -33,6 +33,8 @@ class UpscalerSwinIR(Upscaler):
         self.scalers = scalers
 
     def do_upscale(self, img: Image.Image, model_file: str) -> Image.Image:
+        prepare_free_memory()
+
         current_config = (model_file, shared.opts.SWIN_tile)
 
         if self._cached_model_config == current_config:
diff --git a/modules/dat_model.py b/modules/dat_model.py
index 495d5f49..24c4c968 100644
--- a/modules/dat_model.py
+++ b/modules/dat_model.py
@@ -2,7 +2,7 @@ import os
 
 from modules import modelloader, errors
 from modules.shared import cmd_opts, opts
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 from modules.upscaler_utils import upscale_with_model
 
 
@@ -23,6 +23,7 @@ class UpscalerDAT(Upscaler):
                 self.scalers.append(model)
 
     def do_upscale(self, img, path):
+        prepare_free_memory()
         try:
             info = self.load_model(path)
         except Exception:
diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index 70041ab0..4ade7fce 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -1,6 +1,6 @@
 from modules import modelloader, devices, errors
 from modules.shared import opts
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 from modules.upscaler_utils import upscale_with_model
 
 
@@ -27,6 +27,7 @@ class UpscalerESRGAN(Upscaler):
             self.scalers.append(scaler_data)
 
     def do_upscale(self, img, selected_model):
+        prepare_free_memory()
         try:
             model = self.load_model(selected_model)
         except Exception:
diff --git a/modules/hat_model.py b/modules/hat_model.py
index 7f2abb41..e31dee21 100644
--- a/modules/hat_model.py
+++ b/modules/hat_model.py
@@ -3,7 +3,7 @@ import sys
 
 from modules import modelloader, devices
 from modules.shared import opts
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 from modules.upscaler_utils import upscale_with_model
 
 
@@ -20,6 +20,7 @@ class UpscalerHAT(Upscaler):
             self.scalers.append(scaler_data)
 
     def do_upscale(self, img, selected_model):
+        prepare_free_memory()
         try:
             model = self.load_model(selected_model)
         except Exception as e:
diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index ff9d8ac0..e9314510 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -2,7 +2,7 @@ import os
 
 from modules import modelloader, errors
 from modules.shared import cmd_opts, opts
-from modules.upscaler import Upscaler, UpscalerData
+from modules.upscaler import Upscaler, UpscalerData, prepare_free_memory
 from modules.upscaler_utils import upscale_with_model
 
 
@@ -27,6 +27,8 @@ class UpscalerRealESRGAN(Upscaler):
                 self.scalers.append(scaler)
 
     def do_upscale(self, img, path):
+        prepare_free_memory()
+
         if not self.enable:
             return img