From 924e222004ab54273806c5f2ca7a0e7cfa76ad83 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 18 Jan 2023 23:04:24 +0300
Subject: add option to show/hide warnings removed hiding warnings from LDSR
 fixed/reworked few places that produced warnings

---
 modules/sd_hijack_checkpoint.py | 38 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

(limited to 'modules/sd_hijack_checkpoint.py')

diff --git a/modules/sd_hijack_checkpoint.py b/modules/sd_hijack_checkpoint.py
index 5712972f..2604d969 100644
--- a/modules/sd_hijack_checkpoint.py
+++ b/modules/sd_hijack_checkpoint.py
@@ -1,10 +1,46 @@
 from torch.utils.checkpoint import checkpoint
 
+import ldm.modules.attention
+import ldm.modules.diffusionmodules.openaimodel
+
+
 def BasicTransformerBlock_forward(self, x, context=None):
     return checkpoint(self._forward, x, context)
 
+
 def AttentionBlock_forward(self, x):
     return checkpoint(self._forward, x)
 
+
 def ResBlock_forward(self, x, emb):
-    return checkpoint(self._forward, x, emb)
\ No newline at end of file
+    return checkpoint(self._forward, x, emb)
+
+
+stored = []
+
+
+def add():
+    if len(stored) != 0:
+        return
+
+    stored.extend([
+        ldm.modules.attention.BasicTransformerBlock.forward,
+        ldm.modules.diffusionmodules.openaimodel.ResBlock.forward,
+        ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward
+    ])
+
+    ldm.modules.attention.BasicTransformerBlock.forward = BasicTransformerBlock_forward
+    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = ResBlock_forward
+    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = AttentionBlock_forward
+
+
+def remove():
+    if len(stored) == 0:
+        return
+
+    ldm.modules.attention.BasicTransformerBlock.forward = stored[0]
+    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = stored[1]
+    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = stored[2]
+
+    stored.clear()
+
-- 
cgit v1.2.1