From feb3d7510296e0e16e7e905ef60194fca0f68cc4 Mon Sep 17 00:00:00 2001 From: blepping Date: Fri, 9 May 2025 09:27:31 -0600 Subject: [PATCH] Let optimized_attention handle output reshape for ACE --- comfy/ldm/ace/attention.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/comfy/ldm/ace/attention.py b/comfy/ldm/ace/attention.py index c81801747..f20a01669 100644 --- a/comfy/ldm/ace/attention.py +++ b/comfy/ldm/ace/attention.py @@ -433,11 +433,8 @@ class CustomerAttnProcessor2_0: # the output of sdp = (batch, num_heads, seq_len, head_dim) hidden_states = optimized_attention( - query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True, skip_output_reshape=True, - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) + query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True, + ).to(query.dtype) # linear proj hidden_states = attn.to_out[0](hidden_states)