alternate implementation for unet forward replacement that does not depend on hijack being applied

1 files changed, 8 insertions, 6 deletions

diff --git a/modules/sd_unet.py b/modules/sd_unet.py
index 6a7bc9e2..a771849c 100644
--- a/modules/sd_unet.py
+++ b/modules/sd_unet.py
@@ -5,8 +5,7 @@ from modules import script_callbacks, shared, devices
 unet_options = []

 current_unet_option = None

 current_unet = None

-original_forward = None

-

+original_forward = None  # not used, only left temporarily for compatibility

 

 def list_unets():

     new_unets = script_callbacks.list_unets_callback()

@@ -84,9 +83,12 @@ class SdUnet(torch.nn.Module):
         pass

 

 

-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)

+def create_unet_forward(original_forward):

+    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 original_forward(self, x, timesteps, context, *args, **kwargs)

 

-    return original_forward(self, x, timesteps, context, *args, **kwargs)

+    return UNetModel_forward