wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-10 00:07:33 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Use optimized attention in Dino2AttentionBlock.

Browse Source
This commit is contained in:
Talmaj Marinc 2026-06-08 12:46:58 +02:00
parent 15f4dc401a
commit 8cbdd8f72e
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/image_encoders/dino2.py
View File

@ -53,8 +53,7 @@ class Dino2AttentionBlock(torch.nn.Module):
if rope is not None and pos is not None:
q = rope(q, pos)
k = rope(k, pos)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
out = out.transpose(1, 2).reshape(B, N, C)
out = optimized_attention(q, k, v, h, mask=mask, skip_reshape=True)
return self.output(out)