From 750dd6014a45397979cad42a74634451d0861581 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:27:53 +0800
Subject: Fix potential bugs

---
 modules/devices.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

(limited to 'modules')

diff --git a/modules/devices.py b/modules/devices.py
index dfffaf24..60f7d6d7 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -141,7 +141,12 @@ def manual_cast_forward(target_dtype):
             args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
             kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
 
-        org_dtype = torch_utils.get_param(self).dtype
+        org_dtype = target_dtype
+        for param in self.parameters():
+            if param.dtype != target_dtype:
+                org_dtype = param.dtype
+                break
+
         if org_dtype != target_dtype:
             self.to(target_dtype)
         result = self.org_forward(*args, **kwargs)
@@ -170,7 +175,7 @@ def manual_cast(target_dtype):
             continue
         applied = True
         org_forward = module_type.forward
-        if module_type == torch.nn.MultiheadAttention and has_xpu():
+        if module_type == torch.nn.MultiheadAttention:
             module_type.forward = manual_cast_forward(torch.float32)
         else:
             module_type.forward = manual_cast_forward(target_dtype)
-- 
cgit v1.2.1


From 6e7f0860f7ae4a0ce59f9416fb9b2f3bcab44f1d Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:46:43 +0800
Subject: linting

---
 modules/devices.py | 1 -
 1 file changed, 1 deletion(-)

(limited to 'modules')

diff --git a/modules/devices.py b/modules/devices.py
index 60f7d6d7..8f49f7a4 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -4,7 +4,6 @@ from functools import lru_cache
 
 import torch
 from modules import errors, shared
-from modules import torch_utils
 
 if sys.platform == "darwin":
     from modules import mac_specific
-- 
cgit v1.2.1


From d243e24f539d717b221992e894a5db5a321bf3cd Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:49:45 +0800
Subject: Try to reverse the dtype checking mechanism

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

(limited to 'modules')

diff --git a/modules/devices.py b/modules/devices.py
index 8f49f7a4..f9648e9a 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -4,6 +4,7 @@ from functools import lru_cache
 
 import torch
 from modules import errors, shared
+from modules import torch_utils
 
 if sys.platform == "darwin":
     from modules import mac_specific
@@ -140,11 +141,7 @@ def manual_cast_forward(target_dtype):
             args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
             kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
 
-        org_dtype = target_dtype
-        for param in self.parameters():
-            if param.dtype != target_dtype:
-                org_dtype = param.dtype
-                break
+        org_dtype = torch_utils.get_param(self).dtype
 
         if org_dtype != target_dtype:
             self.to(target_dtype)
-- 
cgit v1.2.1


From f9ba7e648ad5bf7dbdf2b95fa207936179bf784e Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 29 Jan 2024 22:54:12 +0800
Subject: Revert "Try to reverse the dtype checking mechanism"

This reverts commit d243e24f539d717b221992e894a5db5a321bf3cd.
---
 modules/devices.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

(limited to 'modules')

diff --git a/modules/devices.py b/modules/devices.py
index f9648e9a..8f49f7a4 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -4,7 +4,6 @@ from functools import lru_cache
 
 import torch
 from modules import errors, shared
-from modules import torch_utils
 
 if sys.platform == "darwin":
     from modules import mac_specific
@@ -141,7 +140,11 @@ def manual_cast_forward(target_dtype):
             args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
             kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
 
-        org_dtype = torch_utils.get_param(self).dtype
+        org_dtype = target_dtype
+        for param in self.parameters():
+            if param.dtype != target_dtype:
+                org_dtype = param.dtype
+                break
 
         if org_dtype != target_dtype:
             self.to(target_dtype)
-- 
cgit v1.2.1