AUTOMATIC1111/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2026-04-29 18:51:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14689 from AUTOMATIC1111/fix-nested-manual-cast

Browse Source 
Fix nested manual cast
This commit is contained in:
AUTOMATIC1111
2024-01-20 11:43:51 +03:00
committed by GitHub
parent 1dbee391b4 4a66d2fb22
commit c1713bfeac
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
modules/devices.py
View File

@@ -164,7 +164,11 @@ def manual_cast_forward(target_dtype):
@contextlib.contextmanager @contextlib.contextmanager
def manual_cast(target_dtype): def manual_cast(target_dtype):
applied = False
for module_type in patch_module_list: for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
continue
applied = True
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)
@@ -174,8 +178,11 @@ def manual_cast(target_dtype):
try: try:
yield None yield None
finally: finally:
for module_type in patch_module_list: if applied:
module_type.forward = module_type.org_forward 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): def autocast(disable=False):