mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-22 17:23:36 +08:00
Use optimized attention in Dino2AttentionBlock.
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user