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

Avoid early disable

This commit is contained in:
Kohaku-Blueleaf 2024-01-20 16:31:12 +08:00
parent 0181c1f76b
commit 81126027f5
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
modules/devices.py
View File

@ -164,9 +164,11 @@ def manual_cast_forward(target_dtype):
@contextlib.contextmanager
def manual_cast(target_dtype):
applied = False
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
continue
applied = True
org_forward = module_type.forward
if module_type == torch.nn.MultiheadAttention and has_xpu():
module_type.forward = manual_cast_forward(torch.float32)
@ -176,6 +178,8 @@ def manual_cast(target_dtype):
try:
yield None
finally:
if not applied:
return
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward