Merge branch 'dev' into kingljl-patch-memory-leak

1 files changed, 8 insertions, 6 deletions

diff --git a/modules/sd_unet.py b/modules/sd_unet.py
index 5525cfbc..a771849c 100644
--- a/modules/sd_unet.py
+++ b/modules/sd_unet.py
@@ -1,12 +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  # 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 ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs)

+    return UNetModel_forward