fix upscalers disappearing after the user reloads UI

2 files changed, 11 insertions, 22 deletions

diff --git a/modules/modelloader.py b/modules/modelloader.py
index cb85ac4f..a70aa0e3 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 727ebd31..e8f0a63d 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")