Skip to content

S
Stable Diffusion Webui
Commit 0181c1f7 authored by Kohaku-Blueleaf's avatar Kohaku-Blueleaf
Browse files
Options

Fix nested manual cast

parent cb5b335a
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 1 deletion
modules/devices.py
View file @ 0181c1f7
...@@ -165,6 +165,8 @@ def manual_cast_forward(target_dtype): ...@@ -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): ...@@ -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:
if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward module_type.forward = module_type.org_forward
delattr(module_type, "org_forward")
def autocast(disable=False): def autocast(disable=False):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment