From 7e3539df6f4e3979e080ed5d76faa3649c10f76f Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 13 May 2023 20:21:11 +0300
Subject: [PATCH] fix upscalers disappearing after the user reloads UI

---
 modules/modelloader.py | 27 ++++++++++-----------------
 webui.py               |  6 +-----
 2 files changed, 11 insertions(+), 22 deletions(-)

diff --git a/modules/modelloader.py b/modules/modelloader.py
index cb85ac4ff..a70aa0e37 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -117,20 +117,6 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
         pass
 
 
-builtin_upscaler_classes = []
-forbidden_upscaler_classes = set()
-
-
-def list_builtin_upscalers():
-    builtin_upscaler_classes.clear()
-    builtin_upscaler_classes.extend(Upscaler.__subclasses__())
-
-def forbid_loaded_nonbuiltin_upscalers():
-    for cls in Upscaler.__subclasses__():
-        if cls not in builtin_upscaler_classes:
-            forbidden_upscaler_classes.add(cls)
-
-
 def load_upscalers():
     # We can only do this 'magic' method to dynamically load upscalers if they are referenced,
     # so we'll try to import any _model.py files before looking in __subclasses__
@@ -146,10 +132,17 @@ def load_upscalers():
 
     datas = []
     commandline_options = vars(shared.cmd_opts)
-    for cls in Upscaler.__subclasses__():
-        if cls in forbidden_upscaler_classes:
-            continue
 
+    # some of upscaler classes will not go away after reloading their modules, and we'll end
+    # up with two copies of those classes. The newest copy will always be the last in the list,
+    # so we go from end to beginning and ignore duplicates
+    used_classes = {}
+    for cls in reversed(Upscaler.__subclasses__()):
+        classname = str(cls)
+        if classname not in used_classes:
+            used_classes[classname] = cls
+
+    for cls in reversed(used_classes.values()):
         name = cls.__name__
         cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
         scaler = cls(commandline_options.get(cmd_name, None))
diff --git a/webui.py b/webui.py
index 727ebd31d..e8f0a63dd 100644
--- a/webui.py
+++ b/webui.py
@@ -181,14 +181,11 @@ def initialize():
     gfpgan.setup_model(cmd_opts.gfpgan_models_path)
     startup_timer.record("setup gfpgan")
 
-    modelloader.list_builtin_upscalers()
-    startup_timer.record("list builtin upscalers")
-
     modules.scripts.load_scripts()
     startup_timer.record("load scripts")
 
     modelloader.load_upscalers()
-    #startup_timer.record("load upscalers") #Is this necessary? I don't know.
+    startup_timer.record("load upscalers") #Is this necessary? I don't know.
 
     modules.sd_vae.refresh_vae_list()
     startup_timer.record("refresh VAE")
@@ -388,7 +385,6 @@ def webui():
 
         localization.list_localizations(cmd_opts.localizations_dir)
 
-        modelloader.forbid_loaded_nonbuiltin_upscalers()
         modules.scripts.reload_scripts()
         startup_timer.record("load scripts")