Merge 9211284444 into 091b70edda

Signed-off-by: bigcat88 <bigcat88@icloud.com>

AGENTS.md

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

README.md

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

comfy/cli_args.py

@@ -225,6 +225,7 @@ parser.add_argument(
 )
 
 parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.")
+parser.add_argument("--models-directory", type=is_valid_directory, default=None, help="Set the ComfyUI models directory. Overrides the models folder in --base-directory.")
 
 parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.")
 

comfy/ldm/ace/ace_step15.py

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

comfy/ldm/audio/dit.py

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

comfy/ldm/boogu/model.py

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

comfy/ldm/modules/attention.py

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

comfy/ldm/omnigen/omnigen2.py

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

comfy/ops.py

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

comfy/text_encoders/gpt_oss.py

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

comfy/text_encoders/llama.py

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

comfy/text_encoders/qwen35.py

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

comfy_api/latest/_io.py

@@ -1651,6 +1651,15 @@ class Schema:
     Use this for nodes with interactive/operable UI regions that produce intermediate outputs
     (e.g., Image Crop, Painter) rather than final outputs (e.g., Save Image).
     """
+    lazy_outputs: bool=False
+    """When True, cache will invalidate when output connections change, and expected_outputs will be available.
+
+    Use this for nodes that can skip computing outputs that aren't connected downstream.
+    Check `comfy_execution.utils.is_output_needed(i)` inside execute() - False means output i is definitely unused
+    and safe to skip. Only nodes with this flag receive expected_outputs; all others see None.
+
+    Limitation: consumers must exist before this node runs - a subgraph expansion that
+    hand-builds a link to a pre-existing node's already-skipped output reads a stale value."""
 
     def validate(self):
         '''Validate the schema:
@@ -2108,6 +2117,14 @@ class _ComfyNodeBaseInternal(_ComfyNodeInternal):
             cls.GET_SCHEMA()
         return cls._ACCEPT_ALL_INPUTS
 
+    _LAZY_OUTPUTS = None
+    @final
+    @classproperty
+    def LAZY_OUTPUTS(cls):  # noqa
+        if cls._LAZY_OUTPUTS is None:
+            cls.GET_SCHEMA()
+        return cls._LAZY_OUTPUTS
+
     @final
     @classmethod
     def INPUT_TYPES(cls) -> dict[str, dict]:
@@ -2152,6 +2169,8 @@ class _ComfyNodeBaseInternal(_ComfyNodeInternal):
             cls._NOT_IDEMPOTENT = schema.not_idempotent
         if cls._ACCEPT_ALL_INPUTS is None:
             cls._ACCEPT_ALL_INPUTS = schema.accept_all_inputs
+        if cls._LAZY_OUTPUTS is None:
+            cls._LAZY_OUTPUTS = schema.lazy_outputs
 
         if cls._RETURN_TYPES is None:
             output = []

comfy_execution/caching.py

@@ -6,7 +6,7 @@ import time
 import torch
 from typing import Sequence, Mapping, Dict
 from comfy.model_patcher import ModelPatcher
-from comfy_execution.graph import DynamicPrompt
+from comfy_execution.graph import DynamicPrompt, get_expected_outputs_for_node
 from abc import ABC, abstractmethod
 
 import nodes
@@ -116,6 +116,10 @@ class CacheKeySetInputSignature(CacheKeySet):
         signature = [class_type, await self.is_changed_cache.get(node_id)]
         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)
+        # Include expected_outputs in cache key for nodes that opt in via LAZY_OUTPUTS
+        if hasattr(class_def, 'LAZY_OUTPUTS') and class_def.LAZY_OUTPUTS:
+            expected = get_expected_outputs_for_node(dynprompt, node_id)
+            signature.append(("expected_outputs", tuple(sorted(expected))))
         inputs = node["inputs"]
         for key in sorted(inputs.keys()):
             if is_link(inputs[key]):

comfy_execution/graph.py

@@ -18,6 +18,18 @@ class NodeInputError(Exception):
 class NodeNotFoundError(Exception):
     pass
 
+
+def get_expected_outputs_for_node(dynprompt, node_id: str) -> frozenset:
+    """Get the set of output indices that are connected downstream.
+    Returns outputs that MIGHT be used.
+    Outputs NOT in this set are DEFINITELY not used and safe to skip
+    (see Schema.lazy_outputs for the one expansion-related limitation).
+
+    Includes input links and consumers registered via add_output_consumer.
+    """
+    return dynprompt.get_expected_outputs_map().get(node_id, frozenset())
+
+
 class DynamicPrompt:
     def __init__(self, original_prompt):
         # The original prompt provided by the user
@@ -26,6 +38,9 @@ class DynamicPrompt:
         self.ephemeral_prompt = {}
         self.ephemeral_parents = {}
         self.ephemeral_display = {}
+        # Output sockets consumed outside of input links (subgraph expansions)
+        self._external_output_consumers = {}
+        self._expected_outputs_map = None
 
     def get_node(self, node_id):
         if node_id in self.ephemeral_prompt:
@@ -41,6 +56,7 @@ class DynamicPrompt:
         self.ephemeral_prompt[node_id] = node_info
         self.ephemeral_parents[node_id] = parent_id
         self.ephemeral_display[node_id] = display_id
+        self._expected_outputs_map = None
 
     def get_real_node_id(self, node_id):
         while node_id in self.ephemeral_parents:
@@ -58,6 +74,29 @@ class DynamicPrompt:
     def all_node_ids(self):
         return set(self.original_prompt.keys()).union(set(self.ephemeral_prompt.keys()))
 
+    def add_output_consumer(self, node_id, socket):
+        """Record an output socket consumed outside of input links, e.g. a subgraph
+        expansion mapping its parent's output to this node's output."""
+        self._external_output_consumers.setdefault(node_id, set()).add(socket)
+        self._expected_outputs_map = None
+
+    def _build_expected_outputs_map(self):
+        result = {}
+        for node_id in self.all_node_ids():
+            node_data = self.get_node(node_id)
+            for value in node_data.get("inputs", {}).values():
+                if is_link(value):
+                    from_node_id, from_socket = value
+                    result.setdefault(from_node_id, set()).add(from_socket)
+        for node_id, sockets in self._external_output_consumers.items():
+            result.setdefault(node_id, set()).update(sockets)
+        self._expected_outputs_map = {k: frozenset(v) for k, v in result.items()}
+
+    def get_expected_outputs_map(self):
+        if self._expected_outputs_map is None:
+            self._build_expected_outputs_map()
+        return self._expected_outputs_map
+
     def get_original_prompt(self):
         return self.original_prompt
 

comfy_execution/utils.py

@@ -1,23 +1,45 @@
 import contextvars
-from typing import Optional, NamedTuple
+from typing import NamedTuple, FrozenSet
 
 class ExecutionContext(NamedTuple):
     """
     Context information about the currently executing node.
 
     Attributes:
+        prompt_id: The ID of the current prompt execution
         node_id: The ID of the currently executing node
         list_index: The index in a list being processed (for operations on batches/lists)
+        expected_outputs: Set of output indices that might be used downstream.
+                         Outputs NOT in this set are definitely unused (safe to skip).
+                         None means the information is not available.
     """
     prompt_id: str
     node_id: str
-    list_index: Optional[int]
+    list_index: int | None
+    expected_outputs: FrozenSet[int] | None = None
 
-current_executing_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar("current_executing_context", default=None)
+current_executing_context: contextvars.ContextVar[ExecutionContext | None] = contextvars.ContextVar("current_executing_context", default=None)
 
-def get_executing_context() -> Optional[ExecutionContext]:
+def get_executing_context() -> ExecutionContext | None:
     return current_executing_context.get(None)
 
+
+def is_output_needed(output_index: int) -> bool:
+    """Check if an output at the given index is connected downstream.
+
+    Returns True if the output might be used (should be computed).
+    Returns False if the output is definitely not connected (safe to skip).
+
+    Only meaningful for LAZY_OUTPUTS nodes; for all others expected_outputs is
+    None and this always returns True (skipping without the flag would not be
+    reflected in the cache key).
+    """
+    ctx = get_executing_context()
+    if ctx is None or ctx.expected_outputs is None:
+        return True
+    return output_index in ctx.expected_outputs
+
+
 class CurrentNodeContext:
     """
     Context manager for setting the current executing node context.
@@ -25,15 +47,22 @@ class CurrentNodeContext:
     Sets the current_executing_context on enter and resets it on exit.
 
     Example:
-        with CurrentNodeContext(node_id="123", list_index=0):
+        with CurrentNodeContext(prompt_id="abc", node_id="123", list_index=0):
             # Code that should run with the current node context set
             process_image()
     """
-    def __init__(self, prompt_id: str, node_id: str, list_index: Optional[int] = None):
+    def __init__(
+        self,
+        prompt_id: str,
+        node_id: str,
+        list_index: int | None = None,
+        expected_outputs: FrozenSet[int] | None = None,
+    ):
         self.context = ExecutionContext(
-            prompt_id= prompt_id,
-            node_id= node_id,
-            list_index= list_index
+            prompt_id=prompt_id,
+            node_id=node_id,
+            list_index=list_index,
+            expected_outputs=expected_outputs,
         )
         self.token = None
 

execution.py

@@ -35,6 +35,7 @@ from comfy_execution.graph import (
     ExecutionBlocker,
     ExecutionList,
     get_input_info,
+    get_expected_outputs_for_node,
 )
 from comfy_execution.graph_utils import GraphBuilder, is_link
 from comfy_execution.validation import validate_node_input
@@ -237,7 +238,18 @@ async def resolve_map_node_over_list_results(results):
                 raise exc
         return [x.result() if isinstance(x, asyncio.Task) else x for x in results]
 
-async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None, v3_data=None):
+async def _async_map_node_over_list(
+    prompt_id,
+    unique_id,
+    obj,
+    input_data_all,
+    func,
+    allow_interrupt=False,
+    execution_block_cb=None,
+    pre_execute_cb=None,
+    v3_data=None,
+    expected_outputs=None,
+):
     # check if node wants the lists
     input_is_list = getattr(obj, "INPUT_IS_LIST", False)
 
@@ -287,10 +299,12 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
             else:
                 f = getattr(obj, func)
             if inspect.iscoroutinefunction(f):
-                async def async_wrapper(f, prompt_id, unique_id, list_index, args):
-                    with CurrentNodeContext(prompt_id, unique_id, list_index):
+                async def async_wrapper(f, prompt_id, unique_id, list_index, args, expected_outputs):
+                    with CurrentNodeContext(prompt_id, unique_id, list_index, expected_outputs):
                         return await f(**args)
-                task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs))
+                task = asyncio.create_task(
+                    async_wrapper(f, prompt_id, unique_id, index, args=inputs, expected_outputs=expected_outputs)
+                )
                 # Give the task a chance to execute without yielding
                 await asyncio.sleep(0)
                 if task.done():
@@ -299,7 +313,7 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
                 else:
                     results.append(task)
             else:
-                with CurrentNodeContext(prompt_id, unique_id, index):
+                with CurrentNodeContext(prompt_id, unique_id, index, expected_outputs):
                     result = f(**inputs)
                 results.append(result)
         else:
@@ -337,8 +351,17 @@ def merge_result_data(results, obj):
             output.append([o[i] for o in results])
     return output
 
-async def get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=None, pre_execute_cb=None, v3_data=None):
-    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
+async def get_output_data(
+    prompt_id,
+    unique_id,
+    obj,
+    input_data_all,
+    execution_block_cb=None,
+    pre_execute_cb=None,
+    v3_data=None,
+    expected_outputs=None,
+):
+    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data, expected_outputs=expected_outputs)
     has_pending_task = any(isinstance(r, asyncio.Task) and not r.done() for r in return_values)
     if has_pending_task:
         return return_values, {}, False, has_pending_task
@@ -538,8 +561,12 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                 # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)?
                 GraphBuilder.set_default_prefix(unique_id, call_index, 0)
 
+            if getattr(class_def, "LAZY_OUTPUTS", False):
+                expected_outputs = get_expected_outputs_for_node(dynprompt, unique_id)
+            else:
+                expected_outputs = None
             try:
-                output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
+                output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data, expected_outputs=expected_outputs)
             finally:
                 if comfy.memory_management.aimdo_enabled:
                     if args.verbose == "DEBUG":
@@ -596,6 +623,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                         if is_link(node_outputs[i]):
                             from_node_id, from_socket = node_outputs[i][0], node_outputs[i][1]
                             new_output_links.append((from_node_id, from_socket))
+                            dynprompt.add_output_consumer(from_node_id, from_socket)
                     cached_outputs.append((True, node_outputs))
             new_node_ids = set(new_node_ids)
             for cache in caches.all:

folder_paths.py

@@ -17,7 +17,11 @@ if args.base_directory:
 else:
     base_path = os.path.dirname(os.path.realpath(__file__))
 
-models_dir = os.path.join(base_path, "models")
+if args.models_directory:
+    models_dir = os.path.abspath(args.models_directory)
+else:
+    models_dir = os.path.join(base_path, "models")
+
 folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
 folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
 

main.py

@@ -131,6 +131,10 @@ def apply_custom_paths():
     if args.base_directory:
         logging.info(f"Setting base directory to: {folder_paths.base_path}")
 
+    # --models-directory
+    if args.models_directory:
+        logging.info(f"Setting models directory to: {folder_paths.models_dir}")
+
     # --output-directory, --input-directory, --user-directory
     if args.output_directory:
         output_dir = os.path.abspath(args.output_directory)

tests-unit/comfy_test/folder_path_test.py

@@ -163,3 +163,20 @@ def test_base_path_change_clears_old(set_base_dir):
 
     for name in ["controlnet", "diffusion_models", "text_encoders"]:
         assert len(folder_paths.get_folder_paths(name)) == 2
+
+
+def test_models_directory_cli_and_getters(temp_dir):
+    try:
+        with patch.object(sys, 'argv', ["main.py", "--models-directory", temp_dir]):
+            reload(comfy.cli_args)
+            reload(folder_paths)
+
+        assert folder_paths.models_dir == os.path.abspath(temp_dir)
+
+        with pytest.raises(Exception):
+            comfy.cli_args.is_valid_directory(os.path.join(temp_dir, "non_existent_folder_path"))
+    finally:
+        with patch.object(sys, 'argv', ["main.py"]):
+            reload(comfy.cli_args)
+            reload(folder_paths)
+

tests-unit/execution_test/expected_outputs_test.py

@@ -0,0 +1,361 @@
+"""Unit tests for the expected_outputs feature.
+
+This feature allows nodes to know at runtime which outputs are connected downstream,
+enabling them to skip computing outputs that aren't needed.
+"""
+
+from comfy_api.latest import IO
+from comfy_execution.graph import DynamicPrompt, get_expected_outputs_for_node
+from comfy_execution.utils import (
+    CurrentNodeContext,
+    ExecutionContext,
+    get_executing_context,
+    is_output_needed,
+)
+
+
+class TestGetExpectedOutputsForNode:
+    """Tests for get_expected_outputs_for_node() function."""
+
+    def test_single_output_connected(self):
+        """Test node with single output connected to one downstream node."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "ConsumerNode", "inputs": {"image": ["1", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0})
+
+    def test_multiple_outputs_partial_connected(self):
+        """Test node with multiple outputs, only some connected."""
+        prompt = {
+            "1": {"class_type": "MultiOutputNode", "inputs": {}},
+            "2": {"class_type": "ConsumerA", "inputs": {"input": ["1", 0]}},
+            # Output 1 is not connected
+            "3": {"class_type": "ConsumerC", "inputs": {"input": ["1", 2]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0, 2})
+        assert 1 not in expected  # Output 1 is definitely unused
+
+    def test_no_outputs_connected(self):
+        """Test node with no outputs connected."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "OtherNode", "inputs": {}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset()
+
+    def test_same_output_connected_multiple_times(self):
+        """Test same output connected to multiple downstream nodes."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "ConsumerA", "inputs": {"input": ["1", 0]}},
+            "3": {"class_type": "ConsumerB", "inputs": {"input": ["1", 0]}},
+            "4": {"class_type": "ConsumerC", "inputs": {"input": ["1", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0})
+
+    def test_node_not_in_prompt(self):
+        """Test getting expected outputs for a node not in the prompt."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "999")
+        assert expected == frozenset()
+
+    def test_chained_nodes(self):
+        """Test expected outputs in a chain of nodes."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "MiddleNode", "inputs": {"input": ["1", 0]}},
+            "3": {"class_type": "EndNode", "inputs": {"input": ["2", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+
+        # Node 1's output 0 is connected to node 2
+        expected_1 = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected_1 == frozenset({0})
+
+        # Node 2's output 0 is connected to node 3
+        expected_2 = get_expected_outputs_for_node(dynprompt, "2")
+        assert expected_2 == frozenset({0})
+
+        # Node 3 has no downstream connections
+        expected_3 = get_expected_outputs_for_node(dynprompt, "3")
+        assert expected_3 == frozenset()
+
+    def test_complex_graph(self):
+        """Test expected outputs in a complex graph with multiple connections."""
+        prompt = {
+            "1": {"class_type": "MultiOutputNode", "inputs": {}},
+            "2": {"class_type": "ProcessorA", "inputs": {"image": ["1", 0], "mask": ["1", 1]}},
+            "3": {"class_type": "ProcessorB", "inputs": {"data": ["1", 2]}},
+            "4": {"class_type": "Combiner", "inputs": {"a": ["2", 0], "b": ["3", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+
+        # Node 1 has outputs 0, 1, 2 all connected
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0, 1, 2})
+
+    def test_constant_inputs_ignored(self):
+        """Test that constant (non-link) inputs don't affect expected outputs."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {
+                "class_type": "ConsumerNode",
+                "inputs": {
+                    "image": ["1", 0],
+                    "value": 42,
+                    "name": "test",
+                },
+            },
+        }
+        dynprompt = DynamicPrompt(prompt)
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0})
+
+    def test_ephemeral_node_invalidates_cache(self):
+        """Test that adding ephemeral nodes updates expected outputs."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "ConsumerNode", "inputs": {"image": ["1", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+
+        # Initially only output 0 is connected
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0})
+
+        # Add an ephemeral node that connects to output 1
+        dynprompt.add_ephemeral_node(
+            "eph_1",
+            {"class_type": "EphemeralNode", "inputs": {"data": ["1", 1]}},
+            parent_id="2",
+            display_id="2",
+        )
+
+        # Now both outputs 0 and 1 should be expected
+        expected = get_expected_outputs_for_node(dynprompt, "1")
+        assert expected == frozenset({0, 1})
+
+
+class TestExternalOutputConsumers:
+    """Tests for DynamicPrompt.add_output_consumer() — out-of-band consumers
+    (subgraph expansion output mappings) that have no input link in the prompt."""
+
+    def test_external_consumer_only(self):
+        """A socket consumed only externally must appear in expected outputs."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        assert get_expected_outputs_for_node(dynprompt, "1") == frozenset()
+
+        dynprompt.add_output_consumer("1", 1)
+        assert get_expected_outputs_for_node(dynprompt, "1") == frozenset({1})
+
+    def test_external_consumer_merges_with_links(self):
+        """External consumers merge with input-link consumers."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "ConsumerNode", "inputs": {"image": ["1", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        dynprompt.add_output_consumer("1", 2)
+        assert get_expected_outputs_for_node(dynprompt, "1") == frozenset({0, 2})
+
+    def test_external_consumer_invalidates_cached_map(self):
+        """Registering after the map was built must invalidate the cache."""
+        prompt = {
+            "1": {"class_type": "SourceNode", "inputs": {}},
+            "2": {"class_type": "ConsumerNode", "inputs": {"image": ["1", 0]}},
+        }
+        dynprompt = DynamicPrompt(prompt)
+        # Build (and cache) the map first
+        assert get_expected_outputs_for_node(dynprompt, "1") == frozenset({0})
+
+        dynprompt.add_output_consumer("1", 1)
+        assert get_expected_outputs_for_node(dynprompt, "1") == frozenset({0, 1})
+
+
+class TestExecutionContext:
+    """Tests for ExecutionContext with expected_outputs field."""
+
+    def test_context_with_expected_outputs(self):
+        """Test creating ExecutionContext with expected_outputs."""
+        ctx = ExecutionContext(
+            prompt_id="prompt-123", node_id="node-456", list_index=0, expected_outputs=frozenset({0, 2})
+        )
+        assert ctx.prompt_id == "prompt-123"
+        assert ctx.node_id == "node-456"
+        assert ctx.list_index == 0
+        assert ctx.expected_outputs == frozenset({0, 2})
+
+    def test_context_without_expected_outputs(self):
+        """Test ExecutionContext defaults to None for expected_outputs."""
+        ctx = ExecutionContext(prompt_id="prompt-123", node_id="node-456", list_index=0)
+        assert ctx.expected_outputs is None
+
+    def test_context_empty_expected_outputs(self):
+        """Test ExecutionContext with empty expected_outputs set."""
+        ctx = ExecutionContext(
+            prompt_id="prompt-123", node_id="node-456", list_index=None, expected_outputs=frozenset()
+        )
+        assert ctx.expected_outputs == frozenset()
+        assert len(ctx.expected_outputs) == 0
+
+
+class TestCurrentNodeContext:
+    """Tests for CurrentNodeContext context manager with expected_outputs."""
+
+    def test_context_manager_with_expected_outputs(self):
+        """Test CurrentNodeContext sets and resets context correctly."""
+        assert get_executing_context() is None
+
+        with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 1})):
+            ctx = get_executing_context()
+            assert ctx is not None
+            assert ctx.prompt_id == "prompt-1"
+            assert ctx.node_id == "node-1"
+            assert ctx.list_index == 0
+            assert ctx.expected_outputs == frozenset({0, 1})
+
+        assert get_executing_context() is None
+
+    def test_context_manager_without_expected_outputs(self):
+        """Test CurrentNodeContext works without expected_outputs (backwards compatible)."""
+        with CurrentNodeContext("prompt-1", "node-1"):
+            ctx = get_executing_context()
+            assert ctx is not None
+            assert ctx.expected_outputs is None
+
+    def test_nested_context_managers(self):
+        """Test nested CurrentNodeContext managers."""
+        with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0})):
+            ctx1 = get_executing_context()
+            assert ctx1.expected_outputs == frozenset({0})
+
+            with CurrentNodeContext("prompt-1", "node-2", 0, frozenset({1, 2})):
+                ctx2 = get_executing_context()
+                assert ctx2.expected_outputs == frozenset({1, 2})
+                assert ctx2.node_id == "node-2"
+
+            # After inner context exits, should be back to outer context
+            ctx1_again = get_executing_context()
+            assert ctx1_again.expected_outputs == frozenset({0})
+            assert ctx1_again.node_id == "node-1"
+
+    def test_output_check_pattern(self):
+        """Test the typical pattern nodes will use to check expected outputs."""
+        with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 2})):
+            ctx = get_executing_context()
+
+            # Typical usage pattern
+            if ctx and ctx.expected_outputs is not None:
+                should_compute_0 = 0 in ctx.expected_outputs
+                should_compute_1 = 1 in ctx.expected_outputs
+                should_compute_2 = 2 in ctx.expected_outputs
+            else:
+                # Fallback when info not available
+                should_compute_0 = should_compute_1 = should_compute_2 = True
+
+            assert should_compute_0 is True
+            assert should_compute_1 is False  # Not in expected_outputs
+            assert should_compute_2 is True
+
+
+class TestSchemaLazyOutputs:
+    """Tests for lazy_outputs in V3 Schema."""
+
+    def test_schema_lazy_outputs_default(self):
+        """Test that lazy_outputs defaults to False."""
+        schema = IO.Schema(
+            node_id="TestNode",
+            inputs=[],
+            outputs=[IO.Float.Output()],
+        )
+        assert schema.lazy_outputs is False
+
+    def test_schema_lazy_outputs_true(self):
+        """Test setting lazy_outputs to True."""
+        schema = IO.Schema(
+            node_id="TestNode",
+            lazy_outputs=True,
+            inputs=[],
+            outputs=[IO.Float.Output()],
+        )
+        assert schema.lazy_outputs is True
+
+    def test_v3_node_lazy_outputs_property(self):
+        """Test that LAZY_OUTPUTS property works on V3 nodes."""
+
+        class TestNodeWithLazyOutputs(IO.ComfyNode):
+            @classmethod
+            def define_schema(cls):
+                return IO.Schema(
+                    node_id="TestNodeWithLazyOutputs",
+                    lazy_outputs=True,
+                    inputs=[],
+                    outputs=[IO.Float.Output()],
+                )
+
+            @classmethod
+            def execute(cls):
+                return IO.NodeOutput(1.0)
+
+        assert TestNodeWithLazyOutputs.LAZY_OUTPUTS is True
+
+    def test_v3_node_lazy_outputs_default(self):
+        """Test that LAZY_OUTPUTS defaults to False on V3 nodes."""
+
+        class TestNodeWithoutLazyOutputs(IO.ComfyNode):
+            @classmethod
+            def define_schema(cls):
+                return IO.Schema(
+                    node_id="TestNodeWithoutLazyOutputs",
+                    inputs=[],
+                    outputs=[IO.Float.Output()],
+                )
+
+            @classmethod
+            def execute(cls):
+                return IO.NodeOutput(1.0)
+
+        assert TestNodeWithoutLazyOutputs.LAZY_OUTPUTS is False
+
+
+class TestIsOutputNeeded:
+    """Tests for is_output_needed() helper function."""
+
+    def test_output_needed_when_in_expected(self):
+        """Test that output is needed when in expected_outputs."""
+        with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 2})):
+            assert is_output_needed(0) is True
+            assert is_output_needed(2) is True
+
+    def test_output_not_needed_when_not_in_expected(self):
+        """Test that output is not needed when not in expected_outputs."""
+        with CurrentNodeContext("prompt-1", "node-1", 0, frozenset({0, 2})):
+            assert is_output_needed(1) is False
+            assert is_output_needed(3) is False
+
+    def test_output_needed_when_no_context(self):
+        """Test that output is needed when no context."""
+        assert get_executing_context() is None
+        assert is_output_needed(0) is True
+        assert is_output_needed(1) is True
+
+    def test_output_needed_when_expected_outputs_is_none(self):
+        """Test that output is needed when expected_outputs is None."""
+        with CurrentNodeContext("prompt-1", "node-1", 0, None):
+            assert is_output_needed(0) is True
+            assert is_output_needed(1) is True

tests/execution/test_execution.py

@@ -573,6 +573,144 @@ class TestExecution:
         else:
             assert result.did_run(test_node), "The execution should have been re-run"
 
+    def test_expected_outputs_all_connected(self, client: ComfyClient, builder: GraphBuilder):
+        """Test that expected_outputs contains all connected outputs."""
+        g = builder
+        # Create a node with 3 outputs, all connected
+        expected_outputs_node = g.node("TestExpectedOutputs", height=64, width=64)
+
+        # Connect all 3 outputs to preview nodes
+        output0 = g.node("PreviewImage", images=expected_outputs_node.out(0))
+        output1 = g.node("PreviewImage", images=expected_outputs_node.out(1))
+        output2 = g.node("PreviewImage", images=expected_outputs_node.out(2))
+
+        result = client.run(g)
+
+        # All outputs should be white (255) since all are connected
+        images0 = result.get_images(output0)
+        images1 = result.get_images(output1)
+        images2 = result.get_images(output2)
+
+        assert len(images0) == 1, "Should have 1 image for output0"
+        assert len(images1) == 1, "Should have 1 image for output1"
+        assert len(images2) == 1, "Should have 1 image for output2"
+
+        # White pixels = 255, meaning output was in expected_outputs
+        assert numpy.array(images0[0]).min() == 255, "Output 0 should be white (was expected)"
+        assert numpy.array(images1[0]).min() == 255, "Output 1 should be white (was expected)"
+        assert numpy.array(images2[0]).min() == 255, "Output 2 should be white (was expected)"
+
+    def test_expected_outputs_partial_connected(self, client: ComfyClient, builder: GraphBuilder):
+        """Test that expected_outputs only contains connected outputs."""
+        g = builder
+        # Create a node with 3 outputs, only some connected
+        expected_outputs_node = g.node("TestExpectedOutputs", height=64, width=64)
+
+        # Only connect outputs 0 and 2, leave output 1 disconnected
+        output0 = g.node("PreviewImage", images=expected_outputs_node.out(0))
+        # output1 is intentionally not connected
+        output2 = g.node("PreviewImage", images=expected_outputs_node.out(2))
+
+        result = client.run(g)
+
+        # Connected outputs should be white (255)
+        images0 = result.get_images(output0)
+        images2 = result.get_images(output2)
+
+        assert len(images0) == 1, "Should have 1 image for output0"
+        assert len(images2) == 1, "Should have 1 image for output2"
+
+        # White = expected, output 1 is not connected so we can't verify it directly but outputs 0 and 2 should be white
+        assert numpy.array(images0[0]).min() == 255, "Output 0 should be white (was expected)"
+        assert numpy.array(images2[0]).min() == 255, "Output 2 should be white (was expected)"
+
+    def test_expected_outputs_single_connected(self, client: ComfyClient, builder: GraphBuilder):
+        """Test that expected_outputs works with single connected output."""
+        g = builder
+        # Create a node with 3 outputs, only one connected
+        expected_outputs_node = g.node("TestExpectedOutputs", height=64, width=64)
+
+        # Only connect output 1
+        output1 = g.node("PreviewImage", images=expected_outputs_node.out(1))
+
+        result = client.run(g)
+
+        images1 = result.get_images(output1)
+        assert len(images1) == 1, "Should have 1 image for output1"
+
+        # Output 1 should be white (connected), others are not visible in this test
+        assert numpy.array(images1[0]).min() == 255, "Output 1 should be white (was expected)"
+
+    def test_expected_outputs_cache_invalidation(self, client: ComfyClient, builder: GraphBuilder, server):
+        """Test that cache invalidates when output connections change."""
+        g = builder
+        # Use unique dimensions to avoid cache collision with other expected_outputs tests
+        expected_outputs_node = g.node("TestExpectedOutputs", height=32, width=32)
+
+        # First run: only connect output 0
+        output0 = g.node("PreviewImage", images=expected_outputs_node.out(0))
+
+        result1 = client.run(g)
+        assert result1.did_run(expected_outputs_node), "First run should execute the node"
+
+        # Second run: same connections, should be cached
+        result2 = client.run(g)
+        if server["should_cache_results"]:
+            assert not result2.did_run(expected_outputs_node), "Second run should be cached"
+
+        # Third run: add connection to output 2
+        output2 = g.node("PreviewImage", images=expected_outputs_node.out(2))
+
+        result3 = client.run(g)
+        # Because LAZY_OUTPUTS=True, changing connections should invalidate cache
+        if server["should_cache_results"]:
+            assert result3.did_run(expected_outputs_node), "Adding output connection should invalidate cache"
+
+        # Verify both outputs are now white
+        images0 = result3.get_images(output0)
+        images2 = result3.get_images(output2)
+        assert numpy.array(images0[0]).min() == 255, "Output 0 should be white"
+        assert numpy.array(images2[0]).min() == 255, "Output 2 should be white"
+
+    def test_expected_outputs_expansion_output_mapping(self, client: ComfyClient, builder: GraphBuilder):
+        """A socket consumed only via an expansion's parent-output mapping must still
+        be in the inner LAZY_OUTPUTS node's expected_outputs (white, not black)."""
+        g = builder
+        expander = g.node("TestExpectedOutputsExpansion", height=80, width=80)
+        output = g.node("PreviewImage", images=expander.out(0))
+
+        result = client.run(g)
+
+        images = result.get_images(output)
+        assert len(images) == 1, "Should have 1 image"
+        assert numpy.array(images[0]).min() == 255, (
+            "Inner node skipped an output that is consumed via the expansion's "
+            "parent-output mapping (expected white, got black)"
+        )
+
+    def test_expected_outputs_requires_opt_in(self, client: ComfyClient, builder: GraphBuilder, server):
+        """Nodes without LAZY_OUTPUTS must see expected_outputs=None: their cache key
+        ignores topology, so a skipped output would be served stale after rewiring."""
+        g = builder
+        node = g.node("TestExpectedOutputsNotOptedIn", height=96, width=96)
+        output0 = g.node("PreviewImage", images=node.out(0))
+
+        # Only output 0 connected: correct gating -> node sees None, computes all
+        result1 = client.run(g)
+        assert numpy.array(result1.get_images(output0)[0]).min() == 255
+
+        # Connect output 1: key unchanged -> cache hit must still serve correct data
+        output1 = g.node("PreviewImage", images=node.out(1))
+        result2 = client.run(g)
+
+        if server["should_cache_results"]:
+            assert not result2.did_run(node), "Node should be a cache hit (key ignores topology)"
+        images1 = result2.get_images(output1)
+        assert len(images1) == 1, "Should have 1 image for output1"
+        assert numpy.array(images1[0]).min() == 255, (
+            "Non-opted-in node observed expected_outputs and skipped output 1; "
+            "the stale skipped value was then served from cache"
+        )
 
     def test_parallel_sleep_nodes(self, client: ComfyClient, builder: GraphBuilder, skip_timing_checks):
         # Warmup execution to ensure server is fully initialized

tests/execution/testing_nodes/testing-pack/specific_tests.py

@@ -6,6 +6,7 @@ from .tools import VariantSupport
 from comfy_execution.graph_utils import GraphBuilder
 from comfy.comfy_types.node_typing import ComfyNodeABC
 from comfy.comfy_types import IO
+from comfy_execution.utils import get_executing_context, is_output_needed
 
 class TestLazyMixImages:
     @classmethod
@@ -482,6 +483,106 @@ class TestOutputNodeWithSocketOutput:
         result = image * value
         return (result,)
 
+
+class TestExpectedOutputs:
+    """Test node for the expected_outputs feature.
+
+    This node has 3 IMAGE outputs that encode which outputs were expected:
+    - White image (255) if the output was in expected_outputs
+    - Black image (0) if the output was NOT in expected_outputs
+
+    This allows integration tests to verify which outputs were expected by checking pixel values.
+    """
+    LAZY_OUTPUTS = True  # Opt into cache invalidation on output connection changes
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 1024}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
+    RETURN_NAMES = ("output0", "output1", "output2")
+    FUNCTION = "execute"
+    CATEGORY = "_for_testing"
+
+    def execute(self, height, width):
+        # Return white image if expected, black if not
+        # This allows tests to verify which outputs were expected via pixel values
+        white = torch.ones(1, height, width, 3)
+        black = torch.zeros(1, height, width, 3)
+
+        return (
+            white if is_output_needed(0) else black,
+            white if is_output_needed(1) else black,
+            white if is_output_needed(2) else black,
+        )
+
+
+class TestExpectedOutputsExpansion:
+    """Expands into an inner LAZY_OUTPUTS node whose output 1 is consumed ONLY via
+    the parent-output mapping (no input link anywhere). If that mapping is not part
+    of the expected-outputs map, the inner node wrongly skips it -> black not white.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 1024}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+    CATEGORY = "_for_testing"
+
+    def execute(self, height, width):
+        g = GraphBuilder()
+        inner = g.node("TestExpectedOutputs", height=height, width=width)
+        return {"result": (inner.out(1),), "expand": g.finalize()}
+
+
+class TestExpectedOutputsNotOptedIn:
+    """Reads expected_outputs WITHOUT declaring LAZY_OUTPUTS; the executor must pass
+    None (such nodes have no cache-key protection against output rewiring). Outputs
+    are white when the node correctly sees None, otherwise they encode membership.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 1024}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("output0", "output1")
+    FUNCTION = "execute"
+    CATEGORY = "_for_testing"
+
+    def execute(self, height, width):
+        # Raw context access (not is_output_needed): must distinguish None from a set
+        ctx = get_executing_context()
+        expected = ctx.expected_outputs if ctx is not None else None
+
+        white = torch.ones(1, height, width, 3)
+        black = torch.zeros(1, height, width, 3)
+
+        if expected is None:
+            return (white, white.clone())
+        return (
+            white if 0 in expected else black,
+            white if 1 in expected else black,
+        )
+
+
 TEST_NODE_CLASS_MAPPINGS = {
     "TestLazyMixImages": TestLazyMixImages,
     "TestVariadicAverage": TestVariadicAverage,
@@ -498,6 +599,9 @@ TEST_NODE_CLASS_MAPPINGS = {
     "TestSleep": TestSleep,
     "TestParallelSleep": TestParallelSleep,
     "TestOutputNodeWithSocketOutput": TestOutputNodeWithSocketOutput,
+    "TestExpectedOutputs": TestExpectedOutputs,
+    "TestExpectedOutputsExpansion": TestExpectedOutputsExpansion,
+    "TestExpectedOutputsNotOptedIn": TestExpectedOutputsNotOptedIn,
 }
 
 TEST_NODE_DISPLAY_NAME_MAPPINGS = {
@@ -516,4 +620,7 @@ TEST_NODE_DISPLAY_NAME_MAPPINGS = {
     "TestSleep": "Test Sleep",
     "TestParallelSleep": "Test Parallel Sleep",
     "TestOutputNodeWithSocketOutput": "Test Output Node With Socket Output",
+    "TestExpectedOutputs": "Test Expected Outputs",
+    "TestExpectedOutputsExpansion": "Test Expected Outputs Expansion",
+    "TestExpectedOutputsNotOptedIn": "Test Expected Outputs Not Opted In",
 }