Ver código fonte

Optimize model loader

Child classes only get populated to __subclassess__ when they are imported. We don't actually need to import any of them to webui any more, so clean up webUI imports and make sure loader imports children.

Also, fix command line paths not actually being passed to the scalers.
d8ahazard 2 anos atrás
pai
commit
e82ea20299
2 arquivos alterados com 19 adições e 13 exclusões
  1. 16 3
      modules/modelloader.py
  2. 3 10
      webui.py

+ 16 - 3
modules/modelloader.py

@@ -4,7 +4,6 @@ import importlib
 from urllib.parse import urlparse
 from urllib.parse import urlparse
 
 
 from basicsr.utils.download_util import load_file_from_url
 from basicsr.utils.download_util import load_file_from_url
-
 from modules import shared
 from modules import shared
 from modules.upscaler import Upscaler
 from modules.upscaler import Upscaler
 from modules.paths import script_path, models_path
 from modules.paths import script_path, models_path
@@ -120,16 +119,30 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
 
 
 
 
 def load_upscalers():
 def load_upscalers():
+    sd = shared.script_path
+    # 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__
+    modules_dir = os.path.join(sd, "modules")
+    for file in os.listdir(modules_dir):
+        if "_model.py" in file:
+            model_name = file.replace("_model.py", "")
+            full_model = f"modules.{model_name}_model"
+            try:
+                importlib.import_module(full_model)
+            except:
+                pass
     datas = []
     datas = []
+    c_o = vars(shared.cmd_opts)
     for cls in Upscaler.__subclasses__():
     for cls in Upscaler.__subclasses__():
         name = cls.__name__
         name = cls.__name__
         module_name = cls.__module__
         module_name = cls.__module__
         module = importlib.import_module(module_name)
         module = importlib.import_module(module_name)
         class_ = getattr(module, name)
         class_ = getattr(module, name)
-        cmd_name = f"{name.lower().replace('upscaler', '')}-models-path"
+        cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
         opt_string = None
         opt_string = None
         try:
         try:
-            opt_string = shared.opts.__getattr__(cmd_name)
+            if cmd_name in c_o:
+                opt_string = c_o[cmd_name]
         except:
         except:
             pass
             pass
         scaler = class_(opt_string)
         scaler = class_(opt_string)

+ 3 - 10
webui.py

@@ -1,28 +1,21 @@
 import os
 import os
-import threading
-
-from modules import devices
-from modules.paths import script_path
 import signal
 import signal
 import threading
 import threading
-import modules.paths
+
 import modules.codeformer_model as codeformer
 import modules.codeformer_model as codeformer
-import modules.esrgan_model as esrgan
-import modules.bsrgan_model as bsrgan
 import modules.extras
 import modules.extras
 import modules.face_restoration
 import modules.face_restoration
 import modules.gfpgan_model as gfpgan
 import modules.gfpgan_model as gfpgan
 import modules.img2img
 import modules.img2img
-import modules.ldsr_model as ldsr
 import modules.lowvram
 import modules.lowvram
-import modules.realesrgan_model as realesrgan
+import modules.paths
 import modules.scripts
 import modules.scripts
 import modules.sd_hijack
 import modules.sd_hijack
 import modules.sd_models
 import modules.sd_models
 import modules.shared as shared
 import modules.shared as shared
-import modules.swinir_model as swinir
 import modules.txt2img
 import modules.txt2img
 import modules.ui
 import modules.ui
+from modules import devices
 from modules import modelloader
 from modules import modelloader
 from modules.paths import script_path
 from modules.paths import script_path
 from modules.shared import cmd_opts
 from modules.shared import cmd_opts