xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-04 18:27:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Let optimized_attention handle output reshape for ACE

Browse Source
This commit is contained in:
blepping 2025-05-09 09:27:31 -06:00
parent 5284548432
commit feb3d75102
1 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -433,11 +433,8 @@ class CustomerAttnProcessor2_0:
# the output of sdp = (batch, num_heads, seq_len, head_dim) # the output of sdp = (batch, num_heads, seq_len, head_dim)
hidden_states = optimized_attention( hidden_states = optimized_attention(
query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True, skip_output_reshape=True, query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True,
) ).to(query.dtype)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj # linear proj
hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[0](hidden_states)