Skip to content

Commit

Permalink
Slightly lower hunyuan video memory usage.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Jan 16, 2025
1 parent 25683b5 commit 6320d05
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion comfy/ldm/flux/math.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,15 @@
from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management


def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None) -> Tensor:
q, k = apply_rope(q, k, pe)
q_shape = q.shape
k_shape = k.shape

q = q.float().reshape(*q.shape[:-1], -1, 1, 2)
k = k.float().reshape(*k.shape[:-1], -1, 1, 2)
q = (pe[..., 0] * q[..., 0] + pe[..., 1] * q[..., 1]).reshape(*q_shape).type_as(v)
k = (pe[..., 0] * k[..., 0] + pe[..., 1] * k[..., 1]).reshape(*k_shape).type_as(v)

heads = q.shape[1]
x = optimized_attention(q, k, v, heads, skip_reshape=True, mask=mask)
Expand Down

0 comments on commit 6320d05

Please sign in to comment.