From 9d2299ed0bd6c81cae8a7ba4ca22d6a14fb27bef Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Sat, 19 Aug 2023 10:16:27 +0300
Subject: [PATCH] implement undo hijack for SDXL

---
 modules/sd_hijack.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 46652fbd3..592f00551 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -245,7 +245,21 @@ class StableDiffusionModelHijack:
         ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_unet.UNetModel_forward
 
     def undo_hijack(self, m):
-        if type(m.cond_stage_model) == sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords:
+        conditioner = getattr(m, 'conditioner', None)
+        if conditioner:
+            for i in range(len(conditioner.embedders)):
+                embedder = conditioner.embedders[i]
+                if isinstance(embedder, (sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords, sd_hijack_open_clip.FrozenOpenCLIPEmbedder2WithCustomWords)):
+                    embedder.wrapped.model.token_embedding = embedder.wrapped.model.token_embedding.wrapped
+                    conditioner.embedders[i] = embedder.wrapped
+                if isinstance(embedder, sd_hijack_clip.FrozenCLIPEmbedderForSDXLWithCustomWords):
+                    embedder.wrapped.transformer.text_model.embeddings.token_embedding = embedder.wrapped.transformer.text_model.embeddings.token_embedding.wrapped
+                    conditioner.embedders[i] = embedder.wrapped
+
+            if hasattr(m, 'cond_stage_model'):
+                delattr(m, 'cond_stage_model')
+
+        elif type(m.cond_stage_model) == sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords:
             m.cond_stage_model = m.cond_stage_model.wrapped
 
         elif type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords: