diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 08d31080..f93df0a6 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -48,6 +48,11 @@ def apply_optimizations(): undo_optimizations() + if len(optimizers) == 0: + # a script can access the model very early, and optimizations would not be filled by then + current_optimizer = None + return '' + ldm.modules.diffusionmodules.model.nonlinearity = silu ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th @@ -67,8 +72,9 @@ def apply_optimizations(): matching_optimizer = optimizers[0] if matching_optimizer is not None: - print(f"Applying optimization: {matching_optimizer.name}") + print(f"Applying optimization: {matching_optimizer.name}... ", end='') matching_optimizer.apply() + print("done.") current_optimizer = matching_optimizer return current_optimizer.name else: @@ -149,6 +155,13 @@ class StableDiffusionModelHijack: def __init__(self): self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir) + def apply_optimizations(self): + try: + self.optimization_method = apply_optimizations() + except Exception as e: + errors.display(e, "applying cross attention optimization") + undo_optimizations() + def hijack(self, m): if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation: model_embeddings = m.cond_stage_model.roberta.embeddings @@ -168,11 +181,7 @@ class StableDiffusionModelHijack: if m.cond_stage_key == "edit": sd_hijack_unet.hijack_ddpm_edit() - try: - self.optimization_method = apply_optimizations() - except Exception as e: - errors.display(e, "applying cross attention optimization") - undo_optimizations() + self.apply_optimizations() self.clip = m.cond_stage_model diff --git a/webui.py b/webui.py index d4402f55..f0ffbbbf 100644 --- a/webui.py +++ b/webui.py @@ -291,9 +291,20 @@ def initialize_rest(*, reload_script_modules=False): modules.sd_hijack.list_optimizers() startup_timer.record("scripts list_optimizers") - # load model in parallel to other startup stuff - # (when reloading, this does nothing) - Thread(target=lambda: shared.sd_model).start() + def load_model(): + """ + Accesses shared.sd_model property to load model. + After it's available, if it has been loaded before this access by some extension, + its optimization may be None because the list of optimizaers has neet been filled + by that time, so we apply optimization again. + """ + + shared.sd_model # noqa: B018 + + if modules.sd_hijack.current_optimizer is None: + modules.sd_hijack.apply_optimizations() + + Thread(target=load_model).start() shared.reload_hypernetworks() startup_timer.record("reload hypernetworks")