Commit 7e3539df authored by AUTOMATIC's avatar AUTOMATIC
Browse files

fix upscalers disappearing after the user reloads UI

parent 47719935
Loading
Loading
Loading
Loading
+10 −17
Original line number Diff line number Diff line
@@ -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))
+1 −5
Original line number Diff line number Diff line
@@ -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")