Allow model files to contain which attention should be used for each … · code/app-python-comfyui@c194dd0 · GitHub
Skip to content

Commit c194dd0

Browse files
Allow model files to contain which attention should be used for each block. (Comfy-Org#16419)
1 parent 1d61dcc commit c194dd0

5 files changed

Lines changed: 67 additions & 8 deletions

File tree

‎QUANTIZATION.md‎

Lines changed: 21 additions & 0 deletions

‎comfy/ldm/minimax/model.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
import comfy.ops
2626
import comfy.patcher_extension
2727
import comfy.quant_ops
28-
from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention
28+
from comfy.ldm.modules.attention import ComfyAttention, AttentionTensorContainer, optimized_attention
2929

3030
FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
3131
FRAME_RESCALE = 5.0 / 3.0
@@ -158,6 +158,7 @@ def rope_rotation_table(angles, dtype):
158158
class Attention(nn.Module):
159159
def __init__(self, hidden, heads, head_dim, eps, gate_compress=False, dtype=None, device=None, operations=None):
160160
super().__init__()
161+
self.comfy_attention = ComfyAttention()
161162
self.heads = heads
162163
self.head_dim = head_dim
163164
inner = heads * head_dim
@@ -196,7 +197,7 @@ def forward(self, x, rope_freqs=None, transformer_options={}):
196197
q = AttentionTensorContainer(q.transpose(0, 1).unsqueeze(0))
197198
k = AttentionTensorContainer(k.transpose(0, 1).unsqueeze(0))
198199
v = AttentionTensorContainer(v.transpose(0, 1).unsqueeze(0))
199-
out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
200+
out = optimized_attention(q, k, v, self.heads, preferred_attention=self.comfy_attention, mask=None, skip_reshape=True, transformer_options=transformer_options)
200201
return self.out_proj(out.squeeze(0))
201202

202203

‎comfy/ldm/modules/attention.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import math
22
import sys
33
import inspect
4+
import json
45

56
import torch
67
import torch.nn.functional as F
@@ -71,6 +72,34 @@ def get_attention_function(name: str, default: Any=...) -> Union[Callable, None]
7172
return default
7273
return REGISTERED_ATTENTION_FUNCTIONS[name]
7374

75+
76+
class ComfyAttention(nn.Module):
77+
def __init__(self):
78+
super().__init__()
79+
self.config = None
80+
self.function = None
81+
82+
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
83+
self.config = None
84+
self.function = None
85+
metadata = state_dict.pop(prefix + "config", None)
86+
if metadata is not None:
87+
config = json.loads(metadata.numpy().tobytes())
88+
method = config.get("attention")
89+
if method == "comfy_kitchen_int8":
90+
self.config = config
91+
if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE and comfy_kitchen.int8_attention_is_available(model_management.get_torch_device()):
92+
self.function = attention_comfy_kitchen_int8
93+
else:
94+
logging.warning(f"Ignoring unsupported attention method {method!r} for {prefix.rstrip('.')}")
95+
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
96+
97+
def _save_to_state_dict(self, destination, prefix, keep_vars):
98+
super()._save_to_state_dict(destination, prefix, keep_vars)
99+
if self.config is not None:
100+
destination[prefix + "config"] = torch.tensor(list(json.dumps(self.config).encode("utf-8")), dtype=torch.uint8)
101+
102+
74103
from comfy.cli_args import args
75104
import comfy.ops
76105
ops = comfy.ops.disable_weight_init
@@ -171,6 +200,7 @@ def take(self) -> torch.Tensor:
171200
def wrap_attn(func):
172201
@functools.wraps(func)
173202
def wrapper(*args, **kwargs):
203+
preferred_attention = kwargs.pop("preferred_attention", None)
174204
containers = None
175205
if len(args) >= 3 and isinstance(args[0], AttentionTensorContainer):
176206
if not isinstance(args[1], AttentionTensorContainer) or not isinstance(args[2], AttentionTensorContainer):
@@ -191,6 +221,10 @@ def wrapper(*args, **kwargs):
191221
return optimized_attention_override.container_function(*args, **kwargs)
192222
args = tuple(container.take() for container in containers) + args[3:]
193223
return optimized_attention_override(func, *args, **kwargs)
224+
if preferred_attention is not None:
225+
attention = preferred_attention.function
226+
if attention is not None:
227+
return attention(*args, **kwargs)
194228

195229
if containers is not None:
196230
if wrapper.container_function is not None:

‎comfy/ldm/qwen_image21/model.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from comfy.ldm.flux.layers import EmbedND, timestep_embedding
1313
from comfy.ldm.flux.math import apply_rope1
1414
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
1616
from comfy.ldm.wan.model_animate2 import PoseBranchCache
1717

1818

@@ -69,6 +69,7 @@ def forward(self, x):
6969
class Attention(nn.Module):
7070
def __init__(self, dim, heads, dim_head, eps=1e-6, dtype=None, device=None, operations=None):
7171
super().__init__()
72+
self.comfy_attention = ComfyAttention()
7273
self.heads = heads
7374
inner_dim = heads * dim_head
7475
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={}):
103104
q, k = comfy.quant_ops.ck.rms_rope(q, k, pe, q_scale, k_scale, self.norm_q.eps)
104105
comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_stream)
105106
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))
107108

108109

109110
def _split_rows(p):
@@ -164,20 +165,20 @@ def forward(self, x, temb):
164165

165166
def block_causal_attention(segments, transformer_options={}, cache=None, block_index=0, prefix_len=0):
166167
# 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):
168169
if cache is not None:
169170
# K and V stacked on dim 1 so batch stays first and quantized rows are per token and head
170171
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)
172173
for start, end, mask in segments]
173174
return torch.cat(outs, dim=1) if len(outs) > 1 else outs[0]
174175
return attn
175176

176177

177178
def prefix_cached_attention(prefix_k, prefix_v, transformer_options={}):
178179
# 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)
181182
return attn
182183

183184

‎comfy/model_patcher.py‎

Lines changed: 2 additions & 0 deletions

0 commit comments

Comments
 (0)