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: