mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-12 09:27:15 +08:00
Compare commits
64 Commits
dcd1274b60
...
f4619d118f
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f4619d118f | ||
|
|
ffbecfffb9 | ||
|
|
b481bc15af | ||
|
|
6880614319 | ||
|
|
a75bf6f1d8 | ||
|
|
0bad1b06bd | ||
|
|
07c96b7238 | ||
|
|
4bfa533371 | ||
|
|
fa0eaccfcc | ||
|
|
63e08a02fd | ||
|
|
69cbd50aa6 | ||
|
|
6d4f9e86ab | ||
|
|
1548aee40e | ||
|
|
9c210473fc | ||
|
|
c1e9164c63 | ||
|
|
5e9a90186f | ||
|
|
ece906328a | ||
|
|
500ca8e02a | ||
|
|
3143b7981f | ||
|
|
c9b3f81e83 | ||
|
|
5e74e9b3ed | ||
|
|
c702cddf75 | ||
|
|
e13da8104c | ||
|
|
fdcc38b9ea | ||
|
|
c1ce00287c | ||
|
|
6e3bd33665 | ||
|
|
ce05e377a8 | ||
|
|
1a00f7743f | ||
|
|
a6472b1514 | ||
|
|
6158cd5820 | ||
|
|
bff714dda0 | ||
|
|
fce22da313 | ||
|
|
9f9d37bd9a | ||
|
|
088778c35d | ||
|
|
4c5f82971e | ||
|
|
f1d91a4c8c | ||
|
|
dbed5a1b52 | ||
|
|
24fdbb9aca | ||
|
|
a6624a9afd | ||
|
|
0b512198e8 | ||
|
|
9feb26928c | ||
|
|
fadd79ad48 | ||
|
|
77bc7bdd6b | ||
|
|
117afbc1d7 | ||
|
|
064eec2278 | ||
|
|
aceaa5e579 | ||
|
|
763089f681 | ||
|
|
1693dabc8f | ||
|
|
08063d2638 | ||
|
|
e069617e54 | ||
|
|
2bea0ee5d7 | ||
|
|
17863f603a | ||
|
|
31ba844624 | ||
|
|
1451001f64 | ||
|
|
1af99b2e81 | ||
|
|
3568b82b76 | ||
|
|
6728d4d439 | ||
|
|
4b431ffc27 | ||
|
|
880b51ac4f | ||
|
|
4d9516b909 | ||
|
|
39086890e2 | ||
|
|
2adde5a0e1 | ||
|
|
0c1bfad0df | ||
|
|
7d76a4447e |
@ -127,6 +127,8 @@
|
||||
- Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
|
||||
platform, or backend capability detection only when the program has a useful
|
||||
fallback. Prefer specific exception types when changing new code.
|
||||
- If a library version is pinned in `requirements.txt`, do not add code to
|
||||
ComfyUI to handle older versions of that library.
|
||||
- Remove any workarounds for PyTorch versions that ComfyUI no longer officially
|
||||
supports. Deprecated workarounds include catching an exception and rerunning
|
||||
the same op with the input cast to float. If a workaround does not have a
|
||||
|
||||
@ -229,7 +229,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia
|
||||
|
||||
Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
|
||||
|
||||
torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old.
|
||||
torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
|
||||
|
||||
### Instructions:
|
||||
|
||||
|
||||
@ -217,10 +217,7 @@ class AceStepAttention(nn.Module):
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||
|
||||
n_rep = self.num_heads // self.num_kv_heads
|
||||
if n_rep > 1:
|
||||
key_states = key_states.repeat_interleave(n_rep, dim=1)
|
||||
value_states = value_states.repeat_interleave(n_rep, dim=1)
|
||||
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
|
||||
|
||||
attn_bias = None
|
||||
if self.sliding_window is not None and not self.is_cross_attention:
|
||||
@ -244,7 +241,7 @@ class AceStepAttention(nn.Module):
|
||||
else:
|
||||
attn_bias = window_bias
|
||||
|
||||
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False)
|
||||
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False, **gqa_kwargs)
|
||||
attn_output = self.o_proj(attn_output)
|
||||
|
||||
return attn_output
|
||||
|
||||
@ -425,19 +425,16 @@ class Attention(nn.Module):
|
||||
if n == 1 and causal:
|
||||
causal = False
|
||||
|
||||
if h != kv_h:
|
||||
# Repeat interleave kv_heads to match q_heads
|
||||
heads_per_kv_head = h // kv_h
|
||||
k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
|
||||
gqa_kwargs = {"enable_gqa": True} if h != kv_h else {}
|
||||
|
||||
if self.differential:
|
||||
q, q_diff = q.unbind(dim=1)
|
||||
k, k_diff = k.unbind(dim=1)
|
||||
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
|
||||
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
|
||||
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
|
||||
out = out - out_diff
|
||||
else:
|
||||
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
|
||||
|
||||
out = self.to_out(out)
|
||||
|
||||
|
||||
@ -74,11 +74,8 @@ class BooguDoubleStreamProcessor(nn.Module):
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
|
||||
if attn.kv_heads < attn.heads:
|
||||
key = key.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
|
||||
value = value.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
|
||||
|
||||
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
|
||||
gqa_kwargs = {"enable_gqa": True} if attn.kv_heads < attn.heads else {}
|
||||
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
|
||||
|
||||
# Split back to instruction/image, apply per-stream output projections, recombine.
|
||||
instruct_hidden_states = self.instruct_out(hidden_states[:, :L_instruct])
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
import math
|
||||
import sys
|
||||
import inspect
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@ -14,16 +15,16 @@ from .sub_quadratic_attention import efficient_dot_product_attention
|
||||
|
||||
from comfy import model_management
|
||||
|
||||
TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5)
|
||||
|
||||
if model_management.xformers_enabled():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
|
||||
SAGE_ATTENTION_IS_AVAILABLE = False
|
||||
SAGE_ATTENTION_SUPPORTS_MASK = False
|
||||
try:
|
||||
from sageattention import sageattn
|
||||
SAGE_ATTENTION_IS_AVAILABLE = True
|
||||
SAGE_ATTENTION_SUPPORTS_MASK = "attn_mask" in inspect.signature(sageattn).parameters
|
||||
except ImportError as e:
|
||||
if model_management.sage_attention_enabled():
|
||||
if e.name == "sageattention":
|
||||
@ -89,6 +90,44 @@ def default(val, d):
|
||||
return val
|
||||
return d
|
||||
|
||||
def _gqa_repeat_factor(query_heads, key_heads, value_heads):
|
||||
if key_heads != value_heads:
|
||||
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
|
||||
if query_heads == key_heads:
|
||||
return 1
|
||||
if query_heads % key_heads != 0:
|
||||
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
|
||||
return query_heads // key_heads
|
||||
|
||||
def _repeat_kv_for_gqa(k, v, query_heads, head_dim):
|
||||
n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
|
||||
if n_rep > 1:
|
||||
k = k.repeat_interleave(n_rep, dim=head_dim)
|
||||
v = v.repeat_interleave(n_rep, dim=head_dim)
|
||||
return k, v
|
||||
|
||||
def _heads_from_dim(tensor, dim_head, name):
|
||||
inner_dim = tensor.shape[-1]
|
||||
if inner_dim % dim_head != 0:
|
||||
raise ValueError(f"{name} inner dimension {inner_dim} is not divisible by head dimension {dim_head}")
|
||||
return inner_dim // dim_head
|
||||
|
||||
def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_kv=True):
|
||||
q = q.unsqueeze(3).reshape(b, -1, heads, dim_head)
|
||||
if enable_gqa:
|
||||
key_heads = _heads_from_dim(k, dim_head, "Key")
|
||||
value_heads = _heads_from_dim(v, dim_head, "Value")
|
||||
else:
|
||||
key_heads = heads
|
||||
value_heads = heads
|
||||
k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head)
|
||||
v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head)
|
||||
if enable_gqa:
|
||||
_gqa_repeat_factor(heads, key_heads, value_heads)
|
||||
if expand_kv:
|
||||
k, v = _repeat_kv_for_gqa(k, v, heads, -2)
|
||||
return q, k, v
|
||||
|
||||
|
||||
# feedforward
|
||||
class GEGLU(nn.Module):
|
||||
@ -152,28 +191,19 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]:
|
||||
n_rep = q.shape[-3] // k.shape[-3]
|
||||
k = k.repeat_interleave(n_rep, dim=-3)
|
||||
v = v.repeat_interleave(n_rep, dim=-3)
|
||||
|
||||
scale = kwargs.get("scale", dim_head ** -0.5)
|
||||
|
||||
h = heads
|
||||
if skip_reshape:
|
||||
q, k, v = map(
|
||||
if kwargs.get("enable_gqa", False):
|
||||
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
|
||||
q, k, v = map(
|
||||
lambda t: t.reshape(b * heads, -1, dim_head),
|
||||
(q, k, v),
|
||||
)
|
||||
else:
|
||||
q, k, v = map(
|
||||
lambda t: t.unsqueeze(3)
|
||||
.reshape(b, -1, heads, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b * heads, -1, dim_head)
|
||||
.contiguous(),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
|
||||
|
||||
# force cast to fp32 to avoid overflowing
|
||||
if attn_precision == torch.float32:
|
||||
@ -231,13 +261,16 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
|
||||
query = query * (kwargs["scale"] * dim_head ** 0.5)
|
||||
|
||||
if skip_reshape:
|
||||
if kwargs.get("enable_gqa", False):
|
||||
key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3)
|
||||
query = query.reshape(b * heads, -1, dim_head)
|
||||
value = value.reshape(b * heads, -1, dim_head)
|
||||
key = key.reshape(b * heads, -1, dim_head).movedim(1, 2)
|
||||
else:
|
||||
query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
|
||||
value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
|
||||
key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
|
||||
query, key, value = _reshape_qkv_to_heads(query, key, value, b, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
query = query.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
|
||||
value = value.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
|
||||
key = key.permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
|
||||
|
||||
|
||||
dtype = query.dtype
|
||||
@ -304,19 +337,15 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
scale = kwargs.get("scale", dim_head ** -0.5)
|
||||
|
||||
if skip_reshape:
|
||||
q, k, v = map(
|
||||
if kwargs.get("enable_gqa", False):
|
||||
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
|
||||
q, k, v = map(
|
||||
lambda t: t.reshape(b * heads, -1, dim_head),
|
||||
(q, k, v),
|
||||
)
|
||||
else:
|
||||
q, k, v = map(
|
||||
lambda t: t.unsqueeze(3)
|
||||
.reshape(b, -1, heads, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b * heads, -1, dim_head)
|
||||
.contiguous(),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
|
||||
|
||||
r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
|
||||
|
||||
@ -438,7 +467,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
|
||||
disabled_xformers = True
|
||||
|
||||
if disabled_xformers:
|
||||
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
|
||||
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
|
||||
|
||||
if skip_reshape:
|
||||
# b h k d -> b k h d
|
||||
@ -446,13 +475,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
|
||||
lambda t: t.permute(0, 2, 1, 3),
|
||||
(q, k, v),
|
||||
)
|
||||
if kwargs.get("enable_gqa", False):
|
||||
k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2)
|
||||
# actually do the reshaping
|
||||
else:
|
||||
dim_head //= heads
|
||||
q, k, v = map(
|
||||
lambda t: t.reshape(b, -1, heads, dim_head),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
|
||||
if mask is not None:
|
||||
# add a singleton batch dimension
|
||||
@ -474,7 +502,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
|
||||
mask = mask_out[..., :mask.shape[-1]]
|
||||
mask = mask.expand(b, heads, -1, -1)
|
||||
|
||||
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
|
||||
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, scale=kwargs.get("scale", None))
|
||||
|
||||
if skip_output_reshape:
|
||||
out = out.permute(0, 2, 1, 3)
|
||||
@ -498,10 +526,8 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
else:
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
q, k, v = map(
|
||||
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
|
||||
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
|
||||
|
||||
if mask is not None:
|
||||
# add a batch dimension if there isn't already one
|
||||
@ -511,9 +537,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
# Pass through extra SDPA kwargs (scale, enable_gqa) if provided
|
||||
# enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above
|
||||
sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",)
|
||||
sdpa_keys = ("scale", "enable_gqa")
|
||||
sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys}
|
||||
|
||||
if SDP_BATCH_LIMIT >= b:
|
||||
@ -541,20 +565,19 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
|
||||
@wrap_attn
|
||||
def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
||||
if kwargs.get("low_precision_attention", True) is False:
|
||||
if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK):
|
||||
return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
|
||||
|
||||
exception_fallback = False
|
||||
if skip_reshape:
|
||||
b, _, _, dim_head = q.shape
|
||||
tensor_layout = "HND"
|
||||
if kwargs.get("enable_gqa", False):
|
||||
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
|
||||
else:
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
q, k, v = map(
|
||||
lambda t: t.view(b, -1, heads, dim_head),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
tensor_layout = "NHD"
|
||||
|
||||
if mask is not None:
|
||||
@ -565,8 +588,12 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
sage_kwargs = {"is_causal": False, "tensor_layout": tensor_layout, "sm_scale": kwargs.get("scale", None), "smooth_k": False}
|
||||
if mask is not None:
|
||||
sage_kwargs["attn_mask"] = mask
|
||||
|
||||
try:
|
||||
out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
|
||||
out = sageattn(q, k, v, **sage_kwargs)
|
||||
except Exception as e:
|
||||
logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e))
|
||||
exception_fallback = True
|
||||
@ -616,7 +643,6 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
skip_output_reshape=skip_output_reshape,
|
||||
**kwargs
|
||||
)
|
||||
q_s, k_s, v_s = q, k, v
|
||||
N = q.shape[2]
|
||||
dim_head = D
|
||||
else:
|
||||
@ -642,11 +668,15 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
**kwargs
|
||||
)
|
||||
|
||||
if not skip_reshape:
|
||||
q_s, k_s, v_s = map(
|
||||
lambda t: t.view(B, -1, heads, dim_head).permute(0, 2, 1, 3).contiguous(),
|
||||
(q, k, v),
|
||||
)
|
||||
if skip_reshape:
|
||||
q_s = q
|
||||
if kwargs.get("enable_gqa", False):
|
||||
k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
|
||||
else:
|
||||
k_s, v_s = k, v
|
||||
else:
|
||||
q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False))
|
||||
q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s))
|
||||
B, H, L, D = q_s.shape
|
||||
|
||||
try:
|
||||
@ -662,7 +692,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
q, k, v, heads,
|
||||
mask=mask,
|
||||
attn_precision=attn_precision,
|
||||
skip_reshape=False,
|
||||
skip_reshape=skip_reshape,
|
||||
skip_output_reshape=skip_output_reshape,
|
||||
**kwargs
|
||||
)
|
||||
@ -681,19 +711,20 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
try:
|
||||
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
|
||||
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
|
||||
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
|
||||
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
|
||||
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
|
||||
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal, softmax_scale=softmax_scale_arg)
|
||||
|
||||
|
||||
@flash_attn_wrapper.register_fake
|
||||
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False):
|
||||
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0):
|
||||
# Output shape is the same as q
|
||||
return q.new_empty(q.shape)
|
||||
except AttributeError as error:
|
||||
FLASH_ATTN_ERROR = error
|
||||
|
||||
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
|
||||
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
|
||||
assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}"
|
||||
|
||||
@wrap_attn
|
||||
@ -703,10 +734,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
else:
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
q, k, v = map(
|
||||
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
|
||||
(q, k, v),
|
||||
)
|
||||
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
|
||||
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
|
||||
|
||||
if mask is not None:
|
||||
# add a batch dimension if there isn't already one
|
||||
@ -725,10 +754,16 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
v.transpose(1, 2),
|
||||
dropout_p=0.0,
|
||||
causal=False,
|
||||
softmax_scale=kwargs.get("scale", -1.0),
|
||||
).transpose(1, 2)
|
||||
except Exception as e:
|
||||
logging.warning(f"Flash Attention failed, using default SDPA: {e}")
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
sdpa_extra = {}
|
||||
if kwargs.get("enable_gqa", False):
|
||||
sdpa_extra["enable_gqa"] = True
|
||||
if "scale" in kwargs:
|
||||
sdpa_extra["scale"] = kwargs["scale"]
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra)
|
||||
if not skip_output_reshape:
|
||||
out = (
|
||||
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
||||
@ -1209,5 +1244,3 @@ class SpatialVideoTransformer(SpatialTransformer):
|
||||
x = self.proj_out(x)
|
||||
out = x + x_in
|
||||
return out
|
||||
|
||||
|
||||
|
||||
@ -141,11 +141,8 @@ class Attention(nn.Module):
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
|
||||
if self.kv_heads < self.heads:
|
||||
key = key.repeat_interleave(self.heads // self.kv_heads, dim=1)
|
||||
value = value.repeat_interleave(self.heads // self.kv_heads, dim=1)
|
||||
|
||||
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
|
||||
gqa_kwargs = {"enable_gqa": True} if self.kv_heads < self.heads else {}
|
||||
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
|
||||
hidden_states = self.to_out[0](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
@ -190,6 +190,14 @@ def is_wsl():
|
||||
return True
|
||||
return False
|
||||
|
||||
_WSL_SOFT_EMPTY_CACHE_SKIP_LOGGED = False
|
||||
|
||||
def wsl_skip_nonforced_soft_empty_cache():
|
||||
return is_wsl() and os.getenv("COMFYUI_WSL_SOFT_EMPTY_CACHE", "0") != "1"
|
||||
|
||||
def wsl_skip_model_load_synchronize():
|
||||
return is_wsl() and os.getenv("COMFYUI_WSL_MODEL_LOAD_SYNCHRONIZE", "0") != "1"
|
||||
|
||||
def get_torch_device():
|
||||
global directml_enabled
|
||||
global cpu_state
|
||||
@ -941,7 +949,14 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
|
||||
if vram_set_state == VRAMState.NO_VRAM:
|
||||
lowvram_model_memory = 0.1
|
||||
|
||||
model_name = model.model.__class__.__name__ if hasattr(model, "model") else model.__class__.__name__
|
||||
logging.info(
|
||||
f"Loading model {model_name} start: device={torch_dev} "
|
||||
f"vram_state={vram_set_state.name} lowvram_model_memory={lowvram_model_memory} "
|
||||
f"force_full_load={force_full_load} force_patch_weights={force_patch_weights}"
|
||||
)
|
||||
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
logging.info(f"Loading model {model_name} complete")
|
||||
current_loaded_models.insert(0, loaded_model)
|
||||
return
|
||||
|
||||
@ -1961,6 +1976,12 @@ def soft_empty_cache(force=False):
|
||||
elif is_mlu():
|
||||
torch.mlu.empty_cache()
|
||||
elif torch.cuda.is_available():
|
||||
if wsl_skip_nonforced_soft_empty_cache() and not force:
|
||||
global _WSL_SOFT_EMPTY_CACHE_SKIP_LOGGED
|
||||
if not _WSL_SOFT_EMPTY_CACHE_SKIP_LOGGED:
|
||||
logging.info("Skipping non-forced CUDA soft_empty_cache on WSL; set COMFYUI_WSL_SOFT_EMPTY_CACHE=1 to re-enable.")
|
||||
_WSL_SOFT_EMPTY_CACHE_SKIP_LOGGED = True
|
||||
return
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
@ -42,6 +42,8 @@ from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
|
||||
|
||||
import comfy_aimdo.model_vbar
|
||||
|
||||
_WSL_MODEL_LOAD_SYNC_SKIP_LOGGED = False
|
||||
|
||||
def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
|
||||
to = model_options["transformer_options"].copy()
|
||||
|
||||
@ -1006,6 +1008,11 @@ class ModelPatcher:
|
||||
mem_counter += move_weight_functions(m, device_to)
|
||||
|
||||
load_completely.sort(reverse=True)
|
||||
skip_wsl_load_sync = comfy.model_management.is_device_cuda(device_to) and comfy.model_management.wsl_skip_model_load_synchronize()
|
||||
global _WSL_MODEL_LOAD_SYNC_SKIP_LOGGED
|
||||
if skip_wsl_load_sync and len(load_completely) > 0 and not _WSL_MODEL_LOAD_SYNC_SKIP_LOGGED:
|
||||
logging.info("Skipping per-module CUDA synchronize during model load on WSL; set COMFYUI_WSL_MODEL_LOAD_SYNCHRONIZE=1 to re-enable.")
|
||||
_WSL_MODEL_LOAD_SYNC_SKIP_LOGGED = True
|
||||
for x in load_completely:
|
||||
n = x[1]
|
||||
m = x[2]
|
||||
@ -1020,7 +1027,7 @@ class ModelPatcher:
|
||||
key = key_param_name_to_key(n, param)
|
||||
self.unpin_weight(key)
|
||||
self.patch_weight_to_device(key, device_to=device_to)
|
||||
if comfy.model_management.is_device_cuda(device_to):
|
||||
if comfy.model_management.is_device_cuda(device_to) and not skip_wsl_load_sync:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
|
||||
|
||||
@ -174,6 +174,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
|
||||
elif xfer_dest2 is not None:
|
||||
xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False)
|
||||
return
|
||||
else:
|
||||
return
|
||||
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2)
|
||||
|
||||
def handle_pin(m, pin, source, dest, subset="weights", size=None):
|
||||
|
||||
12
comfy/sd.py
12
comfy/sd.py
@ -1198,7 +1198,17 @@ class VAE:
|
||||
else:
|
||||
pixel_samples = pixel_samples.unsqueeze(2)
|
||||
|
||||
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) # TODO: calculate mem required for tile
|
||||
if dims == 2:
|
||||
default_tile_x = 512 if tile_x is None else tile_x
|
||||
default_tile_y = 512 if tile_y is None else tile_y
|
||||
tile_shapes = [
|
||||
(1, pixel_samples.shape[1], min(pixel_samples.shape[2], max(1, default_tile_y)), min(pixel_samples.shape[3], max(1, default_tile_x))),
|
||||
(1, pixel_samples.shape[1], min(pixel_samples.shape[2], max(1, default_tile_y // 2)), min(pixel_samples.shape[3], max(1, default_tile_x * 2))),
|
||||
(1, pixel_samples.shape[1], min(pixel_samples.shape[2], max(1, default_tile_y * 2)), min(pixel_samples.shape[3], max(1, default_tile_x // 2))),
|
||||
]
|
||||
memory_used = max(self.memory_used_encode(shape, self.vae_dtype) for shape in tile_shapes)
|
||||
else:
|
||||
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
|
||||
model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
|
||||
|
||||
args = {}
|
||||
|
||||
@ -12,7 +12,7 @@ import torch.nn.functional as F
|
||||
|
||||
import comfy.ops
|
||||
from comfy import sd1_clip
|
||||
from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
from comfy.text_encoders.llama import RMSNorm, apply_rope
|
||||
|
||||
|
||||
@ -110,10 +110,6 @@ def _attention_with_sinks(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sin
|
||||
putting the sink logit in the mask at that column.
|
||||
"""
|
||||
|
||||
if num_kv_groups > 1 and not TORCH_HAS_GQA:
|
||||
k = k.repeat_interleave(num_kv_groups, dim=1)
|
||||
v = v.repeat_interleave(num_kv_groups, dim=1)
|
||||
|
||||
B, _, S_q, D = q.shape
|
||||
H_kv = k.shape[1]
|
||||
S_kv = k.shape[-2]
|
||||
|
||||
@ -550,10 +550,8 @@ class Attention(nn.Module):
|
||||
xv = xv[:, :, -sliding_window:]
|
||||
attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
|
||||
|
||||
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
|
||||
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
|
||||
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
|
||||
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
|
||||
return self.o_proj(output), present_key_value
|
||||
|
||||
class MLP(nn.Module):
|
||||
|
||||
@ -366,12 +366,8 @@ class GatedAttention(nn.Module):
|
||||
xv = torch.cat((past_value[:, :, :index], xv), dim=2)
|
||||
present_key_value = (xk, xv, index + num_tokens)
|
||||
|
||||
# Expand KV heads for GQA
|
||||
if self.num_heads != self.num_kv_heads:
|
||||
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
|
||||
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
|
||||
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
|
||||
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
|
||||
output = output * gate.sigmoid()
|
||||
|
||||
return self.o_proj(output), present_key_value
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import bisect
|
||||
import gc
|
||||
import itertools
|
||||
import psutil
|
||||
import time
|
||||
@ -17,6 +18,7 @@ NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {}
|
||||
|
||||
|
||||
def include_unique_id_in_input(class_type: str) -> bool:
|
||||
"""Return whether a node class includes UNIQUE_ID among its hidden inputs."""
|
||||
if class_type in NODE_CLASS_CONTAINS_UNIQUE_ID:
|
||||
return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type]
|
||||
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
||||
@ -24,52 +26,412 @@ def include_unique_id_in_input(class_type: str) -> bool:
|
||||
return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type]
|
||||
|
||||
class CacheKeySet(ABC):
|
||||
"""Base helper for building and storing cache keys for prompt nodes."""
|
||||
def __init__(self, dynprompt, node_ids, is_changed_cache):
|
||||
"""Initialize cache-key storage for a dynamic prompt execution pass."""
|
||||
self.keys = {}
|
||||
self.subcache_keys = {}
|
||||
|
||||
@abstractmethod
|
||||
async def add_keys(self, node_ids):
|
||||
"""Populate cache keys for the provided node ids."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def all_node_ids(self):
|
||||
"""Return the set of node ids currently tracked by this key set."""
|
||||
return set(self.keys.keys())
|
||||
|
||||
def get_used_keys(self):
|
||||
"""Return the computed cache keys currently in use."""
|
||||
return self.keys.values()
|
||||
|
||||
def get_used_subcache_keys(self):
|
||||
"""Return the computed subcache keys currently in use."""
|
||||
return self.subcache_keys.values()
|
||||
|
||||
def get_data_key(self, node_id):
|
||||
"""Return the cache key for a node, if present."""
|
||||
return self.keys.get(node_id, None)
|
||||
|
||||
def get_subcache_key(self, node_id):
|
||||
"""Return the subcache key for a node, if present."""
|
||||
return self.subcache_keys.get(node_id, None)
|
||||
|
||||
class Unhashable:
|
||||
def __init__(self):
|
||||
self.value = float("NaN")
|
||||
"""Hashable identity sentinel for values that cannot be represented safely in cache keys."""
|
||||
pass
|
||||
|
||||
def to_hashable(obj):
|
||||
# So that we don't infinitely recurse since frozenset and tuples
|
||||
# are Sequences.
|
||||
if isinstance(obj, (int, float, str, bool, bytes, type(None))):
|
||||
return obj
|
||||
elif isinstance(obj, Mapping):
|
||||
return frozenset([(to_hashable(k), to_hashable(v)) for k, v in sorted(obj.items())])
|
||||
elif isinstance(obj, Sequence):
|
||||
return frozenset(zip(itertools.count(), [to_hashable(i) for i in obj]))
|
||||
else:
|
||||
# TODO - Support other objects like tensors?
|
||||
|
||||
_PRIMITIVE_SIGNATURE_TYPES = (int, float, str, bool, bytes, type(None))
|
||||
_CONTAINER_SIGNATURE_TYPES = (dict, list, tuple, set, frozenset)
|
||||
_MAX_SIGNATURE_DEPTH = 32
|
||||
_MAX_SIGNATURE_CONTAINER_VISITS = 10_000
|
||||
_FAILED_SIGNATURE = object()
|
||||
|
||||
|
||||
def _shallow_is_changed_signature(value):
|
||||
"""Reduce execution-time `is_changed` values through a fail-closed builtin canonicalizer."""
|
||||
value_type = type(value)
|
||||
if value_type in _PRIMITIVE_SIGNATURE_TYPES:
|
||||
return value
|
||||
|
||||
if value_type not in _CONTAINER_SIGNATURE_TYPES:
|
||||
return Unhashable()
|
||||
|
||||
canonical = _signature_to_hashable(value, max_nodes=64)
|
||||
if type(canonical) is Unhashable:
|
||||
return canonical
|
||||
if value_type is list or value_type is tuple:
|
||||
container_tag = "is_changed_list" if value_type is list else "is_changed_tuple"
|
||||
return (container_tag, canonical[1])
|
||||
|
||||
return canonical
|
||||
|
||||
|
||||
def _primitive_signature_sort_key(obj):
|
||||
"""Return a deterministic ordering key for primitive signature values."""
|
||||
obj_type = type(obj)
|
||||
return ("primitive", obj_type.__module__, obj_type.__qualname__, repr(obj))
|
||||
|
||||
|
||||
def _sanitized_sort_key(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=None, memo=None):
|
||||
"""Return a deterministic ordering key for sanitized built-in container content."""
|
||||
if depth >= max_depth:
|
||||
return ("MAX_DEPTH",)
|
||||
|
||||
if active is None:
|
||||
active = set()
|
||||
if memo is None:
|
||||
memo = {}
|
||||
|
||||
obj_type = type(obj)
|
||||
if obj_type is Unhashable:
|
||||
return ("UNHASHABLE",)
|
||||
elif obj_type in _PRIMITIVE_SIGNATURE_TYPES:
|
||||
return (obj_type.__module__, obj_type.__qualname__, repr(obj))
|
||||
elif obj_type not in _CONTAINER_SIGNATURE_TYPES:
|
||||
return (obj_type.__module__, obj_type.__qualname__, "OPAQUE")
|
||||
|
||||
obj_id = id(obj)
|
||||
if obj_id in memo:
|
||||
return memo[obj_id]
|
||||
if obj_id in active:
|
||||
return ("CYCLE",)
|
||||
|
||||
active.add(obj_id)
|
||||
try:
|
||||
if obj_type is dict:
|
||||
items = [
|
||||
(
|
||||
_sanitized_sort_key(k, depth + 1, max_depth, active, memo),
|
||||
_sanitized_sort_key(v, depth + 1, max_depth, active, memo),
|
||||
)
|
||||
for k, v in obj.items()
|
||||
]
|
||||
items.sort()
|
||||
result = ("dict", tuple(items))
|
||||
elif obj_type is list:
|
||||
result = ("list", tuple(_sanitized_sort_key(i, depth + 1, max_depth, active, memo) for i in obj))
|
||||
elif obj_type is tuple:
|
||||
result = ("tuple", tuple(_sanitized_sort_key(i, depth + 1, max_depth, active, memo) for i in obj))
|
||||
elif obj_type is set:
|
||||
result = ("set", tuple(sorted(_sanitized_sort_key(i, depth + 1, max_depth, active, memo) for i in obj)))
|
||||
else:
|
||||
result = ("frozenset", tuple(sorted(_sanitized_sort_key(i, depth + 1, max_depth, active, memo) for i in obj)))
|
||||
finally:
|
||||
active.discard(obj_id)
|
||||
|
||||
memo[obj_id] = result
|
||||
return result
|
||||
|
||||
|
||||
def _signature_to_hashable_impl(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=None, memo=None, budget=None):
|
||||
"""Canonicalize signature inputs directly into their final hashable form."""
|
||||
if depth >= max_depth:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
if active is None:
|
||||
active = set()
|
||||
if memo is None:
|
||||
memo = {}
|
||||
if budget is None:
|
||||
budget = {"remaining": _MAX_SIGNATURE_CONTAINER_VISITS}
|
||||
|
||||
obj_type = type(obj)
|
||||
if obj_type in _PRIMITIVE_SIGNATURE_TYPES:
|
||||
return obj, _primitive_signature_sort_key(obj)
|
||||
if obj_type is Unhashable or obj_type not in _CONTAINER_SIGNATURE_TYPES:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
obj_id = id(obj)
|
||||
if obj_id in memo:
|
||||
return memo[obj_id]
|
||||
if obj_id in active:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
budget["remaining"] -= 1
|
||||
if budget["remaining"] < 0:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
active.add(obj_id)
|
||||
try:
|
||||
if obj_type is dict:
|
||||
try:
|
||||
items = list(obj.items())
|
||||
except RuntimeError:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
ordered_items = []
|
||||
for key, value in items:
|
||||
if type(key) not in _PRIMITIVE_SIGNATURE_TYPES:
|
||||
return _FAILED_SIGNATURE
|
||||
key_result = (key, _primitive_signature_sort_key(key))
|
||||
value_result = _signature_to_hashable_impl(value, depth + 1, max_depth, active, memo, budget)
|
||||
if value_result is _FAILED_SIGNATURE:
|
||||
return _FAILED_SIGNATURE
|
||||
key_value, key_sort = key_result
|
||||
value_value, value_sort = value_result
|
||||
ordered_items.append((key_sort, value_sort, key_value, value_value))
|
||||
|
||||
ordered_items.sort(key=lambda item: (item[0], item[1]))
|
||||
for index in range(1, len(ordered_items)):
|
||||
previous_key_sort = ordered_items[index - 1][0]
|
||||
current_key_sort = ordered_items[index][0]
|
||||
if previous_key_sort == current_key_sort:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
value = ("dict", tuple((key_value, value_value) for _, _, key_value, value_value in ordered_items))
|
||||
sort_key = ("dict", tuple((key_sort, value_sort) for key_sort, value_sort, _, _ in ordered_items))
|
||||
elif obj_type is list or obj_type is tuple:
|
||||
try:
|
||||
items = list(obj)
|
||||
except RuntimeError:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
child_results = []
|
||||
for item in items:
|
||||
child_result = _signature_to_hashable_impl(item, depth + 1, max_depth, active, memo, budget)
|
||||
if child_result is _FAILED_SIGNATURE:
|
||||
return _FAILED_SIGNATURE
|
||||
child_results.append(child_result)
|
||||
|
||||
container_tag = "list" if obj_type is list else "tuple"
|
||||
value = (container_tag, tuple(child for child, _ in child_results))
|
||||
sort_key = (container_tag, tuple(child_sort for _, child_sort in child_results))
|
||||
else:
|
||||
try:
|
||||
items = list(obj)
|
||||
except RuntimeError:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
ordered_items = []
|
||||
for item in items:
|
||||
child_result = _signature_to_hashable_impl(item, depth + 1, max_depth, active, memo, budget)
|
||||
if child_result is _FAILED_SIGNATURE:
|
||||
return _FAILED_SIGNATURE
|
||||
child_value, child_sort = child_result
|
||||
ordered_items.append((child_sort, child_value))
|
||||
|
||||
ordered_items.sort(key=lambda item: item[0])
|
||||
for index in range(1, len(ordered_items)):
|
||||
previous_sort_key, previous_value = ordered_items[index - 1]
|
||||
current_sort_key, current_value = ordered_items[index]
|
||||
if previous_sort_key == current_sort_key and previous_value != current_value:
|
||||
return _FAILED_SIGNATURE
|
||||
|
||||
container_tag = "set" if obj_type is set else "frozenset"
|
||||
value = (container_tag, tuple(child_value for _, child_value in ordered_items))
|
||||
sort_key = (container_tag, tuple(child_sort for child_sort, _ in ordered_items))
|
||||
finally:
|
||||
active.discard(obj_id)
|
||||
|
||||
memo[obj_id] = (value, sort_key)
|
||||
return memo[obj_id]
|
||||
|
||||
|
||||
def _signature_to_hashable(obj, max_nodes=_MAX_SIGNATURE_CONTAINER_VISITS):
|
||||
"""Build the final cache-signature representation in one fail-closed pass."""
|
||||
try:
|
||||
result = _signature_to_hashable_impl(obj, budget={"remaining": max_nodes})
|
||||
except RuntimeError:
|
||||
return Unhashable()
|
||||
if result is _FAILED_SIGNATURE:
|
||||
return Unhashable()
|
||||
return result[0]
|
||||
|
||||
|
||||
def to_hashable(obj, max_nodes=_MAX_SIGNATURE_CONTAINER_VISITS):
|
||||
"""Convert sanitized prompt inputs into a stable hashable representation.
|
||||
|
||||
The input is expected to already be sanitized to plain built-in containers,
|
||||
but this function still fails safe for anything unexpected. Traversal is
|
||||
iterative and memoized so shared built-in substructures do not trigger
|
||||
exponential re-walks during cache-key construction.
|
||||
"""
|
||||
obj_type = type(obj)
|
||||
if obj_type in _PRIMITIVE_SIGNATURE_TYPES or obj_type is Unhashable:
|
||||
return obj
|
||||
if obj_type not in _CONTAINER_SIGNATURE_TYPES:
|
||||
return Unhashable()
|
||||
|
||||
memo = {}
|
||||
active = set()
|
||||
snapshots = {}
|
||||
sort_memo = {}
|
||||
processed = 0
|
||||
# Keep traversal state separate from container snapshots/results.
|
||||
work_stack = [(obj, False)]
|
||||
|
||||
def resolve_value(value):
|
||||
"""Resolve a child value from the completed memo table when available."""
|
||||
value_type = type(value)
|
||||
if value_type in _PRIMITIVE_SIGNATURE_TYPES or value_type is Unhashable:
|
||||
return value
|
||||
return memo.get(id(value), Unhashable())
|
||||
|
||||
def is_failed(value):
|
||||
"""Return whether a resolved child value represents failed canonicalization."""
|
||||
return type(value) is Unhashable
|
||||
|
||||
def resolve_unordered_values(current_items, container_tag):
|
||||
"""Resolve a set-like container or fail closed if ordering is ambiguous."""
|
||||
try:
|
||||
ordered_items = [
|
||||
(_sanitized_sort_key(item, memo=sort_memo), resolve_value(item))
|
||||
for item in current_items
|
||||
]
|
||||
if any(is_failed(value) for _, value in ordered_items):
|
||||
return Unhashable()
|
||||
ordered_items.sort(key=lambda item: item[0])
|
||||
except RuntimeError:
|
||||
return Unhashable()
|
||||
|
||||
for index in range(1, len(ordered_items)):
|
||||
previous_key, previous_value = ordered_items[index - 1]
|
||||
current_key, current_value = ordered_items[index]
|
||||
if previous_key == current_key and previous_value != current_value:
|
||||
return Unhashable()
|
||||
|
||||
return (container_tag, tuple(value for _, value in ordered_items))
|
||||
|
||||
while work_stack:
|
||||
entry = work_stack.pop()
|
||||
if len(entry) == 3:
|
||||
_, current_id, current_type = entry
|
||||
current = None
|
||||
expanded = True
|
||||
else:
|
||||
current, expanded = entry
|
||||
current_type = type(current)
|
||||
current_id = id(current)
|
||||
|
||||
if not expanded and (current_type in _PRIMITIVE_SIGNATURE_TYPES or current_type is Unhashable):
|
||||
continue
|
||||
if not expanded and current_type not in _CONTAINER_SIGNATURE_TYPES:
|
||||
memo[current_id] = Unhashable()
|
||||
continue
|
||||
|
||||
if current_id in memo:
|
||||
continue
|
||||
|
||||
if expanded:
|
||||
active.discard(current_id)
|
||||
try:
|
||||
items = snapshots.pop(current_id, None)
|
||||
if items is None:
|
||||
memo[current_id] = Unhashable()
|
||||
continue
|
||||
|
||||
if current_type is dict:
|
||||
ordered_items = [
|
||||
(_sanitized_sort_key(k, memo=sort_memo), k, resolve_value(v))
|
||||
for k, v in items
|
||||
]
|
||||
if any(type(key) not in _PRIMITIVE_SIGNATURE_TYPES or is_failed(value) for _, key, value in ordered_items):
|
||||
memo[current_id] = Unhashable()
|
||||
continue
|
||||
ordered_items.sort(key=lambda item: item[0])
|
||||
for index in range(1, len(ordered_items)):
|
||||
if ordered_items[index - 1][0] == ordered_items[index][0]:
|
||||
memo[current_id] = Unhashable()
|
||||
break
|
||||
else:
|
||||
memo[current_id] = (
|
||||
"dict",
|
||||
tuple((key, value) for _, key, value in ordered_items),
|
||||
)
|
||||
elif current_type is list:
|
||||
resolved_items = tuple(resolve_value(item) for item in items)
|
||||
if any(is_failed(item) for item in resolved_items):
|
||||
memo[current_id] = Unhashable()
|
||||
else:
|
||||
memo[current_id] = ("list", resolved_items)
|
||||
elif current_type is tuple:
|
||||
resolved_items = tuple(resolve_value(item) for item in items)
|
||||
if any(is_failed(item) for item in resolved_items):
|
||||
memo[current_id] = Unhashable()
|
||||
else:
|
||||
memo[current_id] = ("tuple", resolved_items)
|
||||
elif current_type is set:
|
||||
memo[current_id] = resolve_unordered_values(items, "set")
|
||||
else:
|
||||
memo[current_id] = resolve_unordered_values(items, "frozenset")
|
||||
except RuntimeError:
|
||||
memo[current_id] = Unhashable()
|
||||
continue
|
||||
|
||||
if current_id in active:
|
||||
memo[current_id] = Unhashable()
|
||||
continue
|
||||
|
||||
processed += 1
|
||||
if processed > max_nodes:
|
||||
return Unhashable()
|
||||
|
||||
active.add(current_id)
|
||||
if current_type is dict:
|
||||
try:
|
||||
items = list(current.items())
|
||||
snapshots[current_id] = items
|
||||
except RuntimeError:
|
||||
memo[current_id] = Unhashable()
|
||||
active.discard(current_id)
|
||||
continue
|
||||
for key, value in items:
|
||||
if type(key) not in _PRIMITIVE_SIGNATURE_TYPES:
|
||||
snapshots.pop(current_id, None)
|
||||
memo[current_id] = Unhashable()
|
||||
active.discard(current_id)
|
||||
break
|
||||
else:
|
||||
work_stack.append(("EXPANDED", current_id, current_type))
|
||||
for _, value in reversed(items):
|
||||
work_stack.append((value, False))
|
||||
continue
|
||||
continue
|
||||
else:
|
||||
try:
|
||||
items = list(current)
|
||||
snapshots[current_id] = items
|
||||
except RuntimeError:
|
||||
memo[current_id] = Unhashable()
|
||||
active.discard(current_id)
|
||||
continue
|
||||
work_stack.append(("EXPANDED", current_id, current_type))
|
||||
for item in reversed(items):
|
||||
work_stack.append((item, False))
|
||||
|
||||
return memo.get(id(obj), Unhashable())
|
||||
|
||||
class CacheKeySetID(CacheKeySet):
|
||||
"""Cache-key strategy that keys nodes by node id and class type."""
|
||||
def __init__(self, dynprompt, node_ids, is_changed_cache):
|
||||
"""Initialize identity-based cache keys for the supplied dynamic prompt."""
|
||||
super().__init__(dynprompt, node_ids, is_changed_cache)
|
||||
self.dynprompt = dynprompt
|
||||
|
||||
async def add_keys(self, node_ids):
|
||||
"""Populate identity-based keys for nodes that exist in the dynamic prompt."""
|
||||
for node_id in node_ids:
|
||||
if node_id in self.keys:
|
||||
continue
|
||||
@ -80,15 +442,19 @@ class CacheKeySetID(CacheKeySet):
|
||||
self.subcache_keys[node_id] = (node_id, node["class_type"])
|
||||
|
||||
class CacheKeySetInputSignature(CacheKeySet):
|
||||
"""Cache-key strategy that hashes a node's immediate inputs plus ancestor references."""
|
||||
def __init__(self, dynprompt, node_ids, is_changed_cache):
|
||||
"""Initialize input-signature-based cache keys for the supplied dynamic prompt."""
|
||||
super().__init__(dynprompt, node_ids, is_changed_cache)
|
||||
self.dynprompt = dynprompt
|
||||
self.is_changed_cache = is_changed_cache
|
||||
|
||||
def include_node_id_in_input(self) -> bool:
|
||||
"""Return whether node ids should be included in computed input signatures."""
|
||||
return False
|
||||
|
||||
async def add_keys(self, node_ids):
|
||||
"""Populate input-signature-based keys for nodes in the dynamic prompt."""
|
||||
for node_id in node_ids:
|
||||
if node_id in self.keys:
|
||||
continue
|
||||
@ -99,21 +465,37 @@ class CacheKeySetInputSignature(CacheKeySet):
|
||||
self.subcache_keys[node_id] = (node_id, node["class_type"])
|
||||
|
||||
async def get_node_signature(self, dynprompt, node_id):
|
||||
"""Build the full cache signature for a node and its ordered ancestors."""
|
||||
signature = []
|
||||
ancestors, order_mapping = self.get_ordered_ancestry(dynprompt, node_id)
|
||||
signature.append(await self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
|
||||
immediate = await self.get_immediate_node_signature(dynprompt, node_id, order_mapping)
|
||||
if type(immediate) is Unhashable:
|
||||
return immediate
|
||||
signature.append(immediate)
|
||||
for ancestor_id in ancestors:
|
||||
signature.append(await self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
|
||||
return to_hashable(signature)
|
||||
immediate = await self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping)
|
||||
if type(immediate) is Unhashable:
|
||||
return immediate
|
||||
signature.append(immediate)
|
||||
return tuple(signature)
|
||||
|
||||
async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
|
||||
"""Build the immediate cache-signature fragment for a node.
|
||||
|
||||
Link inputs are reduced to ancestor references here. Non-link values
|
||||
are canonicalized or failed closed before being appended so the final
|
||||
node signature is assembled from already-hashable fragments.
|
||||
"""
|
||||
if not dynprompt.has_node(node_id):
|
||||
# This node doesn't exist -- we can't cache it.
|
||||
return [float("NaN")]
|
||||
return Unhashable()
|
||||
node = dynprompt.get_node(node_id)
|
||||
class_type = node["class_type"]
|
||||
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
||||
signature = [class_type, await self.is_changed_cache.get(node_id)]
|
||||
is_changed_signature = _shallow_is_changed_signature(await self.is_changed_cache.get(node_id))
|
||||
if type(is_changed_signature) is Unhashable:
|
||||
return is_changed_signature
|
||||
signature = [class_type, is_changed_signature]
|
||||
if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type):
|
||||
signature.append(node_id)
|
||||
inputs = node["inputs"]
|
||||
@ -123,18 +505,23 @@ class CacheKeySetInputSignature(CacheKeySet):
|
||||
ancestor_index = ancestor_order_mapping[ancestor_id]
|
||||
signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
|
||||
else:
|
||||
signature.append((key, inputs[key]))
|
||||
return signature
|
||||
value_signature = to_hashable(inputs[key])
|
||||
if type(value_signature) is Unhashable:
|
||||
return value_signature
|
||||
signature.append((key, value_signature))
|
||||
return tuple(signature)
|
||||
|
||||
# This function returns a list of all ancestors of the given node. The order of the list is
|
||||
# deterministic based on which specific inputs the ancestor is connected by.
|
||||
def get_ordered_ancestry(self, dynprompt, node_id):
|
||||
"""Return ancestors in deterministic traversal order and their index mapping."""
|
||||
ancestors = []
|
||||
order_mapping = {}
|
||||
self.get_ordered_ancestry_internal(dynprompt, node_id, ancestors, order_mapping)
|
||||
return ancestors, order_mapping
|
||||
|
||||
def get_ordered_ancestry_internal(self, dynprompt, node_id, ancestors, order_mapping):
|
||||
"""Recursively collect ancestors in input order without revisiting prior nodes."""
|
||||
if not dynprompt.has_node(node_id):
|
||||
return
|
||||
inputs = dynprompt.get_node(node_id)["inputs"]
|
||||
|
||||
@ -1,11 +1,17 @@
|
||||
def is_link(obj):
|
||||
if not isinstance(obj, list):
|
||||
"""Return whether obj is a plain prompt link of the form [node_id, output_index]."""
|
||||
# Prompt links produced by the frontend / GraphBuilder are plain Python
|
||||
# lists in the form [node_id, output_index]. Some custom-node paths can
|
||||
# inject foreign runtime objects into prompt inputs during on-prompt graph
|
||||
# rewriting or subgraph construction. Be strict here so cache signature
|
||||
# building never tries to treat list-like proxy objects as links.
|
||||
if type(obj) is not list:
|
||||
return False
|
||||
if len(obj) != 2:
|
||||
return False
|
||||
if not isinstance(obj[0], str):
|
||||
if type(obj[0]) is not str:
|
||||
return False
|
||||
if not isinstance(obj[1], int) and not isinstance(obj[1], float):
|
||||
if type(obj[1]) is not int:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
473
tests-unit/execution_test/caching_test.py
Normal file
473
tests-unit/execution_test/caching_test.py
Normal file
@ -0,0 +1,473 @@
|
||||
"""Unit tests for cache-signature canonicalization hardening."""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _DummyNode:
|
||||
"""Minimal node stub used to satisfy cache-signature class lookups."""
|
||||
|
||||
@staticmethod
|
||||
def INPUT_TYPES():
|
||||
"""Return a minimal empty input schema for unit tests."""
|
||||
return {"required": {}}
|
||||
|
||||
|
||||
class _FakeDynPrompt:
|
||||
"""Small DynamicPrompt stand-in with only the methods these tests need."""
|
||||
|
||||
def __init__(self, nodes_by_id):
|
||||
"""Store test nodes by id."""
|
||||
self._nodes_by_id = nodes_by_id
|
||||
|
||||
def has_node(self, node_id):
|
||||
"""Return whether the fake prompt contains the requested node."""
|
||||
return node_id in self._nodes_by_id
|
||||
|
||||
def get_node(self, node_id):
|
||||
"""Return the stored node payload for the requested id."""
|
||||
return self._nodes_by_id[node_id]
|
||||
|
||||
|
||||
class _FakeIsChangedCache:
|
||||
"""Async stub for `is_changed` lookups used by cache-key generation."""
|
||||
|
||||
def __init__(self, values):
|
||||
"""Store canned `is_changed` responses keyed by node id."""
|
||||
self._values = values
|
||||
|
||||
async def get(self, node_id):
|
||||
"""Return the canned `is_changed` value for a node."""
|
||||
return self._values[node_id]
|
||||
|
||||
|
||||
class _OpaqueValue:
|
||||
"""Hashable opaque object used to exercise fail-closed unordered hashing paths."""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def caching_module(monkeypatch):
|
||||
"""Import `comfy_execution.caching` with lightweight stub dependencies."""
|
||||
torch_module = types.ModuleType("torch")
|
||||
psutil_module = types.ModuleType("psutil")
|
||||
nodes_module = types.ModuleType("nodes")
|
||||
nodes_module.NODE_CLASS_MAPPINGS = {}
|
||||
graph_module = types.ModuleType("comfy_execution.graph")
|
||||
|
||||
class DynamicPrompt:
|
||||
"""Placeholder graph type so the caching module can import cleanly."""
|
||||
|
||||
pass
|
||||
|
||||
graph_module.DynamicPrompt = DynamicPrompt
|
||||
|
||||
monkeypatch.setitem(sys.modules, "torch", torch_module)
|
||||
monkeypatch.setitem(sys.modules, "psutil", psutil_module)
|
||||
monkeypatch.setitem(sys.modules, "nodes", nodes_module)
|
||||
monkeypatch.setitem(sys.modules, "comfy_execution.graph", graph_module)
|
||||
monkeypatch.delitem(sys.modules, "comfy_execution.caching", raising=False)
|
||||
|
||||
module = importlib.import_module("comfy_execution.caching")
|
||||
module = importlib.reload(module)
|
||||
return module, nodes_module
|
||||
|
||||
|
||||
def test_signature_to_hashable_handles_shared_builtin_substructures(caching_module):
|
||||
"""Shared built-in substructures should canonicalize without collapsing to Unhashable."""
|
||||
caching, _ = caching_module
|
||||
shared = [{"value": 1}, {"value": 2}]
|
||||
|
||||
signature = caching._signature_to_hashable([shared, shared])
|
||||
|
||||
assert signature[0] == "list"
|
||||
assert signature[1][0] == signature[1][1]
|
||||
assert signature[1][0][0] == "list"
|
||||
assert signature[1][0][1][0] == ("dict", (("value", 1),))
|
||||
assert signature[1][0][1][1] == ("dict", (("value", 2),))
|
||||
|
||||
|
||||
def test_signature_to_hashable_fails_closed_on_opaque_values(caching_module):
|
||||
"""Opaque values should collapse the full signature to Unhashable immediately."""
|
||||
caching, _ = caching_module
|
||||
|
||||
signature = caching._signature_to_hashable(["safe", object()])
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_signature_to_hashable_stops_descending_after_failure(caching_module, monkeypatch):
|
||||
"""Once canonicalization fails, later recursive descent should stop immediately."""
|
||||
caching, _ = caching_module
|
||||
original = caching._signature_to_hashable_impl
|
||||
marker = object()
|
||||
marker_seen = False
|
||||
|
||||
def tracking_canonicalize(obj, *args, **kwargs):
|
||||
"""Track whether recursion reaches the nested marker after failure."""
|
||||
nonlocal marker_seen
|
||||
if obj is marker:
|
||||
marker_seen = True
|
||||
return original(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_signature_to_hashable_impl", tracking_canonicalize)
|
||||
|
||||
signature = caching._signature_to_hashable([object(), [marker]])
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
assert marker_seen is False
|
||||
|
||||
|
||||
def test_signature_to_hashable_snapshots_list_before_recursing(caching_module, monkeypatch):
|
||||
"""List canonicalization should read a point-in-time snapshot before recursive descent."""
|
||||
caching, _ = caching_module
|
||||
original = caching._signature_to_hashable_impl
|
||||
marker = ("marker",)
|
||||
values = [marker, 2]
|
||||
|
||||
def mutating_canonicalize(obj, *args, **kwargs):
|
||||
"""Mutate the live list during recursion to verify snapshot-based traversal."""
|
||||
if obj is marker:
|
||||
values[1] = 3
|
||||
return original(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_signature_to_hashable_impl", mutating_canonicalize)
|
||||
|
||||
signature = caching._signature_to_hashable(values)
|
||||
|
||||
assert signature == ("list", (("tuple", ("marker",)), 2))
|
||||
assert values[1] == 3
|
||||
|
||||
|
||||
def test_signature_to_hashable_snapshots_dict_before_recursing(caching_module, monkeypatch):
|
||||
"""Dict canonicalization should read a point-in-time snapshot before recursive descent."""
|
||||
caching, _ = caching_module
|
||||
original = caching._signature_to_hashable_impl
|
||||
marker = ("marker",)
|
||||
values = {"first": marker, "second": 2}
|
||||
|
||||
def mutating_canonicalize(obj, *args, **kwargs):
|
||||
"""Mutate the live dict during recursion to verify snapshot-based traversal."""
|
||||
if obj is marker:
|
||||
values["second"] = 3
|
||||
return original(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_signature_to_hashable_impl", mutating_canonicalize)
|
||||
|
||||
signature = caching._signature_to_hashable(values)
|
||||
|
||||
assert signature == ("dict", (("first", ("tuple", ("marker",))), ("second", 2)))
|
||||
assert values["second"] == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"container_factory",
|
||||
[
|
||||
lambda marker: [marker],
|
||||
lambda marker: (marker,),
|
||||
lambda marker: {marker},
|
||||
lambda marker: frozenset({marker}),
|
||||
lambda marker: {"key": marker},
|
||||
],
|
||||
)
|
||||
def test_signature_to_hashable_fails_closed_on_runtimeerror(caching_module, monkeypatch, container_factory):
|
||||
"""Traversal RuntimeError should degrade canonicalization to Unhashable."""
|
||||
caching, _ = caching_module
|
||||
original = caching._signature_to_hashable_impl
|
||||
marker = object()
|
||||
|
||||
def raising_canonicalize(obj, *args, **kwargs):
|
||||
"""Raise a traversal RuntimeError for the marker value and delegate otherwise."""
|
||||
if obj is marker:
|
||||
raise RuntimeError("container changed during iteration")
|
||||
return original(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_signature_to_hashable_impl", raising_canonicalize)
|
||||
|
||||
signature = caching._signature_to_hashable(container_factory(marker))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_to_hashable_handles_shared_builtin_substructures(caching_module):
|
||||
"""The legacy helper should still hash sanitized built-ins stably when used directly."""
|
||||
caching, _ = caching_module
|
||||
shared = [{"value": 1}, {"value": 2}]
|
||||
|
||||
sanitized = [shared, shared]
|
||||
hashable = caching.to_hashable(sanitized)
|
||||
|
||||
assert hashable[0] == "list"
|
||||
assert hashable[1][0] == hashable[1][1]
|
||||
assert hashable[1][0][0] == "list"
|
||||
|
||||
|
||||
def test_to_hashable_uses_parent_snapshot_during_expanded_phase(caching_module, monkeypatch):
|
||||
"""Expanded-phase assembly should not reread a live parent container after snapshotting."""
|
||||
caching, _ = caching_module
|
||||
original_sort_key = caching._sanitized_sort_key
|
||||
outer = [{"marker"}, 2]
|
||||
|
||||
def mutating_sort_key(obj, *args, **kwargs):
|
||||
"""Mutate the live parent while a child container is being canonicalized."""
|
||||
if obj == "marker":
|
||||
outer[1] = 3
|
||||
return original_sort_key(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_sanitized_sort_key", mutating_sort_key)
|
||||
|
||||
hashable = caching.to_hashable(outer)
|
||||
|
||||
assert hashable == ("list", (("set", ("marker",)), 2))
|
||||
assert outer[1] == 3
|
||||
|
||||
|
||||
def test_to_hashable_fails_closed_for_ordered_container_with_opaque_child(caching_module):
|
||||
"""Ordered containers should fail closed when a child cannot be canonicalized."""
|
||||
caching, _ = caching_module
|
||||
|
||||
result = caching.to_hashable([object()])
|
||||
|
||||
assert isinstance(result, caching.Unhashable)
|
||||
|
||||
|
||||
def test_to_hashable_canonicalizes_dict_insertion_order(caching_module):
|
||||
"""Dicts with the same content should hash identically regardless of insertion order."""
|
||||
caching, _ = caching_module
|
||||
|
||||
first = {"b": 2, "a": 1}
|
||||
second = {"a": 1, "b": 2}
|
||||
|
||||
assert caching.to_hashable(first) == ("dict", (("a", 1), ("b", 2)))
|
||||
assert caching.to_hashable(first) == caching.to_hashable(second)
|
||||
|
||||
|
||||
def test_to_hashable_fails_closed_for_opaque_dict_key(caching_module):
|
||||
"""Opaque dict keys should fail closed instead of being traversed during hashing."""
|
||||
caching, _ = caching_module
|
||||
|
||||
hashable = caching.to_hashable({_OpaqueValue(): 1})
|
||||
|
||||
assert isinstance(hashable, caching.Unhashable)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"container_factory",
|
||||
[
|
||||
set,
|
||||
frozenset,
|
||||
],
|
||||
)
|
||||
def test_to_hashable_fails_closed_on_runtimeerror(caching_module, monkeypatch, container_factory):
|
||||
"""Traversal RuntimeError should degrade unordered hash conversion to Unhashable."""
|
||||
caching, _ = caching_module
|
||||
|
||||
def raising_sort_key(obj, *args, **kwargs):
|
||||
"""Raise a traversal RuntimeError while unordered values are canonicalized."""
|
||||
raise RuntimeError("container changed during iteration")
|
||||
|
||||
monkeypatch.setattr(caching, "_sanitized_sort_key", raising_sort_key)
|
||||
|
||||
hashable = caching.to_hashable(container_factory({"value"}))
|
||||
|
||||
assert isinstance(hashable, caching.Unhashable)
|
||||
|
||||
|
||||
def test_to_hashable_fails_closed_for_ambiguous_dict_ordering(caching_module, monkeypatch):
|
||||
"""Ambiguous dict key ordering should fail closed instead of using insertion order."""
|
||||
caching, _ = caching_module
|
||||
original_sort_key = caching._sanitized_sort_key
|
||||
ambiguous = {"a": 1, "b": 1}
|
||||
|
||||
def colliding_sort_key(obj, *args, **kwargs):
|
||||
"""Force two distinct primitive keys to share the same ordering key."""
|
||||
if obj == "a" or obj == "b":
|
||||
return ("COLLIDE",)
|
||||
return original_sort_key(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_sanitized_sort_key", colliding_sort_key)
|
||||
|
||||
hashable = caching.to_hashable(ambiguous)
|
||||
|
||||
assert isinstance(hashable, caching.Unhashable)
|
||||
|
||||
|
||||
def test_signature_to_hashable_fails_closed_for_ambiguous_dict_ordering(caching_module, monkeypatch):
|
||||
"""Ambiguous dict sort ties should fail closed instead of depending on input order."""
|
||||
caching, _ = caching_module
|
||||
original_sort_key = caching._primitive_signature_sort_key
|
||||
ambiguous = {"a": 1, "b": 1}
|
||||
|
||||
def colliding_sort_key(obj):
|
||||
"""Force two distinct primitive keys to share the same ordering key."""
|
||||
if obj == "a" or obj == "b":
|
||||
return ("COLLIDE",)
|
||||
return original_sort_key(obj)
|
||||
|
||||
monkeypatch.setattr(caching, "_primitive_signature_sort_key", colliding_sort_key)
|
||||
|
||||
sanitized = caching._signature_to_hashable(ambiguous)
|
||||
|
||||
assert isinstance(sanitized, caching.Unhashable)
|
||||
|
||||
|
||||
def test_signature_to_hashable_fails_closed_for_opaque_dict_key(caching_module):
|
||||
"""Opaque dict keys should fail closed instead of being recursively canonicalized."""
|
||||
caching, _ = caching_module
|
||||
|
||||
sanitized = caching._signature_to_hashable({_OpaqueValue(): 1})
|
||||
|
||||
assert isinstance(sanitized, caching.Unhashable)
|
||||
|
||||
|
||||
def test_signature_to_hashable_fails_closed_on_dict_key_sort_collisions_even_with_distinct_values(caching_module, monkeypatch):
|
||||
"""Different values must not mask dict key-sort collisions during canonicalization."""
|
||||
caching, _ = caching_module
|
||||
original_sort_key = caching._primitive_signature_sort_key
|
||||
|
||||
def colliding_sort_key(obj):
|
||||
"""Force two distinct primitive keys to share the same ordering key."""
|
||||
if obj == "a" or obj == "b":
|
||||
return ("COLLIDE",)
|
||||
return original_sort_key(obj)
|
||||
|
||||
monkeypatch.setattr(caching, "_primitive_signature_sort_key", colliding_sort_key)
|
||||
|
||||
sanitized = caching._signature_to_hashable({"a": 1, "b": 2})
|
||||
|
||||
assert isinstance(sanitized, caching.Unhashable)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"container_factory",
|
||||
[
|
||||
set,
|
||||
frozenset,
|
||||
],
|
||||
)
|
||||
def test_to_hashable_fails_closed_for_ambiguous_unordered_values(caching_module, monkeypatch, container_factory):
|
||||
"""Ambiguous unordered values should fail closed instead of depending on iteration order."""
|
||||
caching, _ = caching_module
|
||||
original_sort_key = caching._sanitized_sort_key
|
||||
container = container_factory({"a", "b"})
|
||||
|
||||
def colliding_sort_key(obj, *args, **kwargs):
|
||||
"""Force two distinct primitive values to share the same ordering key."""
|
||||
if obj == "a" or obj == "b":
|
||||
return ("COLLIDE",)
|
||||
return original_sort_key(obj, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(caching, "_sanitized_sort_key", colliding_sort_key)
|
||||
|
||||
hashable = caching.to_hashable(container)
|
||||
|
||||
assert isinstance(hashable, caching.Unhashable)
|
||||
|
||||
|
||||
def test_get_node_signature_returns_top_level_unhashable_for_tainted_signature(caching_module, monkeypatch):
|
||||
"""Tainted full signatures should fail closed before `to_hashable()` runs."""
|
||||
caching, nodes_module = caching_module
|
||||
monkeypatch.setitem(nodes_module.NODE_CLASS_MAPPINGS, "UnitTestNode", _DummyNode)
|
||||
monkeypatch.setattr(
|
||||
caching,
|
||||
"to_hashable",
|
||||
lambda *_args, **_kwargs: pytest.fail("to_hashable should not run for tainted signatures"),
|
||||
)
|
||||
|
||||
is_changed_value = []
|
||||
is_changed_value.append(is_changed_value)
|
||||
|
||||
dynprompt = _FakeDynPrompt(
|
||||
{
|
||||
"node": {
|
||||
"class_type": "UnitTestNode",
|
||||
"inputs": {"value": 5},
|
||||
}
|
||||
}
|
||||
)
|
||||
key_set = caching.CacheKeySetInputSignature(
|
||||
dynprompt,
|
||||
["node"],
|
||||
_FakeIsChangedCache({"node": is_changed_value}),
|
||||
)
|
||||
|
||||
signature = asyncio.run(key_set.get_node_signature(dynprompt, "node"))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_accepts_primitive_lists(caching_module):
|
||||
"""Primitive-only `is_changed` lists should stay hashable without deep descent."""
|
||||
caching, _ = caching_module
|
||||
|
||||
sanitized = caching._shallow_is_changed_signature([1, "two", None, True])
|
||||
|
||||
assert sanitized == ("is_changed_list", (1, "two", None, True))
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_accepts_structured_builtin_fingerprint_lists(caching_module):
|
||||
"""Structured built-in `is_changed` fingerprints should remain representable."""
|
||||
caching, _ = caching_module
|
||||
|
||||
sanitized = caching._shallow_is_changed_signature([("seed", 42), {"cfg": 8}])
|
||||
|
||||
assert sanitized == (
|
||||
"is_changed_list",
|
||||
(
|
||||
("tuple", ("seed", 42)),
|
||||
("dict", (("cfg", 8),)),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_fails_closed_for_opaque_payload(caching_module):
|
||||
"""Opaque `is_changed` payloads should still fail closed."""
|
||||
caching, _ = caching_module
|
||||
|
||||
sanitized = caching._shallow_is_changed_signature([_OpaqueValue()])
|
||||
|
||||
assert isinstance(sanitized, caching.Unhashable)
|
||||
|
||||
|
||||
def test_get_immediate_node_signature_fails_closed_for_unhashable_is_changed(caching_module, monkeypatch):
|
||||
"""Recursive `is_changed` payloads should fail the full fragment closed."""
|
||||
caching, nodes_module = caching_module
|
||||
monkeypatch.setitem(nodes_module.NODE_CLASS_MAPPINGS, "UnitTestNode", _DummyNode)
|
||||
|
||||
is_changed_value = []
|
||||
is_changed_value.append(is_changed_value)
|
||||
dynprompt = _FakeDynPrompt(
|
||||
{
|
||||
"node": {
|
||||
"class_type": "UnitTestNode",
|
||||
"inputs": {"value": 5},
|
||||
}
|
||||
}
|
||||
)
|
||||
key_set = caching.CacheKeySetInputSignature(
|
||||
dynprompt,
|
||||
["node"],
|
||||
_FakeIsChangedCache({"node": is_changed_value}),
|
||||
)
|
||||
|
||||
signature = asyncio.run(key_set.get_immediate_node_signature(dynprompt, "node", {}))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_get_immediate_node_signature_fails_closed_for_missing_node(caching_module):
|
||||
"""Missing nodes should return the fail-closed sentinel instead of a NaN tuple."""
|
||||
caching, _ = caching_module
|
||||
dynprompt = _FakeDynPrompt({})
|
||||
key_set = caching.CacheKeySetInputSignature(
|
||||
dynprompt,
|
||||
[],
|
||||
_FakeIsChangedCache({}),
|
||||
)
|
||||
|
||||
signature = asyncio.run(key_set.get_immediate_node_signature(dynprompt, "missing", {}))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
242
tests/execution/test_caching.py
Normal file
242
tests/execution/test_caching.py
Normal file
@ -0,0 +1,242 @@
|
||||
import asyncio
|
||||
|
||||
from comfy_execution import caching
|
||||
|
||||
|
||||
class _StubDynPrompt:
|
||||
def __init__(self, nodes):
|
||||
self._nodes = nodes
|
||||
|
||||
def has_node(self, node_id):
|
||||
return node_id in self._nodes
|
||||
|
||||
def get_node(self, node_id):
|
||||
return self._nodes[node_id]
|
||||
|
||||
|
||||
class _StubIsChangedCache:
|
||||
async def get(self, node_id):
|
||||
return None
|
||||
|
||||
|
||||
class _StubNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {}}
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_keeps_primitive_only_list_shallow():
|
||||
assert caching._shallow_is_changed_signature([1, "two", None, True]) == (
|
||||
"is_changed_list",
|
||||
(1, "two", None, True),
|
||||
)
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_keeps_primitive_only_tuple_shallow():
|
||||
assert caching._shallow_is_changed_signature((1, "two", None, True)) == (
|
||||
"is_changed_tuple",
|
||||
(1, "two", None, True),
|
||||
)
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_keeps_structured_builtin_fingerprint_list():
|
||||
assert caching._shallow_is_changed_signature([("seed", 42), {"cfg": 8}]) == (
|
||||
"is_changed_list",
|
||||
(
|
||||
("tuple", ("seed", 42)),
|
||||
("dict", (("cfg", 8),)),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_shallow_is_changed_signature_does_not_use_to_hashable(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
caching,
|
||||
"to_hashable",
|
||||
lambda *_args, **_kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("is_changed signature must not deep-canonicalize")
|
||||
),
|
||||
)
|
||||
|
||||
signature = caching._shallow_is_changed_signature([("seed", 42), {"cfg": 8}])
|
||||
|
||||
assert signature == (
|
||||
"is_changed_list",
|
||||
(
|
||||
("tuple", ("seed", 42)),
|
||||
("dict", (("cfg", 8),)),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_get_immediate_node_signature_canonicalizes_non_link_inputs(monkeypatch):
|
||||
live_value = [1, {"nested": [2, 3]}]
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": live_value},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_immediate_node_signature(dynprompt, "1", {}))
|
||||
|
||||
assert signature == (
|
||||
"TestCacheNode",
|
||||
None,
|
||||
("value", ("list", (1, ("dict", (("nested", ("list", (2, 3))),))))),
|
||||
)
|
||||
|
||||
|
||||
def test_to_hashable_walks_dicts_without_rebinding_traversal_stack():
|
||||
live_value = {
|
||||
"outer": {"nested": [2, 3]},
|
||||
"items": [{"leaf": 4}],
|
||||
}
|
||||
|
||||
assert caching.to_hashable(live_value) == (
|
||||
"dict",
|
||||
(
|
||||
("items", ("list", (("dict", (("leaf", 4),)),))),
|
||||
("outer", ("dict", (("nested", ("list", (2, 3))),))),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_get_immediate_node_signature_fails_closed_for_opaque_non_link_input(monkeypatch):
|
||||
class OpaqueRuntimeValue:
|
||||
pass
|
||||
|
||||
live_value = OpaqueRuntimeValue()
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": live_value},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_immediate_node_signature(dynprompt, "1", {}))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_get_node_signature_propagates_unhashable_immediate_fragment(monkeypatch):
|
||||
class OpaqueRuntimeValue:
|
||||
pass
|
||||
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": OpaqueRuntimeValue()},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_node_signature(dynprompt, "1"))
|
||||
|
||||
assert isinstance(signature, caching.Unhashable)
|
||||
|
||||
|
||||
def test_get_node_signature_never_visits_raw_non_link_input(monkeypatch):
|
||||
live_value = [1, 2, 3]
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": live_value},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
monkeypatch.setattr(
|
||||
caching,
|
||||
"_signature_to_hashable",
|
||||
lambda *_args, **_kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("outer signature canonicalizer should not run")
|
||||
),
|
||||
)
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_node_signature(dynprompt, "1"))
|
||||
|
||||
assert isinstance(signature, tuple)
|
||||
|
||||
|
||||
def test_get_node_signature_keeps_deep_canonicalized_input_fragment(monkeypatch):
|
||||
live_value = 1
|
||||
for _ in range(8):
|
||||
live_value = [live_value]
|
||||
expected = caching.to_hashable(live_value)
|
||||
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": live_value},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_node_signature(dynprompt, "1"))
|
||||
|
||||
assert isinstance(signature, tuple)
|
||||
assert signature[0][2][0] == "value"
|
||||
assert signature[0][2][1] == expected
|
||||
|
||||
|
||||
def test_get_node_signature_keeps_large_precanonicalized_fragment(monkeypatch):
|
||||
live_value = object()
|
||||
canonical_fragment = ("tuple", tuple(("list", (index, index + 1)) for index in range(256)))
|
||||
dynprompt = _StubDynPrompt(
|
||||
{
|
||||
"1": {
|
||||
"class_type": "TestCacheNode",
|
||||
"inputs": {"value": live_value},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
|
||||
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
|
||||
monkeypatch.setattr(
|
||||
caching,
|
||||
"to_hashable",
|
||||
lambda value, max_nodes=caching._MAX_SIGNATURE_CONTAINER_VISITS: (
|
||||
canonical_fragment if value is live_value else caching.Unhashable()
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
caching,
|
||||
"_signature_to_hashable",
|
||||
lambda *_args, **_kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("outer signature canonicalizer should not run")
|
||||
),
|
||||
)
|
||||
|
||||
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
|
||||
signature = asyncio.run(keyset.get_node_signature(dynprompt, "1"))
|
||||
|
||||
assert isinstance(signature, tuple)
|
||||
assert signature[0][2] == ("value", canonical_fragment)
|
||||
Loading…
Reference in New Issue
Block a user