From 6a7bc35db845179a26e62534f3d4b789151e52fe Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 9 Jan 2024 13:46:52 -0500 Subject: [PATCH] Use basic attention implementation for small inputs on old pytorch. --- comfy/ldm/modules/attention.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 309240d..fd8888d 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -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