|
12 | 12 | from comfy.ldm.flux.layers import EmbedND, timestep_embedding |
13 | 13 | from comfy.ldm.flux.math import apply_rope1 |
14 | 14 | from comfy.ldm.lightricks.model import TimestepEmbedding |
15 | | -from comfy.ldm.modules.attention import optimized_attention |
| 15 | +from comfy.ldm.modules.attention import ComfyAttention, optimized_attention |
16 | 16 | from comfy.ldm.wan.model_animate2 import PoseBranchCache |
17 | 17 |
|
18 | 18 |
|
@@ -69,6 +69,7 @@ def forward(self, x): |
69 | 69 | class Attention(nn.Module): |
70 | 70 | def __init__(self, dim, heads, dim_head, eps=1e-6, dtype=None, device=None, operations=None): |
71 | 71 | super().__init__() |
| 72 | + self.comfy_attention = ComfyAttention() |
72 | 73 | self.heads = heads |
73 | 74 | inner_dim = heads * dim_head |
74 | 75 | self.to_q = operations.Linear(dim, inner_dim, bias=False, dtype=dtype, device=device) |
@@ -103,7 +104,7 @@ def forward(self, x, pe, attn_fn, prefix_len, transformer_options={}): |
103 | 104 | q, k = comfy.quant_ops.ck.rms_rope(q, k, pe, q_scale, k_scale, self.norm_q.eps) |
104 | 105 | comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_stream) |
105 | 106 | comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_stream) |
106 | | - return self.to_out[0](attn_fn(q, k, v, self.heads)) |
| 107 | + return self.to_out[0](attn_fn(q, k, v, self.heads, preferred_attention=self.comfy_attention)) |
107 | 108 |
|
108 | 109 |
|
109 | 110 | def _split_rows(p): |
@@ -164,20 +165,20 @@ def forward(self, x, temb): |
164 | 165 |
|
165 | 166 | def block_causal_attention(segments, transformer_options={}, cache=None, block_index=0, prefix_len=0): |
166 | 167 | # segments: (start, end, mask); text segments get a causal mask, image blocks attend to everything before their end |
167 | | - def attn(q, k, v, heads): |
| 168 | + def attn(q, k, v, heads, preferred_attention=None): |
168 | 169 | if cache is not None: |
169 | 170 | # K and V stacked on dim 1 so batch stays first and quantized rows are per token and head |
170 | 171 | cache.put(block_index, torch.stack([k[:, :prefix_len], v[:, :prefix_len]], dim=1)) |
171 | | - outs = [optimized_attention(q[:, start:end].flatten(2), k[:, :end].flatten(2), v[:, :end].flatten(2), heads, mask=mask, transformer_options=transformer_options) |
| 172 | + outs = [optimized_attention(q[:, start:end].flatten(2), k[:, :end].flatten(2), v[:, :end].flatten(2), heads, mask=mask, transformer_options=transformer_options, preferred_attention=preferred_attention) |
172 | 173 | for start, end, mask in segments] |
173 | 174 | return torch.cat(outs, dim=1) if len(outs) > 1 else outs[0] |
174 | 175 | return attn |
175 | 176 |
|
176 | 177 |
|
177 | 178 | def prefix_cached_attention(prefix_k, prefix_v, transformer_options={}): |
178 | 179 | # target-only queries: block-causal reduces to full attention over [cached prefix, target] |
179 | | - def attn(q, k, v, heads): |
180 | | - return optimized_attention(q.flatten(2), torch.cat([prefix_k, k], dim=1).flatten(2), torch.cat([prefix_v, v], dim=1).flatten(2), heads, transformer_options=transformer_options) |
| 180 | + def attn(q, k, v, heads, preferred_attention=None): |
| 181 | + return optimized_attention(q.flatten(2), torch.cat([prefix_k, k], dim=1).flatten(2), torch.cat([prefix_v, v], dim=1).flatten(2), heads, transformer_options=transformer_options, preferred_attention=preferred_attention) |
181 | 182 | return attn |
182 | 183 |
|
183 | 184 |
|
|
0 commit comments