Compare commits

...

64 Commits

Author SHA1 Message Date
xmarre
f4619d118f
Merge a75bf6f1d8 into ffbecfffb9 2026-07-08 00:18:04 -04:00
comfyanonymous
ffbecfffb9
Fix crash when using UNetSelfAttentionMultiply (#14823)
Some checks are pending
Detect Unreviewed Merge / detect (push) Waiting to run
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
2026-07-07 21:17:31 -07:00
comfyanonymous
b481bc15af
Support gqa on all attention backends, drop support for pytorch 2.4 (#14772) 2026-07-07 22:57:52 -04:00
comfyanonymous
6880614319
Update AGENTS.md (#14819) 2026-07-07 18:36:13 -07:00
xmarre
a75bf6f1d8
Merge branch 'Comfy-Org:master' into master 2026-06-22 16:14:14 +02:00
xmarre
0bad1b06bd Restore ComfyUI README 2026-06-22 15:55:15 +02:00
xmarre
07c96b7238 Update README with performance-oriented WSL setup 2026-06-22 15:50:02 +02:00
xmarre
4bfa533371
Merge branch 'Comfy-Org:master' into master 2026-06-21 16:57:46 +02:00
xmarre
fa0eaccfcc
Merge pull request #4 from xmarre/codex/wsl-model-load-guard
Guard WSL CUDA sync during model load
2026-05-30 17:56:37 +02:00
xmarre
63e08a02fd Guard WSL CUDA sync during model load 2026-05-30 17:43:20 +02:00
xmarre
69cbd50aa6
Merge branch 'Comfy-Org:master' into master 2026-05-30 14:22:43 +02:00
xmarre
6d4f9e86ab
Merge branch 'Comfy-Org:master' into master 2026-04-18 09:20:41 +02:00
xmarre
1548aee40e
Merge pull request #3 from xmarre/codex/vae-encode-tiled-admission-fix
Fix 2D tiled VAE encode memory admission estimation
2026-04-16 12:59:16 +02:00
xmarre
9c210473fc Fix tiled VAE encode memory admission estimate 2026-04-16 12:49:49 +02:00
xmarre
c1e9164c63
Merge branch 'master' into master 2026-04-16 10:07:30 +02:00
xmarre
5e9a90186f
Merge branch 'master' into master 2026-04-14 20:11:29 +02:00
xmarre
ece906328a
Merge branch 'master' into master 2026-04-02 18:55:32 +02:00
xmarre
500ca8e02a
Merge branch 'Comfy-Org:master' into master 2026-03-25 17:45:49 +01:00
xmarre
3143b7981f
Merge branch 'Comfy-Org:master' into master 2026-03-23 02:13:08 +01:00
xmarre
c9b3f81e83
Merge branch 'master' into master 2026-03-18 14:06:06 +01:00
xmarre
5e74e9b3ed
Merge pull request #1 from xmarre/codex/fix-cache-signature-shallow-check
Enforce shallow is_changed signature handling
2026-03-18 13:29:34 +01:00
xmarre
c702cddf75 Fix shallow is_changed logic 2026-03-18 13:15:04 +01:00
xmarre
e13da8104c Fix shallow is_changed handling 2026-03-18 12:26:30 +01:00
xmarre
fdcc38b9ea Return Unhashable on missing node 2026-03-17 07:48:14 +01:00
xmarre
c1ce00287c Stop requeueing live containers 2026-03-16 19:21:24 +01:00
xmarre
6e3bd33665 Prevent dict key canonicalization 2026-03-16 17:06:09 +01:00
xmarre
ce05e377a8 Stop canonicalizing dict keys 2026-03-16 16:48:42 +01:00
xmarre
1a00f7743f Stop traversing dict keys 2026-03-16 16:10:01 +01:00
xmarre
a6472b1514 Fix to_hashable traversal stack handling 2026-03-16 15:34:15 +01:00
xmarre
6158cd5820 Prevent redundant signature rewalk 2026-03-16 13:31:02 +01:00
xmarre
bff714dda0 Fix non-link input cache signature 2026-03-16 10:13:04 +01:00
xmarre
fce22da313 Prevent signature traversal of raw 2026-03-16 09:29:00 +01:00
xmarre
9f9d37bd9a
Merge branch 'master' into master 2026-03-16 09:07:29 +01:00
xmarre
088778c35d Stop canonicalizing is_changed 2026-03-15 17:06:20 +01:00
xmarre
4c5f82971e Restrict is_changed canonicalization 2026-03-15 16:44:25 +01:00
xmarre
f1d91a4c8c Prevent canonicalizing is_changed 2026-03-15 16:14:23 +01:00
xmarre
dbed5a1b52 Replace sanitize and hash passes 2026-03-15 07:39:10 +01:00
xmarre
24fdbb9aca Replace sanitize hash two pass 2026-03-15 07:30:18 +01:00
xmarre
a6624a9afd Unify signature sanitize and hash 2026-03-15 07:09:24 +01:00
xmarre
0b512198e8 Adopt single-pass signature hashing 2026-03-15 05:41:39 +01:00
xmarre
9feb26928c Change signature cache to bail early 2026-03-15 04:31:32 +01:00
xmarre
fadd79ad48 Fix nondeterministic set signing 2026-03-15 03:29:59 +01:00
xmarre
77bc7bdd6b Merge branch 'master' of https://github.com/xmarre/ComfyUI 2026-03-15 02:56:09 +01:00
xmarre
117afbc1d7 Add docstrings and harden signature 2026-03-15 02:55:39 +01:00
xmarre
064eec2278
Merge branch 'master' into master 2026-03-15 02:32:56 +01:00
xmarre
aceaa5e579 fail closed on ambiguous container ordering in cache signatures 2026-03-15 02:32:25 +01:00
xmarre
763089f681
Merge branch 'master' into master 2026-03-15 01:48:10 +01:00
xmarre
1693dabc8f
Merge branch 'master' into master 2026-03-15 00:28:34 +01:00
xmarre
08063d2638
Merge branch 'Comfy-Org:master' into master 2026-03-14 23:38:46 +01:00
xmarre
e069617e54
Merge branch 'Comfy-Org:master' into master 2026-03-14 21:27:17 +01:00
xmarre
2bea0ee5d7 Simplify Unhashable sentinel implementation 2026-03-14 12:42:04 +01:00
xmarre
17863f603a Add comprehensive docstrings for cache key helpers 2026-03-14 12:26:27 +01:00
xmarre
31ba844624 Add cycle detection to signature input sanitization 2026-03-14 12:04:31 +01:00
xmarre
1451001f64 Add docstrings for cache signature hardening helpers 2026-03-14 10:57:45 +01:00
xmarre
1af99b2e81 Update caching hash recursion 2026-03-14 10:31:07 +01:00
xmarre
3568b82b76 Revert "Add missing docstrings"
This reverts commit 4b431ffc27.
2026-03-14 10:11:35 +01:00
xmarre
6728d4d439 Revert "Harden to_hashable against cycles"
This reverts commit 880b51ac4f.
2026-03-14 10:11:04 +01:00
xmarre
4b431ffc27 Add missing docstrings 2026-03-14 09:57:22 +01:00
xmarre
880b51ac4f Harden to_hashable against cycles 2026-03-14 09:46:27 +01:00
xmarre
4d9516b909 Fix caching sanitization logic 2026-03-14 07:06:39 +01:00
xmarre
39086890e2 Fix sanitize_signature_input 2026-03-14 06:56:49 +01:00
xmarre
2adde5a0e1 Keep container types in sanitizer 2026-03-14 06:36:06 +01:00
xmarre
0c1bfad0df
Merge branch 'Comfy-Org:master' into master 2026-03-14 06:13:25 +01:00
xmarre
7d76a4447e Sanitize execution cache inputs 2026-03-14 02:36:40 +01:00
18 changed files with 1289 additions and 128 deletions

View File

@ -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

View File

@ -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:

View File

@ -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

View File

@ -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)

View File

@ -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])

View File

@ -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

View File

@ -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

View File

@ -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()

View File

@ -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))

View File

@ -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):

View File

@ -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 = {}

View File

@ -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]

View File

@ -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):

View File

@ -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

View File

@ -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"]

View File

@ -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

View 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)

View 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)