From 0181c1f76b97162c42401f1e6286ae73d8aa6033 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Fri, 19 Jan 2024 00:14:03 +0800
Subject: [PATCH 1/3] Fix nested manual cast

---
 modules/devices.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/modules/devices.py b/modules/devices.py
index 0321d12c6..370286297 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -165,6 +165,8 @@ def manual_cast_forward(target_dtype):
 @contextlib.contextmanager
 def manual_cast(target_dtype):
     for module_type in patch_module_list:
+        if hasattr(module_type, "org_forward"):
+            continue
         org_forward = module_type.forward
         if module_type == torch.nn.MultiheadAttention and has_xpu():
             module_type.forward = manual_cast_forward(torch.float32)
@@ -175,7 +177,9 @@ def manual_cast(target_dtype):
         yield None
     finally:
         for module_type in patch_module_list:
-            module_type.forward = module_type.org_forward
+            if hasattr(module_type, "org_forward"):
+                module_type.forward = module_type.org_forward
+                delattr(module_type, "org_forward")
 
 
 def autocast(disable=False):

From 81126027f5226e7ee58e1a99194eb9ec7b8ec6e7 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sat, 20 Jan 2024 16:31:12 +0800
Subject: [PATCH 2/3] Avoid early disable

---
 modules/devices.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modules/devices.py b/modules/devices.py
index 370286297..3bde16990 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -164,9 +164,11 @@ def manual_cast_forward(target_dtype):
 
 @contextlib.contextmanager
 def manual_cast(target_dtype):
+    applied = False
     for module_type in patch_module_list:
         if hasattr(module_type, "org_forward"):
             continue
+        applied = True
         org_forward = module_type.forward
         if module_type == torch.nn.MultiheadAttention and has_xpu():
             module_type.forward = manual_cast_forward(torch.float32)
@@ -176,6 +178,8 @@ def manual_cast(target_dtype):
     try:
         yield None
     finally:
+        if not applied:
+            return
         for module_type in patch_module_list:
             if hasattr(module_type, "org_forward"):
                 module_type.forward = module_type.org_forward

From 4a66d2fb228584bb38dc22db6a3e657561834c7a Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sat, 20 Jan 2024 16:33:59 +0800
Subject: [PATCH 3/3] Avoid exceptions to be silenced

---
 modules/devices.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/modules/devices.py b/modules/devices.py
index 3bde16990..dfffaf24f 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -178,12 +178,11 @@ def manual_cast(target_dtype):
     try:
         yield None
     finally:
-        if not applied:
-            return
-        for module_type in patch_module_list:
-            if hasattr(module_type, "org_forward"):
-                module_type.forward = module_type.org_forward
-                delattr(module_type, "org_forward")
+        if applied:
+            for module_type in patch_module_list:
+                if hasattr(module_type, "org_forward"):
+                    module_type.forward = module_type.org_forward
+                    delattr(module_type, "org_forward")
 
 
 def autocast(disable=False):