Merge branch 'master' into fix-flash-attn-memory-estimate

comfy/model_prefetch.py

@@ -1,4 +1,5 @@
 import torch
+import warnings
 import weakref
 
 import comfy_aimdo.model_vbar
@@ -28,6 +29,18 @@ def cleanup_prefetched_modules(module, comfy_modules):
         comfy_aimdo.model_vbar.vbar_unpin(module._v_block)
         del module._v_block_faulted
 
+def _drop_graph(module):
+    graph = getattr(module, "_comfy_graph", None)
+    if graph is None:
+        return
+    # reset() through the bound method surfaces the allocator's benign
+    # "uncaptured free of a captured allocation" as catchable Python warnings;
+    # a plain del frees from the C++ dealloc path and spams stderr instead
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        graph["graph"].reset()
+    del module._comfy_graph
+
 def cleanup_prefetch_queues():
     global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS
 
@@ -41,7 +54,7 @@ def cleanup_prefetch_queues():
                 cleanup_prefetched_modules(prefetched_module, comfy_modules)
     PREFETCH_QUEUES = []
     for module in GRAPH_MODULES:
-        del module._comfy_graph
+        _drop_graph(module)
     GRAPH_MODULES.clear()
     GRAPH_WARMED_MODULES.clear()
     GRAPH_CAPTURE_STREAMS = {}
@@ -117,6 +130,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
                 if signature is not None:
                     module._v_block_faulted = True
             if signature is not None:
+                _drop_graph(module)
                 graph = torch.cuda.CUDAGraph()
                 if generator is not None:
                     graph.register_generator_state(generator)

comfy/ops.py

@@ -1639,7 +1639,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                             self.norm_type, self.scale_grad_by_freq, self.sparse)
                     target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
                     x = x.to(dtype=target_dtype)
-                    if scale is not None and scale != 1.0:
+                    if scale is not None:
                         x = x * scale.to(dtype=target_dtype)
                     return x
 

comfy/text_encoders/gemma4.py

@@ -6,13 +6,16 @@ import numpy as np
 from tokenizers import Tokenizer
 from dataclasses import dataclass
 import math
+import re
 
 from comfy import sd1_clip
 import comfy.model_management
+import comfy.model_prefetch
 import comfy.ops
+import comfy.quant_ops
 from comfy.ldm.modules.attention import optimized_attention_for_device
 from comfy.rmsnorm import rms_norm
-from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, _make_scaled_embedding
+from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, FixedKV, _make_scaled_embedding
 
 
 # Intentional minor divergences from transformers -reference implementation:
@@ -109,7 +112,28 @@ class Gemma4_12B_Config(Gemma4Config):
     suppress_tokens = [258883, 258882]
 
 
-# unfused RoPE as addcmul_ RoPE diverges from reference code
+class RingKV(FixedKV):
+    # sliding-window ring: writes wrap at capacity, validity saturates
+    def prepare(self, num_tokens):
+        capacity = self.key.shape[2]
+        self.position.fill_(self.index % capacity)
+        self.seqlen.fill_(min(self.index + num_tokens, capacity))
+
+
+def _fixed_kv_decode_mask(mask, cache, min_val):
+    capacity = cache.key.shape[2]
+    valid = min(cache.index + 1, capacity)
+    output = mask.new_full((*mask.shape[:-1], capacity), min_val)
+    if isinstance(cache, RingKV):
+        positions = torch.arange(cache.index + 1 - valid, cache.index + 1, device=mask.device) % capacity
+        output.index_copy_(-1, positions, mask[..., -valid:])
+    else:
+        output[..., :valid] = mask[..., :valid]
+    return output
+
+
+# unfused RoPE as addcmul_ RoPE diverges from reference code (vision only; text
+# layers use the kitchen split-half kernel, bitwise-equal to this with bf16 freqs)
 def _apply_rotary_pos_emb(x, freqs_cis):
     cos, sin = freqs_cis[0], freqs_cis[1]
     half = x.shape[-1] // 2
@@ -140,6 +164,23 @@ class Gemma4Attention(nn.Module):
         if config.k_norm == "gemma3":
             self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps, device=device, dtype=dtype)
 
+    def _decode_attention(self, xq, cache, bias):
+        if bias is None:
+            # eager decode: slice the cache to the valid length (python-side index,
+            # no mask needed; a full ring is order-invariant under softmax)
+            n = min(cache.index + 1, cache.key.shape[2])
+            gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
+            attention = optimized_attention_for_device(xq.device, mask=False, small_input=True)
+            return attention(xq, cache.key[:, :, :n], cache.value[:, :, :n], self.num_heads, skip_reshape=True, scale=1.0, **gqa_kwargs)
+        # graph capture: fixed-length masked attention over the full capacity, explicit
+        # math (SDPA leaves its fast path on broadcast-bias + GQA and costs ~0.5ms/layer)
+        batch_size = xq.shape[0]
+        groups = self.num_heads // self.num_kv_heads
+        q = xq.reshape(batch_size, self.num_kv_heads, groups, self.head_dim)
+        scores = q @ cache.key.transpose(-1, -2) + bias
+        probs = torch.softmax(scores.float(), dim=-1).to(xq.dtype)
+        return (probs @ cache.value).reshape(batch_size, 1, self.inner_size)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -156,10 +197,16 @@ class Gemma4Attention(nn.Module):
         if self.q_norm is not None:
             xq = self.q_norm(xq)
 
+        if isinstance(shared_kv, FixedKV):
+            # decode on a KV-shared layer: attend the source layer's fixed cache
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
+            output = self._decode_attention(xq, shared_kv, attention_mask)
+            return self.o_proj(output), None, None
+
         if shared_kv is not None:
             xk, xv = shared_kv
             # Apply RoPE to Q only (K already has RoPE from source layer)
-            xq = _apply_rotary_pos_emb(xq, freqs_cis)
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
             present_key_value = None
             shareable_kv = None
         else:
@@ -173,11 +220,39 @@ class Gemma4Attention(nn.Module):
             xv = rms_norm(xv)
             xk = xk.transpose(1, 2)
             xv = xv.transpose(1, 2)
-            xq = _apply_rotary_pos_emb(xq, freqs_cis)
-            xk = _apply_rotary_pos_emb(xk, freqs_cis)
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
+            xk = comfy.quant_ops.ck.apply_rope_split_half1(xk, freqs_cis)
 
             present_key_value = None
-            if past_key_value is not None:
+            fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None
+            if fixed_cache is not None:
+                if seq_length == 1:
+                    # CUDA-graphable decode: write at the device-side ring/linear position
+                    fixed_cache.key.index_copy_(2, fixed_cache.position, xk)
+                    fixed_cache.value.index_copy_(2, fixed_cache.position, xv)
+                    output = self._decode_attention(xq, fixed_cache, attention_mask)
+                    return self.o_proj(output), fixed_cache, None
+
+                # prefill: attend the local sequence, persist the tail into the cache
+                capacity = fixed_cache.key.shape[2]
+                index = fixed_cache.index
+                if index + seq_length <= capacity:
+                    fixed_cache.key[:, :, index:index + seq_length] = xk
+                    fixed_cache.value[:, :, index:index + seq_length] = xv
+                    if index > 0:
+                        xk = fixed_cache.key[:, :, :index + seq_length]
+                        xv = fixed_cache.value[:, :, :index + seq_length]
+                elif index == 0:
+                    # prefill longer than the sliding ring: attend the full local K/V
+                    # (per-query windows come from the prefill sliding mask), cache only
+                    # the last `capacity` keys at their wrapped slots (position % capacity)
+                    slots = torch.arange(seq_length - capacity, seq_length, device=xk.device) % capacity
+                    fixed_cache.key.index_copy_(2, slots, xk[:, :, -capacity:])
+                    fixed_cache.value.index_copy_(2, slots, xv[:, :, -capacity:])
+                else:
+                    raise RuntimeError("gemma4: chunked prefill past the sliding window is not supported")
+                present_key_value = fixed_cache
+            elif past_key_value is not None:
                 cumulative_len = 0
                 if len(past_key_value) > 0:
                     past_key, past_value, cumulative_len = past_key_value
@@ -245,6 +320,7 @@ class TransformerBlockGemma4(nn.Module):
         self.register_buffer("layer_scalar", torch.empty(1, device=device, dtype=dtype))
 
     def forward(self, x, attention_mask=None, freqs_cis=None, past_key_value=None, per_layer_input=None, shared_kv=None):
+        output = x
         sliding_window = None
         if self.sliding_attention:
             sliding_window = self.sliding_attention
@@ -281,7 +357,8 @@ class TransformerBlockGemma4(nn.Module):
             x = self.post_per_layer_input_norm(x)
             x = residual + x
 
-        x = x * comfy.ops.cast_to_input(self.layer_scalar, x)
+        # in-place into the input buffer so CUDA-graph replays land in the static x
+        x = torch.mul(x, comfy.ops.cast_to_input(self.layer_scalar, x), out=output)
 
         return x, present_key_value, shareable_kv
 
@@ -290,6 +367,9 @@ class Gemma4Transformer(nn.Module):
     def __init__(self, config, device=None, dtype=None, ops=None):
         super().__init__()
         self.config = config
+        self.fixed_kv = True
+        self.prefetch_dynamic_vbars = True
+        self.graph_dynamic_vbar_blocks = True
 
         self.embed_tokens = _make_scaled_embedding(ops, config.vocab_size, config.hidden_size, config.hidden_size ** 0.5, device, dtype)
 
@@ -298,6 +378,19 @@ class Gemma4Transformer(nn.Module):
             for i in range(config.num_hidden_layers)
         ])
 
+        # KV-shared layers never run k_proj/v_proj/k_norm: their never-resolved vbar
+        # signatures would block layer graph capture, so prefetch only what executes
+        first_kv_shared = config.num_hidden_layers - config.num_kv_shared_layers if config.num_kv_shared_layers > 0 else config.num_hidden_layers
+        self._prefetch_units = []
+        for i, layer in enumerate(self.layers):
+            if i >= first_kv_shared:
+                dead = {layer.self_attn.k_proj, layer.self_attn.v_proj, layer.self_attn.k_norm}
+                self._prefetch_units.append([
+                    m for m in layer.modules() if next(m.children(), None) is None and m not in dead
+                ])
+            else:
+                self._prefetch_units.append(layer)
+
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, device=device, dtype=dtype) if config.final_norm else None
 
         # Precompute RoPE inv_freq on CPU to match reference code's exact value
@@ -311,6 +404,9 @@ class Gemma4Transformer(nn.Module):
         sliding_inv = 1.0 / (config.rope_theta[1] ** (torch.arange(0, config.head_dim, 2).float() / config.head_dim))
         self.register_buffer("_sliding_inv_freq", sliding_inv, persistent=False)
 
+        if config.suppress_tokens:
+            self.register_buffer("_suppress_tokens", torch.tensor(config.suppress_tokens, dtype=torch.long), persistent=False)
+
         # Per-layer input mechanism
         self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
         if self.hidden_size_per_layer_input:
@@ -322,19 +418,26 @@ class Gemma4Transformer(nn.Module):
                 self.hidden_size_per_layer_input, eps=config.rms_norm_eps,
                 device=device, dtype=dtype)
 
+    def get_dynamic_vram__units(self):
+        return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
+
     def get_past_len(self, past_key_values):
         for kv in past_key_values:
+            if isinstance(kv, FixedKV):
+                return kv.index
             if len(kv) >= 3:
                 return kv[2]
         return 0
 
     def _freqs_from_inv(self, inv_freq, position_ids, device, dtype):
-        """Compute cos/sin from stored inv_freq"""
+        """Compute per-pair 2x2 rotation matrices [B, 1, S, d/2, 2, 2] from stored inv_freq"""
         inv_exp = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(device)
         pos_exp = position_ids[:, None, :].float()
         freqs = (inv_exp @ pos_exp).transpose(1, 2)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        return emb.cos().unsqueeze(1).to(dtype), emb.sin().unsqueeze(1).to(dtype)
+        cos, sin = freqs.cos(), freqs.sin()
+        mat = torch.stack((torch.stack((cos, -sin), dim=-1),
+                           torch.stack((sin, cos), dim=-1)), dim=-2)
+        return mat.unsqueeze(1).to(dtype)
 
     def compute_freqs_cis(self, position_ids, device, dtype=None):
         global_freqs = self._freqs_from_inv(self._global_inv_freq, position_ids, device, dtype)
@@ -401,6 +504,71 @@ class Gemma4Transformer(nn.Module):
         first_kv_shared = self.config.num_hidden_layers - num_kv_shared if num_kv_shared > 0 else self.config.num_hidden_layers
         shared_sliding_kv = None  # KV from last non-shared sliding layer
         shared_global_kv = None   # KV from last non-shared global layer
+        share_source = {}
+        if num_kv_shared > 0:
+            for i in range(first_kv_shared):
+                share_source[bool(self.layers[i].sliding_attention)] = i
+
+        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(
+            list(self._prefetch_units), x.device,
+            {"prefetch_dynamic_vbars": self.prefetch_dynamic_vbars and past_key_values is not None})
+
+        fixed_kv = (past_key_values is not None and len(past_key_values) > 0
+                    and isinstance(past_key_values[0], FixedKV))
+        decode = fixed_kv and seq_len == 1
+        # mirror the conditions under which prefetch_queue_pop can actually capture, so
+        # eager fallbacks keep the sliced decode path instead of the full-capacity one
+        enable_graph = (decode and mask is None and self.graph_dynamic_vbar_blocks
+                        and prefetch_queue is not None
+                        and hasattr(self.layers[0], "_v_block")
+                        and not comfy.model_management.args.disable_cuda_graphs
+                        and comfy.model_management.is_device_cuda(x.device))
+        decode_bias = None
+        decode_masks = None
+        if decode:
+            prepared = set()
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV) and id(kv.position) not in prepared:
+                    kv.prepare(seq_len)
+                    prepared.add(id(kv.position))
+            if mask is not None:
+                decode_masks = {}
+                for kv in past_key_values:
+                    if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
+                        decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
+        if enable_graph:
+            # static buffers + per-capacity attention biases: layer graphs replay against
+            # stable storage, refreshed eagerly each step
+            capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
+            state_key = (x.shape, x.dtype, x.device, tuple(t.shape for t in freqs_cis), capacities,
+                         None if per_layer_inputs is None else per_layer_inputs.shape)
+            state = getattr(self, "_comfy_cross_step_state", None)
+            if state is None or state["key"] != state_key:
+                state = {"key": state_key,
+                         "x": torch.empty_like(x),
+                         "freqs_cis": [torch.empty_like(t) for t in freqs_cis],
+                         "bias": {c: torch.empty((1, 1, 1, c), dtype=x.dtype, device=x.device) for c in capacities},
+                         "per_layer": None if per_layer_inputs is None else torch.empty_like(per_layer_inputs),
+                         "bias_valid": -1}
+                self._comfy_cross_step_state = state
+                comfy.model_management._register_cross_step(self)
+            state["x"].copy_(x)
+            for source, target in zip(freqs_cis, state["freqs_cis"]):
+                target.copy_(source)
+            x = state["x"]
+            freqs_cis = state["freqs_cis"]
+            if per_layer_inputs is not None:
+                state["per_layer"].copy_(per_layer_inputs)
+                per_layer_inputs = state["per_layer"]
+            valid = past_len + 1
+            for capacity, bias in state["bias"].items():
+                if state["bias_valid"] != past_len:
+                    bias.fill_(min_val)
+                    bias[..., :min(valid, capacity)] = 0
+                elif past_len < capacity:
+                    bias[..., past_len:valid] = 0
+            state["bias_valid"] = valid
+            decode_bias = state["bias"]
 
         intermediate = None
         all_intermediate = None
@@ -429,12 +597,36 @@ class Gemma4Transformer(nn.Module):
 
             is_sliding = hasattr(layer, 'sliding_attention') and layer.sliding_attention
             if i >= first_kv_shared and num_kv_shared > 0:
-                shared = shared_sliding_kv if is_sliding else shared_global_kv
-                if shared is not None:
-                    layer_kwargs['shared_kv'] = shared
+                if decode:
+                    layer_kwargs['shared_kv'] = past_key_values[share_source[bool(is_sliding)]]
+                else:
+                    shared = shared_sliding_kv if is_sliding else shared_global_kv
+                    if shared is not None:
+                        layer_kwargs['shared_kv'] = shared
 
-            x, current_kv, shareable_kv = layer(x=x, attention_mask=mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
+            if enable_graph:
+                bias_cache = layer_kwargs.get('shared_kv', past_kv)
+                layer_mask = decode_bias[bias_cache.key.shape[2]]
+            elif decode:
+                bias_cache = layer_kwargs.get('shared_kv', past_kv)
+                layer_mask = None if decode_masks is None else decode_masks[id(bias_cache.position)]
+            else:
+                layer_mask = mask
 
+            result = []
+
+            def core():
+                nonlocal x
+                x, current_kv, shareable_kv = layer(x=x, attention_mask=layer_mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
+                result.append((current_kv, shareable_kv))
+
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
+
+            if result:
+                current_kv, shareable_kv = result[0]
+            else:
+                # graph replay: the cache already holds this step's write
+                current_kv, shareable_kv = past_kv, None
             next_key_values.append(current_kv if current_kv is not None else ())
 
             # Only track the last sliding/global before the sharing boundary
@@ -447,6 +639,14 @@ class Gemma4Transformer(nn.Module):
             if i == intermediate_output:
                 intermediate = x.clone()
 
+        if prefetch_queue is not None:
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
+
+        if fixed_kv:
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV):
+                    kv.advance(seq_len)
+
         if self.norm is not None:
             x = self.norm(x)
 
@@ -481,14 +681,37 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module):
         if cap:
             logits = cap * torch.tanh(logits / cap)
         if self.model.config.suppress_tokens:
-            logits[..., self.model.config.suppress_tokens] = torch.finfo(logits.dtype).min
+            logits.index_fill_(-1, self.model._suppress_tokens, torch.finfo(logits.dtype).min)
         return logits
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
-        past_key_values = []
-        for _ in range(self.model.config.num_hidden_layers):
-            past_key_values.append(())
-        return past_key_values
+        cfg = self.model.config
+        num_layers = cfg.num_hidden_layers
+        if not self.model.fixed_kv:
+            return [() for _ in range(num_layers)]
+        first_shared = num_layers - cfg.num_kv_shared_layers if cfg.num_kv_shared_layers > 0 else num_layers
+        # position/seqlen device tensors are shared per cache geometry and filled once per step
+        trackers = {}
+        caches = []
+        for i in range(num_layers):
+            if i >= first_shared:
+                caches.append(())
+                continue
+            sliding = cfg.sliding_attention[i % len(cfg.sliding_attention)] if cfg.sliding_attention else False
+            head_dim = cfg.head_dim if sliding else cfg.global_head_dim
+            k_eq_v = cfg.attention_k_eq_v and not sliding
+            kv_heads = cfg.num_global_key_value_heads if k_eq_v else cfg.num_key_value_heads
+            length = min(sliding, max_cache_len) if sliding else max_cache_len
+            cache_cls = RingKV if sliding else FixedKV
+            tracker = trackers.get((cache_cls, length))
+            if tracker is None:
+                tracker = (torch.empty((1,), device=device, dtype=torch.int64),
+                           torch.empty((batch,), device=device, dtype=torch.int32))
+                trackers[(cache_cls, length)] = tracker
+            # zero-init: decode attends full capacity with masked tails, 0*0 stays finite
+            key = torch.zeros((batch, kv_heads, length, head_dim), device=device, dtype=execution_dtype)
+            caches.append(cache_cls(key, torch.zeros_like(key), 0, tracker[0], tracker[1]))
+        return caches
 
     def preprocess_embed(self, embed, device):
         if embed["type"] == "image":
@@ -1402,11 +1625,13 @@ class Gemma4SDTokenizer(Gemma4_Tokenizer, sd1_clip.SDTokenizer):
 
     def decode(self, token_ids, **kwargs):
         text = super().decode(token_ids, skip_special_tokens=False)
-        # Translate thinking channel markers to standard <think>/</think> tags
+        # Only a close that ends a thought channel becomes </think>: generation primed with
+        # another channel leaves its opener in the prompt, so its close is not reasoning.
+        text = re.sub(r"<\|channel>thought\n(.*?)<channel\|>", r"<think>\n\1</think>", text, flags=re.DOTALL)
         text = text.replace("<|channel>thought\n", "<think>\n")
-        text = text.replace("<channel|>", "</think>")
         # Strip remaining special tokens
-        text = text.replace("<turn|>", "").replace("<eos>", "").strip()
+        text = re.sub(r"<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?|<turn\|>", "", text)
+        text = text.replace("<eos>", "").strip()
         return text
 
 
@@ -1497,6 +1722,9 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(device=device, dtype=dtype, name="gemma4", clip_model=clip_model, model_options=model_options)
+
+        def get_dynamic_vram__units(self):
+            return getattr(self, self.clip).transformer.model.get_dynamic_vram__units()
     return Gemma4TEModel_
 
 

comfy_api_nodes/apis/bytedance.py

@@ -116,6 +116,7 @@ class Seedance2TaskCreationRequest(BaseModel):
     seed: int | None = Field(None, ge=0, le=2147483647)
     watermark: bool | None = Field(None)
     output_format: str | None = Field(None)
+    omni_reference_task_type: str | None = Field(None, description="One of: auto, reference, edit, extend.")
 
 
 class TaskCreationResponse(BaseModel):

comfy_api_nodes/nodes_bytedance.py

@@ -56,6 +56,8 @@ from comfy_api_nodes.util import (
     ApiEndpoint,
     audio_bytes_to_audio_input,
     audio_input_to_mp3,
+    bytesio_to_image_tensor,
+    download_url_as_bytesio,
     download_url_to_image_tensor,
     download_url_to_video_output,
     downscale_image_tensor_by_max_side,
@@ -1315,7 +1317,9 @@ class ByteDanceSeedreamLayerSeparationNode(IO.ComfyNode):
             left, top, rect_w, rect_h = spec["left"], spec["top"], spec["rect_w"], spec["rect_h"]
             async with semaphore:
                 try:
-                    rgba = (await download_url_to_image_tensor(str(item["url"])))[0]
+                    # the layer math below needs the alpha channel, and ByteDance encodes
+                    # alpha-less images as plain RGB (the base plate is one), so force RGBA
+                    rgba = bytesio_to_image_tensor(await download_url_as_bytesio(str(item["url"])), mode="RGBA")[0]
                 except ProcessingInterrupted:
                     raise
                 except Exception as exc:
@@ -2069,7 +2073,7 @@ def _seedance2_text_inputs(resolutions: list[str], default_ratio: str = "16:9"):
     ]
 
 
-def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool = False):
+def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool = False, with_task_type: bool = False):
     return [
         IO.String.Input(
             "prompt",
@@ -2124,6 +2128,29 @@ def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool =
             if with_video_editing
             else []
         ),
+        *(
+            [
+                IO.Combo.Input(
+                    "task_type",
+                    options=["auto", "reference", "edit", "extend"],
+                    default="auto",
+                    tooltip="What to do with the reference media. Every value except auto is "
+                    "validated when the task is submitted, so mismatched settings fail before "
+                    "generation starts. auto: the model infers the task from the prompt and "
+                    "inputs, and settings that conflict with its reading fail only after "
+                    "generation has started. reference: generate a new video guided by the "
+                    "reference images, videos, and audio. edit: change a connected reference "
+                    "video (add, remove, replace); the output keeps the source clip's own length "
+                    "and aspect ratio, and the duration and ratio widgets are ignored. extend: "
+                    "continue a connected reference video forward or backward; the prompt should "
+                    "say 'extend forward', 'extend backward', or 'continue', the aspect ratio "
+                    "follows the source clip, and the output contains only the newly generated "
+                    "segment of the duration you set, not the source clip.",
+                )
+            ]
+            if with_task_type
+            else []
+        ),
         IO.Combo.Input(
             "output_format",
             options=["mp4"],
@@ -2133,9 +2160,9 @@ def _seedance25_text_inputs(with_ratio: bool = True, with_video_editing: bool =
     ]
 
 
-def _seedance25_reference_inputs():
+def _seedance25_reference_inputs(with_video_editing: bool = False, with_task_type: bool = False):
     return [
-        *_seedance25_text_inputs(with_video_editing=True),
+        *_seedance25_text_inputs(with_video_editing=with_video_editing, with_task_type=with_task_type),
         IO.Autogrow.Input(
             "reference_images",
             template=IO.Autogrow.TemplateNames(
@@ -2196,17 +2223,23 @@ def _seedance2_build_request(
     watermark: bool,
     ratio: str,
 ) -> Seedance2TaskCreationRequest:
-    video_editing = bool(model.get("video_editing"))
+    task_type = model.get("task_type", "auto")
+    duration = model["duration"]
+    if model.get("video_editing") or task_type == "edit":
+        ratio, duration = "adaptive", -1
+    elif task_type == "extend":
+        ratio = "adaptive"
     return Seedance2TaskCreationRequest(
         model=model_id,
         content=content,
         generate_audio=model["generate_audio"],
         resolution=model["resolution"],
-        ratio="adaptive" if video_editing else ratio,
-        duration=-1 if video_editing else model["duration"],
+        ratio=ratio,
+        duration=duration,
         seed=seed,
         watermark=watermark,
         output_format=model.get("output_format"),
+        omni_reference_task_type=None if task_type == "auto" else task_type,
     )
 
 
@@ -2216,7 +2249,7 @@ _SEEDANCE2_PRICE_EXPR_TEMPLATE = """
   $res := $lookup(widgets, "model.resolution");
   $ratio := $lookup(widgets, "model.ratio");
   $dur := $lookup(widgets, "model.duration");
-  $auto := $lookup(widgets, "model.video_editing") = true;
+  $auto := __IS_EDIT__;
   $hasVideo := __HAS_VIDEO__;
   $ready := $type($m) = "string" and $type($res) = "string" and ($auto or $type($dur) = "number");
   $ready ? (
@@ -2261,6 +2294,7 @@ _SEEDANCE2_PRICE_EXPR_TEMPLATE = """
 
 _SEEDANCE_AUDIO_POLICY_CODE = "OutputAudioSensitiveContentDetected.PolicyViolation"
 _SEEDANCE_TASK_TYPE_CONSTRAINT_CODE = "InvalidParameter.TaskTypeConstraint"
+_SEEDANCE_TASK_TYPE_MISMATCH_CODE = "InvalidParameter.TaskTypeMismatch"
 
 
 async def _seedance2_poll_video_task(
@@ -2269,6 +2303,7 @@ async def _seedance2_poll_video_task(
     model_id: str,
     resolution: str,
     has_video_input: bool,
+    task_type: str | None = None,
 ) -> TaskStatusResponse:
     try:
         return await poll_op(
@@ -2289,19 +2324,48 @@ async def _seedance2_poll_video_task(
                 "to get a silent video, or adjust the prompt and try again."
             ) from exc
         if _SEEDANCE_TASK_TYPE_CONSTRAINT_CODE in str(exc):
+            if task_type is None:
+                raise ValueError(
+                    "Seedance read this prompt as editing the reference video, and an edit always "
+                    "takes its duration and aspect ratio from that video. Enable video_editing on "
+                    "this node and run again, or reword the prompt so it describes a new video "
+                    "rather than a change to the reference one."
+                ) from exc
+            if task_type == "edit":
+                raise ValueError(
+                    "The request does not satisfy the 'edit' constraints: the clip being edited "
+                    "must be 4 to 30 seconds long."
+                ) from exc
+            if task_type == "extend":
+                raise ValueError(
+                    "The request does not satisfy the 'extend' constraints: the clip being "
+                    "extended must be 1.9 to 30 seconds long."
+                ) from exc
             raise ValueError(
-                "Seedance read this prompt as editing the reference video, and an edit always "
-                "takes its duration and aspect ratio from that video. Enable video_editing on "
-                "this node and run again, or reword the prompt so it describes a new video "
-                "rather than a change to the reference one."
+                "Seedance decided from the prompt that this task's duration or aspect ratio "
+                "must come from the reference video, and the current settings conflict with "
+                "that. Set task_type to the task you mean ('edit' or 'extend') and run again, "
+                "or reword the prompt so it describes a new video rather than a change to the "
+                "reference one."
+            ) from exc
+        if _SEEDANCE_TASK_TYPE_MISMATCH_CODE in str(exc):
+            raise ValueError(
+                f"Seedance read this prompt as a different task than the selected task_type "
+                f"'{task_type}'. Reword the prompt so it matches: an extend prompt should say "
+                "'extend forward', 'extend backward', or 'continue'; an edit prompt should use "
+                "words like add, remove, replace, or change. Or set task_type to auto."
             ) from exc
         raise
 
 
-def _seedance2_price_badge(with_reference_videos: bool) -> IO.PriceBadge:
+def _seedance2_price_badge(with_reference_videos: bool, legacy_video_editing: bool = False) -> IO.PriceBadge:
     widgets = ["model", "model.resolution", "model.ratio", "model.duration"]
+    if legacy_video_editing:
+        is_edit = '$lookup(widgets, "model.video_editing") = true'
+    else:
+        is_edit = '$lookup(widgets, "model.task_type") = "edit"'
     if with_reference_videos:
-        widgets.append("model.video_editing")
+        widgets.append("model.video_editing" if legacy_video_editing else "model.task_type")
     has_video = (
         '$exists(inputGroups) and $lookup(inputGroups, "model.reference_videos") > 0'
         if with_reference_videos
@@ -2312,7 +2376,7 @@ def _seedance2_price_badge(with_reference_videos: bool) -> IO.PriceBadge:
             widgets=widgets,
             input_groups=["model.reference_videos"] if with_reference_videos else [],
         ),
-        expr=_SEEDANCE2_PRICE_EXPR_TEMPLATE.replace("__HAS_VIDEO__", has_video),
+        expr=_SEEDANCE2_PRICE_EXPR_TEMPLATE.replace("__HAS_VIDEO__", has_video).replace("__IS_EDIT__", is_edit),
     )
 
 
@@ -2662,12 +2726,12 @@ def _seedance2_reference_inputs(resolutions: list[str], default_ratio: str = "16
     ]
 
 
-class ByteDance2ReferenceNode(IO.ComfyNode):
+class ByteDance2ReferenceNodeV2(IO.ComfyNode):
 
     @classmethod
     def define_schema(cls):
         return IO.Schema(
-            node_id="ByteDance2ReferenceNode",
+            node_id="ByteDance2ReferenceNodeV2",
             display_name="ByteDance Seedance 2.5 Reference to Video",
             category="partner/video/ByteDance",
             description="Generate, edit, or extend video using Seedance 2.5 or 2.0 with reference "
@@ -2676,7 +2740,7 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
-                        IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs()),
+                        IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs(with_task_type=True)),
                         IO.DynamicCombo.Option(
                             "Seedance 2.0",
                             _seedance2_reference_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
@@ -2761,6 +2825,13 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
                 f"(videos={len(reference_videos)}, video assets={len(reference_video_assets)}). "
                 f"Maximum is {limits['max_videos']}."
             )
+        task_type = model.get("task_type")
+        if task_type in ("edit", "extend") and total_videos == 0:
+            raise ValueError(
+                f"A '{task_type}' task needs at least one reference video. Connect the video "
+                f"you want to {'change' if task_type == 'edit' else 'continue'}, or set "
+                "task_type to 'reference' to generate a new video from the references you have."
+            )
         total_audios = len(reference_audios) + len(reference_audio_assets)
         if total_audios > limits["max_audios"]:
             raise ValueError(
@@ -2893,11 +2964,78 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
             response_model=TaskCreationResponse,
         )
         response = await _seedance2_poll_video_task(
-            cls, initial_response.id, model_id, model["resolution"], has_video_input=has_video_input
+            cls,
+            initial_response.id,
+            model_id,
+            model["resolution"],
+            has_video_input=has_video_input,
+            task_type=task_type,
         )
         return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
 
 
+class ByteDance2ReferenceNode(ByteDance2ReferenceNodeV2):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="ByteDance2ReferenceNode",
+            display_name="ByteDance Seedance 2.5 Reference to Video (Legacy)",
+            category="partner/video/ByteDance",
+            description="Generate, edit, or extend video using Seedance 2.5 or 2.0 with reference "
+            "images, videos, and audio. Supports multimodal reference, video editing, and video extension.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option("Seedance 2.5", _seedance25_reference_inputs(with_video_editing=True)),
+                        IO.DynamicCombo.Option(
+                            "Seedance 2.0",
+                            _seedance2_reference_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
+                        ),
+                        IO.DynamicCombo.Option(
+                            "Seedance 2.0 Fast",
+                            _seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
+                        ),
+                        IO.DynamicCombo.Option(
+                            "Seedance 2.0 Mini",
+                            _seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
+                        ),
+                    ],
+                    tooltip=SEEDANCE_MODEL_TOOLTIP,
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed controls whether the node should re-run; "
+                    "results are non-deterministic regardless of seed.",
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add a watermark to the video.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Video.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            is_deprecated=True,
+            price_badge=_seedance2_price_badge(with_reference_videos=True, legacy_video_editing=True),
+        )
+
+
 async def process_video_task(
     cls: type[IO.ComfyNode],
     payload: Text2VideoTaskCreationRequest | Image2VideoTaskCreationRequest,
@@ -3413,6 +3551,7 @@ class ByteDanceExtension(ComfyExtension):
             ByteDance2TextToVideoNode,
             ByteDance2FirstLastFrameNode,
             ByteDance2ReferenceNode,
+            ByteDance2ReferenceNodeV2,
             ByteDanceCreateImageAsset,
             ByteDanceCreateVideoAsset,
             ByteDanceSeedAudioNode,

comfy_api_nodes/nodes_gemini.py

@@ -43,6 +43,7 @@ from comfy_api_nodes.util import (
     download_url_to_image_tensor,
     download_url_to_video_output,
     get_number_of_images,
+    pad_images_to_common_channels,
     sync_op,
     tensor_to_base64_string,
     upload_audio_to_comfyapi,
@@ -233,8 +234,8 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
                 "Try rephrasing your prompt or changing the response modality to 'IMAGE+TEXT' "
                 "to see the model's reasoning."
             )
-        return torch.zeros((1, 1024, 1024, 4))
-    return torch.cat(image_tensors, dim=0)
+        return torch.zeros((1, 1024, 1024, 3))
+    return torch.cat(pad_images_to_common_channels(image_tensors), dim=0)
 
 
 def get_text_from_interaction(interaction: GeminiInteraction) -> str:

comfy_api_nodes/nodes_recraft.py

@@ -27,6 +27,7 @@ from comfy_api_nodes.util import (
     ApiEndpoint,
     bytesio_to_image_tensor,
     download_url_as_bytesio,
+    pad_images_to_common_channels,
     resize_mask_to_image,
     sync_op,
     tensor_to_bytesio,
@@ -621,7 +622,7 @@ class RecraftImageToImageNode(IO.ComfyNode):
                 images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftImageInpaintingNode(IO.ComfyNode):
@@ -723,7 +724,7 @@ class RecraftImageInpaintingNode(IO.ComfyNode):
                 images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftTextToVectorNode(IO.ComfyNode):
@@ -954,7 +955,7 @@ class RecraftReplaceBackgroundNode(IO.ComfyNode):
             images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftRemoveBackgroundNode(IO.ComfyNode):
@@ -995,7 +996,7 @@ class RecraftRemoveBackgroundNode(IO.ComfyNode):
                 image=image[i],
                 path="/proxy/recraft/images/removeBackground",
             )
-            images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
+            images.append(torch.cat([bytesio_to_image_tensor(x, mode="RGBA") for x in sub_bytes], dim=0))
             pbar.update(1)
 
         images_tensor = torch.cat(images, dim=0)
@@ -1047,7 +1048,7 @@ class RecraftCrispUpscaleNode(IO.ComfyNode):
             images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode):

comfy_api_nodes/util/__init__.py

@@ -18,6 +18,7 @@ from .conversions import (
     downscale_image_tensor_by_max_side,
     downscale_video_to_max_pixels,
     image_tensor_pair_to_batch,
+    pad_images_to_common_channels,
     pil_to_bytesio,
     resize_mask_to_image,
     tensor_to_base64_string,
@@ -92,6 +93,7 @@ __all__ = [
     "downscale_image_tensor_by_max_side",
     "downscale_video_to_max_pixels",
     "image_tensor_pair_to_batch",
+    "pad_images_to_common_channels",
     "pil_to_bytesio",
     "resize_mask_to_image",
     "tensor_to_base64_string",

comfy_api_nodes/util/conversions.py

@@ -16,12 +16,14 @@ from comfy_api.latest import Input, InputImpl, Types
 from ._helpers import mimetype_to_extension
 
 
-def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor:
+def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str | None = None) -> torch.Tensor:
     """Converts image data from BytesIO to a torch.Tensor.
 
     Args:
         image_bytesio: BytesIO object containing the image data.
-        mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA").
+        mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). Defaults
+            to RGBA when the decoded image carries transparency and RGB when it
+            does not, so an API that returns no alpha does not get an opaque one.
 
     Returns:
         A torch.Tensor representing the image (1, H, W, C).
@@ -31,6 +33,8 @@ def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch
         ValueError: If the specified mode is invalid.
     """
     image = Image.open(image_bytesio)
+    if mode is None:
+        mode = "RGBA" if "A" in image.getbands() or "transparency" in image.info else "RGB"
     image = image.convert(mode)
     image_array = np.array(image).astype(np.float32) / 255.0
     return torch.from_numpy(image_array).unsqueeze(0)
@@ -53,6 +57,17 @@ def image_tensor_pair_to_batch(image1: torch.Tensor, image2: torch.Tensor) -> to
     return torch.cat((image1, image2), dim=0)
 
 
+def pad_images_to_common_channels(images: list[torch.Tensor]) -> list[torch.Tensor]:
+    """Pads [B, H, W, C] image tensors with opaque alpha so they all share the largest channel count."""
+    channels = max(image.shape[-1] for image in images)
+    return [
+        torch.nn.functional.pad(image, (0, channels - image.shape[-1]), value=1.0)
+        if image.shape[-1] < channels
+        else image
+        for image in images
+    ]
+
+
 def tensor_to_bytesio(
     image: torch.Tensor,
     *,

comfy_extras/nodes_textgen.py

@@ -256,12 +256,10 @@ class TextGenerateLTX2Prompt(TextGenerate):
 
         out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
 
-        text = out.args[0]
-        text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
-        if "</think>" in text:  # unclosed/truncated reasoning: keep what follows the last close
-            text = text.rsplit("</think>", 1)[-1]
-        text = re.sub(r"</?think>|<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?", "", text).strip()
-        return io.NodeOutput(text)
+        # Drop reasoning, including a block left unclosed by max_length. Both system prompts ask
+        # for the original prompt back when there is nothing to give; empty conditions on nothing.
+        text = re.sub(r"<think>.*?(?:</think>|$)", "", out.args[0], flags=re.DOTALL).strip()
+        return io.NodeOutput(text or prompt)
 
 
 class TextgenExtension(ComfyExtension):

main.py

@@ -248,7 +248,17 @@ import hook_breaker_ac10a0
 import comfy.memory_management
 import comfy.model_patcher
 
-if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia()):
+
+def dynamic_vram_supported():
+    if comfy.model_management.is_nvidia():
+        return True
+    if comfy.model_management.is_amd():
+        if comfy.model_management.rocm_version >= (7, 14):
+            return True
+    return False
+
+
+if args.enable_dynamic_vram or (enables_dynamic_vram() and dynamic_vram_supported()):
     if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
         logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
     else:

openapi.yaml

@@ -35,6 +35,10 @@ components:
                     description: Timestamp when the asset was last accessed
                     format: date-time
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 metadata:
                     additionalProperties: true
                     description: System-managed metadata from download sources (HuggingFace, CivitAI, etc.) - read-only, not user-modifiable
@@ -165,6 +169,10 @@ components:
                     format: uuid
                     nullable: true
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 mime_type:
                     description: Updated MIME type of the asset
                     type: string
@@ -188,6 +196,31 @@ components:
                 - id
                 - updated_at
             type: object
+        ChurnkeyAuthResponse:
+            description: |
+                Credentials the Churnkey embed requires to launch the cancel flow.
+                `auth_hash` is hex-encoded HMAC-SHA256 of `customer_id` signed with the
+                server's CHURNKEY_HMAC_SECRET; it is bound to that single customer ID
+                and must not be reused for other customers.
+            properties:
+                auth_hash:
+                    description: Hex-encoded HMAC-SHA256(customer_id, CHURNKEY_HMAC_SECRET)
+                    type: string
+                customer_id:
+                    description: Stripe customer ID for the workspace
+                    type: string
+                mode:
+                    description: Churnkey environment matching the configured app
+                    enum:
+                        - live
+                        - test
+                        - sandbox
+                    type: string
+            required:
+                - customer_id
+                - auth_hash
+                - mode
+            type: object
         CreateWorkflowRequest:
             description: Request body for creating a new saved workflow.
             properties:
@@ -511,6 +544,25 @@ components:
             required:
                 - history
             type: object
+        JobAssetsResponse:
+            description: Paginated list of the assets produced by a single job.
+            properties:
+                assets:
+                    description: The job's output assets for the requested page (empty when the job produced none)
+                    items:
+                        $ref: '#/components/schemas/JobOutputAsset'
+                    type: array
+                job_id:
+                    description: ID of the job these assets belong to
+                    format: uuid
+                    type: string
+                pagination:
+                    $ref: '#/components/schemas/PaginationInfo'
+            required:
+                - job_id
+                - assets
+                - pagination
+            type: object
         JobCancelResponse:
             description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
             properties:
@@ -565,6 +617,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -597,6 +652,13 @@ components:
                 workflow_id:
                     description: UUID identifying the workflow graph definition
                     type: string
+                workflow_version_id:
+                    description: |
+                        UUID of the cloud workflow version this job is pinned to, if the
+                        submission carried one (see PromptRequest's workflow_version_id).
+                        Absent for jobs submitted without that association, including
+                        every job submitted through the public API v2 today.
+                    type: string
                 workspace_id:
                     description: |
                         ID of the workspace that owns this job. A successful (200)
@@ -645,6 +707,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only present for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -662,6 +727,56 @@ components:
                 - status
                 - create_time
             type: object
+        JobOutputAsset:
+            description: |
+                An asset produced by a job, enriched with the per-output node context
+                (`node_id`, `output_key`, `output_index`) correlated from the job's
+                execution outputs by content hash. The node-context fields are null
+                when the asset cannot be matched to an output entry.
+            properties:
+                created_at:
+                    description: Timestamp when the asset was created
+                    format: date-time
+                    type: string
+                hash:
+                    description: Blake3 hash of the asset content.
+                    pattern: ^blake3:[a-f0-9]{64}$
+                    type: string
+                id:
+                    description: Unique identifier for the asset
+                    format: uuid
+                    type: string
+                mime_type:
+                    description: MIME type of the asset
+                    type: string
+                name:
+                    description: Name of the asset file
+                    type: string
+                node_id:
+                    description: ID of the workflow node that produced this asset, if known
+                    nullable: true
+                    type: string
+                output_index:
+                    description: Zero-based index of this asset within the node's output slot, if known
+                    nullable: true
+                    type: integer
+                output_key:
+                    description: Output slot key under the producing node (e.g. "images"), if known
+                    nullable: true
+                    type: string
+                preview_url:
+                    description: Relative URL for asset preview/thumbnail
+                    format: uri-reference
+                    type: string
+                size:
+                    description: Size of the asset in bytes
+                    format: int64
+                    type: integer
+            required:
+                - id
+                - name
+                - created_at
+            type: object
         JobStatusResponse:
             description: Job status information
             properties:
@@ -1522,7 +1637,11 @@ paths:
             operationId: listAssets
             parameters:
                 - deprecated: true
-                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. Combining it
+                    with `tags_all`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: include_tags
@@ -1532,7 +1651,11 @@ paths:
                     type: array
                   style: form
                 - deprecated: true
-                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. Combining it
+                    with `tags_none`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: exclude_tags
@@ -1541,7 +1664,14 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have ALL of these tags
+                - description: |
+                    Filter assets that have ALL of these tags. Tag values are opaque
+                    byte-strings compared exactly and case-sensitively; unknown tags
+                    are not an error — they simply match nothing. Replaces the
+                    deprecated `include_tags`. Sending both spellings, listing the
+                    same tag here and in `tags_none`, or exceeding 100 tags per list
+                    (counted after removing empty values and duplicates) returns 400
+                    `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_all
@@ -1550,7 +1680,14 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have AT LEAST ONE of these tags
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins;
+                    overlap with `tags_none` is allowed and leaves a dead term).
+                    Supplying a positive tag filter (`tags_any`, `tags_all`, or
+                    `include_tags`) replaces the default category filter that is
+                    otherwise applied. Lists over 100 tags (counted after removing
+                    empty values and duplicates) return 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_any
@@ -1559,7 +1696,11 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the
+                    deprecated `exclude_tags`. Sending both spellings, or exceeding
+                    100 tags per list (counted after removing empty values and
+                    duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_none
@@ -2342,7 +2483,11 @@ paths:
             operationId: getAssetTagHistogram
             parameters:
                 - deprecated: true
-                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: include_tags
@@ -2352,7 +2497,11 @@ paths:
                     type: array
                   style: form
                 - deprecated: true
-                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: exclude_tags
@@ -2361,7 +2510,10 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have ALL of these tags
+                - description: |
+                    Filter assets that have ALL of these tags. Replaces the deprecated
+                    `include_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_all
@@ -2370,7 +2522,11 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have AT LEAST ONE of these tags
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins).
+                    The same combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_any
@@ -2379,7 +2535,10 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the deprecated
+                    `exclude_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_none
@@ -2440,6 +2599,49 @@ paths:
             summary: Get tag histogram for filtered assets
             tags:
                 - file
+    /api/billing/churnkey/auth:
+        get:
+            description: |
+                Returns the Stripe customer identifier and a server-signed
+                HMAC-SHA256 of the customer ID, used to launch the Churnkey-hosted
+                cancellation flow embed.
+            operationId: getChurnkeyAuth
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ChurnkeyAuthResponse'
+                    description: Success
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Workspace has no Stripe customer (never subscribed)
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+                "503":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Churnkey is not configured on the server
+            security:
+                - BearerAuth: []
+            summary: Get Churnkey HMAC auth credentials
+            tags:
+                - billing
     /api/embeddings:
         get:
             description: Returns the list of text-encoder embeddings available on disk.
@@ -2460,9 +2662,10 @@ paths:
                 Returns a list of model folders available in the system.
                 This is an experimental endpoint that replaces the legacy /models endpoint.
                 Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
-                Once the model_type migration is active the names are model_type folder_names
-                (e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
-                its directory path.
+                The folder vocabulary is resolved per request from the caller's identity: where the
+                model_type migration is active for that caller the names are model_type folder_names
+                (e.g. `ultralytics_bbox`), and a folder with no folder_name mapping is returned by its
+                directory path. An authenticated response can therefore differ from an anonymous one.
             operationId: getModelFolders
             responses:
                 "200":
@@ -2479,7 +2682,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get available model folders
             tags:
                 - file
@@ -2488,6 +2694,10 @@ paths:
             description: |
                 Returns a list of models available in the specified folder.
                 This is an experimental endpoint that provides enhanced model information.
+                Accepted folder identifiers are those returned by /api/experiment/models for the same
+                caller. That vocabulary is request-scoped, so list folders and fetch a folder's models
+                with the same credentials — a name obtained anonymously may not resolve when
+                authenticated, and vice versa.
             operationId: getModelsInFolder
             parameters:
                 - description: The folder name to list models from
@@ -2518,7 +2728,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get models in a specific folder
             tags:
                 - file
@@ -3155,6 +3368,74 @@ paths:
             summary: Get full job details
             tags:
                 - workflow
+    /api/jobs/{job_id}/assets:
+        get:
+            description: |
+                Retrieve a paginated list of the assets produced by a specific job,
+                enriched with the per-output node context (`node_id`, `output_key`,
+                `output_index`) correlated from the job's execution outputs by content
+                hash. Unlike `GET /api/assets?job_ids={id}`, this endpoint is scoped to a
+                single job and carries node-level placement, making it suited to job
+                output views rather than the general asset browser. Returns an empty
+                `assets` array for jobs that produced no assets.
+            operationId: getJobAssets
+            parameters:
+                - description: Job identifier (UUID)
+                  in: path
+                  name: job_id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+                - description: Maximum number of assets to return (1-500)
+                  in: query
+                  name: limit
+                  schema:
+                    default: 20
+                    maximum: 500
+                    minimum: 1
+                    type: integer
+                - description: Number of assets to skip for pagination
+                  in: query
+                  name: offset
+                  schema:
+                    default: 0
+                    minimum: 0
+                    type: integer
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/JobAssetsResponse'
+                    description: Success - Job assets returned
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Invalid request parameters
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized - Authentication required
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Job not found or does not belong to the user
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+            summary: List a job's output assets
+            tags:
+                - workflow
     /api/jobs/{job_id}/cancel:
         post:
             description: |
@@ -3360,6 +3641,12 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
                     description: Payment required - Insufficient credits
+                "403":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/PromptErrorResponse'
+                    description: Workspace governance policy blocks one or more partner providers (error.type PARTNER_NODE_DISABLED; error.class_types lists the offending nodes, error.providers the disabled providers)
                 "413":
                     content:
                         application/json:
@@ -3371,7 +3658,7 @@ paths:
                         application/json:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
-                    description: Payment required - User has not paid
+                    description: 'Retryable backpressure. Two distinct causes, disambiguated by the body''s `error.type`, NOT by parsing `error.message`: `PAYMENT_REQUIRED` / `FREE_TIER_UNAVAILABLE` / `FREE_TIER_EXHAUSTED` / `PARTNER_NODE_PAYMENT_REQUIRED` (a billing gate - retrying without paying never succeeds), or `QUEUE_LIMIT` (this workspace''s bounded job queue is full - retrying after some queued jobs complete will succeed).'
                 "500":
                     content:
                         application/json:
@@ -5210,6 +5497,8 @@ tags:
       name: user
     - description: Background task management
       name: task
+    - description: Workspace billing and subscription management
+      name: billing
     - description: Workflow storage and version management
       name: workflows
     - description: Job queue state and control

requirements.txt

@@ -1,6 +1,6 @@
-comfyui-frontend-package==1.48.7
+comfyui-frontend-package==1.49.6
 comfyui-workflow-templates==0.11.41
-comfyui-embedded-docs==0.5.9
+comfyui-embedded-docs==0.5.10
 torch
 torchsde
 torchvision

tests-unit/comfy_api_nodes_test/gemini_response_images_test.py

@@ -0,0 +1,57 @@
+import asyncio
+import base64
+from io import BytesIO
+
+import torch
+from PIL import Image
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+from comfy_api_nodes.apis.gemini import (  # noqa: E402
+    GeminiCandidate,
+    GeminiContent,
+    GeminiGenerateContentResponse,
+    GeminiInlineData,
+    GeminiPart,
+)
+from comfy_api_nodes.nodes_gemini import get_image_from_response  # noqa: E402
+
+
+def image_part(mode, color):
+    buffer = BytesIO()
+    Image.new(mode, (4, 4), color).save(buffer, format="PNG")
+    return GeminiPart(
+        inlineData=GeminiInlineData(
+            data=base64.b64encode(buffer.getvalue()).decode(),
+            mimeType="image/png",
+        )
+    )
+
+
+def response(*parts):
+    return GeminiGenerateContentResponse(
+        candidates=[GeminiCandidate(content=GeminiContent(parts=list(parts), role="model"))]
+    )
+
+
+def test_rgb_only_response_stays_three_channels():
+    out = asyncio.run(get_image_from_response(response(image_part("RGB", (10, 20, 30)))))
+    assert out.shape == (1, 4, 4, 3)
+
+
+def test_mixed_rgb_and_rgba_parts_are_padded_to_the_same_width():
+    out = asyncio.run(
+        get_image_from_response(
+            response(
+                image_part("RGB", (10, 20, 30)),
+                image_part("RGBA", (10, 20, 30, 0)),
+            )
+        )
+    )
+    assert out.shape == (2, 4, 4, 4)
+    # the part that had no alpha is padded opaque, the transparent one is preserved
+    assert out[0, ..., 3].min() == 1.0
+    assert out[1, ..., 3].max() == 0.0

tests-unit/comfy_api_nodes_test/image_conversions_test.py

@@ -0,0 +1,80 @@
+from io import BytesIO
+
+import pytest
+import torch
+from PIL import Image
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+from comfy_api_nodes.util.conversions import bytesio_to_image_tensor, pad_images_to_common_channels  # noqa: E402
+
+
+def encode(image: Image.Image, image_format: str = "PNG") -> BytesIO:
+    buffer = BytesIO()
+    image.save(buffer, format=image_format)
+    buffer.seek(0)
+    return buffer
+
+
+def test_rgb_png_stays_three_channels():
+    tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30))))
+    assert tensor.shape == (1, 4, 4, 3)
+
+
+def test_jpeg_stays_three_channels():
+    tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30)), "JPEG"))
+    assert tensor.shape == (1, 4, 4, 3)
+
+
+def test_grayscale_is_expanded_to_rgb():
+    tensor = bytesio_to_image_tensor(encode(Image.new("L", (4, 4), 128)))
+    assert tensor.shape == (1, 4, 4, 3)
+
+
+def test_rgba_png_keeps_its_alpha():
+    tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 0))))
+    assert tensor.shape == (1, 4, 4, 4)
+    assert tensor[..., 3].max() == 0.0
+
+
+def test_palette_png_with_transparency_keeps_its_alpha():
+    image = Image.new("P", (4, 4), 1)
+    image.putpalette([0, 0, 0, 255, 255, 255])
+    image.info["transparency"] = 0
+    image.putpixel((0, 0), 0)
+    tensor = bytesio_to_image_tensor(encode(image))
+    assert tensor.shape == (1, 4, 4, 4)
+    assert tensor[0, 0, 0, 3] == 0.0
+    assert tensor[0, 1, 1, 3] == 1.0
+
+
+@pytest.mark.parametrize("mode,channels", [("RGB", 3), ("RGBA", 4)])
+def test_explicit_mode_is_respected(mode, channels):
+    tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 128))), mode=mode)
+    assert tensor.shape == (1, 4, 4, channels)
+
+
+def test_pad_mixed_channels_concatenates():
+    rgb = torch.rand(1, 4, 4, 3)
+    rgba = torch.rand(2, 4, 4, 4)
+    padded = pad_images_to_common_channels([rgb, rgba])
+    result = torch.cat(padded, dim=0)
+    assert result.shape == (3, 4, 4, 4)
+
+
+def test_pad_adds_opaque_alpha_and_keeps_rgb_values():
+    rgb = torch.rand(1, 4, 4, 3)
+    rgba = torch.rand(1, 4, 4, 4)
+    padded_rgb, padded_rgba = pad_images_to_common_channels([rgb, rgba])
+    assert torch.equal(padded_rgb[..., :3], rgb)
+    assert padded_rgb[..., 3].min() == 1.0
+    assert padded_rgba is rgba
+
+
+def test_pad_leaves_homogeneous_channels_unchanged():
+    images = [torch.rand(1, 4, 4, 3), torch.rand(2, 4, 4, 3)]
+    padded = pad_images_to_common_channels(images)
+    assert all(p is i for p, i in zip(padded, images))