298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-03-13 14:21:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use basic attention implementation for small inputs on old pytorch.

Browse Source
This commit is contained in:
comfyanonymous 2024-01-09 13:46:52 -05:00
parent b3b5ddb07a
commit 6a7bc35db8
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy/ldm/modules/attention.py
View File

@ -351,8 +351,11 @@ else:
optimized_attention_masked = optimized_attention
def optimized_attention_for_device(device, mask=False, small_input=False):
if small_input and model_management.pytorch_attention_enabled():
return attention_pytorch #TODO: need to confirm but this is probably slightly faster for small inputs in all cases
if small_input:
if model_management.pytorch_attention_enabled():
return attention_pytorch #TODO: need to confirm but this is probably slightly faster for small inputs in all cases
else:
return attention_basic
if device == torch.device("cpu"):
return attention_sub_quad