1 files changed, 12 insertions, 5 deletions
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 592f0055..22a1eb5c 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -2,7 +2,7 @@ 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
@@ -10,6 +10,7 @@ from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hija
 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()
@@ -239,10 +242,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)
@@ -279,7 +285,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: