Merge branch 'dev' into master

1 files changed, 14 insertions, 7 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 0a7e5135..4fff418d 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -2,14 +2,15 @@ import torch
 from torch.nn.functional import silu

 from types import MethodType

 

-from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet

+from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet, patches

 from modules.hypernetworks import hypernetwork

 from modules.shared import cmd_opts

-from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr

+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr, xlmr_m18

 

 import ldm.modules.attention

 import ldm.modules.diffusionmodules.model

 import ldm.modules.diffusionmodules.openaimodel

+import ldm.models.diffusion.ddpm

 import ldm.models.diffusion.ddim

 import ldm.models.diffusion.plms

 import ldm.modules.encoders.modules

@@ -37,6 +38,8 @@ ldm.models.diffusion.ddpm.print = shared.ldm_print
 optimizers = []

 current_optimizer: sd_hijack_optimizations.SdOptimization = None

 

+ldm_original_forward = patches.patch(__file__, ldm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward)

+sgm_original_forward = patches.patch(__file__, sgm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward)

 

 def list_optimizers():

     new_optimizers = script_callbacks.list_optimizers_callback()

@@ -220,7 +223,7 @@ class StableDiffusionModelHijack:
             else:

                 m.cond_stage_model = conditioner

 

-        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

+        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation or type(m.cond_stage_model) == xlmr_m18.BertSeriesModelWithTransformation:

             model_embeddings = m.cond_stage_model.roberta.embeddings

             model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.word_embeddings, self)

             m.cond_stage_model = sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords(m.cond_stage_model, self)

@@ -251,10 +254,13 @@ class StableDiffusionModelHijack:
 

         self.layers = flatten(m)

 

-        if not hasattr(ldm.modules.diffusionmodules.openaimodel, 'copy_of_UNetModel_forward_for_webui'):

-            ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui = ldm.modules.diffusionmodules.openaimodel.UNetModel.forward

+        if isinstance(m, ldm.models.diffusion.ddpm.LatentDiffusion):

+            sd_unet.original_forward = ldm_original_forward

+        elif isinstance(m, sgm.models.diffusion.DiffusionEngine):

+            sd_unet.original_forward = sgm_original_forward

+        else:

+            sd_unet.original_forward = None

 

-        ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_unet.UNetModel_forward

 

     def undo_hijack(self, m):

         conditioner = getattr(m, 'conditioner', None)

@@ -291,7 +297,8 @@ class StableDiffusionModelHijack:
         self.layers = None

         self.clip = None

 

-        ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui

+        sd_unet.original_forward = None

+

 

     def apply_circular(self, enable):

         if self.circular_enabled == enable: