OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-03-03 12:34:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Avoid early disable

Browse Source
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 @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"): if hasattr(module_type, "org_forward"):
continue 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)
@ -176,6 +178,8 @@ def manual_cast(target_dtype):
try: try:
yield None yield None
finally: finally:
if not applied:
return
for module_type in patch_module_list: for module_type in patch_module_list:
if hasattr(module_type, "org_forward"): if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward module_type.forward = module_type.org_forward