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