Merge 9a5e6233f4 into 091b70edda

* fix(Video): don't crash on videos with undecodable audio streams

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

* Update comfy_api_nodes/util/upload_helpers.py

---------

Signed-off-by: bigcat88 <bigcat88@icloud.com>
Co-authored-by: Alexis Rolland <alexisrolland@hotmail.com>

.github/workflows/cla.yml

@@ -0,0 +1,91 @@
+name: CLA Assistant
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_target:
+    types: [opened, synchronize, closed]
+
+permissions:
+  actions: write
+  contents: read # 'read' is enough because signatures live in a REMOTE repo
+  pull-requests: write
+  statuses: write
+
+jobs:
+  cla-assistant:
+    runs-on: ubuntu-latest
+    steps:
+      # The CLA action normally requires every commit author in a PR to sign.
+      # We only want the PR author to sign, so we allowlist all other committers
+      # by computing them from the PR's commits and excluding the PR author.
+      - name: Build author-only allowlist
+        id: allowlist
+        if: >
+          github.event_name == 'pull_request_target' ||
+          (github.event_name == 'issue_comment' && github.event.issue.pull_request && (
+            github.event.comment.body == 'recheck' ||
+            github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
+          ))
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
+          PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
+          BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
+        run: |
+          others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
+            --jq '.[] | (.author.login // empty), (.committer.login // empty)' \
+            | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
+          if [ -n "$others" ]; then
+            echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
+          else
+            echo "allowlist=${BASE_ALLOWLIST}" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: CLA Assistant
+        # Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
+        # IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
+        if: >
+          github.event_name == 'pull_request_target' ||
+          (github.event_name == 'issue_comment' && github.event.issue.pull_request && (
+            github.event.comment.body == 'recheck' ||
+            github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
+          ))
+        uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          # PAT required to write to the centralized signatures repo.
+          PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
+        with:
+          # Where the CLA document lives (shown to contributors)
+          path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
+
+          # Centralized signature storage
+          remote-organization-name: comfy-org
+          remote-repository-name: comfy-cla
+          path-to-signatures: signatures/cla.json
+          branch: main
+
+          # Only the PR author must sign: bots plus every non-author committer
+          # are allowlisted via the "Build author-only allowlist" step above.
+          # *[bot] is a catch-all for any GitHub App bot account.
+          allowlist: ${{ steps.allowlist.outputs.allowlist }}
+
+          # Custom PR comment messages
+          custom-notsigned-prcomment: |
+            🎉 Thank you for your contribution, we really appreciate it! 🎉
+
+            Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
+
+            - Confirm that you own your contribution.
+            - Keep the right to reuse your own code.
+            - Grant us a copyright license to include and share it within our projects.
+
+            CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
+
+            ✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
+
+          custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
+
+          custom-allsigned-prcomment: |
+            ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.

AGENTS.md

@@ -127,6 +127,8 @@
 - Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
   platform, or backend capability detection only when the program has a useful
   fallback. Prefer specific exception types when changing new code.
+- If a library version is pinned in `requirements.txt`, do not add code to
+  ComfyUI to handle older versions of that library.
 - Remove any workarounds for PyTorch versions that ComfyUI no longer officially
   supports. Deprecated workarounds include catching an exception and rerunning
   the same op with the input cast to float. If a workaround does not have a

CLAUDE.md

@@ -1 +0,0 @@
-AGENTS.md

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

@@ -782,15 +782,23 @@ def resolve_areas_and_cond_masks_multidim(conditions, dims, device):
             mask = c['mask']
             mask = mask.to(device=device)
             modified = c.copy()
-            if len(mask.shape) == len(dims):
+            # Normalize mask to [batch, *spatial_dims]
+            target_ndim = len(dims) + 1
+            while mask.ndim > target_ndim and mask.shape[0] == 1:
+                mask = mask.squeeze(0)
+            while mask.ndim < target_ndim:
                 mask = mask.unsqueeze(0)
             if mask.shape[1:] != dims:
-                if mask.ndim < 4:
+                if len(dims) == 1:
+                    mask = torch.nn.functional.interpolate(
+                        mask.unsqueeze(1), size=dims[0],
+                        mode='linear', align_corners=False).squeeze(1)
+                elif mask.ndim < 4:
                     mask = comfy.utils.common_upscale(mask.unsqueeze(1), dims[-1], dims[-2], 'bilinear', 'none').squeeze(1)
                 else:
                     mask = comfy.utils.common_upscale(mask, dims[-1], dims[-2], 'bilinear', 'none')
 
-            if modified.get("set_area_to_bounds", False): #TODO: handle dim != 2
+            if modified.get("set_area_to_bounds", False) and len(dims) == 2:
                 bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
                 boxes, is_empty = get_mask_aabb(bounds)
                 if is_empty[0]:

comfy/sd.py

@@ -468,6 +468,9 @@ class CLIP:
     def decode(self, token_ids, skip_special_tokens=True):
         return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
 
+    def is_dynamic(self):
+        return self.patcher.is_dynamic()
+
 class VAE:
     def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
         if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
@@ -1251,6 +1254,8 @@ class VAE:
         except:
             return None
 
+    def is_dynamic(self):
+        return self.patcher.is_dynamic()
 
 class StyleModel:
     def __init__(self, model, device="cpu"):

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

@@ -1313,6 +1313,7 @@ def reshape_mask(input_mask, output_shape):
     dims = len(output_shape) - 2
 
     if dims == 1:
+        input_mask = input_mask.reshape((-1, 1, input_mask.shape[-1]))
         scale_mode = "linear"
 
     if dims == 2:

comfy_api/latest/_input_impl/video_types.py

@@ -281,11 +281,18 @@ class VideoFromFile(VideoInput):
         video_done = False
         audio_done = True
 
-        if len(container.streams.audio):
-            audio_stream = container.streams.audio[-1]
+        # 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:
@@ -457,10 +464,13 @@ class VideoFromFile(VideoInput):
                         else:
                             output_container.metadata[key] = json.dumps(value)
 
-                # Add streams to the new container
+                # Add streams to the new container. Streams with no codec context cannot be used as an output template.
                 stream_map = {}
                 for stream in streams:
                     if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)):
+                        if stream.codec_context is None:
+                            logging.warning("Skipping %s stream %d with unsupported codec", stream.type, stream.index)
+                            continue
                         out_stream = output_container.add_stream_from_template(template=stream, opaque=True)
                         stream_map[stream] = out_stream
 

comfy_api_nodes/util/upload_helpers.py

@@ -158,7 +158,14 @@ async def upload_video_to_comfyapi(
 
     # Convert VideoInput to BytesIO using specified container/codec
     video_bytes_io = BytesIO()
-    video.save_to(video_bytes_io, format=container, codec=codec)
+    try:
+        video.save_to(video_bytes_io, format=container, codec=codec)
+    except Exception as e:
+        raise ValueError(
+            f"Could not convert the input video to {container.value.upper()} for upload; "
+            f"the file may be corrupted or use an unsupported codec. "
+            f"Try re-exporting it as MP4 (H.264). Original error: {e}"
+        ) from e
     video_bytes_io.seek(0)
 
     return await upload_file_to_comfyapi(cls, video_bytes_io, filename, upload_mime_type, wait_label)

comfy_execution/caching.py

@@ -503,6 +503,21 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
 
 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
 
+
+def all_outputs_dynamic(outputs):
+    if outputs is None:
+        return False
+
+    for output in outputs:
+        if isinstance(output, (list, tuple)):
+            if not all_outputs_dynamic(output):
+                return False
+        elif not hasattr(output, "is_dynamic") or not output.is_dynamic():
+            return False
+
+    return True
+
+
 class RAMPressureCache(LRUCache):
 
     def __init__(self, key_class, enable_providers=False):
@@ -533,7 +548,11 @@ class RAMPressureCache(LRUCache):
         for key, cache_entry in self.cache.items():
             if not free_active and self.used_generation[key] == self.generation:
                 continue
-            oom_score =  RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
+
+            if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation:
+                continue
+
+            oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
 
             ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
             def scan_list_for_ram_usage(outputs):

comfy_extras/nodes_text_overlay.py

@@ -0,0 +1,150 @@
+import numpy as np
+import torch
+from PIL import Image as PILImage, ImageColor, ImageDraw, ImageFont
+from typing_extensions import override
+
+from comfy_api.latest import ComfyExtension, IO
+
+
+class TextOverlay(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="TextOverlay",
+            display_name="Draw Text Overlay",
+            category="text",
+            description="Draw text overlay on an image or batch of images.",
+            search_aliases=["text", "label", "caption", "subtitle", "watermark", "title", "addlabel", "overlay"],
+            inputs=[
+                IO.Image.Input("images"),
+                IO.String.Input("text", multiline=True, default=""),
+                IO.Float.Input("font_size", default=5.0, min=0.5, max=50.0, step=0.5, tooltip="Font size as a percentage of the image height."),
+                IO.Color.Input("color", default="#ffffff", tooltip="Color of the text."),
+                IO.Combo.Input("position", options=["top", "bottom"], default="top"),
+                IO.Combo.Input("align", options=["left", "center", "right"], default="left"),
+                IO.Boolean.Input("outline", default=True, tooltip="Draw a black outline around the text."),
+            ],
+            outputs=[IO.Image.Output(display_name="images")],
+        )
+
+    @classmethod
+    def execute(cls, images, text, font_size, color, position, align, outline) -> IO.NodeOutput:
+        if text.strip() == "":
+            return IO.NodeOutput(images)
+
+        text = text.replace("\\n", "\n").replace("\\t", "\t")
+
+        text_rgba = cls.parse_color_to_rgba(color)
+        outline_rgba = (0, 0, 0, 255) if outline else (0, 0, 0, 0)
+
+        # Render the overlay once and composite it across all frames in the batch
+        height = images.shape[1]
+        width = images.shape[2]
+        overlay_rgb, overlay_alpha = cls.render_overlay_text(width, height, text, position, align, font_size, text_rgba, outline_rgba)
+        overlay_rgb = overlay_rgb.to(device=images.device, dtype=images.dtype)
+        overlay_alpha = overlay_alpha.to(device=images.device, dtype=images.dtype)
+
+        result = images * (1.0 - overlay_alpha) + overlay_rgb * overlay_alpha
+        return IO.NodeOutput(result)
+
+    @staticmethod
+    def parse_color_to_rgba(color_string):
+        parsed = ImageColor.getrgb(color_string)
+
+        if len(parsed) == 3:
+            return (*parsed, 255)
+
+        return parsed
+
+    @classmethod
+    def render_overlay_text(cls, width, height, text, position, align, font_size, text_rgba, outline_rgba):
+        line_spacing = 1.2
+        margin_percent = 1.0
+        min_font_percent = 2.0
+        min_font_pixels = 10
+        outline_thickness_factor = 0.04
+
+        # Draw onto a transparent layer so the result can be alpha-composited over any frame.
+        layer = PILImage.new("RGBA", (width, height), (0, 0, 0, 0))
+        draw = ImageDraw.Draw(layer)
+
+        margin = int(round(margin_percent / 100.0 * min(width, height)))
+        max_width = max(1, width - 2 * margin)
+        max_height = max(1, height - 2 * margin)
+
+        # Font scales with resolution, then shrinks to fit the height.
+        size = max(1, int(round(font_size / 100.0 * height)))
+        floor = min(size, max(min_font_pixels, int(round(min_font_percent / 100.0 * height))))
+
+        while True:
+            font = ImageFont.load_default(size=size)
+            stroke = max(1, int(round(size * outline_thickness_factor))) if outline_rgba[3] > 0 else 0
+            block = "\n".join(cls.wrap_text(text, font, max_width))
+            # convert line spacing to pixel spacing
+            single = draw.textbbox((0, 0), "Ay", font=font, stroke_width=stroke)
+            double = draw.multiline_textbbox((0, 0), "Ay\nAy", font=font, spacing=0, stroke_width=stroke)
+            natural_advance = (double[3] - double[1]) - (single[3] - single[1])
+            pixel_spacing = int(round(size * line_spacing - natural_advance))
+            box = draw.multiline_textbbox((0, 0), block, font=font, spacing=pixel_spacing, stroke_width=stroke)
+            block_height = box[3] - box[1]
+
+            if block_height <= max_height or size <= floor:
+                break
+
+            size = max(floor, int(size * 0.9))
+
+        anchor_h, x = {"left": ("l", margin), "center": ("m", width / 2), "right": ("r", width - margin)}[align]
+
+        # Offset y so the rendered text sits flush against the margin
+        if position == "bottom":
+            y = height - margin - box[3]
+        else:
+            y = margin - box[1]
+
+        draw.multiline_text((x, y), block, font=font, fill=text_rgba, anchor=anchor_h + "a",
+                            align=align, spacing=pixel_spacing, stroke_width=stroke, stroke_fill=outline_rgba)
+
+        overlay = np.array(layer).astype(np.float32) / 255.0
+        overlay_rgb = torch.from_numpy(overlay[:, :, :3])
+        overlay_alpha = torch.from_numpy(overlay[:, :, 3:4])
+        return overlay_rgb, overlay_alpha
+
+    @staticmethod
+    def wrap_text(text, font, max_width):
+        lines = []
+        for raw_line in text.split("\n"):
+            words = raw_line.split()
+            if not words:
+                lines.append("")
+                continue
+            current = ""
+            # Break the line into words and split words that are too long
+            for word in words:
+                while font.getlength(word) > max_width and len(word) > 1:
+                    cut = 1
+                    while cut < len(word) and font.getlength(word[:cut + 1]) <= max_width:
+                        cut += 1
+                    if current:
+                        lines.append(current)
+                        current = ""
+                    lines.append(word[:cut])
+                    word = word[cut:]
+                candidate = word if not current else current + " " + word
+                if not current or font.getlength(candidate) <= max_width:
+                    current = candidate
+                else:
+                    lines.append(current)
+                    current = word
+            if current:
+                lines.append(current)
+        return lines
+
+
+class TextOverlayExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [TextOverlay]
+
+
+async def comfy_entrypoint() -> TextOverlayExtension:
+    return TextOverlayExtension()

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)

nodes.py

@@ -2478,6 +2478,7 @@ async def init_builtin_extra_nodes():
         "nodes_glsl.py",
         "nodes_lora_debug.py",
         "nodes_textgen.py",
+        "nodes_text_overlay.py",
         "nodes_color.py",
         "nodes_toolkit.py",
         "nodes_replacements.py",

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/comfy_test/mask_test.py

@@ -0,0 +1,336 @@
+"""
+Tests for mask handling across arbitrary spatial dimensions.
+
+Covers resolve_areas_and_cond_masks_multidim (conditioning masks) and
+reshape_mask (denoise masks) for 1D (audio), 2D (image), and 3D (video)
+spatial dims, including edge cases around batch size, spurious unsqueeze
+from upstream nodes, and size mismatches.
+"""
+
+import torch
+
+from comfy.samplers import resolve_areas_and_cond_masks_multidim
+from comfy.utils import reshape_mask
+
+
+def make_cond(mask):
+    """Create a minimal conditioning dict with a mask."""
+    return {"mask": mask, "model_conds": {}}
+
+
+def run_resolve(mask, dims, device="cpu"):
+    """Run resolve on a single condition and return the resolved mask."""
+    conds = [make_cond(mask)]
+    resolve_areas_and_cond_masks_multidim(conds, dims, device)
+    return conds[0]["mask"]
+
+
+# ============================================================
+# 1D spatial dims (audio models like AceStep v1.5)
+# dims = (length,), expected mask output: [batch, length]
+# ============================================================
+
+class Test1DSpatial:
+    """Tests for 1D spatial models (e.g. audio with noise shape [B, C, L])."""
+
+    def test_correct_shape_same_length(self):
+        """[B, L] mask with matching length — should pass through unchanged."""
+        mask = torch.ones(2, 100)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (2, 100)
+
+    def test_correct_shape_resize(self):
+        """[B, L] mask with different length — should resize via linear interp."""
+        mask = torch.ones(1, 50)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (1, 100)
+
+    def test_bare_spatial_mask(self):
+        """[L] mask (no batch) — should get batch dim added."""
+        mask = torch.ones(50)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (1, 100)
+
+    def test_spurious_unsqueeze_from_hooks(self):
+        """[1, B, L] mask (from set_mask_for_conditioning unsqueezing a [B, L] mask)
+        — should squeeze back to [B, L]."""
+        # Simulates: mask is [B, L], hooks.py does unsqueeze(0) -> [1, B, L]
+        mask = torch.ones(1, 2, 100)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (2, 100)
+
+    def test_spurious_unsqueeze_batch1(self):
+        """[1, 1, L] mask (batch=1, hooks added extra dim) — should become [1, L]."""
+        mask = torch.ones(1, 1, 50)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (1, 100)
+
+    def test_batch_gt1_same_length(self):
+        """[B, L] mask with batch=4 and matching length — no changes needed."""
+        mask = torch.rand(4, 100)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (4, 100)
+        torch.testing.assert_close(result, mask)
+
+    def test_batch_gt1_resize(self):
+        """[B, L] mask with batch=4 and different length — should resize each batch."""
+        mask = torch.rand(4, 50)
+        result = run_resolve(mask, dims=(100,))
+        assert result.shape == (4, 100)
+
+    def test_values_preserved_no_resize(self):
+        """Mask values should be preserved when no resize is needed."""
+        mask = torch.tensor([[0.0, 0.5, 1.0]])
+        result = run_resolve(mask, dims=(3,))
+        torch.testing.assert_close(result, mask)
+
+    def test_linear_interpolation_values(self):
+        """Check that linear interpolation produces sensible values."""
+        mask = torch.tensor([[0.0, 1.0]])  # [1, 2]
+        result = run_resolve(mask, dims=(5,))
+        assert result.shape == (1, 5)
+        # Should interpolate from 0 to 1
+        assert result[0, 0].item() < result[0, -1].item()
+
+    def test_set_area_to_bounds_skipped_for_1d(self):
+        """set_area_to_bounds should be skipped for 1D (no crash)."""
+        mask = torch.zeros(1, 100)
+        mask[0, 10:50] = 1.0
+        conds = [{"mask": mask, "model_conds": {}, "set_area_to_bounds": True}]
+        resolve_areas_and_cond_masks_multidim(conds, (100,), "cpu")
+        assert "area" not in conds[0]
+
+
+# ============================================================
+# 2D spatial dims (image models) — regression tests
+# dims = (H, W), expected mask output: [batch, H, W]
+# ============================================================
+
+class Test2DSpatial:
+    """Regression tests for standard 2D image models."""
+
+    def test_correct_shape_same_size(self):
+        """[B, H, W] mask matching dims — pass through."""
+        mask = torch.ones(1, 64, 64)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (1, 64, 64)
+
+    def test_bare_spatial_mask(self):
+        """[H, W] mask — should get batch dim added."""
+        mask = torch.ones(64, 64)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (1, 64, 64)
+
+    def test_resize_different_resolution(self):
+        """[B, H1, W1] mask with different size than dims — should bilinear resize."""
+        mask = torch.ones(1, 32, 32)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (1, 64, 64)
+
+    def test_4d_mask(self):
+        """[B, C, H, W] mask (4D) — should resize via common_upscale 4D path."""
+        mask = torch.ones(1, 1, 32, 32)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (1, 64, 64)
+
+    def test_batch_gt1(self):
+        """[B, H, W] mask with batch > 1."""
+        mask = torch.rand(4, 64, 64)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (4, 64, 64)
+
+    def test_batch_gt1_resize(self):
+        """[B, H, W] mask with batch > 1 and different resolution."""
+        mask = torch.rand(4, 32, 32)
+        result = run_resolve(mask, dims=(64, 64))
+        assert result.shape == (4, 64, 64)
+
+    def test_set_area_to_bounds(self):
+        """set_area_to_bounds should work for 2D masks."""
+        mask = torch.zeros(1, 64, 64)
+        mask[0, 10:20, 10:30] = 1.0
+        conds = [{"mask": mask, "model_conds": {}, "set_area_to_bounds": True}]
+        resolve_areas_and_cond_masks_multidim(conds, (64, 64), "cpu")
+        assert "area" in conds[0]
+
+    def test_non_square_resize(self):
+        """[B, H1, W1] mask resized to non-square dims."""
+        mask = torch.ones(1, 16, 32)
+        result = run_resolve(mask, dims=(64, 128))
+        assert result.shape == (1, 64, 128)
+
+
+# ============================================================
+# 3D spatial dims (video models)
+# dims = (T, H, W), expected mask output: [batch, T, H, W]
+# ============================================================
+
+class Test3DSpatial:
+    """Tests for 3D spatial models (e.g. video with noise shape [B, C, T, H, W])."""
+
+    def test_correct_shape_same_size(self):
+        """[B, T, H, W] mask matching dims — pass through."""
+        mask = torch.ones(1, 8, 64, 64)
+        result = run_resolve(mask, dims=(8, 64, 64))
+        assert result.shape == (1, 8, 64, 64)
+
+    def test_bare_spatial_mask(self):
+        """[T, H, W] mask — should get batch dim added."""
+        mask = torch.ones(8, 64, 64)
+        result = run_resolve(mask, dims=(8, 64, 64))
+        assert result.shape == (1, 8, 64, 64)
+
+    def test_resize_hw(self):
+        """[B, T, H1, W1] mask with different H, W — should resize last 2 dims."""
+        mask = torch.ones(1, 8, 32, 32)
+        result = run_resolve(mask, dims=(8, 64, 64))
+        assert result.shape == (1, 8, 64, 64)
+
+    def test_set_area_to_bounds_skipped_for_3d(self):
+        """set_area_to_bounds should be skipped for 3D (no crash)."""
+        mask = torch.zeros(1, 8, 64, 64)
+        mask[0, :, 10:20, 10:30] = 1.0
+        conds = [{"mask": mask, "model_conds": {}, "set_area_to_bounds": True}]
+        resolve_areas_and_cond_masks_multidim(conds, (8, 64, 64), "cpu")
+        assert "area" not in conds[0]
+
+
+class TestNoMask:
+    """Conditions without masks should pass through untouched."""
+
+    def test_no_mask_key(self):
+        """Condition with no mask key — untouched."""
+        conds = [{"model_conds": {}}]
+        resolve_areas_and_cond_masks_multidim(conds, (64, 64), "cpu")
+        assert "mask" not in conds[0]
+
+    def test_empty_conditions(self):
+        """Empty conditions list — no crash."""
+        conds = []
+        resolve_areas_and_cond_masks_multidim(conds, (64, 64), "cpu")
+        assert len(conds) == 0
+
+
+# ============================================================
+# Area resolution (percentage-based)
+# ============================================================
+
+class TestAreaResolution:
+    """Test that percentage-based area resolution works for different dims."""
+
+    def test_percentage_area_2d(self):
+        """Percentage area for 2D should resolve to pixel coords."""
+        conds = [{"area": ("percentage", 0.5, 0.5, 0.25, 0.25), "model_conds": {}}]
+        resolve_areas_and_cond_masks_multidim(conds, (64, 64), "cpu")
+        area = conds[0]["area"]
+        assert area == (32, 32, 16, 16)
+
+    def test_percentage_area_1d(self):
+        """Percentage area for 1D should resolve to frame coords."""
+        conds = [{"area": ("percentage", 0.5, 0.25), "model_conds": {}}]
+        resolve_areas_and_cond_masks_multidim(conds, (100,), "cpu")
+        area = conds[0]["area"]
+        assert area == (50, 25)
+
+
+# ============================================================
+# reshape_mask — mask reshaping for F.interpolate
+# ============================================================
+
+class TestReshapeMask1D:
+    """Tests for reshape_mask with 1D output (e.g. audio with noise shape [B, C, L])."""
+
+    def test_4d_input_same_length(self):
+        """[1, 1, 1, L] input (typical from pipeline) — should reshape and expand channels."""
+        mask = torch.ones(1, 1, 1, 100)
+        result = reshape_mask(mask, torch.Size([1, 64, 100]))
+        assert result.shape == (1, 64, 100)
+
+    def test_4d_input_resize(self):
+        """[1, 1, 1, L1] input resized to different length."""
+        mask = torch.ones(1, 1, 1, 50)
+        result = reshape_mask(mask, torch.Size([1, 64, 100]))
+        assert result.shape == (1, 64, 100)
+
+    def test_3d_input(self):
+        """[1, 1, L] input — should work directly."""
+        mask = torch.ones(1, 1, 100)
+        result = reshape_mask(mask, torch.Size([1, 64, 100]))
+        assert result.shape == (1, 64, 100)
+
+    def test_2d_input(self):
+        """[B, L] input — should reshape to [B, 1, L]."""
+        mask = torch.ones(1, 50)
+        result = reshape_mask(mask, torch.Size([1, 64, 100]))
+        assert result.shape == (1, 64, 100)
+
+    def test_1d_input(self):
+        """[L] input — should reshape to [1, 1, L]."""
+        mask = torch.ones(50)
+        result = reshape_mask(mask, torch.Size([1, 64, 100]))
+        assert result.shape == (1, 64, 100)
+
+    def test_channel_repeat(self):
+        """Mask with 1 channel should repeat to match output channels."""
+        mask = torch.full((1, 1, 1, 100), 0.5)
+        result = reshape_mask(mask, torch.Size([1, 32, 100]))
+        assert result.shape == (1, 32, 100)
+        torch.testing.assert_close(result, torch.full_like(result, 0.5))
+
+    def test_batch_repeat(self):
+        """Single-batch mask should repeat to match output batch size."""
+        mask = torch.full((1, 1, 1, 100), 0.7)
+        result = reshape_mask(mask, torch.Size([4, 64, 100]))
+        assert result.shape == (4, 64, 100)
+
+    def test_values_preserved_no_resize(self):
+        """Values should be preserved when no resize is needed."""
+        values = torch.tensor([[[0.0, 0.5, 1.0]]])  # [1, 1, 3]
+        result = reshape_mask(values, torch.Size([1, 1, 3]))
+        torch.testing.assert_close(result, values)
+
+    def test_interpolation_values(self):
+        """Linear interpolation should produce sensible intermediate values."""
+        mask = torch.tensor([[[[0.0, 1.0]]]])  # [1, 1, 1, 2]
+        result = reshape_mask(mask, torch.Size([1, 1, 4]))
+        assert result.shape == (1, 1, 4)
+        # Should interpolate from 0 to 1
+        assert result[0, 0, 0].item() < result[0, 0, -1].item()
+
+
+class TestReshapeMask2D:
+    """Regression tests for reshape_mask with 2D output (image models)."""
+
+    def test_standard_resize(self):
+        """[1, 1, H, W] mask resized to different resolution."""
+        mask = torch.ones(1, 1, 32, 32)
+        result = reshape_mask(mask, torch.Size([1, 4, 64, 64]))
+        assert result.shape == (1, 4, 64, 64)
+
+    def test_same_size(self):
+        """[1, 1, H, W] mask with matching size — no resize needed."""
+        mask = torch.rand(1, 1, 64, 64)
+        result = reshape_mask(mask, torch.Size([1, 4, 64, 64]))
+        assert result.shape == (1, 4, 64, 64)
+
+    def test_3d_input(self):
+        """[B, H, W] input — should reshape to [B, 1, H, W]."""
+        mask = torch.ones(1, 32, 32)
+        result = reshape_mask(mask, torch.Size([1, 4, 64, 64]))
+        assert result.shape == (1, 4, 64, 64)
+
+
+class TestReshapeMask3D:
+    """Regression tests for reshape_mask with 3D output (video models)."""
+
+    def test_standard_resize(self):
+        """[1, 1, T, H, W] mask resized to different resolution."""
+        mask = torch.ones(1, 1, 8, 32, 32)
+        result = reshape_mask(mask, torch.Size([1, 4, 8, 64, 64]))
+        assert result.shape == (1, 4, 8, 64, 64)
+
+    def test_4d_input(self):
+        """[B, T, H, W] input — should reshape to [1, 1, T, H, W]."""
+        mask = torch.ones(1, 8, 32, 32)
+        result = reshape_mask(mask, torch.Size([1, 4, 8, 64, 64]))
+        assert result.shape == (1, 4, 8, 64, 64)