From c715ef04d1edb1a112a602639ed3bb292fdeb0e2 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 29 Sep 2022 15:40:28 +0300
Subject: fix for incorrect model weight loading for #814

---
 modules/sd_hijack.py | 9 +++++++++
 1 file changed, 9 insertions(+)

(limited to 'modules/sd_hijack.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 2848a251..5945b7c2 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -245,6 +245,7 @@ class StableDiffusionModelHijack:
 
         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 @@ class StableDiffusionModelHijack:
 
         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
-- 
cgit v1.2.3