Add a way for nodes to add pre attn patches to flux model. (#12861)

This commit is contained in:
comfyanonymous
2026-03-09 20:37:58 -07:00
committed by GitHub
parent 740d998c9c
commit c4fb0271cd
3 changed files with 17 additions and 2 deletions

View File

@@ -31,6 +31,8 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
def _apply_rope1(x: Tensor, freqs_cis: Tensor):
x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
if x_.shape[2] != 1 and freqs_cis.shape[2] != 1 and x_.shape[2] != freqs_cis.shape[2]:
freqs_cis = freqs_cis[:, :, :x_.shape[2]]
x_out = freqs_cis[..., 0] * x_[..., 0]
x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])