Mirrors
/
stable-diffusion-webui
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix nested manual cast

This commit is contained in:
Kohaku-Blueleaf 2024-01-19 00:14:03 +08:00
parent cb5b335acd
commit 0181c1f76b
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
modules/devices.py
View File

@ -165,6 +165,8 @@ def manual_cast_forward(target_dtype):
@contextlib.contextmanager @contextlib.contextmanager
def manual_cast(target_dtype): def manual_cast(target_dtype):
for module_type in patch_module_list: for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
continue
org_forward = module_type.forward org_forward = module_type.forward
if module_type == torch.nn.MultiheadAttention and has_xpu(): if module_type == torch.nn.MultiheadAttention and has_xpu():
module_type.forward = manual_cast_forward(torch.float32) module_type.forward = manual_cast_forward(torch.float32)
@ -175,7 +177,9 @@ def manual_cast(target_dtype):
yield None yield None
finally: finally:
for module_type in patch_module_list: 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): def autocast(disable=False):