Try to fix the model reloading issue some people have.

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.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.
+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.
 
 ### Instructions:
 

comfy/ldm/ace/ace_step15.py

@@ -217,7 +217,10 @@ class AceStepAttention(nn.Module):
                 cos, sin = position_embeddings
                 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
-        gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
+        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)
 
         attn_bias = None
         if self.sliding_window is not None and not self.is_cross_attention:
@@ -241,7 +244,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, **gqa_kwargs)
+        attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False)
         attn_output = self.o_proj(attn_output)
 
         return attn_output

comfy/ldm/audio/dit.py

@@ -425,16 +425,19 @@ class Attention(nn.Module):
         if n == 1 and causal:
             causal = False
 
-        gqa_kwargs = {"enable_gqa": True} if h != kv_h else {}
+        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))
 
         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, **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      = 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 = out - out_diff
         else:
-            out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
+            out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
 
         out = self.to_out(out)
 

comfy/ldm/boogu/model.py

@@ -74,8 +74,11 @@ class BooguDoubleStreamProcessor(nn.Module):
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
 
-        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)
+        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)
 
         # 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,6 +1,5 @@
 import math
 import sys
-import inspect
 
 import torch
 import torch.nn.functional as F
@@ -15,16 +14,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":
@@ -90,44 +89,6 @@ 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):
@@ -191,19 +152,28 @@ 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:
-        if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
-        q, k, v = map(
+         q, k, v = map(
             lambda t: t.reshape(b * heads, -1, dim_head),
             (q, k, v),
         )
     else:
-        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))
+        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),
+        )
 
     # force cast to fp32 to avoid overflowing
     if attn_precision == torch.float32:
@@ -261,16 +231,13 @@ 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, 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)
+        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)
 
 
     dtype = query.dtype
@@ -337,15 +304,19 @@ 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:
-        if kwargs.get("enable_gqa", False):
-            k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
-        q, k, v = map(
+         q, k, v = map(
             lambda t: t.reshape(b * heads, -1, dim_head),
             (q, k, v),
         )
     else:
-        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))
+        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),
+        )
 
     r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
 
@@ -467,7 +438,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, skip_output_reshape=skip_output_reshape, **kwargs)
+        return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
 
     if skip_reshape:
         # b h k d -> b k h d
@@ -475,12 +446,13 @@ 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 = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
+        q, k, v = map(
+            lambda t: t.reshape(b, -1, heads, dim_head),
+            (q, k, v),
+        )
 
     if mask is not None:
         # add a singleton batch dimension
@@ -502,7 +474,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, scale=kwargs.get("scale", None))
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
 
     if skip_output_reshape:
         out = out.permute(0, 2, 1, 3)
@@ -526,8 +498,10 @@ 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 = _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))
+        q, k, v = map(
+            lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
+            (q, k, v),
+        )
 
     if mask is not None:
         # add a batch dimension if there isn't already one
@@ -537,7 +511,9 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
         if mask.ndim == 3:
             mask = mask.unsqueeze(1)
 
-    sdpa_keys = ("scale", "enable_gqa")
+    # 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_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys}
 
     if SDP_BATCH_LIMIT >= b:
@@ -565,19 +541,20 @@ 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 or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK):
+    if kwargs.get("low_precision_attention", True) is False:
         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 = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
+        q, k, v = map(
+            lambda t: t.view(b, -1, heads, dim_head),
+            (q, k, v),
+        )
         tensor_layout = "NHD"
 
     if mask is not None:
@@ -588,12 +565,8 @@ 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, **sage_kwargs)
+        out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
     except Exception as e:
         logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e))
         exception_fallback = True
@@ -643,6 +616,7 @@ 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:
@@ -668,15 +642,11 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
                 **kwargs
             )
 
-    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))
+    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),
+        )
         B, H, L, D = q_s.shape
 
     try:
@@ -692,7 +662,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=skip_reshape,
+                skip_reshape=False,
                 skip_output_reshape=skip_output_reshape,
                 **kwargs
             )
@@ -711,20 +681,19 @@ 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, 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)
+                    dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
+        return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
 
 
     @flash_attn_wrapper.register_fake
-    def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0):
+    def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False):
         # 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, softmax_scale: float = -1.0) -> torch.Tensor:
+                    dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
         assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}"
 
 @wrap_attn
@@ -734,8 +703,10 @@ 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 = _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))
+        q, k, v = map(
+            lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
+            (q, k, v),
+        )
 
     if mask is not None:
         # add a batch dimension if there isn't already one
@@ -754,16 +725,10 @@ 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}")
-        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)
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
     if not skip_output_reshape:
         out = (
             out.transpose(1, 2).reshape(b, -1, heads * dim_head)
@@ -1244,3 +1209,5 @@ class SpatialVideoTransformer(SpatialTransformer):
             x = self.proj_out(x)
         out = x + x_in
         return out
+
+

comfy/ldm/omnigen/omnigen2.py

@@ -141,8 +141,11 @@ class Attention(nn.Module):
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
 
-        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)
+        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)
         hidden_states = self.to_out[0](hidden_states)
         return hidden_states
 

comfy/model_management.py

@@ -616,6 +616,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
 #Freeing registerables on pressure does imply a GPU sync, so go big on
 #the hysteresis so each expensive sync gives us back a good chunk.
 REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
+WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0
+WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2
 
 def module_size(module):
     module_mem = 0
@@ -642,6 +644,15 @@ def free_pins(size, evict_active=False):
             size -= freed
     return freed_total
 
+def should_free_pins_for_ram_pressure(shortfall):
+    if shortfall <= 0:
+        return False
+    if not WINDOWS:
+        return True
+    if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
+        return True
+    return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
+
 def ensure_pin_budget(size, evict_active=False):
     if args.high_ram:
         return True

comfy/ops.py

@@ -174,8 +174,6 @@ 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 optimized_attention_for_device
+from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device
 from comfy.text_encoders.llama import RMSNorm, apply_rope
 
 
@@ -110,6 +110,10 @@ 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,8 +550,10 @@ class Attention(nn.Module):
                 xv = xv[:, :, -sliding_window:]
                 attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
 
-        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)
+        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)
         return self.o_proj(output), present_key_value
 
 class MLP(nn.Module):

comfy/text_encoders/qwen35.py

@@ -366,8 +366,12 @@ class GatedAttention(nn.Module):
                     xv = torch.cat((past_value[:, :, :index], xv), dim=2)
                 present_key_value = (xk, xv, index + num_tokens)
 
-        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)
+        # 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)
         output = output * gate.sigmoid()
 
         return self.o_proj(output), present_key_value

comfy_api/latest/_input_impl/video_types.py

@@ -1,6 +1,5 @@
 from av.container import InputContainer
 from av.subtitles.stream import SubtitleStream
-from av.video.reformatter import ColorRange
 from fractions import Fraction
 from typing import Optional
 from .._input import AudioInput, VideoInput
@@ -10,7 +9,6 @@ import itertools
 import json
 import numpy as np
 import math
-import os
 import torch
 from .._util import VideoContainer, VideoCodec, VideoComponents
 import logging
@@ -60,57 +58,6 @@ def video_stream_bit_depth(stream) -> int:
     return max(component.bits for component in stream.format.components)
 
 
-def last_decodable_audio_stream(container: InputContainer):
-    """Streams FFmpeg has no decoder for have no codec context, and decoding their
-    packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
-    stream = next(
-        (s for s in reversed(container.streams.audio) if s.codec_context is not None),
-        None,
-    )
-    if stream is None and len(container.streams.audio):
-        logging.warning("No decodable audio stream found in video; ignoring audio.")
-    return stream
-
-
-def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
-    """Containers probed only up to a window (mpegts) leave audio codec parameters unset when
-    audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
-    Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
-    for i, packet in enumerate(container.demux(audio_stream)):
-        try:
-            frames = packet.decode()
-        except av.error.FFmpegError:
-            return 0, 0
-        if frames:
-            return frames[0].sample_rate, frames[0].layout.nb_channels
-        if i >= max_packets:
-            break
-    return 0, 0
-
-
-def write_output_metadata(container: InputContainer, output, metadata: dict | None):
-    """Copy the source container's metadata, then overlay the caller's tags."""
-    for key, value in container.metadata.items():
-        if metadata is None or key not in metadata:
-            output.metadata[key] = value
-    if metadata is not None:
-        for key, value in metadata.items():
-            output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
-
-
-def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
-    if format != VideoContainer.AUTO and format != VideoContainer.MP4:
-        raise ValueError("Only MP4 format is supported for now")
-    if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
-        raise ValueError("Only H264 codec is supported for now")
-    open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
-    if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
-        open_kwargs["format"] = format.value
-    elif isinstance(path, io.BytesIO):
-        open_kwargs["format"] = "mp4"  # no file extension to infer the format from
-    return open_kwargs
-
-
 class VideoFromFile(VideoInput):
     """
     Class representing video input from a file.
@@ -245,10 +192,13 @@ class VideoFromFile(VideoInput):
                     return estimated_frames
 
             # 3. Last resort: decode frames and count them (streaming)
-            start_time, duration = self.get_active_trim_window()
+            if self.__start_time < 0:
+                start_time = max(self._get_raw_duration() + self.__start_time, 0)
+            else:
+                start_time = self.__start_time
             frame_count = 1
             start_pts = int(start_time / video_stream.time_base)
-            end_pts = int((start_time + duration) / video_stream.time_base)
+            end_pts = int((start_time + self.__duration) / video_stream.time_base)
             container.seek(start_pts, stream=video_stream)
             frame_iterator = (
                 container.decode(video_stream)
@@ -303,14 +253,17 @@ class VideoFromFile(VideoInput):
 
     def get_components_internal(self, container: InputContainer) -> VideoComponents:
         video_stream = self._get_first_video_stream(container)
-        start_time, duration = self.get_active_trim_window()
+        if self.__start_time < 0:
+            start_time = max(self._get_raw_duration() + self.__start_time, 0)
+        else:
+            start_time = self.__start_time
 
         # Get video frames
         frames = []
         audio_frames = []
         alphas = None
         start_pts = int(start_time / video_stream.time_base)
-        end_pts = int((start_time + duration) / video_stream.time_base)
+        end_pts = int((start_time + self.__duration) / video_stream.time_base)
 
         if start_pts != 0:
             container.seek(start_pts, stream=video_stream)
@@ -328,11 +281,18 @@ class VideoFromFile(VideoInput):
         video_done = False
         audio_done = True
 
-        audio_stream = last_decodable_audio_stream(container)
+        # Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
+        # and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
+        audio_stream = next(
+            (s for s in reversed(container.streams.audio) if s.codec_context is not None),
+            None,
+        )
         if audio_stream is not None:
             streams += [audio_stream]
             resampler = av.audio.resampler.AudioResampler(format='fltp')
             audio_done = False
+        elif len(container.streams.audio):
+            logging.warning("No decodable audio stream found in video; ignoring audio.")
 
         for packet in container.demux(*streams):
             if video_done and audio_done:
@@ -345,7 +305,7 @@ class VideoFromFile(VideoInput):
                     for frame in packet.decode():
                         if frame.pts < start_pts:
                             continue
-                        if duration and frame.pts >= end_pts:
+                        if self.__duration and frame.pts >= end_pts:
                             video_done = True
                             break
 
@@ -412,7 +372,7 @@ class VideoFromFile(VideoInput):
                     map(resampler.resample, packet.decode())
                 )
                 for frame in aframes:
-                    if duration and frame.time > start_time + duration:
+                    if self.__duration and frame.time > start_time + self.__duration:
                         audio_done = True
                         break
 
@@ -434,8 +394,8 @@ class VideoFromFile(VideoInput):
 
         if len(audio_frames) > 0:
             audio_data = np.concatenate(audio_frames, axis=1)  # shape: (channels, total_samples)
-            if duration:
-                audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
+            if self.__duration:
+                audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
 
             audio_tensor = torch.from_numpy(audio_data).unsqueeze(0)  # shape: (1, channels, total_samples)
             audio = AudioInput({
@@ -481,14 +441,28 @@ class VideoFromFile(VideoInput):
             if not reuse_streams:
                 if bit_depth is None:
                     bit_depth = source_bit_depth
-                return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
+                components = self.get_components_internal(container)
+                video = VideoFromComponents(components)
+                return video.save_to(
+                    path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
+                )
 
             streams = container.streams
 
             open_kwargs = get_open_write_kwargs(path, container_format, format)
             with av.open(path, **open_kwargs) as output_container:
-                # Add metadata before writing any streams
-                write_output_metadata(container, output_container, metadata)
+                # Copy over the original metadata
+                for key, value in container.metadata.items():
+                    if metadata is None or key not in metadata:
+                        output_container.metadata[key] = value
+
+                # Add our new metadata
+                if metadata is not None:
+                    for key, value in metadata.items():
+                        if isinstance(value, str):
+                            output_container.metadata[key] = value
+                        else:
+                            output_container.metadata[key] = json.dumps(value)
 
                 # Add streams to the new container. Streams with no codec context cannot be used as an output template.
                 stream_map = {}
@@ -506,254 +480,6 @@ class VideoFromFile(VideoInput):
                         packet.stream = stream_map[packet.stream]
                         output_container.mux(packet)
 
-    def _save_transcoded(
-        self,
-        container: InputContainer,
-        path: str | io.BytesIO,
-        format: VideoContainer,
-        codec: VideoCodec,
-        metadata: dict | None,
-        bit_depth: int,
-    ):
-        """Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
-        open_kwargs = mp4_output_open_kwargs(path, format, codec)
-        video_stream = self._get_first_video_stream(container)
-        start_time, duration = self.get_active_trim_window()
-        start_pts = int(start_time / video_stream.time_base)
-        end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
-        if start_pts != 0:
-            container.seek(start_pts, stream=video_stream)
-
-        audio_stream = last_decodable_audio_stream(container)
-        pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
-        rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
-
-        resampler = None
-        sample_rate = 0
-        audio_time_base = None
-        duration_cap = None
-        if audio_stream is not None:
-            sample_rate = audio_stream.codec_context.sample_rate
-            channels = audio_stream.codec_context.channels
-            if not sample_rate:
-                sample_rate, channels = probe_audio_params(container, audio_stream)
-                container.seek(start_pts, stream=video_stream)
-                if sample_rate:
-                    audio_stream.codec_context.flush_buffers()
-                else:
-                    logging.warning("Audio stream parameters could not be determined; ignoring audio.")
-                    audio_stream = None
-        if audio_stream is not None:
-            audio_time_base = Fraction(1, sample_rate)
-            layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
-            resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
-            if duration:
-                duration_cap = math.ceil(duration * sample_rate)
-
-        streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
-        pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
-        video_done = False
-        audio_done = audio_stream is None
-        video_pts_offset = None
-        last_video_pts = None
-        last_video_end = None
-        # rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
-        video_frame_durations = {}
-        source_size = None
-        rotation_k = 0
-        rotation_filter = None
-        audio_started = False
-        samples_written = 0
-        pending_audio = []
-        # The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
-        # and never seeking back keeps webm/mkv leading audio intact.
-        output = None
-        out_video = None
-        out_audio = None
-
-        def audio_frame_from_ndarray(nd_planar):
-            frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
-            frame.sample_rate = sample_rate
-            return frame
-
-        def drain_audio(final=False):
-            # Audio may cover the pts span of the video written so far, capped by the requested duration
-            nonlocal samples_written, audio_done
-            if last_video_end is None:
-                cap = 0
-            else:
-                cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
-            if duration_cap is not None:
-                cap = min(cap, duration_cap)
-            while pending_audio and not audio_done:
-                frame = pending_audio[0]
-                if samples_written + frame.samples <= cap:
-                    frame.pts = samples_written
-                    frame.time_base = audio_time_base
-                    output.mux(out_audio.encode(frame))
-                    samples_written += frame.samples
-                    pending_audio.pop(0)
-                    continue
-                if final:
-                    keep = frame.to_ndarray()[..., :cap - samples_written]
-                    if keep.shape[-1] > 0:
-                        tail = audio_frame_from_ndarray(keep)
-                        tail.pts = samples_written
-                        tail.time_base = audio_time_base
-                        output.mux(out_audio.encode(tail))
-                        samples_written += keep.shape[-1]
-                    pending_audio.clear()
-                break
-            if duration_cap is not None and samples_written >= duration_cap:
-                audio_done = True
-            return cap
-
-        try:
-            for packet in container.demux(*streams):
-                if video_done and audio_done:
-                    break
-
-                if packet.stream == video_stream and not video_done:
-                    try:
-                        frames = packet.decode()
-                    except av.error.InvalidDataError:
-                        logging.info("pyav decode error")
-                        continue
-                    for frame in frames:
-                        if frame.pts is not None and frame.pts < start_pts:
-                            continue
-                        if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
-                            video_done = True
-                            if last_video_pts is not None:
-                                # the source continues past the window: hold the last kept frame to the window end
-                                last_video_end = max(last_video_end, end_pts - video_pts_offset)
-                            break
-                        # the source's true display duration of this frame; average_rate is not a
-                        # frame duration (sparse/VFR sources), so it is only the fallback
-                        frame_duration = frame.duration if frame.duration else pts_step
-                        if end_pts is not None and frame.pts is not None:
-                            frame_duration = min(frame_duration, end_pts - frame.pts)
-                        if output is None:
-                            rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
-                            if rotation_k % 2:
-                                out_width, out_height = frame.height, frame.width
-                            else:
-                                out_width, out_height = frame.width, frame.height
-                            if out_width % 2 or out_height % 2:
-                                raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
-                            source_size = (frame.width, frame.height)
-                            output = av.open(path, **open_kwargs)
-                            # Add metadata before writing any streams
-                            write_output_metadata(container, output, metadata)
-                            out_video = output.add_stream("h264", rate=rate)
-                            # no B-frames: reordering makes mp4 sample durations follow decode order,
-                            # so irregular-VFR spans and trim windows land wrong
-                            out_video.codec_context.max_b_frames = 0
-                            out_video.width = out_width
-                            out_video.height = out_height
-                            out_video.pix_fmt = pix_fmt
-                            # source pts pass through (rebased to 0), so variable frame rate survives
-                            out_video.codec_context.time_base = video_stream.time_base
-                            if audio_stream is not None:
-                                out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
-                        if (frame.width, frame.height) != source_size:
-                            # encoding would silently rescale the new geometry into the old one
-                            raise ValueError(
-                                f"Video resolution changes mid-stream "
-                                f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
-                            )
-                        if rotation_k:
-                            if rotation_filter is None:
-                                g = av.filter.Graph()
-                                g_src = g.add_buffer(width=frame.width, height=frame.height,
-                                                     format=frame.format.name, time_base=video_stream.time_base)
-                                tail = g_src
-                                for filter_name, filter_args in {1: [("transpose", "cclock")],
-                                                                 2: [("hflip", None), ("vflip", None)],
-                                                                 3: [("transpose", "clock")]}[rotation_k]:
-                                    step = g.add(filter_name, filter_args)
-                                    tail.link_to(step)
-                                    tail = step
-                                g_sink = g.add("buffersink")
-                                tail.link_to(g_sink)
-                                g.configure()
-                                rotation_filter = (g_src, g_sink)
-                            rotation_filter[0].push(frame)
-                            frame = rotation_filter[1].pull()
-                        if frame.color_range == ColorRange.JPEG:
-                            # compress full-range sources (yuvj/MJPEG) to limited range
-                            frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
-                        else:
-                            frame = frame.reformat(format=pix_fmt)
-                        if frame.pts is not None:
-                            if video_pts_offset is None:
-                                video_pts_offset = frame.pts
-                            frame.pts -= video_pts_offset
-                        if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
-                            # broken sources emit missing/backward timestamps mid-stream, which the
-                            # muxer rejects; nudge them forward by one nominal frame interval
-                            frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
-                        last_video_pts = frame.pts
-                        last_video_end = frame.pts + frame_duration
-                        video_frame_durations[frame.pts] = frame_duration
-                        # the decoded pict_type would force x264's frame types (intra-only
-                        # sources like MJPEG/ProRes would come out all-keyframe)
-                        frame.pict_type = 0
-                        for out_packet in out_video.encode(frame):
-                            out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
-                            output.mux(out_packet)
-                        drain_audio()
-
-                elif packet.stream == audio_stream and not audio_done:
-                    for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
-                        if not audio_started:
-                            if resampled.pts is None:
-                                frame_start = 0.0
-                            else:
-                                # passthrough frames keep the source stream's time base
-                                tb = resampled.time_base if resampled.time_base else audio_time_base
-                                frame_start = float(resampled.pts * tb)
-                            to_skip = max(0, int((start_time - frame_start) * sample_rate))
-                            if to_skip >= resampled.samples:
-                                continue
-                            audio_started = True
-                            if to_skip:
-                                pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
-                                continue
-                        pending_audio.append(resampled)
-                        if video_done:
-                            # the video window is complete so the cap is final, but containers
-                            # that interleave audio behind video (fragmented mp4) still owe most
-                            # of it: stop only once the demuxed audio covers the cap
-                            cap = drain_audio()
-                            if pending_audio or samples_written >= cap:
-                                drain_audio(final=True)
-                                audio_done = True
-                                break
-
-            if output is None:
-                raise ValueError(f"No decodable video frames found in file '{self.__file}'")
-            if out_audio is not None and not audio_done:
-                drain_audio(final=True)
-            window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
-            for out_packet in out_video.encode(None):
-                duration = video_frame_durations.pop(out_packet.pts, 0)
-                if out_packet.pts == last_video_pts:
-                    duration = max(duration, window_fill)
-                out_packet.duration = duration
-                output.mux(out_packet)
-            if out_audio is not None:
-                output.mux(out_audio.encode(None))
-        except BaseException:
-            if output is not None:
-                output.close()
-                if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
-                    os.remove(path)
-            raise
-        else:
-            if output is not None:
-                output.close()
-
     def _get_first_video_stream(self, container: InputContainer):
         if len(container.streams.video):
             return container.streams.video[0]
@@ -801,12 +527,22 @@ class VideoFromComponents(VideoInput):
         bit_depth: int | None = None,
     ):
         """Save the video to a file path or BytesIO buffer."""
-        open_kwargs = mp4_output_open_kwargs(path, format, codec)
+        if format != VideoContainer.AUTO and format != VideoContainer.MP4:
+            raise ValueError("Only MP4 format is supported for now")
+        if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
+            raise ValueError("Only H264 codec is supported for now")
         # None means "use the depth this video was created with" (CreateVideo's choice).
         if bit_depth is None:
             bit_depth = self.__bit_depth
         is_10bit = bit_depth >= 10
-        with av.open(path, **open_kwargs) as output:
+        extra_kwargs = {}
+        if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
+            extra_kwargs["format"] = format.value
+        elif isinstance(path, io.BytesIO):
+            # BytesIO has no file extension, so av.open can't infer the format.
+            # Default to mp4 since that's the only supported format anyway.
+            extra_kwargs["format"] = "mp4"
+        with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
             # Add metadata before writing any streams
             if metadata is not None:
                 for key, value in metadata.items():

comfy_execution/caching.py

@@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
 
 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
 
+RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
+
 
 def all_outputs_dynamic(outputs):
     if outputs is None:
@@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs):
 
     return True
 
-
 class RAMPressureCache(LRUCache):
 
     def __init__(self, key_class, enable_providers=False):
@@ -539,9 +540,9 @@ class RAMPressureCache(LRUCache):
         self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
         super().set_local(node_id, value)
 
-    def ram_release(self, target, free_active=False):
+    def ram_release(self, target, free_active=False, min_entry_size=0):
         if psutil.virtual_memory().available >= target:
-            return
+            return 0
 
         clean_list = []
 
@@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache):
             oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
 
             ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
+            oom_ram_usage = ram_usage
             def scan_list_for_ram_usage(outputs):
-                nonlocal ram_usage
+                nonlocal ram_usage, oom_ram_usage
                 if outputs is None:
                     return
                 for output in outputs:
@@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache):
                         scan_list_for_ram_usage(output)
                     elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
                         ram_usage += output.numel() * output.element_size()
+                        oom_ram_usage += output.numel() * output.element_size()
                     elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
                         #old ModelPatchers are the first to go
-                        ram_usage = 1e30
+                        oom_ram_usage = 1e30
             scan_list_for_ram_usage(cache_entry.outputs)
 
-            oom_score *= ram_usage
+            if ram_usage < min_entry_size:
+                continue
+
+            oom_score *= oom_ram_usage
             #In the case where we have no information on the node ram usage at all,
             #break OOM score ties on the last touch timestamp (pure LRU)
-            bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
+            bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
 
+        freed = 0
         while psutil.virtual_memory().available < target and clean_list:
-            _, _, key = clean_list.pop()
+            _, _, key, ram_usage = clean_list.pop()
             del self.cache[key]
             self.used_generation.pop(key, None)
             self.timestamps.pop(key, None)
             self.children.pop(key, None)
+            freed += ram_usage
+        return freed

execution.py

@@ -29,6 +29,7 @@ from comfy_execution.caching import (
     HierarchicalCache,
     LRUCache,
     RAMPressureCache,
+    RAM_CACHE_LARGE_INTERMEDIATE,
 )
 from comfy_execution.graph import (
     DynamicPrompt,
@@ -794,12 +795,16 @@ class PromptExecutor:
                     if self.cache_type == CacheType.RAM_PRESSURE:
                         ram_release_callback(ram_inactive_headroom)
                         ram_shortfall = ram_headroom - psutil.virtual_memory().available
-                        freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
-                        if freed < ram_shortfall:
-                            if freed > 64 * (1024 ** 2):
-                                # AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
-                                time.sleep(0.05)
-                            ram_release_callback(ram_headroom, free_active=True)
+                        if ram_shortfall > 0:
+                            freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE)
+                            ram_shortfall -= freed
+                        if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall):
+                            freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
+                            if freed < ram_shortfall:
+                                if freed > 64 * (1024 ** 2):
+                                    # AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
+                                    time.sleep(0.05)
+                                ram_release_callback(ram_headroom, free_active=True)
                 else:
                     # Only execute when the while-loop ends without break
                     # Send cached UI for intermediate output nodes that weren't executed

tests-unit/comfy_api_test/video_types_test.py

@@ -2,12 +2,11 @@ import pytest
 import torch
 import tempfile
 import os
-import sys
 import av
 import io
 from fractions import Fraction
 from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
-from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
+from comfy_api.util.video_types import VideoComponents
 from comfy_api.input.basic_types import AudioInput
 from av.error import InvalidDataError
 
@@ -238,386 +237,3 @@ def test_duration_consistency(video_components):
     manual_duration = float(components.images.shape[0] / components.frame_rate)
 
     assert duration == pytest.approx(manual_duration)
-
-
-def create_transcode_source(
-    width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False,
-    container_format="mov", audio_codec="pcm_s16le",
-):
-    """Create a temp video that save_to must transcode (mpeg4 video, so codec != h264).
-
-    ``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists
-    (``codec_context is None``), like the APAC track in iPhone spatial-audio recordings.
-    ``rotation`` patches a 90-degree display matrix into the video track header.
-    """
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format=container_format) as container:
-        video_stream = container.add_stream("mpeg4", rate=fps)
-        video_stream.width = width
-        video_stream.height = height
-        video_stream.pix_fmt = "yuv420p"
-        audio = []
-        for _ in range(audio_streams + undecodable_audio):
-            stream = container.add_stream(audio_codec, rate=44100)
-            stream.sample_rate = 44100
-            audio.append(stream)
-
-        for i in range(frames):
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(),
-                format="rgb24",
-            )
-            container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
-        # write audio in 1024-sample frames, like real decoders produce, so the
-        # per-frame skip/cap logic in the transcode path actually runs
-        for stream in audio:
-            for offset in range(0, 44100 * frames // fps, 1024):
-                n = min(1024, 44100 * frames // fps - offset)
-                audio_frame = av.AudioFrame.from_ndarray(
-                    torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono"
-                )
-                audio_frame.sample_rate = 44100
-                audio_frame.pts = offset
-                container.mux(stream.encode(audio_frame))
-        for stream in [video_stream, *audio]:
-            container.mux(stream.encode(None))
-
-    data = bytearray(buffer.getvalue())
-    end = len(data)
-    for _ in range(undecodable_audio):
-        end = data.rindex(b"sowt", 0, end)
-        data[end:end + 4] = b"Xpac"
-    if rotation:
-        # the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd
-        # inside moov = video track (search from moov so mdat bytes can't false-match)
-        matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40
-        values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30]
-        data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values)
-
-    tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False)
-    tmp.write(bytes(data))
-    tmp.close()
-    return tmp.name
-
-
-def transcode_and_probe(video):
-    buffer = io.BytesIO()
-    video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264)
-    buffer.seek(0)
-    with av.open(buffer) as container:
-        video_stream = container.streams.video[0]
-        audio_stream = container.streams.audio[0] if container.streams.audio else None
-        frames = 0
-        first_pts = None
-        for packet in container.demux(video_stream):
-            for frame in packet.decode():
-                if first_pts is None:
-                    first_pts = frame.pts
-                frames += 1
-        return {
-            "codec": video_stream.codec_context.name,
-            "width": video_stream.codec_context.width,
-            "height": video_stream.codec_context.height,
-            "frames": frames,
-            "first_pts": first_pts,
-            "video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None,
-            "audio_seconds": float(audio_stream.duration * audio_stream.time_base)
-            if audio_stream and audio_stream.duration else None,
-            "audio_codecs": [s.codec_context.name for s in container.streams.audio],
-        }
-
-
-def test_save_to_transcode_streams_without_buffering_frames():
-    """Transcoding must not decode the whole video into memory first (~2 GiB for this source)"""
-    resource = pytest.importorskip("resource")  # no getrusage on Windows
-    rss_scale = 1 if sys.platform == "darwin" else 1024  # ru_maxrss: bytes on macOS, KiB elsewhere
-    # ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured
-    # delta and quietly defang this canary, so keep this source the biggest thing in the suite
-    file_path = create_transcode_source(width=640, height=480, frames=300)
-    try:
-        rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale
-        result = transcode_and_probe(VideoFromFile(file_path))
-        rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before
-
-        assert result["codec"] == "h264"
-        assert result["frames"] == 300
-        assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)"
-    finally:
-        os.unlink(file_path)
-
-
-def test_save_to_transcode_honors_trim_window():
-    """start_time/duration trim applies to both video and audio on the streaming path"""
-    file_path = create_transcode_source(frames=90)  # 3s @ 30fps
-    try:
-        result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
-        assert result["frames"] == pytest.approx(30, abs=2)
-        assert result["first_pts"] == 0  # trimmed output is rebased to start at zero
-        assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
-        assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
-    finally:
-        os.unlink(file_path)
-
-
-def test_save_to_transcode_keeps_audio_of_sparse_video():
-    """Audio that runs ahead of a sparse video track (slideshows, timelapses) must be
-    kept in full — it is only clamped to the video's end, never to the video cursor."""
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mp4") as container:
-        video_stream = container.add_stream("mpeg4", rate=30)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
-        for t in (0, 30, 60):  # 3 frames spread over 60 seconds
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24"
-            ).reformat(format="yuv420p")
-            frame.pts = t * 15360
-            frame.time_base = Fraction(1, 15360)
-            container.mux(video_stream.encode(frame))
-        container.mux(video_stream.encode(None))
-        for offset in range(0, 48000 * 60, 1024):
-            n = min(1024, 48000 * 60 - offset)
-            audio_frame = av.AudioFrame.from_ndarray(
-                torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
-            )
-            audio_frame.sample_rate = 48000
-            audio_frame.pts = offset
-            audio_frame.time_base = Fraction(1, 48000)
-            container.mux(audio_stream.encode(audio_frame))
-        container.mux(audio_stream.encode(None))
-
-    buffer.seek(0)
-    result = transcode_and_probe(VideoFromFile(buffer))
-    assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0)
-
-
-def test_save_to_transcode_vfr_audio_covers_video_span():
-    """A trim window in the sparse region of a VFR file keeps audio for the true pts span
-    of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the
-    average is dominated by the dense region (and can be plain wrong on MediaRecorder files)."""
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mp4") as container:
-        video_stream = container.add_stream("mpeg4", rate=30)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
-        # 10 frames inside the first second, then one every 1.25 s
-        for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]):
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24"
-            ).reformat(format="yuv420p")
-            frame.pts = int(t * 15360)
-            frame.time_base = Fraction(1, 15360)
-            container.mux(video_stream.encode(frame))
-        container.mux(video_stream.encode(None))
-        for offset in range(0, 48000 * 6, 1024):
-            n = min(1024, 48000 * 6 - offset)
-            audio_frame = av.AudioFrame.from_ndarray(
-                torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
-            )
-            audio_frame.sample_rate = 48000
-            audio_frame.pts = offset
-            audio_frame.time_base = Fraction(1, 48000)
-            container.mux(audio_stream.encode(audio_frame))
-        container.mux(audio_stream.encode(None))
-
-    buffer.seek(0)
-    result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5))
-    # kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval
-    assert result["frames"] == 4
-    assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45)
-
-
-def test_save_to_transcode_trims_audio_in_stream_time_base_units():
-    """Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing
-    must convert through the frame's time base instead of assuming sample units. AAC audio,
-    because it decodes straight to the encoder's format and hits the resampler passthrough
-    that keeps the source time base on the frames."""
-    file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac")
-    try:
-        result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
-        assert result["audio_codecs"] == ["aac"]
-        assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
-        assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
-    finally:
-        os.unlink(file_path)
-
-
-def test_save_to_transcode_learns_unprobed_audio_params():
-    """mpegts is only probed a few seconds deep at open, so an audio stream whose first
-    packet comes later (live captures where audio kicks in late) still has sample_rate 0
-    when the transcode starts; the parameters must be learned from the stream itself."""
-    sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mpegts") as container:
-        video_stream = container.add_stream("mpeg4", rate=fps)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
-        for i in range(video_seconds * fps):
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
-            )
-            container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
-        for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024):
-            n = min(1024, (video_seconds - audio_start) * sample_rate - offset)
-            audio_frame = av.AudioFrame.from_ndarray(
-                torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
-            )
-            audio_frame.sample_rate = sample_rate
-            audio_frame.pts = audio_start * sample_rate + offset
-            container.mux(audio_stream.encode(audio_frame))
-        for stream in (video_stream, audio_stream):
-            container.mux(stream.encode(None))
-
-    buffer.seek(0)
-    with av.open(buffer) as container:
-        # the scenario requires unprobed parameters; if a future FFmpeg probes deeper,
-        # push audio_start/video_seconds further out to restore it
-        assert container.streams.audio[0].codec_context.sample_rate == 0
-    result = transcode_and_probe(VideoFromFile(buffer))
-    assert result["frames"] == video_seconds * fps
-    assert result["audio_codecs"] == ["aac"]
-    assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
-
-
-def test_save_to_transcode_trimmed_fragmented_mp4_keeps_audio():
-    """Fragmented mp4 (MediaRecorder, DASH/HLS-derived files) delivers audio well behind
-    video, so when the trim window's last video frame arrives the audio demuxed so far
-    does not cover the window yet; the transcode must keep demuxing audio until it does
-    instead of finalizing on the first audio frame it sees afterwards."""
-    sample_rate, fps, seconds = 48000, 30, 6
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mp4", options={"movflags": "frag_keyframe+empty_moov"}) as container:
-        video_stream = container.add_stream("h264", rate=fps)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
-        next_audio_pts = 0
-        for i in range(seconds * fps):
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
-            )
-            container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
-            while next_audio_pts / sample_rate <= i / fps:  # feed audio alongside, like a live pipeline
-                audio_frame = av.AudioFrame.from_ndarray(
-                    torch.zeros(1, 1024, dtype=torch.float32).numpy(), format="fltp", layout="mono"
-                )
-                audio_frame.sample_rate = sample_rate
-                audio_frame.pts = next_audio_pts
-                container.mux(audio_stream.encode(audio_frame))
-                next_audio_pts += 1024
-        for stream in (video_stream, audio_stream):
-            container.mux(stream.encode(None))
-
-    result = transcode_and_probe(VideoFromFile(buffer, start_time=0.5, duration=1.0))
-    assert result["video_seconds"] == pytest.approx(1.0, abs=0.05)
-    assert result["audio_seconds"] == pytest.approx(1.0, abs=0.05)
-
-
-def test_save_to_transcode_sparse_video_keeps_true_duration():
-    """average_rate is not a frame duration: a 3-frame video spanning 60 s averages
-    0.05 fps, and padding the last frame with 1/average_rate used to extend the
-    output — and the audio kept with it — about 20 s past the source span."""
-    sample_rate = 48000
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mp4") as container:
-        video_stream = container.add_stream("mpeg4", rate=30)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
-        for i, second in enumerate((0, 30, 60)):
-            frame = av.VideoFrame.from_ndarray(
-                torch.full((64, 64, 3), i * 80, dtype=torch.uint8).numpy(), format="rgb24"
-            ).reformat(format="yuv420p")
-            frame.pts = second * 30
-            frame.time_base = Fraction(1, 30)
-            container.mux(video_stream.encode(frame))
-        for offset in range(0, 90 * sample_rate, 1024):
-            n = min(1024, 90 * sample_rate - offset)
-            audio_frame = av.AudioFrame.from_ndarray(
-                torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
-            )
-            audio_frame.sample_rate = sample_rate
-            audio_frame.pts = offset
-            container.mux(audio_stream.encode(audio_frame))
-        for stream in (video_stream, audio_stream):
-            container.mux(stream.encode(None))
-
-    result = transcode_and_probe(VideoFromFile(buffer))
-    assert result["frames"] == 3
-    # the last frame keeps its true stts duration (1/30 s), not 1/average_rate (~20 s)
-    assert result["video_seconds"] == pytest.approx(60.03, abs=0.05)
-    assert result["audio_seconds"] == pytest.approx(60.03, abs=0.1)
-
-    trimmed = transcode_and_probe(VideoFromFile(buffer, duration=45))
-    assert trimmed["frames"] == 2
-    # a kept frame whose source duration crosses the window end is clamped to it
-    assert trimmed["video_seconds"] == pytest.approx(45.0, abs=0.05)
-    assert trimmed["audio_seconds"] == pytest.approx(45.0, abs=0.1)
-
-
-def test_save_to_transcode_irregular_vfr_keeps_span():
-    """B-frames reorder packets, and mp4 sample durations follow decode order: the dts
-    timeline ends before the pts timeline, so an irregular-VFR source's tail holds fell
-    out of the container (this 20.23 s span used to come out as 15.27 s, and the 10 s
-    trim as 6.03 s). The transcode encodes without B-frames so every sample keeps its
-    true display duration."""
-    durations = [1, 1, 60, 1, 1, 120, 1, 180, 1, 1, 150, 90]  # 1/30 s ticks, span 20.2333 s
-    generator = torch.Generator().manual_seed(7)
-    buffer = io.BytesIO()
-    with av.open(buffer, mode="w", format="mp4") as container:
-        video_stream = container.add_stream("mpeg4", rate=30)
-        video_stream.width = video_stream.height = 64
-        video_stream.pix_fmt = "yuv420p"
-        pts = 0
-        for duration in durations:
-            # textured frames, so an encoder with default settings has B-frames to gain from
-            frame = av.VideoFrame.from_ndarray(
-                torch.randint(0, 255, (64, 64, 3), generator=generator, dtype=torch.uint8).numpy(),
-                format="rgb24",
-            ).reformat(format="yuv420p")
-            frame.pts = pts
-            frame.time_base = Fraction(1, 30)
-            pts += duration
-            for packet in video_stream.encode(frame):
-                packet.duration = duration  # exact stts in the source
-                container.mux(packet)
-        container.mux(video_stream.encode(None))
-
-    result = transcode_and_probe(VideoFromFile(buffer))
-    assert result["frames"] == len(durations)
-    assert result["video_seconds"] == pytest.approx(sum(durations) / 30, abs=0.05)
-
-    trimmed = transcode_and_probe(VideoFromFile(buffer, duration=10))
-    assert trimmed["frames"] == 8  # frames at 12.167 s+ fall outside the window
-    assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05)
-
-
-def test_save_to_transcode_bakes_rotation():
-    """A 90-degree display-matrix rotation swaps the output dimensions (portrait video)"""
-    file_path = create_transcode_source(width=64, height=32, rotation=True)
-    try:
-        result = transcode_and_probe(VideoFromFile(file_path))
-        assert (result["width"], result["height"]) == (32, 64)
-        assert result["frames"] == 30
-    finally:
-        os.unlink(file_path)
-
-
-def test_save_to_transcode_skips_undecodable_audio():
-    """Streaming transcode keeps the decodable audio track and drops undecodable ones;
-    with no decodable audio at all the output is video-only instead of crashing."""
-    mixed = all_bad = None
-    try:
-        mixed = create_transcode_source(audio_streams=1, undecodable_audio=1)
-        all_bad = create_transcode_source(audio_streams=0, undecodable_audio=2)
-        result = transcode_and_probe(VideoFromFile(mixed))
-        assert result["audio_codecs"] == ["aac"]
-        assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
-        assert transcode_and_probe(VideoFromFile(all_bad))["audio_codecs"] == []
-    finally:
-        for path in (mixed, all_bad):
-            if path:
-                os.unlink(path)