wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-12-18 02:23:06 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

Add optimized to get_attention_function

Browse Source
This commit is contained in:
Jedrzej Kosinski 2025-08-29 21:48:36 -07:00
parent d553073a1e
commit cb959f9669
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -51,7 +51,9 @@ def register_attention_function(name: str, func: Callable):
logging.warning(f"Attention function {name} already registered, skipping registration.")
def get_attention_function(name: str, default: Any=...) -> Union[Callable, None]:
if name not in REGISTERED_ATTENTION_FUNCTIONS:
if name == "optimized":
return optimized_attention
elif name not in REGISTERED_ATTENTION_FUNCTIONS:
if default is ...:
raise KeyError(f"Attention function {name} not found.")
else:
@ -62,7 +64,7 @@ def _register_core_attention_functions():
"""
Register attention functions exposed by core ComfyUI.
"""
# NOTE: attention_basic is purposely not registered, as it is not used in code
# NOTE: attention_basic is purposely not registered, as it should not be used
if SAGE_ATTENTION_IS_AVAILABLE:
register_attention_function("sage", attention_sage)
if FLASH_ATTENTION_IS_AVAILABLE: