1 files changed, 28 insertions, 23 deletions
diff --git a/webui.py b/webui.py
index c70a11c7..b8cccd54 100644
--- a/webui.py
+++ b/webui.py
@@ -1,37 +1,38 @@
 import os
 import threading
 
+from modules import devices
 from modules.paths import script_path
-
 import signal
-
-from modules.shared import opts, cmd_opts, state
-import modules.shared as shared
-import modules.ui
-import modules.scripts
-import modules.sd_hijack
-import modules.codeformer_model
-import modules.gfpgan_model
-import modules.face_restoration
-import modules.realesrgan_model as realesrgan
+import threading
+import modules.paths
+import modules.codeformer_model as codeformer
 import modules.esrgan_model as esrgan
-import modules.ldsr_model as ldsr
+import modules.bsrgan_model as bsrgan
 import modules.extras
-import modules.lowvram
-import modules.txt2img
+import modules.face_restoration
+import modules.gfpgan_model as gfpgan
 import modules.img2img
-import modules.swinir as swinir
+import modules.ldsr_model as ldsr
+import modules.lowvram
+import modules.realesrgan_model as realesrgan
+import modules.scripts
+import modules.sd_hijack
 import modules.sd_models
+import modules.shared as shared
+import modules.swinir_model as swinir
+import modules.txt2img
+import modules.ui
+from modules import modelloader
+from modules.paths import script_path
+from modules.shared import cmd_opts
 
-
-modules.codeformer_model.setup_codeformer()
-modules.gfpgan_model.setup_gfpgan()
+modelloader.cleanup_models()
+modules.sd_models.setup_model(cmd_opts.ckpt_dir)
+codeformer.setup_model(cmd_opts.codeformer_models_path)
+gfpgan.setup_model(cmd_opts.gfpgan_models_path)
 shared.face_restorers.append(modules.face_restoration.FaceRestoration())
-
-esrgan.load_models(cmd_opts.esrgan_models_path)
-swinir.load_models(cmd_opts.swinir_models_path)
-realesrgan.setup_realesrgan()
-ldsr.add_lsdr()
+modelloader.load_upscalers()
 queue_lock = threading.Lock()
 
 
@@ -47,6 +48,8 @@ def wrap_queued_call(func):
 
 def wrap_gradio_gpu_call(func):
     def f(*args, **kwargs):
+        devices.torch_gc()
+
         shared.state.sampling_step = 0
         shared.state.job_count = -1
         shared.state.job_no = 0
@@ -62,6 +65,8 @@ def wrap_gradio_gpu_call(func):
         shared.state.job = ""
         shared.state.job_count = 0
 
+        devices.torch_gc()
+
         return res
 
     return modules.ui.wrap_gradio_call(f)