Merge 296b7c7b6d into b08debceca

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

.coderabbit.yaml

@@ -4,12 +4,12 @@ early_access: false
 tone_instructions: "Only comment on issues introduced by this PR's changes. Do not flag pre-existing problems in moved, re-indented, or reformatted code."
 
 reviews:
-  profile: "chill"
-  request_changes_workflow: false
+  profile: "assertive"
+  request_changes_workflow: true
   high_level_summary: false
   poem: false
   review_status: false
-  review_details: false
+  review_details: true
   commit_status: true
   collapse_walkthrough: true
   changed_files_summary: false
@@ -39,6 +39,14 @@ reviews:
     - path: "**"
       instructions: |
         IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
+        Treat AGENTS.md as mandatory repository policy, not optional style guidance.
+        Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
+        In particular, enforce architecture boundaries, dtype/device/memory rules,
+        interface contracts, import style, no unnecessary try/except blocks, no inline
+        imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
+        Prefer direct findings over suggestions when a rule is violated. Only ignore
+        AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
+        in the PR.
         Do NOT flag pre-existing issues in code that was merely moved, re-indented,
         de-indented, or reformatted without logic changes. If code appears in the diff
         only due to whitespace or structural reformatting (e.g., removing a `with:` block),
@@ -123,5 +131,10 @@ chat:
 
 knowledge_base:
   opt_out: false
+  code_guidelines:
+    enabled: true
+    filePatterns:
+      - files: "AGENTS.md"
+        applyTo: "**"
   learnings:
     scope: "auto"

AGENTS.md

@@ -171,6 +171,9 @@
 - Reuse existing model classes, blocks, ops, and helper modules when appropriate.
   Before implementing a new version of a model component, search the existing
   model code for a class or helper that already provides the behavior.
+- Model detection code that inspects linear weight shapes should only use the
+  first dimension. The second dimension may be half the original size for
+  NVFP4 or other 4-bit quantized models.
 - Avoid adding `einops` usage in core inference code. Use native torch tensor
   ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
   `unsqueeze`, and `squeeze` instead.

CLAUDE.md

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

comfy/ldm/kandinsky5/model.py

@@ -6,6 +6,12 @@ import comfy.ldm.common_dit
 from comfy.ldm.modules.attention import optimized_attention
 from comfy.ldm.flux.math import apply_rope1
 from comfy.ldm.flux.layers import EmbedND
+from comfy.ldm.kandinsky5.utils_nabla import (
+    fractal_flatten,
+    fractal_unflatten,
+    fast_sta_nabla,
+    nabla,
+)
 
 def attention(q, k, v, heads, transformer_options={}):
     return optimized_attention(
@@ -116,14 +122,17 @@ class SelfAttention(nn.Module):
         result = proj_fn(x).view(*x.shape[:-1], self.num_heads, -1)
         return apply_rope1(norm_fn(result), freqs)
 
-    def _forward(self, x, freqs, transformer_options={}):
+    def _forward(self, x, freqs, sparse_params=None, transformer_options={}):
         q = self._compute_qk(x, freqs, self.to_query, self.query_norm)
         k = self._compute_qk(x, freqs, self.to_key, self.key_norm)
         v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1)
-        out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
+        if sparse_params is None:
+            out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
+        else:
+            out = nabla(q, k, v, sparse_params)
         return self.out_layer(out)
 
-    def _forward_chunked(self, x, freqs, transformer_options={}):
+    def _forward_chunked(self, x, freqs, sparse_params=None, transformer_options={}):
         def process_chunks(proj_fn, norm_fn):
             x_chunks = torch.chunk(x, self.num_chunks, dim=1)
             freqs_chunks = torch.chunk(freqs, self.num_chunks, dim=1)
@@ -135,14 +144,17 @@ class SelfAttention(nn.Module):
         q = process_chunks(self.to_query, self.query_norm)
         k = process_chunks(self.to_key, self.key_norm)
         v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1)
-        out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
+        if sparse_params is None:
+            out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
+        else:
+            out = nabla(q, k, v, sparse_params)
         return self.out_layer(out)
 
-    def forward(self, x, freqs, transformer_options={}):
+    def forward(self, x, freqs, sparse_params=None, transformer_options={}):
         if x.shape[1] > 8192:
-            return self._forward_chunked(x, freqs, transformer_options=transformer_options)
+            return self._forward_chunked(x, freqs, sparse_params=sparse_params, transformer_options=transformer_options)
         else:
-            return self._forward(x, freqs, transformer_options=transformer_options)
+            return self._forward(x, freqs, sparse_params=sparse_params, transformer_options=transformer_options)
 
 
 class CrossAttention(SelfAttention):
@@ -251,12 +263,12 @@ class TransformerDecoderBlock(nn.Module):
         self.feed_forward_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
         self.feed_forward = FeedForward(model_dim, ff_dim, operation_settings=operation_settings)
 
-    def forward(self, visual_embed, text_embed, time_embed, freqs, transformer_options={}):
+    def forward(self, visual_embed, text_embed, time_embed, freqs, sparse_params=None, transformer_options={}):
         self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1)
         # self attention
         shift, scale, gate = get_shift_scale_gate(self_attn_params)
         visual_out = apply_scale_shift_norm(self.self_attention_norm, visual_embed, scale, shift)
-        visual_out = self.self_attention(visual_out, freqs, transformer_options=transformer_options)
+        visual_out = self.self_attention(visual_out, freqs, sparse_params=sparse_params, transformer_options=transformer_options)
         visual_embed = apply_gate_sum(visual_embed, visual_out, gate)
         # cross attention
         shift, scale, gate = get_shift_scale_gate(cross_attn_params)
@@ -369,21 +381,82 @@ class Kandinsky5(nn.Module):
 
         visual_embed = self.visual_embeddings(x)
         visual_shape = visual_embed.shape[:-1]
-        visual_embed = visual_embed.flatten(1, -2)
 
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.visual_transformer_blocks)
         transformer_options["block_type"] = "double"
+
+        B, _, T, H, W = x.shape
+        NABLA_THR = 31 # long (10 sec) generation
+        if T > NABLA_THR:
+            assert self.patch_size[0] == 1
+
+            # pro video model uses lower P at higher resolutions
+            P = 0.7 if self.model_dim == 4096 and H * W >= 14080 else 0.9
+
+            freqs = freqs.view(freqs.shape[0], *visual_shape[1:], *freqs.shape[2:])
+            visual_embed, freqs = fractal_flatten(visual_embed, freqs, visual_shape[1:])
+            pt, ph, pw = self.patch_size
+            T, H, W = T // pt, H // ph, W // pw
+
+            wT, wW, wH = 11, 3, 3
+            sta_mask = fast_sta_nabla(T, H // 8, W // 8, wT, wH, wW, device=x.device)
+
+            sparse_params = dict(
+                sta_mask=sta_mask.unsqueeze_(0).unsqueeze_(0),
+                attention_type="nabla",
+                to_fractal=True,
+                P=P,
+                wT=wT, wW=wW, wH=wH,
+                add_sta=True,
+                visual_shape=(T, H, W),
+                method="topcdf",
+            )
+        else:
+            sparse_params = None
+            visual_embed = visual_embed.flatten(1, -2)
+
         for i, block in enumerate(self.visual_transformer_blocks):
             transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
-                    return block(x=args["x"], context=args["context"], time_embed=args["time_embed"], freqs=args["freqs"], transformer_options=args.get("transformer_options"))
-                visual_embed = blocks_replace[("double_block", i)]({"x": visual_embed, "context": context, "time_embed": time_embed, "freqs": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})["x"]
+                    return block(
+                        x=args["x"],
+                        context=args["context"],
+                        time_embed=args["time_embed"],
+                        freqs=args["freqs"],
+                        sparse_params=args.get("sparse_params"),
+                        transformer_options=args.get("transformer_options"),
+                    )
+                visual_embed = blocks_replace[("double_block", i)](
+                    {
+                        "x": visual_embed,
+                        "context": context,
+                        "time_embed": time_embed,
+                        "freqs": freqs,
+                        "sparse_params": sparse_params,
+                        "transformer_options": transformer_options,
+                    },
+                    {"original_block": block_wrap},
+                )["x"]
             else:
-                visual_embed = block(visual_embed, context, time_embed, freqs=freqs, transformer_options=transformer_options)
+                visual_embed = block(
+                    visual_embed,
+                    context,
+                    time_embed,
+                    freqs=freqs,
+                    sparse_params=sparse_params,
+                    transformer_options=transformer_options,
+                )
+
+        if T > NABLA_THR:
+            visual_embed = fractal_unflatten(
+                visual_embed,
+                visual_shape[1:],
+            )
+        else:
+            visual_embed = visual_embed.reshape(*visual_shape, -1)
 
-        visual_embed = visual_embed.reshape(*visual_shape, -1)
         return self.out_layer(visual_embed, time_embed)
 
     def _forward(self, x, timestep, context, y, time_dim_replace=None, transformer_options={}, **kwargs):

comfy/ldm/kandinsky5/utils_nabla.py

@@ -0,0 +1,146 @@
+import math
+
+import torch
+from torch import Tensor
+from torch.nn.attention.flex_attention import BlockMask, flex_attention
+
+
+def fractal_flatten(x, rope, shape):
+    pixel_size = 8
+    x = local_patching(x, shape, (1, pixel_size, pixel_size), dim=1)
+    rope = local_patching(rope, shape, (1, pixel_size, pixel_size), dim=1)
+    x = x.flatten(1, 2)
+    rope = rope.flatten(1, 2)
+    return x, rope
+
+
+def fractal_unflatten(x, shape):
+    pixel_size = 8
+    x = x.reshape(x.shape[0], -1, pixel_size**2, x.shape[-1])
+    x = local_merge(x, shape, (1, pixel_size, pixel_size), dim=1)
+    return x
+
+def local_patching(x, shape, group_size, dim=0):
+    duration, height, width = shape
+    g1, g2, g3 = group_size
+    x = x.reshape(
+        *x.shape[:dim],
+        duration // g1,
+        g1,
+        height // g2,
+        g2,
+        width // g3,
+        g3,
+        *x.shape[dim + 3 :]
+    )
+    x = x.permute(
+        *range(len(x.shape[:dim])),
+        dim,
+        dim + 2,
+        dim + 4,
+        dim + 1,
+        dim + 3,
+        dim + 5,
+        *range(dim + 6, len(x.shape))
+    )
+    x = x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3)
+    return x
+
+
+def local_merge(x, shape, group_size, dim=0):
+    duration, height, width = shape
+    g1, g2, g3 = group_size
+    x = x.reshape(
+        *x.shape[:dim],
+        duration // g1,
+        height // g2,
+        width // g3,
+        g1,
+        g2,
+        g3,
+        *x.shape[dim + 2 :]
+    )
+    x = x.permute(
+        *range(len(x.shape[:dim])),
+        dim,
+        dim + 3,
+        dim + 1,
+        dim + 4,
+        dim + 2,
+        dim + 5,
+        *range(dim + 6, len(x.shape))
+    )
+    x = x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3)
+    return x
+
+def fast_sta_nabla(T: int, H: int, W: int, wT: int = 3, wH: int = 3, wW: int = 3, device="cuda") -> Tensor:
+    l = torch.Tensor([T, H, W]).amax()
+    r = torch.arange(0, l, 1, dtype=torch.int16, device=device)
+    mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs()
+    sta_t, sta_h, sta_w = (
+        mat[:T, :T].flatten(),
+        mat[:H, :H].flatten(),
+        mat[:W, :W].flatten(),
+    )
+    sta_t = sta_t <= wT // 2
+    sta_h = sta_h <= wH // 2
+    sta_w = sta_w <= wW // 2
+    sta_hw = (
+        (sta_h.unsqueeze(1) * sta_w.unsqueeze(0))
+        .reshape(H, H, W, W)
+        .transpose(1, 2)
+        .flatten()
+    )
+    sta = (
+        (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0))
+        .reshape(T, T, H * W, H * W)
+        .transpose(1, 2)
+    )
+    return sta.reshape(T * H * W, T * H * W)
+
+def nablaT_v2(q: Tensor, k: Tensor, sta: Tensor, thr: float = 0.9) -> BlockMask:
+    # Map estimation
+    B, h, S, D = q.shape
+    s1 = S // 64
+    qa = q.reshape(B, h, s1, 64, D).mean(-2)
+    ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1)
+    map = qa @ ka
+
+    map = torch.softmax(map / math.sqrt(D), dim=-1)
+    # Map binarization
+    vals, inds = map.sort(-1)
+    cvals = vals.cumsum_(-1)
+    mask = (cvals >= 1 - thr).int()
+    mask = mask.gather(-1, inds.argsort(-1))
+    mask = torch.logical_or(mask, sta)
+
+    # BlockMask creation
+    kv_nb = mask.sum(-1).to(torch.int32)
+    kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32)
+    return BlockMask.from_kv_blocks(
+        torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None
+    )
+
+@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True)
+def nabla(query, key, value, sparse_params=None):
+    query = query.transpose(1, 2).contiguous()
+    key = key.transpose(1, 2).contiguous()
+    value = value.transpose(1, 2).contiguous()
+    block_mask = nablaT_v2(
+        query,
+        key,
+        sparse_params["sta_mask"],
+        thr=sparse_params["P"],
+    )
+    out = (
+        flex_attention(
+            query,
+            key,
+            value,
+            block_mask=block_mask
+        )
+        .transpose(1, 2)
+        .contiguous()
+    )
+    out = out.flatten(-2, -1)
+    return out

comfy/sd1_clip.py

@@ -543,18 +543,24 @@ class SDTokenizer:
     def _try_get_embedding(self, embedding_name:str):
         '''
         Takes a potential embedding name and tries to retrieve it.
-        Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
+        Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
         '''
         split_embed = embedding_name.split()
         embedding_name = split_embed[0]
         leftover = ' '.join(split_embed[1:])
+
+        match = re.search(r'[<\[]', embedding_name)
+        if match is not None:
+            leftover = embedding_name[match.start():] + (" " + leftover if leftover else "")
+            embedding_name = embedding_name[:match.start()]
+
         embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
         if embed is None:
             stripped = embedding_name.strip(',')
             if len(stripped) < len(embedding_name):
                 embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
-                return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
-        return (embed, leftover)
+                return (embed, embedding_name, "{} {}".format(embedding_name[len(stripped):], leftover))
+        return (embed, embedding_name, leftover)
 
     def pad_tokens(self, tokens, amount):
         if self.pad_left:
@@ -585,7 +591,7 @@ class SDTokenizer:
         tokens = []
         for weighted_segment, weight in parsed_weights:
             to_tokenize = unescape_important(weighted_segment)
-            split = re.split(' {0}|\n{0}'.format(self.embedding_identifier), to_tokenize)
+            split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize)
             to_tokenize = [split[0]]
             for i in range(1, len(split)):
                 to_tokenize.append("{}{}".format(self.embedding_identifier, split[i]))
@@ -595,7 +601,7 @@ class SDTokenizer:
                 # if we find an embedding, deal with the embedding
                 if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
                     embedding_name = word[len(self.embedding_identifier):].strip('\n')
-                    embed, leftover = self._try_get_embedding(embedding_name)
+                    embed, embedding_name, leftover = self._try_get_embedding(embedding_name)
                     if embed is None:
                         logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
                     else:

comfy/text_encoders/llama.py

@@ -937,22 +937,41 @@ class BaseGenerate:
             return torch.argmax(logits, dim=-1, keepdim=True)
 
         # Sampling mode
-        if repetition_penalty != 1.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
-
-        if presence_penalty is not None and presence_penalty != 0.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] -= presence_penalty
+        if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
+            token_ids = torch.tensor(list(set(token_history)), device=logits.device)
+            token_logits = logits[:, token_ids]
+            if repetition_penalty != 1.0:
+                token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
+            if presence_penalty is not None and presence_penalty != 0.0:
+                token_logits = token_logits - presence_penalty
+            logits[:, token_ids] = token_logits
 
         if temperature != 1.0:
             logits = logits / temperature
 
         if top_k > 0:
-            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
+            top_k = min(top_k, logits.shape[-1])
+            logits, top_indices = torch.topk(logits, top_k)
+
+            if min_p > 0.0:
+                probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
+                top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
+                min_threshold = min_p * top_probs
+                indices_to_remove = probs_before_filter < min_threshold
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 0] = False
+                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+                indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1, generator=generator)
+            return top_indices.gather(1, next_token)
 
         if min_p > 0.0:
             probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)

comfy_api_nodes/nodes_bytedance.py

@@ -2611,7 +2611,7 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
         return IO.Schema(
             node_id="ByteDanceSeedAudio",
             display_name="ByteDance Seed Audio 1.0",
-            category="api node/audio/ByteDance",
+            category="partner/audio/ByteDance",
             description=(
                 "Generate speech, music, sound effects and multi-speaker dialogue from a single prompt "
                 "with ByteDance Seed Audio 1.0. Describe the voice(s), emotion, ambience, background music "

comfy_api_nodes/util/request_logger.py

@@ -9,6 +9,7 @@ from typing import Any
 import folder_paths
 
 logger = logging.getLogger(__name__)
+_SENSITIVE_HEADERS = {"authorization", "x-api-key"}
 
 
 def get_log_directory():
@@ -73,6 +74,10 @@ def _format_data_for_logging(data: Any) -> str:
     return str(data)
 
 
+def _redact_headers(headers: dict) -> dict:
+    return {k: ("***" if k.lower() in _SENSITIVE_HEADERS else v) for k, v in headers.items()}
+
+
 def log_request_response(
     operation_id: str,
     request_method: str,
@@ -101,7 +106,7 @@ def log_request_response(
         log_content.append(f"Method: {request_method}")
         log_content.append(f"URL: {request_url}")
         if request_headers:
-            log_content.append(f"Headers:\n{_format_data_for_logging(request_headers)}")
+            log_content.append(f"Headers:\n{_format_data_for_logging(_redact_headers(request_headers))}")
         if request_params:
             log_content.append(f"Params:\n{_format_data_for_logging(request_params)}")
         if request_data is not None:

comfy_extras/nodes_color.py

@@ -16,23 +16,30 @@ class ColorToRGBInt(io.ComfyNode):
             ],
             outputs=[
                 io.Int.Output(display_name="rgb_int"),
-                io.Color.Output(display_name="hex")
+                io.Color.Output(display_name="hex"),
+                io.Float.Output(display_name="alpha"),
             ],
         )
 
     @classmethod
     def execute(cls, color: str) -> io.NodeOutput:
-        # expect format #RRGGBB
-        if len(color) != 7 or color[0] != "#":
-            raise ValueError("Color must be in format #RRGGBB")
+        # expect format #RRGGBB or #RRGGBBAA
+        if len(color) not in (7, 9) or color[0] != "#":
+            raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA")
         try:
             int(color[1:], 16)
         except ValueError:
-            raise ValueError("Color must be in format #RRGGBB") from None
+            raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA") from None
+
+        alpha = 1.0
+        if len(color) == 9:
+            alpha = int(color[7:9], 16) / 255.0
+            color = color[:7]
+
         r, g, b = hex_to_rgb(color)
 
         rgb_int = r * 256 * 256 + g * 256 + b
-        return io.NodeOutput(rgb_int, color)
+        return io.NodeOutput(rgb_int, color, alpha)
 
 
 class ColorExtension(ComfyExtension):

comfy_extras/nodes_kandinsky5.py

@@ -34,6 +34,9 @@ class Kandinsky5ImageToVideo(io.ComfyNode):
 
     @classmethod
     def execute(cls, positive, negative, vae, width, height, length, batch_size, start_image=None) -> io.NodeOutput:
+        if length > 121: # 10 sec generation, for nabla
+            height = 128 * round(height / 128)
+            width = 128 * round(width / 128)
         latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
         cond_latent_out = {}
         if start_image is not None:

requirements.txt

@@ -1,6 +1,6 @@
 comfyui-frontend-package==1.45.20
 comfyui-workflow-templates==0.11.2
-comfyui-embedded-docs==0.5.6
+comfyui-embedded-docs==0.5.7
 torch
 torchsde
 torchvision