Merge f9bc8cc781 into 091b70edda

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

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/chroma/model.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 from torch import Tensor, nn
-from einops import rearrange, repeat
+from einops import rearrange
 import comfy.patcher_extension
 import comfy.ldm.common_dit
 
@@ -267,29 +267,54 @@ class Chroma(nn.Module):
             img = self.final_layer(img, vec=final_mod)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
+    def forward(self, x, timestep, context, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
-        ).execute(x, timestep, context, guidance, control, transformer_options, **kwargs)
+        ).execute(x, timestep, context, guidance, ref_latents, control, transformer_options, **kwargs)
 
-    def _forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
-        bs, c, h, w = x.shape
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
-
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size)
+    def _forward(self, x, timestep, context, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
+        bs, c, h_orig, w_orig = x.shape
 
+        h_len = ((h_orig + (self.patch_size // 2)) // self.patch_size)
+        w_len = ((w_orig + (self.patch_size // 2)) // self.patch_size)
+        img, img_ids = comfy.ldm.common_dit.process_img(x, patch_size=self.patch_size, transformer_options=transformer_options)
         if img.ndim != 3 or context.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
+        img_tokens = img.shape[1]
+        if ref_latents is not None:
+            h = 0
+            w = 0
+            index = 0
+            ref_latents_method = kwargs.get("ref_latents_method", "offset")
+            for ref in ref_latents:
+                if ref_latents_method == "index":
+                    index += 1
+                    h_offset = 0
+                    w_offset = 0
+                elif ref_latents_method == "uxo":
+                    index = 0
+                    h_offset = h_len * self.patch_size + h
+                    w_offset = w_len * self.patch_size + w
+                    h += ref.shape[-2]
+                    w += ref.shape[-1]
+                else:
+                    index = 1
+                    h_offset = 0
+                    w_offset = 0
+                    if ref.shape[-2] + h > ref.shape[-1] + w:
+                        w_offset = w
+                    else:
+                        h_offset = h
+                    h = max(h, ref.shape[-2] + h_offset)
+                    w = max(w, ref.shape[-1] + w_offset)
 
-        h_len = ((h + (self.patch_size // 2)) // self.patch_size)
-        w_len = ((w + (self.patch_size // 2)) // self.patch_size)
-        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
-        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+                kontext, kontext_ids = comfy.ldm.common_dit.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, patch_size=self.patch_size, transformer_options=transformer_options)
+                img = torch.cat([img, kontext], dim=1)
+                img_ids = torch.cat([img_ids, kontext_ids], dim=1)
 
         txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
         out = self.forward_orig(img, img_ids, context, txt_ids, timestep, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
-        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h,:w]
+        out = out[:, :img_tokens]
+        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h_orig,:w_orig]

comfy/ldm/common_dit.py

@@ -1,4 +1,5 @@
 import torch
+from einops import rearrange, repeat
 import comfy.rmsnorm
 
 
@@ -14,3 +15,32 @@ def pad_to_patch_size(img, patch_size=(2, 2), padding_mode="circular"):
 
 
 rms_norm = comfy.rmsnorm.rms_norm
+
+def process_img(x, index=0, h_offset=0, w_offset=0, patch_size=2, transformer_options={}, num_axes=3):
+    bs, c, h, w = x.shape
+    x = pad_to_patch_size(x, (patch_size, patch_size))
+
+    img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+    h_len = ((h + (patch_size // 2)) // patch_size)
+    w_len = ((w + (patch_size // 2)) // patch_size)
+
+    h_offset = ((h_offset + (patch_size // 2)) // patch_size)
+    w_offset = ((w_offset + (patch_size // 2)) // patch_size)
+
+    steps_h = h_len
+    steps_w = w_len
+
+    rope_options = transformer_options.get("rope_options", None)
+    if rope_options is not None:
+        h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
+        w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
+
+        index += rope_options.get("shift_t", 0.0)
+        h_offset += rope_options.get("shift_y", 0.0)
+        w_offset += rope_options.get("shift_x", 0.0)
+
+    img_ids = torch.zeros((steps_h, steps_w, num_axes), device=x.device, dtype=x.dtype)
+    img_ids[:, :, 0] = img_ids[:, :, 1] + index
+    img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=steps_h, device=x.device, dtype=x.dtype).unsqueeze(1)
+    img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=steps_w, device=x.device, dtype=x.dtype).unsqueeze(0)
+    return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)

comfy/ldm/flux/model.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 from torch import Tensor, nn
-from einops import rearrange, repeat
+from einops import rearrange
 import comfy.ldm.common_dit
 import comfy.patcher_extension
 
@@ -311,36 +311,6 @@ class Flux(nn.Module):
         img = self.final_layer(img, vec_orig, **extra_kwargs)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def process_img(self, x, index=0, h_offset=0, w_offset=0, transformer_options={}):
-        bs, c, h, w = x.shape
-        patch_size = self.patch_size
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
-
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
-        h_len = ((h + (patch_size // 2)) // patch_size)
-        w_len = ((w + (patch_size // 2)) // patch_size)
-
-        h_offset = ((h_offset + (patch_size // 2)) // patch_size)
-        w_offset = ((w_offset + (patch_size // 2)) // patch_size)
-
-        steps_h = h_len
-        steps_w = w_len
-
-        rope_options = transformer_options.get("rope_options", None)
-        if rope_options is not None:
-            h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
-            w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
-
-            index += rope_options.get("shift_t", 0.0)
-            h_offset += rope_options.get("shift_y", 0.0)
-            w_offset += rope_options.get("shift_x", 0.0)
-
-        img_ids = torch.zeros((steps_h, steps_w, len(self.params.axes_dim)), device=x.device, dtype=torch.float32)
-        img_ids[:, :, 0] = img_ids[:, :, 1] + index
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=steps_h, device=x.device, dtype=torch.float32).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=steps_w, device=x.device, dtype=torch.float32).unsqueeze(0)
-        return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)
-
     def forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
@@ -354,7 +324,8 @@ class Flux(nn.Module):
 
         h_len = ((h_orig + (patch_size // 2)) // patch_size)
         w_len = ((w_orig + (patch_size // 2)) // patch_size)
-        img, img_ids = self.process_img(x, transformer_options=transformer_options)
+        num_axes = len(self.params.axes_dim)
+        img, img_ids = comfy.ldm.common_dit.process_img(x, patch_size=patch_size, transformer_options=transformer_options, num_axes=num_axes)
         img_tokens = img.shape[1]
         timestep_zero_index = None
         if ref_latents is not None:
@@ -386,7 +357,7 @@ class Flux(nn.Module):
                     h = max(h, ref.shape[-2] + h_offset)
                     w = max(w, ref.shape[-1] + w_offset)
 
-                kontext, kontext_ids = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, transformer_options=transformer_options)
+                kontext, kontext_ids = comfy.ldm.common_dit.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, patch_size=patch_size, transformer_options=transformer_options, num_axes=num_axes)
                 img = torch.cat([img, kontext], dim=1)
                 img_ids = torch.cat([img_ids, kontext_ids], dim=1)
                 ref_num_tokens.append(kontext.shape[1])

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

folder_paths.py

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

main.py

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

tests-unit/comfy_test/folder_path_test.py

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