From 59544321aa019d71d220b1da1eec703aa44fa8eb Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 11 Sep 2023 21:17:28 +0300
Subject: [PATCH] initial work on sd_unet for SDXL

---
 modules/sd_hijack.py | 17 ++++++++++++-----
 modules/sd_unet.py   |  4 ++--
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 592f00551f1..22a1eb5ca7d 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -2,7 +2,7 @@
 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 @@
 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 @@
 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 @@ def flatten(el):
 
         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 @@ def undo_hijack(self, m):
         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:
diff --git a/modules/sd_unet.py b/modules/sd_unet.py
index 5525cfbc3a0..6a7bc9e2646 100644
--- a/modules/sd_unet.py
+++ b/modules/sd_unet.py
@@ -1,11 +1,11 @@
 import torch.nn
-import ldm.modules.diffusionmodules.openaimodel
 
 from modules import script_callbacks, shared, devices
 
 unet_options = []
 current_unet_option = None
 current_unet = None
+original_forward = None
 
 
 def list_unets():
@@ -88,5 +88,5 @@ def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs):
     if current_unet is not None:
         return current_unet.forward(x, timesteps, context, *args, **kwargs)
 
-    return ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs)
+    return original_forward(self, x, timesteps, context, *args, **kwargs)