Merge b656217514 into 091b70edda

- parse_ids_filter: one shared parser/validator for the ids query param, used
  by the /api/jobs handler AND its tests (no more hand-copied wiring that can
  drift from — and silently outlive a regression in — the shipped handler)
- present-but-empty ids (?ids=, ?ids=,,) is now a zero-match filter, not a
  silent 'return the entire job history'
- bounded history lookup when an ids filter is present: a batch poll costs
  O(requested ids), not O(total history)
- dedupe ids so the max-count cap bounds distinct values, not repeats
- .get('id') instead of j['id'] so a job missing its id degrades to no-match
  rather than a 500

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_execution/jobs.py

@@ -31,6 +31,12 @@ class JobStatus:
     ALL = [PENDING, IN_PROGRESS, COMPLETED, FAILED, CANCELLED]
 
 
+# Maximum number of (distinct) ids accepted by the `ids` filter on the jobs
+# listing. Caps request size; the bounded id-lookup in get_all_jobs then keeps
+# a batch-poll request at O(requested ids), not O(total history).
+MAX_JOB_IDS_FILTER = 100
+
+
 def validate_job_id(value) -> str:
     """Validate a client-supplied job (prompt) id.
 
@@ -50,6 +56,56 @@ def validate_job_id(value) -> str:
     return value
 
 
+class JobIdsFilterError(ValueError):
+    """Raised when the ``ids`` query-param value is malformed.
+
+    Carries an HTTP-ready ``payload`` dict so the caller can return it verbatim
+    with a 400 without re-deriving the message.
+    """
+
+    def __init__(self, payload: dict):
+        self.payload = payload
+        super().__init__(payload.get("error", "invalid ids"))
+
+
+def parse_ids_filter(ids_param: Optional[str]) -> Optional[list[str]]:
+    """Parse the ``ids`` query-param value into a filter list.
+
+    Single source of truth for ``ids`` parsing/validation, shared by the HTTP
+    handler and its tests so the two cannot drift.
+
+    Returns:
+        - ``None`` when the param is absent (``ids_param is None``) -> no filter.
+        - A de-duplicated list when present. An empty/blank value (``?ids=``,
+          ``?ids=,,``) yields ``[]``, which ``get_all_jobs`` treats as a
+          zero-match filter -- NOT "return everything".
+
+    Raises:
+        JobIdsFilterError: more than ``MAX_JOB_IDS_FILTER`` distinct ids, or any
+            id not in canonical UUID form. ``.payload`` is a 400-ready dict.
+    """
+    if ids_param is None:
+        return None
+    # De-dupe up front: a repeated id must not count toward the cap or be
+    # looked up twice. dict.fromkeys keeps first-seen order.
+    ids_filter = list(dict.fromkeys(i.strip() for i in ids_param.split(',') if i.strip()))
+    if len(ids_filter) > MAX_JOB_IDS_FILTER:
+        raise JobIdsFilterError(
+            {"error": f"ids must contain at most {MAX_JOB_IDS_FILTER} values"}
+        )
+    invalid_ids = []
+    for jid in ids_filter:
+        try:
+            validate_job_id(jid)
+        except (ValueError, AttributeError):
+            invalid_ids.append(jid)
+    if invalid_ids:
+        raise JobIdsFilterError(
+            {"error": "ids contains invalid id(s)", "invalid_ids": invalid_ids}
+        )
+    return ids_filter
+
+
 # Media types that can be previewed in the frontend
 PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'})
 
@@ -362,6 +418,7 @@ def get_all_jobs(
     history: dict,
     status_filter: Optional[list[str]] = None,
     workflow_id: Optional[str] = None,
+    ids: Optional[list[str]] = None,
     sort_by: str = "created_at",
     sort_order: str = "desc",
     limit: Optional[int] = None,
@@ -376,6 +433,8 @@ def get_all_jobs(
         history: Dict of history items keyed by prompt_id
         status_filter: List of statuses to include (from JobStatus.ALL)
         workflow_id: Filter by workflow ID
+        ids: Restrict the result to these job ids. None = no filter; a present
+            list (including empty) restricts to that set, so [] = zero matches
         sort_by: Field to sort by ('created_at', 'execution_duration')
         sort_order: 'asc' or 'desc'
         limit: Maximum number of items to return
@@ -389,6 +448,10 @@ def get_all_jobs(
     if status_filter is None:
         status_filter = JobStatus.ALL
 
+    # None => no id filter; a present list (including empty) restricts to that
+    # set (empty => zero matches).
+    id_set = set(ids) if ids is not None else None
+
     if JobStatus.IN_PROGRESS in status_filter:
         for item in running:
             jobs.append(normalize_queue_item(item, JobStatus.IN_PROGRESS))
@@ -400,14 +463,30 @@ def get_all_jobs(
     history_statuses = {JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED}
     requested_history_statuses = history_statuses & set(status_filter)
     if requested_history_statuses:
-        for prompt_id, history_item in history.items():
-            job = normalize_history_item(prompt_id, history_item)
-            if job.get('status') in requested_history_statuses:
-                jobs.append(job)
+        if id_set is not None:
+            # Batch-poll fast path: history is keyed by id, so look up only the
+            # requested ids instead of normalizing the whole (unbounded) history.
+            for prompt_id in id_set:
+                history_item = history.get(prompt_id)
+                if history_item is None:
+                    continue
+                job = normalize_history_item(prompt_id, history_item)
+                if job.get('status') in requested_history_statuses:
+                    jobs.append(job)
+        else:
+            for prompt_id, history_item in history.items():
+                job = normalize_history_item(prompt_id, history_item)
+                if job.get('status') in requested_history_statuses:
+                    jobs.append(job)
 
     if workflow_id:
         jobs = [j for j in jobs if j.get('workflow_id') == workflow_id]
 
+    if id_set is not None:
+        # `.get('id')` (not `j['id']`): prune_dict can drop a None id, and a
+        # job missing its id should degrade to "no match", not raise KeyError.
+        jobs = [j for j in jobs if j.get('id') in id_set]
+
     jobs = apply_sorting(jobs, sort_by, sort_order)
 
     total_count = len(jobs)

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)

server.py

@@ -16,6 +16,8 @@ from comfy_execution.jobs import (
     cancel_job,
     CANCEL_PENDING,
     CANCEL_RUNNING,
+    parse_ids_filter,
+    JobIdsFilterError,
 )
 import uuid
 import urllib
@@ -807,6 +809,7 @@ class PromptServer():
             Query parameters:
                 status: Filter by status (comma-separated): pending, in_progress, completed, failed
                 workflow_id: Filter by workflow ID
+                ids: Filter by job id (comma-separated UUIDs, max 100)
                 sort_by: Sort field: created_at (default), execution_duration
                 sort_order: Sort direction: asc, desc (default)
                 limit: Max items to return (positive integer)
@@ -816,6 +819,7 @@ class PromptServer():
 
             status_param = query.get('status')
             workflow_id = query.get('workflow_id')
+            ids_param = query.get('ids')
             sort_by = query.get('sort_by', 'created_at').lower()
             sort_order = query.get('sort_order', 'desc').lower()
 
@@ -829,6 +833,16 @@ class PromptServer():
                         status=400
                     )
 
+            # Optional batch filter: narrow the result to a known set of job ids
+            # (e.g. polling a submitted batch in one request). Parsing/validation
+            # lives in parse_ids_filter so this handler and its tests share one
+            # implementation. Absent => no filter; present-but-empty (`?ids=`,
+            # `?ids=,,`) => zero matches, not "everything".
+            try:
+                ids_filter = parse_ids_filter(ids_param)
+            except JobIdsFilterError as e:
+                return web.json_response(e.payload, status=400)
+
             if sort_by not in {'created_at', 'execution_duration'}:
                 return web.json_response(
                     {"error": "sort_by must be 'created_at' or 'execution_duration'"},
@@ -880,6 +894,7 @@ class PromptServer():
                 running, queued, history,
                 status_filter=status_filter,
                 workflow_id=workflow_id,
+                ids=ids_filter,
                 sort_by=sort_by,
                 sort_order=sort_order,
                 limit=limit,

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/jobs_list_test/jobs_list_test.py

@@ -0,0 +1,277 @@
+"""Tests for the ``ids`` batch filter on the jobs listing endpoint.
+
+Covers both layers:
+
+* the pure ``comfy_execution.jobs.get_all_jobs`` filtering logic (the ``ids``
+  argument narrows the result, composes with ``status_filter``, and silently
+  ignores ids that match nothing), and
+
+* the HTTP contract of ``GET /api/jobs`` for the ``ids`` query parameter
+  (a valid set narrows the response, an oversized set or a malformed id is
+  rejected with 400).
+
+The HTTP layer is exercised against a small aiohttp app whose handler calls the
+SAME ``parse_ids_filter`` that ``server.py`` uses (no hand-copied wiring, so it
+cannot drift), driven by a fake queue. This keeps the test free of the heavy
+ComfyUI runtime (torch, nodes, ...) while still testing the real parsing
+contract.
+"""
+
+import pytest
+from aiohttp import web
+
+from comfy_execution.jobs import (
+    JobStatus,
+    JobIdsFilterError,
+    MAX_JOB_IDS_FILTER,
+    get_all_jobs,
+    parse_ids_filter,
+)
+
+# Canonical UUID ids (the endpoint validates UUID format).
+_UUID_A = "aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa"
+_UUID_B = "bbbbbbbb-bbbb-4bbb-bbbb-bbbbbbbbbbbb"
+_UUID_C = "cccccccc-cccc-4ccc-cccc-cccccccccccc"
+_UUID_MISSING = "ffffffff-ffff-4fff-ffff-ffffffffffff"
+
+
+def make_queue_item(prompt_id, priority=0):
+    """Build a queue tuple shaped like the real ones (5 elements, id at index 1)."""
+    return (priority, prompt_id, {}, {}, [])
+
+
+def make_history_item(status_str="success"):
+    """Build a history item dict shaped like the real ones."""
+    return {
+        "prompt": (0, "", {}, {}, []),
+        "status": {"status_str": status_str, "messages": []},
+        "outputs": {},
+    }
+
+
+# ---------------------------------------------------------------------------
+# Pure get_all_jobs filtering logic
+# ---------------------------------------------------------------------------
+
+
+def test_ids_filter_returns_only_requested():
+    running = [make_queue_item(_UUID_A)]
+    queued = [make_queue_item(_UUID_B)]
+    history = {_UUID_C: make_history_item()}
+
+    jobs, total = get_all_jobs(running, queued, history, ids=[_UUID_A, _UUID_C])
+
+    returned = {j["id"] for j in jobs}
+    assert returned == {_UUID_A, _UUID_C}
+    assert total == 2
+    assert _UUID_B not in returned
+
+
+def test_ids_filter_absent_returns_all():
+    running = [make_queue_item(_UUID_A)]
+    queued = [make_queue_item(_UUID_B)]
+    history = {_UUID_C: make_history_item()}
+
+    jobs, total = get_all_jobs(running, queued, history)
+
+    assert {j["id"] for j in jobs} == {_UUID_A, _UUID_B, _UUID_C}
+    assert total == 3
+
+
+def test_ids_filter_empty_list_returns_none():
+    """A present-but-empty ids list is a zero-match filter, not "no filter".
+
+    ``None`` means "no id filter"; ``[]`` means "restrict to nothing".
+    """
+    running = [make_queue_item(_UUID_A)]
+    queued = [make_queue_item(_UUID_B)]
+
+    jobs, total = get_all_jobs(running, queued, {}, ids=[])
+
+    assert jobs == []
+    assert total == 0
+
+
+def test_ids_filter_unknown_id_silently_absent():
+    """An id that matches nothing is simply not present (no error)."""
+    running = [make_queue_item(_UUID_A)]
+
+    jobs, total = get_all_jobs(running, [], {}, ids=[_UUID_A, _UUID_MISSING])
+
+    assert {j["id"] for j in jobs} == {_UUID_A}
+    assert total == 1
+
+
+def test_ids_filter_composes_with_status():
+    """ids only narrows; it composes with the status filter."""
+    running = [make_queue_item(_UUID_A)]
+    queued = [make_queue_item(_UUID_B)]
+    history = {_UUID_C: make_history_item()}
+
+    # Request A and C by id, but restrict to in_progress only -> just A.
+    jobs, total = get_all_jobs(
+        running, queued, history,
+        status_filter=[JobStatus.IN_PROGRESS],
+        ids=[_UUID_A, _UUID_C],
+    )
+
+    assert {j["id"] for j in jobs} == {_UUID_A}
+    assert total == 1
+
+
+# ---------------------------------------------------------------------------
+# parse_ids_filter -- the shared parsing/validation (server.py + these tests)
+# ---------------------------------------------------------------------------
+
+
+def test_parse_ids_absent_is_none():
+    assert parse_ids_filter(None) is None
+
+
+def test_parse_ids_present_but_empty_is_empty_list():
+    # `?ids=` and `?ids=,,` parse to [] -> zero-match filter, not None.
+    assert parse_ids_filter("") == []
+    assert parse_ids_filter(",,") == []
+
+
+def test_parse_ids_dedupes_preserving_order():
+    assert parse_ids_filter(f"{_UUID_A},{_UUID_B},{_UUID_A}") == [_UUID_A, _UUID_B]
+
+
+def test_parse_ids_cap_counts_distinct_not_duplicates():
+    # A small distinct set repeated far past the cap is still under it.
+    repeated = ",".join([_UUID_A, _UUID_B] * MAX_JOB_IDS_FILTER)
+    assert parse_ids_filter(repeated) == [_UUID_A, _UUID_B]
+    # But more than MAX distinct ids is rejected.
+    distinct = ",".join(
+        f"{i:08d}-0000-4000-8000-000000000000" for i in range(MAX_JOB_IDS_FILTER + 1)
+    )
+    with pytest.raises(JobIdsFilterError):
+        parse_ids_filter(distinct)
+
+
+def test_parse_ids_invalid_raises_with_payload():
+    with pytest.raises(JobIdsFilterError) as exc:
+        parse_ids_filter(f"{_UUID_A},not-a-uuid")
+    assert "not-a-uuid" in exc.value.payload["invalid_ids"]
+
+
+# ---------------------------------------------------------------------------
+# HTTP contract for the ids query parameter
+# ---------------------------------------------------------------------------
+
+
+class FakePromptQueue:
+    """Minimal stand-in exposing the accessors get_jobs uses."""
+
+    def __init__(self, running=None, queued=None, history=None):
+        self._running = list(running or [])
+        self._queued = list(queued or [])
+        self._history = dict(history or {})
+
+    def get_current_queue_volatile(self):
+        return (list(self._running), list(self._queued))
+
+    def get_history(self):
+        return dict(self._history)
+
+
+def make_app(prompt_queue):
+    """Build an aiohttp app whose handler calls the REAL parse_ids_filter.
+
+    No hand-copied parsing wiring, so this test cannot stay green while the
+    shipped parsing in server.py regresses -- both go through parse_ids_filter.
+    """
+
+    async def get_jobs(request):
+        try:
+            ids_filter = parse_ids_filter(request.rel_url.query.get('ids'))
+        except JobIdsFilterError as e:
+            return web.json_response(e.payload, status=400)
+
+        running, queued = prompt_queue.get_current_queue_volatile()
+        history = prompt_queue.get_history()
+
+        jobs, total = get_all_jobs(running, queued, history, ids=ids_filter)
+
+        return web.json_response({
+            'jobs': jobs,
+            'pagination': {'total': total},
+        })
+
+    app = web.Application()
+    app.router.add_get('/api/jobs', get_jobs)
+    return app
+
+
+@pytest.fixture
+def queue():
+    return FakePromptQueue(
+        running=[make_queue_item(_UUID_A)],
+        queued=[make_queue_item(_UUID_B)],
+        history={_UUID_C: make_history_item()},
+    )
+
+
+@pytest.mark.asyncio
+async def test_http_ids_filter_narrows(aiohttp_client, queue):
+    client = await aiohttp_client(make_app(queue))
+
+    resp = await client.get(f"/api/jobs?ids={_UUID_A},{_UUID_C}")
+    assert resp.status == 200
+    body = await resp.json()
+    assert {j["id"] for j in body["jobs"]} == {_UUID_A, _UUID_C}
+
+
+@pytest.mark.asyncio
+async def test_http_ids_unknown_id_is_not_an_error(aiohttp_client, queue):
+    client = await aiohttp_client(make_app(queue))
+
+    resp = await client.get(f"/api/jobs?ids={_UUID_A},{_UUID_MISSING}")
+    assert resp.status == 200
+    body = await resp.json()
+    assert {j["id"] for j in body["jobs"]} == {_UUID_A}
+
+
+@pytest.mark.asyncio
+async def test_http_ids_over_limit_returns_400(aiohttp_client, queue):
+    client = await aiohttp_client(make_app(queue))
+
+    # Distinct ids past the cap. (Repeats of one id are de-duped and would NOT
+    # trip the cap -- see test_parse_ids_cap_counts_distinct_not_duplicates.)
+    too_many = ",".join(
+        f"{i:08d}-0000-4000-8000-000000000000" for i in range(MAX_JOB_IDS_FILTER + 1)
+    )
+    resp = await client.get(f"/api/jobs?ids={too_many}")
+    assert resp.status == 400
+
+
+@pytest.mark.asyncio
+async def test_http_ids_invalid_id_returns_400(aiohttp_client, queue):
+    client = await aiohttp_client(make_app(queue))
+
+    resp = await client.get(f"/api/jobs?ids={_UUID_A},not-a-uuid")
+    assert resp.status == 400
+    body = await resp.json()
+    assert "not-a-uuid" in body["invalid_ids"]
+
+
+@pytest.mark.asyncio
+async def test_http_ids_absent_returns_all(aiohttp_client, queue):
+    client = await aiohttp_client(make_app(queue))
+
+    resp = await client.get("/api/jobs")
+    assert resp.status == 200
+    body = await resp.json()
+    assert {j["id"] for j in body["jobs"]} == {_UUID_A, _UUID_B, _UUID_C}
+
+
+@pytest.mark.asyncio
+async def test_http_ids_present_but_empty_returns_none(aiohttp_client, queue):
+    """`?ids=` (present but empty) is a zero-match filter, not "return all"."""
+    client = await aiohttp_client(make_app(queue))
+
+    resp = await client.get("/api/jobs?ids=")
+    assert resp.status == 200
+    body = await resp.json()
+    assert body["jobs"] == []