diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 2848a251d6b..5945b7c23fe 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -245,6 +245,7 @@ def hijack(self, m): model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self) m.cond_stage_model = FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self) + self.clip = m.cond_stage_model if cmd_opts.opt_split_attention_v1: @@ -263,6 +264,14 @@ def flatten(el): self.layers = flatten(m) + def undo_hijack(self, m): + if type(m.cond_stage_model) == FrozenCLIPEmbedderWithCustomWords: + m.cond_stage_model = m.cond_stage_model.wrapped + + model_embeddings = m.cond_stage_model.transformer.text_model.embeddings + if type(model_embeddings.token_embedding) == EmbeddingsWithFixes: + model_embeddings.token_embedding = model_embeddings.token_embedding.wrapped + def apply_circular(self, enable): if self.circular_enabled == enable: return diff --git a/modules/sd_models.py b/modules/sd_models.py index 7a5edceda3b..eb21e498a35 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -137,7 +137,7 @@ def load_model(): def reload_model_weights(sd_model, info=None): - from modules import lowvram, devices + from modules import lowvram, devices, sd_hijack checkpoint_info = info or select_checkpoint() if sd_model.sd_model_checkpint == checkpoint_info.filename: @@ -148,8 +148,12 @@ def reload_model_weights(sd_model, info=None): else: sd_model.to(devices.cpu) + sd_hijack.model_hijack.undo_hijack(sd_model) + load_model_weights(sd_model, checkpoint_info.filename, checkpoint_info.hash) + sd_hijack.model_hijack.hijack(sd_model) + if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: sd_model.to(devices.device)