Merge branch 'master' into automation/comfyui-frontend-bump

comfy/audio_encoders/mert2.py

@@ -163,8 +163,9 @@ class MERT2(nn.Module):
 
     def position_embeddings(self, x):
         inverse = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2, device=x.device, dtype=torch.float32) / self.head_dim))
-        angles = torch.arange(x.shape[1], device=x.device, dtype=torch.float32)[:, None] * inverse
-        cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
+        # Match the released encoder's autocast outer product before sin/cos.
+        angles = torch.arange(x.shape[1], device=x.device, dtype=torch.float32).to(x.dtype)[:, None] * inverse.to(x.dtype)
+        cos, sin = angles.cos(), angles.sin()
         return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(1, 1, x.shape[1], -1, 2, 2).float()
 
     def forward(self, mel, layer_weight, output_hidden_states=False):

comfy/model_prefetch.py

@@ -46,6 +46,22 @@ class _PauseMallocGraph:
 def pause_malloc_graph(sync=False):
     return _PauseMallocGraph(sync)
 
+class _MallocGraphScope:
+    def __init__(self, device):
+        self.device = device
+
+    def __enter__(self):
+        malloc_graph_begin(self.device)
+
+    def __exit__(self, exc_type, *args):
+        if exc_type is None:
+            malloc_graph_end()
+        else:
+            cleanup_malloc_graph()
+
+def malloc_graph_scope(device):
+    return _MallocGraphScope(device)
+
 def malloc_graph_begin(device):
     global MALLOC_GRAPH_USED
     if not malloc_graph_enabled(device):
@@ -86,6 +102,24 @@ def cleanup_malloc_graph():
         MALLOC_GRAPH_ROGUES += graph.rogue_count
         del graph
 
+def pin_modules(comfy_modules, device, dtype=None):
+    registerable_size = 0
+    for s in comfy_modules:
+        registerable_size += comfy.memory_management.vram_aligned_size([s.weight, s.bias])
+        for param_key in ("weight", "bias"):
+            lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
+            if lowvram_fn is not None:
+                registerable_size += lowvram_fn.memory_required()
+
+    offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True)
+    if not (comfy_modules and comfy_modules[0]._pin_state["fast_disk"]):
+        comfy.model_management.ensure_pin_registerable(registerable_size)
+    comfy.model_management.sync_stream(device, offload_stream)
+    if fully_faulted and dtype is not None:
+        for comfy_module in comfy_modules:
+            comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False)
+    return offload_stream, fully_faulted
+
 def cleanup_prefetched_modules(module, comfy_modules):
     for s in comfy_modules:
         prefetch = getattr(s, "_prefetch", None)
@@ -201,21 +235,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
                 if hasattr(s, "_v"):
                     comfy_modules.append(s)
 
-        registerable_size = 0
-        for s in comfy_modules:
-            registerable_size += comfy.memory_management.vram_aligned_size([s.weight, s.bias])
-            for param_key in ("weight", "bias"):
-                lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
-                if lowvram_fn is not None:
-                    registerable_size += lowvram_fn.memory_required()
-
-        offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True)
-        if not (comfy_modules and comfy_modules[0]._pin_state["fast_disk"]):
-            comfy.model_management.ensure_pin_registerable(registerable_size)
-        comfy.model_management.sync_stream(device, offload_stream)
-        if fully_faulted and dtype is not None:
-            for comfy_module in comfy_modules:
-                comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False)
+        offload_stream, fully_faulted = pin_modules(comfy_modules, device, dtype)
         queue[0] = (offload_stream, (module, comfy_modules))
 
     if core is not None:

comfy/sd.py

@@ -469,7 +469,7 @@ class CLIP:
     def get_key_patches(self):
         return self.patcher.get_key_patches()
 
-    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None, presence_penalty=0.0):
+    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None, presence_penalty=0.0, mtp=True):
         self.cond_stage_model.reset_clip_options()
 
         self.load_model(tokens)
@@ -478,7 +478,7 @@ class CLIP:
         self.cond_stage_model.set_clip_options({"execution_device": device})
 
         with model_management.cuda_device_context(device), comfy.ops.use_quantized_matmul(self.cond_stage_model, device):
-            return self.cond_stage_model.generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed, presence_penalty=presence_penalty)
+            return self.cond_stage_model.generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed, presence_penalty=presence_penalty, mtp=mtp)
 
     def decode(self, token_ids, skip_special_tokens=True):
         return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
@@ -1910,7 +1910,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
         elif te_model in (TEModel.QWEN35_08B, TEModel.QWEN35_2B, TEModel.QWEN35_4B, TEModel.QWEN35_9B, TEModel.QWEN35_27B):
             clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
             qwen35_type = {TEModel.QWEN35_08B: "qwen35_08b", TEModel.QWEN35_2B: "qwen35_2b", TEModel.QWEN35_4B: "qwen35_4b", TEModel.QWEN35_9B: "qwen35_9b", TEModel.QWEN35_27B: "qwen35_27b"}[te_model]
-            clip_target.clip = comfy.text_encoders.qwen35.te(**llama_detect(clip_data), model_type=qwen35_type)
+            clip_target.clip = comfy.text_encoders.qwen35.te(**llama_detect(clip_data), model_type=qwen35_type, mtp="mtp.fc.weight" in clip_data[0])
             clip_target.tokenizer = comfy.text_encoders.qwen35.tokenizer(model_type=qwen35_type)
         elif te_model in (TEModel.QWEN3VL_4B, TEModel.QWEN3VL_8B):
             if clip_type == CLIPType.IDEOGRAM4 and te_model == TEModel.QWEN3VL_8B:  # Ideogram4 reuses the full Qwen3-VL-8B (13-layer tap for conditioning + multimodal generate).

comfy/sd1_clip.py

@@ -308,7 +308,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
     def load_sd(self, sd):
         return self.transformer.load_state_dict(sd, strict=False, assign=getattr(self, "can_assign_sd", False))
 
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0, mtp=True):
         if isinstance(tokens, dict):
             tokens_only = next(iter(tokens.values())) # todo: get this better?
         else:
@@ -746,5 +746,11 @@ class SD1ClipModel(torch.nn.Module):
     def load_sd(self, sd):
         return getattr(self, self.clip).load_sd(sd)
 
-    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None, presence_penalty=0.0):
-        return getattr(self, self.clip).generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed, presence_penalty=presence_penalty)
+    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None, presence_penalty=0.0, mtp=True):
+        return getattr(self, self.clip).generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed, presence_penalty=presence_penalty, mtp=mtp)
+
+    def get_dynamic_vram__units(self):
+        # forward to the inner transformer so ModelPatcher can register vbar units (graph decode)
+        model = getattr(getattr(getattr(self, self.clip), "transformer", None), "model", None)
+        get_units = getattr(model, "get_dynamic_vram__units", None)
+        return get_units() if get_units is not None else ([], [])

comfy/text_encoders/gemma4.py

@@ -1664,7 +1664,7 @@ class Gemma4Model(sd1_clip.SDClipModel):
         self.dtypes.add(dtype)
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=self.model_class, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0, mtp=True):
         if isinstance(tokens, dict):
             tokens = next(iter(tokens.values()))
         tokens_only = [[t[0] for t in b] for b in tokens]
@@ -1708,9 +1708,6 @@ 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/text_encoders/llama.py

@@ -32,6 +32,72 @@ class FixedKV:
     def advance(self, num_tokens):
         self.index += num_tokens
 
+    def rollback(self, discard=1):
+        # drop the rejected verify tail; stale slots and bias are overwritten by the next write/prepare
+        self.index -= discard
+
+@dataclass
+class FixedKVBias(FixedKV):
+    # full-capacity decode bias [1, 1, rows, capacity], last `seq` rows serve the queries; shared across layers
+    bias: torch.Tensor = None
+    tracker: dict = None
+
+    def prepare(self, num_tokens):
+        if self.tracker["step"] == (self.index, num_tokens):
+            return
+        self.tracker["step"] = (self.index, num_tokens)
+        i = self.index
+        rows = self.bias.shape[-2]
+        if num_tokens <= rows:
+            torch.arange(i, i + rows, out=self.position)
+            window = self.tracker[num_tokens]
+            start = i - rows + 1
+            skip = max(-start, 0)
+            end = min(i + rows, self.bias.shape[-1])
+            self.bias[..., :, start + skip:end] = window[:, skip:end - start]
+        else:
+            self.bias[..., :, i:i + num_tokens] = 0
+
+    @staticmethod
+    def shared(capacity, device, dtype):
+        # all layers advance in lockstep, so the bias caches share one position/bias/tracker
+        rows = 6
+        position = torch.empty((rows,), device=device, dtype=torch.int64)
+        bias = torch.full((1, 1, rows, capacity), torch.finfo(dtype).min, device=device, dtype=dtype)
+        tracker = {"step": -1}
+        # window templates per decode width: row r serves query j = r - (rows - n) and may see slots <= index + j
+        for n in range(1, rows + 1):
+            window = torch.full((rows, 2 * rows - 1), torch.finfo(dtype).min, device=device, dtype=dtype)
+            for r in range(rows):
+                window[r, :rows + max(r - (rows - n), 0)] = 0
+            tracker[n] = window
+        return position, bias, tracker
+
+    @classmethod
+    def zeros(cls, batch, kv_heads, capacity, head_dim, device, dtype, shared):
+        # zero-init: decode attends full capacity with masked tails, 0*0 stays finite
+        key = torch.zeros((batch, kv_heads, capacity, head_dim), device=device, dtype=dtype)
+        return cls(key, torch.zeros_like(key), 0, shared[0], None, shared[1], shared[2])
+
+    def append(self, xk, xv):
+        seq = xk.shape[2]
+        self.key[:, :, self.index:self.index + seq] = xk
+        self.value[:, :, self.index:self.index + seq] = xv
+        return self.key[:, :, :self.index + seq], self.value[:, :, :self.index + seq]
+
+    def decode(self, xq, xk, xv, num_kv_heads):
+        # CUDA-graphable: device-side write position, masked attention over the full capacity
+        batch_size, num_heads, seq, head_dim = xq.shape
+        self.key.index_copy_(2, self.position[:seq], xk)
+        self.value.index_copy_(2, self.position[:seq], xv)
+        groups = num_heads // num_kv_heads
+        q = xq.reshape(batch_size, num_kv_heads, groups, seq, head_dim) * head_dim ** -0.5
+        bias = self.bias[..., self.bias.shape[-2] - seq:, :].unsqueeze(1)
+        scores = (q @ self.key.transpose(-1, -2).unsqueeze(2)).add_(bias)
+        probs = torch.softmax(scores, dim=-1, dtype=torch.float32).to(xq.dtype)
+        out = probs @ self.value.unsqueeze(2)
+        return out.permute(0, 3, 1, 2, 4).reshape(batch_size, seq, num_heads * head_dim)
+
 @dataclass
 class Llama2Config:
     vocab_size: int = 128320
@@ -278,6 +344,9 @@ class Qwen3VL_8BConfig(Qwen3_8BConfig):
     rope_theta: float = 5000000.0
     rope_dims = [24, 20, 20]
     interleaved_mrope = True
+    fixed_kv: bool = True
+    graph_dynamic_vbar_blocks = True
+    prefetch_dynamic_vbars = True
 
 @dataclass
 class Qwen3VL_4BConfig(Qwen3VL_8BConfig):
@@ -776,7 +845,8 @@ class Llama2_(nn.Module):
         super().__init__()
         self.config = config
         self.fixed_kv = getattr(config, "fixed_kv", False)
-        self.graph_dynamic_vbar_blocks = False
+        self.graph_dynamic_vbar_blocks = getattr(config, "graph_dynamic_vbar_blocks", False)
+        self.prefetch_dynamic_vbars = getattr(config, "prefetch_dynamic_vbars", False)
         self.vocab_size = config.vocab_size
 
         if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
@@ -808,19 +878,25 @@ class Llama2_(nn.Module):
 
     def init_kv_cache(self, batch, capacity, device, dtype):
         caches = []
-        fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device)
+        flash = getattr(comfy_kitchen, "flash_attention_decode_is_available", None)
+        flash_kv = self.fixed_kv and flash is not None and flash(device)
         for _ in range(self.config.num_hidden_layers):
-            if fixed_kv:
+            if flash_kv:
                 key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype)
                 value = torch.empty_like(key)
-                position = torch.empty((batch,), device=device, dtype=torch.int64)
+                pos = torch.empty((batch,), device=device, dtype=torch.int64)
                 seqlen = torch.zeros((batch,), device=device, dtype=torch.int32)
-                caches.append(FixedKV(key, value, 0, position, seqlen))
+                caches.append(FixedKV(key, value, 0, pos, seqlen))
             else:
                 key = torch.empty((batch, self.config.num_key_value_heads, capacity, self.config.head_dim), device=device, dtype=dtype)
                 caches.append((key, torch.empty_like(key), 0))
         return caches
 
+    def init_decode_buffers(self, batch, device, dtype):
+        hidden = torch.empty((batch, 1, self.config.hidden_size), device=device, dtype=dtype)
+        positions = torch.zeros((1, 1), device=device, dtype=torch.int64)
+        return hidden, rope_matrix(self.compute_freqs_cis(positions, device))
+
     def compute_freqs_cis(self, position_ids, device):
         return precompute_freqs_cis(self.config.head_dim,
                                     position_ids,
@@ -856,7 +932,8 @@ class Llama2_(nn.Module):
             mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, seq_len, attention_mask.shape[-1])
             mask = mask.masked_fill(mask.to(torch.bool), torch.finfo(x.dtype).min / 4)
 
-        if seq_len > 1:
+        spec_decode = fixed_kv and any(isinstance(kv, FixedKVBias) for kv in past_key_values) and 2 <= seq_len <= 6 and past_len > 0 and attention_mask is None
+        if seq_len > 1 and not spec_decode:  # spec verify: the staircase decode bias is causal
             causal_mask = torch.empty(past_len + seq_len, past_len + seq_len, dtype=x.dtype, device=x.device).fill_(torch.finfo(x.dtype).min / 4).triu_(1)
             if mask is not None:
                 mask += causal_mask
@@ -865,7 +942,7 @@ class Llama2_(nn.Module):
 
         optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
 
-        enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv_decode
+        enable_graph = self.graph_dynamic_vbar_blocks and (fixed_kv_decode or spec_decode)
         if enable_graph:
             if decode_buffers is None:
                 x = x.clone()
@@ -889,7 +966,7 @@ class Llama2_(nn.Module):
             elif intermediate_output < 0:
                 intermediate_output = len(self.layers) + intermediate_output
 
-        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": getattr(self, "prefetch_dynamic_vbars", False)})
+        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": self.prefetch_dynamic_vbars and past_key_values is not None})
         next_key_values = list(past_key_values) if past_key_values is not None else []
         for i, layer in enumerate(self.layers):
             if all_intermediate is not None:
@@ -1001,6 +1078,18 @@ class BaseLlama:
     def forward(self, input_ids, *args, **kwargs):
         return self.model(input_ids, *args, **kwargs)
 
+def penalty_active(repetition_penalty, presence_penalty):
+    return repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)
+
+
+def apply_penalty(logits, repetition_penalty, presence_penalty):
+    if repetition_penalty != 1.0:
+        logits = torch.where(logits < 0, logits * repetition_penalty, logits / repetition_penalty)
+    if presence_penalty is not None and presence_penalty != 0.0:
+        logits = logits - presence_penalty
+    return logits
+
+
 class BaseGenerate:
     def logits(self, x):
         input = x[:, -1:]
@@ -1010,7 +1099,7 @@ class BaseGenerate:
             module = self.model.embed_tokens
 
         if not module.comfy_cast_weights:
-            return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
+            return torch.nn.functional.linear(input, module.weight.to(x), None)
         with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
             return torch.nn.functional.linear(input, weight, None)
 
@@ -1045,7 +1134,14 @@ class BaseGenerate:
         next_pos = int(position_ids[:, -1].max()) + 1 if position_ids is not None else None
 
         compile_allocations = self.model.graph_dynamic_vbar_blocks and comfy.model_prefetch.malloc_graph_enabled(device)
+        decode_buffers = None
+        if compile_allocations and not comfy.model_management.args.disable_cuda_graphs:
+            init_decode_buffers = getattr(self.model, "init_decode_buffers", None)
+            if init_decode_buffers is not None:
+                decode_buffers = init_decode_buffers(embeds.shape[0], device, execution_dtype)
         decode_tokens = torch.empty((embeds.shape[0], 1), dtype=torch.long, device=device)
+        penalize = penalty_active(repetition_penalty, presence_penalty)
+        penalty_mask = None
 
         # Generation loop
         current_input_ids = initial_input_ids
@@ -1059,14 +1155,23 @@ class BaseGenerate:
 
             # DeepStack visual features are injected on the prefill only; gemma4's forward lacks these kwargs.
             extra = {}
+            if decode_buffers is not None:
+                extra["decode_buffers"] = decode_buffers
             if step == 0 and deepstack_embeds is not None:
                 extra["deepstack_embeds"] = deepstack_embeds
                 extra["visual_pos_masks"] = visual_pos_masks
             x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra, embeds_info=(embeds_info if step == 0 else None))
             logits = self.logits(x)[:, -1]
-            next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample, presence_penalty=presence_penalty)
+            if penalty_mask is None and do_sample and penalize:
+                # allocated on the (unbracketed) first step; later steps only index_fill_ it
+                penalty_mask = torch.zeros((logits.shape[-1],), dtype=torch.bool, device=device)
+                if len(initial_tokens) > 0:
+                    penalty_mask.index_fill_(0, torch.tensor(initial_tokens, device=device), True)
+            next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, [], generator, do_sample=do_sample, presence_penalty=presence_penalty, penalty_mask=penalty_mask)
 
             decode_tokens.copy_(next_token)
+            if penalty_mask is not None:
+                penalty_mask.index_fill_(0, decode_tokens[0], True)
             del next_token, logits, x, embeds, position_ids
             if step > 0 and compile_allocations:
                 comfy.model_prefetch.malloc_graph_end()
@@ -1083,24 +1188,19 @@ class BaseGenerate:
 
         return generated_token_ids
 
-    def sample_token(self, logits, temperature, top_k, top_p, min_p, repetition_penalty, token_history, generator, do_sample=True, presence_penalty=0.0):
-
-        if not do_sample or temperature == 0.0:
-            return torch.argmax(logits, dim=-1, keepdim=True)
-
-        # Sampling mode
-        if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
+    def processed_probs(self, logits, temperature, top_k, top_p, min_p, repetition_penalty, token_history, presence_penalty=0.0, penalty_mask=None):
+        # returns (probs, indices); penalty_mask [vocab] bool stands in for token_history
+        if penalty_active(repetition_penalty, presence_penalty):
+            if penalty_mask is not None:
+                logits = torch.where(penalty_mask.unsqueeze(0), apply_penalty(logits, repetition_penalty, presence_penalty), logits)
+            elif len(token_history) > 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
+                logits[:, token_ids] = apply_penalty(logits[:, token_ids], repetition_penalty, presence_penalty)
 
         if temperature != 1.0:
             logits = logits / temperature
 
+        top_indices = None
         if top_k > 0:
             top_k = min(top_k, logits.shape[-1])
             logits, top_indices = torch.topk(logits, top_k)
@@ -1121,29 +1221,18 @@ class BaseGenerate:
             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)
+        return torch.nn.functional.softmax(logits, dim=-1), top_indices
+
+    def sample_token(self, logits, temperature, top_k, top_p, min_p, repetition_penalty, token_history, generator, do_sample=True, presence_penalty=0.0, penalty_mask=None):
+
+        if not do_sample or temperature == 0.0:
+            return torch.argmax(logits, dim=-1, keepdim=True)
+
+        probs, top_indices = self.processed_probs(logits, temperature, top_k, top_p, min_p, repetition_penalty, token_history, presence_penalty=presence_penalty, penalty_mask=penalty_mask)
         next_token = torch.multinomial(probs, num_samples=1, generator=generator)
+        if top_indices is not None:
             return top_indices.gather(1, next_token)
-
-        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)
-
-        return torch.multinomial(probs, num_samples=1, generator=generator)
+        return next_token
 
 class BaseQwen3:
     def logits(self, x):

comfy/text_encoders/lt.py

@@ -102,7 +102,7 @@ class Gemma3_12BModel(sd1_clip.SDClipModel):
         self.dtypes.add(dtype)
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_12B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty):
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty, mtp=True):
         tokens_only = [[t[0] for t in b] for b in tokens]
         embeds, _, _, _ = self.process_tokens(tokens_only, self.execution_device)
         return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, stop_tokens=[106], presence_penalty=presence_penalty)  # 106 is <end_of_turn>
@@ -205,7 +205,7 @@ class LTXAVTEModel(torch.nn.Module):
 
         return out.to(device=out_device, dtype=torch.float), pooled, extra
 
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty):
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty, mtp=True):
         return self.gemma3_12b.generate(tokens[self.text_encoder_key], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty)
 
     def load_sd(self, sd):

comfy/text_encoders/qwen35.py

@@ -2,14 +2,53 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from dataclasses import dataclass, field
+from tqdm import tqdm
+import contextlib
 import os
+import warnings
 
 import comfy.model_management
+import comfy.model_prefetch
+import comfy.ops
+import comfy_kitchen
+from comfy.quant_ops import QuantizedTensor
 from comfy.ldm.modules.attention import optimized_attention_for_device
 from comfy import sd1_clip
 import comfy.text_encoders.qwen_vl
 
-from .llama import BaseLlama, BaseGenerate, Llama2_, MLP, RMSNorm, apply_rope
+from .llama import BaseLlama, BaseGenerate, FixedKV, FixedKVBias, Llama2_, MLP, RMSNorm, apply_penalty, apply_rope, penalty_active, precompute_freqs_cis, rope_matrix
+
+
+@dataclass
+class LinearKV(FixedKV):
+    # DeltaNet state on the FixedKV interface: key=conv_state, value=recurrent_state (fp32)
+    g_decay: torch.Tensor = None
+    dt_bias: torch.Tensor = None
+    snapshots: list = None  # [(recurrent, conv)] taken after step 1, 2, ... of the last verify
+    last_seq: int = 1
+    snap_backing: torch.Tensor = None
+    conv_snap_backing: torch.Tensor = None
+    norm_weight: torch.Tensor = None
+
+    def prepare(self, num_tokens):
+        pass
+
+    def rollback(self, discard=1):
+        # discard the rejected tail: restore the snapshot taken after the last kept token
+        rec, conv = self.snapshots[self.last_seq - discard - 1]
+        self.recurrent_state.copy_(rec)
+        self.conv_state.copy_(conv)
+        self.index -= discard
+
+    @property
+    def conv_state(self):
+        return self.key
+
+    @property
+    def recurrent_state(self):
+        return self.value
+
+
 
 
 def _qwen35_layer_types(n):
@@ -48,6 +87,7 @@ class Qwen35Config:
     transformer_type: str = "qwen35_2b"
     rope_dims: list = None
     rope_scale: float = None
+    mtp: bool = False
 
 QWEN35_VISION_DEFAULTS = dict(hidden_size=1024, num_heads=16, intermediate_size=4096, depth=24, patch_size=16, temporal_patch_size=2, in_channels=3, spatial_merge_size=2, num_position_embeddings=2304)
 
@@ -135,18 +175,6 @@ def torch_chunk_gated_delta_rule(query, key, value, g, beta, chunk_size=64, init
     return core_attn_out, last_recurrent_state
 
 
-def torch_causal_conv1d_update(x, conv_state, weight, bias=None):
-    # conv_state: [B, channels, kernel_size-1], x: [B, channels, 1]
-    # weight: [channels, kernel_size]
-    state_len = conv_state.shape[-1]
-    combined = torch.cat([conv_state, x], dim=-1).to(weight.dtype)  # [B, channels, kernel_size]
-    conv_state.copy_(combined[:, :, -state_len:])
-    out = (combined * weight).sum(dim=-1, keepdim=True)  # [B, channels, 1]
-    if bias is not None:
-        out = out + bias.unsqueeze(0).unsqueeze(-1)
-    return F.silu(out).to(x.dtype)
-
-
 # GatedDeltaNet - Linear Attention Layer
 
 class GatedDeltaNet(nn.Module):
@@ -185,25 +213,55 @@ class GatedDeltaNet(nn.Module):
 
         use_recurrent = (
             past_key_value is not None
-            and past_key_value[2] > 0
-            and seq_len == 1
+            and past_key_value.index > 0
+            and seq_len <= 6
         )
 
+        fused_available = getattr(comfy_kitchen, "gated_delta_decode_is_available", None)
+        use_fused = (use_recurrent and fused_available is not None and fused_available(x.device, self.key_head_dim, self.value_head_dim)
+                     and (seq_len == 1 or past_key_value.snap_backing is not None))
+
         # Projections (shared)
-        mixed_qkv = self.in_proj_qkv(x).transpose(1, 2)  # [B, conv_dim, seq_len]
+        proj = self.in_proj_qkv(x)  # [B, seq_len, conv_dim]
         z = self.in_proj_z(x)
+
+        if use_fused:
+            # decode: kitchen conv step, then gates + delta rule + gated norm in one kernel
+            if seq_len > 1:
+                past_key_value.last_seq = seq_len
+            with comfy.ops.CastBiasWeightContext(self.conv1d, proj, offloadable=True) as (conv_weight, conv_bias):
+                conv_out = comfy_kitchen.deltanet_conv_step(proj, past_key_value.conv_state, conv_weight, conv_bias,
+                                                            past_key_value.conv_snap_backing[:seq_len - 1] if seq_len > 1 else None)
+            with comfy.ops.CastBiasWeightContext(self.in_proj_a, x, offloadable=True) as (w_a, _), \
+                 comfy.ops.CastBiasWeightContext(self.in_proj_b, x, offloadable=True) as (w_b, _):
+                if isinstance(w_a, QuantizedTensor):
+                    w_a = w_a.dequantize()
+                if isinstance(w_b, QuantizedTensor):
+                    w_b = w_b.dequantize()
+                core_attn_out = comfy_kitchen.gated_delta_decode_fused(
+                    conv_out, x, w_a, w_b, past_key_value.dt_bias, past_key_value.g_decay, past_key_value.recurrent_state,
+                    self.key_dim, self.num_key_heads, self.key_head_dim ** -0.5, z, past_key_value.norm_weight, self.norm.eps,
+                    past_key_value.snap_backing[:seq_len - 1] if seq_len > 1 else None)
+            return self.out_proj(core_attn_out.reshape(batch_size, seq_len, -1)), past_key_value
+
+        mixed_qkv = proj.transpose(1, 2)  # [B, conv_dim, seq_len]
         b = self.in_proj_b(x)
         a = self.in_proj_a(x)
 
         # Conv1d
         if use_recurrent:
-            recurrent_state, conv_state, step_index = past_key_value
-            conv_weight = comfy.model_management.cast_to_device(self.conv1d.weight, mixed_qkv.device, mixed_qkv.dtype).squeeze(1)
-            conv_bias = comfy.model_management.cast_to_device(self.conv1d.bias, mixed_qkv.device, mixed_qkv.dtype) if self.conv1d.bias is not None else None
-            mixed_qkv = torch_causal_conv1d_update(mixed_qkv, conv_state, conv_weight, conv_bias)
+            # decode: exact-width causal window, weight resolved via the vbar-aware context
+            combined = torch.cat([past_key_value.conv_state, mixed_qkv], dim=-1)
+            if seq_len > 1:
+                past_key_value.last_seq = seq_len
+                for s in range(seq_len - 1):
+                    past_key_value.snapshots[s][1].copy_(combined[:, :, 1 + s:1 + s + self.conv_kernel_size - 1])
+            past_key_value.conv_state.copy_(combined[:, :, seq_len:])
+            with comfy.ops.CastBiasWeightContext(self.conv1d, combined, offloadable=True) as (conv_weight, conv_bias):
+                mixed_qkv = F.silu(F.conv1d(combined, conv_weight, conv_bias, groups=self.conv1d.groups))
         else:
             if past_key_value is not None:
-                recurrent_state, conv_state, step_index = past_key_value
+                conv_state = past_key_value.conv_state
                 conv_state_init = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 conv_state.copy_(conv_state_init[:, :, -conv_state.shape[-1]:])
             mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])
@@ -211,47 +269,51 @@ class GatedDeltaNet(nn.Module):
         # Split QKV and compute beta/g
         mixed_qkv = mixed_qkv.transpose(1, 2)  # [B, seq_len, conv_dim]
         query, key, value = mixed_qkv.split([self.key_dim, self.key_dim, self.value_dim], dim=-1)
-        beta = b.sigmoid()
-        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float())
-
-        # Delta rule
         if use_recurrent:
-            # single-token path: work in [B, heads, dim] without seq dim
-            query = query.reshape(batch_size, self.num_key_heads, self.key_head_dim)
-            key = key.reshape(batch_size, self.num_key_heads, self.key_head_dim)
-            value = value.reshape(batch_size, self.num_value_heads, self.value_head_dim)
-
-            if self.num_value_heads != self.num_key_heads:
-                rep = self.num_value_heads // self.num_key_heads
-                query = query.repeat_interleave(rep, dim=1)
-                key = key.repeat_interleave(rep, dim=1)
-
-            scale = self.key_head_dim ** -0.5
-            q = F.normalize(query.float(), dim=-1) * scale
-            k = F.normalize(key.float(), dim=-1)
-            v = value.float()
-            beta_t = beta.reshape(batch_size, -1)
-            g_t = g.reshape(batch_size, -1).exp()
-
-            # In-place state update: [B, heads, k_dim, v_dim]
-            recurrent_state.mul_(g_t[:, :, None, None])
-            kv_mem = torch.einsum('bhk,bhkv->bhv', k, recurrent_state)
-            delta = (v - kv_mem) * beta_t[:, :, None]
-            recurrent_state.add_(k.unsqueeze(-1) * delta.unsqueeze(-2))
-            core_attn_out = torch.einsum('bhk,bhkv->bhv', q, recurrent_state)
-
-            core_attn_out = core_attn_out.to(x.dtype).unsqueeze(1)
-            present_key_value = (recurrent_state, conv_state, step_index + 1)
+            g_decay, dt_bias = past_key_value.g_decay, past_key_value.dt_bias
         else:
-            query = query.reshape(batch_size, seq_len, -1, self.key_head_dim)
-            key = key.reshape(batch_size, seq_len, -1, self.key_head_dim)
-            value = value.reshape(batch_size, seq_len, -1, self.value_head_dim)
+            g_decay = -comfy.model_management.cast_to_device(self.A_log, x.device, torch.float32).exp()
+            dt_bias = comfy.model_management.cast_to_device(self.dt_bias, x.device, torch.float32)
+            if past_key_value is not None:
+                past_key_value.g_decay = g_decay
+                past_key_value.dt_bias = dt_bias
+                past_key_value.norm_weight = comfy.model_management.cast_to(self.norm.weight, dtype=x.dtype, device=x.device)
 
+        beta = b.sigmoid()
+        g = g_decay * F.softplus(a.float() + dt_bias)
+        query = query.reshape(batch_size, seq_len, self.num_key_heads, self.key_head_dim)
+        key = key.reshape(batch_size, seq_len, self.num_key_heads, self.key_head_dim)
+        value = value.reshape(batch_size, seq_len, self.num_value_heads, self.value_head_dim)
         if self.num_value_heads != self.num_key_heads:
             rep = self.num_value_heads // self.num_key_heads
             query = query.repeat_interleave(rep, dim=2)
             key = key.repeat_interleave(rep, dim=2)
 
+        # Delta rule
+        if use_recurrent:
+            scale = self.key_head_dim ** -0.5
+            q = F.normalize(query.float(), dim=-1) * scale
+            k = F.normalize(key.float(), dim=-1)
+            v = value.float()
+            beta_t = beta.reshape(batch_size, seq_len, -1)
+            g_t = g.reshape(batch_size, seq_len, -1).exp()
+
+            # In-place state update: [B, heads, k_dim, v_dim]
+            recurrent_state = past_key_value.recurrent_state
+            outs = []
+            for s in range(seq_len):
+                recurrent_state.mul_(g_t[:, s, :, None, None])
+                kv_mem = torch.einsum('bhk,bhkv->bhv', k[:, s], recurrent_state)
+                delta = (v[:, s] - kv_mem) * beta_t[:, s, :, None]
+                # rank-1 update via baddbmm_: no materialized [B, H, D, D] outer-product temp
+                recurrent_state.view(-1, self.key_head_dim, self.value_head_dim).baddbmm_(
+                    k[:, s].reshape(-1, self.key_head_dim, 1), delta.reshape(-1, 1, self.value_head_dim))
+                outs.append(torch.einsum('bhk,bhkv->bhv', q[:, s], recurrent_state))
+                if seq_len > 1 and s < seq_len - 1:
+                    past_key_value.snapshots[s][0].copy_(recurrent_state)
+            core_attn_out = torch.stack(outs, dim=1).to(x.dtype)
+            present_key_value = past_key_value
+        else:
             core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query, key, value, g=g, beta=beta,
                 initial_state=None,
@@ -261,8 +323,8 @@ class GatedDeltaNet(nn.Module):
             present_key_value = None
             if past_key_value is not None:
                 if last_recurrent_state is not None:
-                    recurrent_state.copy_(last_recurrent_state.to(recurrent_state.dtype))
-                present_key_value = (recurrent_state, conv_state, step_index + seq_len)
+                    past_key_value.recurrent_state.copy_(last_recurrent_state.to(past_key_value.recurrent_state.dtype))
+                present_key_value = past_key_value
 
         # Gated norm + output projection (shared)
         core_attn_out = self.norm(core_attn_out.reshape(-1, self.value_head_dim), z.reshape(-1, self.value_head_dim))
@@ -271,29 +333,6 @@ class GatedDeltaNet(nn.Module):
 
 
 # GatedAttention - Full Attention with output gating
-def precompute_partial_rope(head_dim, rotary_dim, position_ids, theta, device=None, mrope_section=None):
-    """Compute RoPE frequencies for partial rotary embeddings."""
-    theta_numerator = torch.arange(0, rotary_dim, 2, device=device).float()
-    inv_freq = 1.0 / (theta ** (theta_numerator / rotary_dim))
-
-    inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-    position_ids_expanded = position_ids[:, None, :].float()
-    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-    emb = torch.cat((freqs, freqs), dim=-1)
-    cos = emb.cos()
-    sin = emb.sin()
-
-    if mrope_section is not None and position_ids.shape[0] == 3:
-        mrope_section_2 = [s * 2 for s in mrope_section]
-        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section_2, dim=-1))], dim=-1).unsqueeze(0)
-        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section_2, dim=-1))], dim=-1).unsqueeze(0)
-
-    cos = cos.unsqueeze(1)
-    sin = sin.unsqueeze(1)
-    sin_split = sin.shape[-1] // 2
-    return (cos, sin[..., :sin_split], -sin[..., sin_split:])
-
-
 def apply_partial_rope(xq, xk, freqs_cis, rotary_dim):
     """Apply RoPE to only the first rotary_dim dimensions."""
     xq_rot = xq[..., :rotary_dim]
@@ -350,26 +389,16 @@ class GatedAttention(nn.Module):
         xq, xk = apply_partial_rope(xq, xk, freqs_cis, self.rotary_dim)
 
         # KV cache
-        present_key_value = None
-        if past_key_value is not None:
-            past_key, past_value, index = past_key_value
-            num_tokens = xk.shape[2]
-            if past_key.shape[2] >= (index + num_tokens):
-                past_key[:, :, index:index + num_tokens] = xk
-                past_value[:, :, index:index + num_tokens] = xv
-                xk = past_key[:, :, :index + num_tokens]
-                xv = past_value[:, :, :index + num_tokens]
-                present_key_value = (past_key, past_value, index + num_tokens)
+        present_key_value = past_key_value
+        if past_key_value is not None and seq_length <= 6 and attention_mask is None:
+            output = past_key_value.decode(xq, xk, xv, self.num_kv_heads)
         else:
-                if index > 0:
-                    xk = torch.cat((past_key[:, :, :index], xk), dim=2)
-                    xv = torch.cat((past_value[:, :, :index], xv), dim=2)
-                present_key_value = (xk, xv, index + num_tokens)
-
+            if past_key_value is not None:
+                xk, xv = past_key_value.append(xk, xv)
             gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
             output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
-        output = output * gate.sigmoid()
 
+        output = output * gate.sigmoid()
         return self.o_proj(output), present_key_value
 
 
@@ -387,13 +416,15 @@ class Qwen35TransformerBlock(nn.Module):
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
 
     def forward(self, x, attention_mask=None, freqs_cis=None, optimized_attention=None, past_key_value=None):
+        output = x
         if self.layer_type == "linear_attention":
             h, present_key_value = self.linear_attn(self.input_layernorm(x), attention_mask=attention_mask, past_key_value=past_key_value)
         else:
             h, present_key_value = self.self_attn(self.input_layernorm(x), attention_mask=attention_mask, freqs_cis=freqs_cis, optimized_attention=optimized_attention, past_key_value=past_key_value)
 
-        x = x + h
-        x = x + self.mlp(self.post_attention_layernorm(x))
+        # in-place into the input buffer so CUDA-graph replays land in the static x
+        x = torch.add(x, h, out=output)
+        x = torch.add(x, self.mlp(self.post_attention_layernorm(x)), out=output)
         return x, present_key_value
 
 
@@ -402,6 +433,8 @@ class Qwen35Transformer(Llama2_):
     def __init__(self, config, device=None, dtype=None, ops=None):
         nn.Module.__init__(self)
         self.config = config
+        self.prefetch_dynamic_vbars = True
+        self.graph_dynamic_vbar_blocks = True
         self.vocab_size = config.vocab_size
         self.embed_tokens = ops.Embedding(config.vocab_size, config.hidden_size, device=device, dtype=dtype)
         self.layers = nn.ModuleList([
@@ -417,21 +450,10 @@ class Qwen35Transformer(Llama2_):
         if config.lm_head:
             self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
 
-    def get_past_len(self, past_key_values):
-        for i, layer in enumerate(self.layers):
-            if layer.layer_type == "full_attention":
-                if len(past_key_values) > i:
-                    return past_key_values[i][2]
-                break
-        return 0
-
     def compute_freqs_cis(self, position_ids, device):
         rotary_dim = int(self.config.head_dim * self.config.partial_rotary_factor)
-        return precompute_partial_rope(
-            self.config.head_dim, rotary_dim, position_ids,
-            self.config.rope_theta, device=device,
-            mrope_section=self.config.mrope_section,
-        )
+        return precompute_freqs_cis(rotary_dim, position_ids, self.config.rope_theta,
+                                    rope_dims=self.config.mrope_section, interleaved_mrope=True, device=device)
 
 
 # Vision Encoder
@@ -671,6 +693,23 @@ class Qwen35VisionModel(nn.Module):
             return merged, deepstack_features
         return merged
 
+class MTPHead(nn.Module):
+    # MTP draft head: fc(cat[norm(embed), norm(hidden)]) -> attention block (own KV) -> norm
+    def __init__(self, config, device=None, dtype=None, ops=None):
+        super().__init__()
+        self.fc = ops.Linear(config.hidden_size * 2, config.hidden_size, bias=False, device=device, dtype=dtype)
+        self.pre_fc_norm_embedding = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
+        self.pre_fc_norm_hidden = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
+        self.layers = nn.ModuleList([Qwen35TransformerBlock(config, index=3, device=device, dtype=dtype, ops=ops)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
+
+    def forward(self, embeds, hidden, freqs_cis, past_key_value):
+        x = self.fc(torch.cat([self.pre_fc_norm_embedding(embeds), self.pre_fc_norm_hidden(hidden)], dim=-1))
+        attention = optimized_attention_for_device(x.device, mask=False, small_input=True)
+        x, _ = self.layers[0](x, attention_mask=None, freqs_cis=freqs_cis, optimized_attention=attention, past_key_value=past_key_value)
+        return self.norm(x), x  # (for logits, pre-norm hidden for recursive drafting)
+
+
 # Model Wrapper
 class Qwen35(BaseLlama, BaseGenerate, torch.nn.Module):
     model_type = "qwen35_2b"
@@ -680,6 +719,9 @@ class Qwen35(BaseLlama, BaseGenerate, torch.nn.Module):
         config = _make_config(self.model_type, config_dict)
         self.num_layers = config.num_hidden_layers
         self.model = Qwen35Transformer(config, device=device, dtype=dtype, ops=operations)
+        self.mtp = None
+        if config.mtp:
+            self.mtp = MTPHead(config, device=device, dtype=dtype, ops=operations)
         vision_overrides = QWEN35_MODELS.get(self.model_type, {}).get("vision", {})
         vision_config = {**QWEN35_VISION_DEFAULTS, **vision_overrides, "out_hidden_size": config.hidden_size}
         self.visual = Qwen35VisionModel(vision_config, device=device, dtype=dtype, ops=operations)
@@ -687,7 +729,8 @@ class Qwen35(BaseLlama, BaseGenerate, torch.nn.Module):
 
     def preprocess_embed(self, embed, device):
         if embed["type"] == "image":
-            image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(embed["data"], patch_size=16)
+            # Qwen3.5 normalizes to [-1, 1] (mean/std 0.5), same as Qwen3-VL.
+            image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(embed["data"], patch_size=16, image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
             return self.visual(image.to(device, dtype=torch.float32), grid), grid
         return None, None
 
@@ -695,9 +738,271 @@ class Qwen35(BaseLlama, BaseGenerate, torch.nn.Module):
         position_ids = comfy.text_encoders.qwen_vl.qwen2vl_mrope_position_ids(embeds_info, embeds.shape[1], embeds.device)
         return super().forward(x, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=dtype, position_ids=position_ids, past_key_values=past_key_values)
 
+    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, **kwargs):
+        mtp = kwargs.pop("mtp", True)
+        if self.mtp is None or not mtp or kwargs.get("position_ids") is not None or kwargs.get("initial_input_ids") is not None:
+            return super().generate(embeds=embeds, do_sample=do_sample, max_length=max_length, temperature=temperature,
+                                    top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty,
+                                    seed=seed, stop_tokens=stop_tokens, **kwargs)
+        sampling = None
+        if do_sample and temperature != 0.0:
+            sampling = {"temperature": temperature, "top_k": top_k, "top_p": top_p, "min_p": min_p,
+                        "repetition_penalty": repetition_penalty,
+                        "presence_penalty": kwargs.get("presence_penalty", 0.0) or 0.0,
+                        "seed": seed if seed is not None else 42}
+        fixed_depth = None if mtp is True else max(2, min(5, int(mtp)))
+        return self._generate_mtp(embeds, max_length, stop_tokens, sampling=sampling, fixed_depth=fixed_depth)
+
+    def _generate_mtp(self, embeds, max_length, stop_tokens, sampling=None, fixed_depth=None):
+        device = embeds.device
+        cfg = self.model.config
+        if stop_tokens is None:
+            stop_tokens = cfg.stop_tokens
+        dt = torch.bfloat16 if comfy.model_management.should_use_bf16(device) else torch.float32
+        embeds = embeds.to(dt)
+        if embeds.ndim == 2:
+            embeds = embeds.unsqueeze(0)
+        # greedy drafts 3 deep (5 after the probe); sampled stays at 2
+        depth = fixed_depth if fixed_depth is not None else (3 if sampling is None else 2)
+        cap = embeds.shape[1] + max_length + 7
+        pkv = self.init_kv_cache(embeds.shape[0], cap, device, dt)
+        # repair window: drafting ahead plus a near-full rollback
+        mtp_kv = FixedKVBias.zeros(embeds.shape[0], cfg.num_key_value_heads, cap, cfg.head_dim, device, dt,
+                                   FixedKVBias.shared(cap, device, dt))
+        head = self.model.lm_head if hasattr(self.model, "lm_head") else self.model.embed_tokens
+
+        def verify_logits(x):
+            if not head.comfy_cast_weights:
+                return F.linear(x, head.weight.to(x), None)
+            with comfy.ops.CastBiasWeightContext(head, x, offloadable=True) as (w, _bias):
+                return F.linear(x, w)
+
+        # the draft graph bakes these weights' addresses: keep them resident for the generate
+        hot = list({id(m): m for m in [head, self.model.embed_tokens, *self.mtp.modules()]}.values())
+        pinned = [m for m in hot if hasattr(m, "_v")]
+
+        generator = None
+        if sampling is not None:
+            generator = torch.Generator(device=device).manual_seed(sampling["seed"])
+        penalized = sampling is not None and penalty_active(sampling["repetition_penalty"], sampling["presence_penalty"])
+
+        # rope table once per generate, sliced per draft
+        ftab = rope_matrix(self.model.compute_freqs_cis(torch.arange(cap, device=device, dtype=torch.float).unsqueeze(0), device))
+
+        def freqs_at(p, n=1):
+            return ftab[:, :, p:p + n]
+
+        # cross-step state lives in static carriers: nothing allocated inside a step may outlive it
+        nt_buf = torch.empty((embeds.shape[0], 1), device=device, dtype=torch.long)
+        h_buf = torch.empty((embeds.shape[0], 1, cfg.hidden_size), device=device, dtype=dt)
+        x, _, _ = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=pkv)
+        lg0 = self.logits(x)[:, -1]
+        if sampling is None:
+            nt_buf.copy_(lg0.argmax(dim=-1, keepdim=True))
+        else:
+            nt_buf.copy_(self.sample_token(lg0, sampling["temperature"], sampling["top_k"], sampling["top_p"], sampling["min_p"],
+                                           sampling["repetition_penalty"], [], generator, presence_penalty=sampling["presence_penalty"]))
+        pen_mask = torch.zeros((lg0.shape[-1],), device=device, dtype=torch.bool) if penalized else None
+        h_buf.copy_(x[:, -1:, :])
+        del x, lg0
+        ids = [nt_buf[0].item()]
+        if penalized:
+            pen_mask.index_fill_(0, nt_buf.reshape(-1), True)
+        pos = embeds.shape[1]
+        progress = comfy.utils.ProgressBar(max_length)
+        console = tqdm(total=max_length, desc="Generating tokens", initial=1)
+
+        def update_progress(n):
+            progress.update(n)
+            console.update(n)
+
+        verify_buffers = None
+        snapshot_bytes = sum(kv.recurrent_state.numel() * 4 + kv.conv_state.numel() * kv.conv_state.element_size()
+                             for kv in pkv if isinstance(kv, LinearKV))
+
+        def set_depth(d):
+            nonlocal depth, verify_buffers
+            depth = d
+            for kv in pkv:
+                if isinstance(kv, LinearKV):
+                    # snapshot views share one backing slab so the fused kernel can write them
+                    kv.snap_backing = torch.empty((d,) + tuple(kv.recurrent_state.shape), device=device, dtype=torch.float32)
+                    kv.conv_snap_backing = torch.empty((d,) + tuple(kv.conv_state.shape), device=device, dtype=kv.conv_state.dtype)
+                    kv.snapshots = [(kv.snap_backing[s], kv.conv_snap_backing[s]) for s in range(d)]
+            # static hidden/rope buffers: graphed layers bake their input addresses
+            verify_buffers = (torch.empty((embeds.shape[0], d + 1, cfg.hidden_size), device=device, dtype=dt), freqs_at(pos, d + 1).clone())
+
+        set_depth(depth)
+        use_graph = (device.type == "cuda"
+                     and comfy.model_management.NUM_STREAMS > 0
+                     and not comfy.model_management.args.disable_cuda_graphs)
+        compile_allocations = use_graph and self.model.graph_dynamic_vbar_blocks and comfy.model_prefetch.malloc_graph_enabled(device)
+        draft_state = {}
+
+        def drop_draft_graph():
+            # free inside torch API calls so the allocator's benign notices stay catchable
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                for t in (draft_state.get("d"), draft_state.get("r"), *draft_state.get("keep", ())):
+                    if t is not None:
+                        t.set_()
+                g = draft_state.pop("graph", None)
+                if g is not None:
+                    g.reset()
+                draft_state.clear()
+
+        def draft_capture():
+            # captured outside the compiler bracket; its static buffers live for the generate
+            ds = draft_state
+            mtp_kv.prepare(1)
+            ds["tok"] = nt_buf.clone()
+            ds["hid"] = h_buf.clone()
+            ds["f"] = freqs_at(pos).clone()
+            side = torch.cuda.Stream()
+            side.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(side):
+                for _ in range(2):
+                    n1, r1 = self.mtp(self.model.embed_tokens(ds["tok"]).to(dt), ds["hid"], ds["f"], mtp_kv)
+                    self.logits(n1)[:, -1].argmax(dim=-1, keepdim=True)
+            torch.cuda.current_stream().wait_stream(side)
+            del n1, r1  # freed before the capture, not shadowed inside it
+            g = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g, capture_error_mode="thread_local"):
+                n1, r1 = self.mtp(self.model.embed_tokens(ds["tok"]).to(dt), ds["hid"], ds["f"], mtp_kv)
+                lg1 = self.logits(n1)
+                ds["d"] = lg1[:, -1].argmax(dim=-1, keepdim=True)
+                ds["r"] = r1
+                ds["keep"] = (n1, lg1)  # captured allocations must outlive the graph
+            ds["graph"] = g
+
+        def draft(token, hidden, p):
+            # one drafted token: mtp head + lm_head argmax, graph-replayed on cuda
+            mtp_kv.prepare(1)
+            f = freqs_at(p)
+            if not use_graph:
+                n1, r1 = self.mtp(self.model.embed_tokens(token).to(dt), hidden, f, mtp_kv)
+                mtp_kv.advance(1)
+                return self.logits(n1)[:, -1].argmax(dim=-1, keepdim=True), r1
+            ds = draft_state
+            ds["tok"].copy_(token)
+            ds["hid"].copy_(hidden)
+            ds["f"].copy_(f)
+            ds["graph"].replay()
+            mtp_kv.advance(1)
+            return ds["d"], ds["r"]
+
+        def verify_sample(lg, drafts):
+            # accept draft i w.p. p_i(draft), else sample the residual; all depth+1 columns as one batch
+            s = sampling
+            rows = lg[0].float()
+            if penalized:
+                # committed tokens penalize every column, each draft the columns after it
+                mask = pen_mask.unsqueeze(0).repeat(depth + 1, 1)
+                for c, d in enumerate(drafts):
+                    mask[c + 1:].index_fill_(1, d.reshape(-1), True)
+                rows = torch.where(mask, apply_penalty(rows, s["repetition_penalty"], s["presence_penalty"]), rows)
+            probs, idx = self.processed_probs(rows, s["temperature"], s["top_k"], s["top_p"], s["min_p"], 1.0, [])
+            dr = torch.cat(drafts, dim=1).reshape(-1, 1)
+            if idx is None:
+                p_draft = probs[:depth].gather(1, dr).reshape(-1)
+                w = probs.clone()
+                w[:depth].scatter_(1, dr, 0.0)
+            else:
+                hit = idx[:depth] == dr
+                p_draft = (probs[:depth] * hit).sum(1)
+                w = probs.clone()
+                w[:depth].masked_fill_(hit, 0.0)
+            # residual: p with the draft zeroed (the rescue term only matters when discarded anyway)
+            w = w + (w.sum(1, keepdim=True) == 0) * probs
+            tok = torch.multinomial(w, num_samples=1, generator=generator)
+            corr = tok if idx is None else idx.gather(1, tok)
+            u = torch.rand(depth, device=device, generator=generator)
+            accepted = (u < p_draft).long().cumprod(0).sum()
+            return dr, corr, accepted
+
+        def step():
+            # scoped so every temporary dies before the compiler bracket closes
+            nonlocal pos
+            drafts = []
+            tok_in, hid_in = nt_buf, h_buf
+            for k in range(depth):
+                dk, rk = draft(tok_in, hid_in, pos + k)
+                if k < depth - 1:
+                    dk = dk.clone()  # later replays overwrite the static output
+                drafts.append(dk)
+                tok_in, hid_in = dk, rk
+            ev = self.model.embed_tokens(torch.cat([nt_buf] + drafts, dim=1)).to(dt)
+            x, _, _ = self.model.forward(None, embeds=ev, attention_mask=None, past_key_values=pkv, decode_buffers=verify_buffers)
+            # all verify positions in one lm_head GEMV, accept decided GPU-side, one sync
+            lg = verify_logits(x)
+            if sampling is None:
+                toks = lg.argmax(dim=-1)
+                vals = torch.cat([toks[0]] + [d[0] for d in drafts]).tolist()
+                t, dr = vals[:depth + 1], vals[depth + 1:]
+                accepts = 0
+                while accepts < depth and t[accepts] == dr[accepts]:
+                    accepts += 1
+                next_toks = tuple(toks[:, i:i + 1] for i in range(depth + 1))
+                commit = tuple(dr[:accepts]) + (t[accepts],)
+            else:
+                dr, corr, accepted = verify_sample(lg, drafts)
+                vals = torch.cat([dr[:, 0], corr[:, 0], accepted.reshape(1)]).tolist()
+                dr, cv, accepts = vals[:depth], vals[depth:2 * depth + 1], vals[-1]
+                next_toks = tuple(corr[i:i + 1] for i in range(depth + 1))
+                commit = tuple(dr[:accepts]) + (cv[accepts],)
+            if accepts < depth:
+                for kv in pkv:
+                    kv.rollback(depth - accepts)
+            if accepts < depth - 1:
+                mtp_kv.rollback(depth - 1 - accepts)  # mtp entries fed by a rejected draft token
+            nt_buf.copy_(next_toks[accepts])
+            h_buf.copy_(x[:, accepts:accepts + 1, :])
+            if penalized:
+                for d in drafts[:accepts]:
+                    pen_mask.index_fill_(0, d.reshape(-1), True)
+                pen_mask.index_fill_(0, nt_buf.reshape(-1), True)
+            pos += accepts + 1
+            return accepts, commit
+
+        probe = None if fixed_depth is not None else [0, 0]  # steps, accepted drafts
+        try:
+            if pinned:
+                comfy.model_prefetch.pin_modules(pinned, device, dt)
+            if use_graph and len(ids) < max_length and ids[-1] not in stop_tokens:
+                draft_capture()
+            while len(ids) < max_length and ids[-1] not in stop_tokens:
+                with (comfy.model_prefetch.malloc_graph_scope(device) if compile_allocations else contextlib.nullcontext()):
+                    accepts, commit = step()
+                commit = list(commit[:max_length - len(ids)])
+                stop = next((i for i, t in enumerate(commit) if t in stop_tokens), None)
+                if stop is not None:
+                    del commit[stop + 1:]
+                ids.extend(commit)
+                update_progress(len(commit))
+                if probe is not None:
+                    probe[0] += 1
+                    probe[1] += accepts
+                    if probe[0] == 32:
+                        # deepen once acceptance sustains it and a recapture round can amortize
+                        if (sampling is None and max_length - len(ids) > 512 and 1 + probe[1] / probe[0] >= 2.2
+                                and 5 * snapshot_bytes < comfy.model_management.get_free_memory(device)):
+                            comfy.model_prefetch.cleanup_prefetch_queues()
+                            drop_draft_graph()
+                            set_depth(5)
+                            if use_graph:
+                                draft_capture()
+                        probe = None
+        finally:
+            console.close()
+            drop_draft_graph()
+            if pinned:
+                comfy.model_prefetch.cleanup_prefetched_modules(None, pinned)
+        return ids
+
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
         model_config = self.model.config
         past_key_values = []
+        shared = FixedKVBias.shared(max_cache_len, device, execution_dtype)
         for i in range(model_config.num_hidden_layers):
             if model_config.layer_types[i] == "linear_attention":
                 recurrent_state = torch.zeros(
@@ -709,13 +1014,9 @@ class Qwen35(BaseLlama, BaseGenerate, torch.nn.Module):
                     [batch, conv_dim, model_config.conv_kernel_size - 1],
                     device=device, dtype=execution_dtype
                 )
-                past_key_values.append((recurrent_state, conv_state, 0))
+                past_key_values.append(LinearKV(conv_state, recurrent_state, 0, None, None))
             else:
-                past_key_values.append((
-                    torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype),
-                    torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype),
-                    0
-                ))
+                past_key_values.append(FixedKVBias.zeros(batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim, device, execution_dtype, shared))
         return past_key_values
 
 # Tokenizer and Text Encoder Wrappers
@@ -777,19 +1078,28 @@ class Qwen35ImageTokenizer(sd1_clip.SD1Tokenizer):
 
 
 class Qwen35ClipModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}, model_type="qwen35_2b"):
+    def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}, model_type="qwen35_2b", mtp=False):
         class Qwen35_(Qwen35):
             pass
         Qwen35_.model_type = model_type
 
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={"mtp": True} if mtp else {},
             dtype=dtype, special_tokens={"pad": 248044}, layer_norm_hidden_state=False,
             model_class=Qwen35_, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0, mtp=True):
+        if isinstance(tokens, dict):
+            tokens = next(iter(tokens.values()))
+        tokens_only = [[t[0] for t in b] for b in tokens]
+        embeds, _, _, embeds_info = self.process_tokens(tokens_only, self.execution_device)
+        position_ids = comfy.text_encoders.qwen_vl.qwen2vl_mrope_position_ids(embeds_info, embeds.shape[1], embeds.device)
+        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed,
+                                         presence_penalty=presence_penalty, position_ids=position_ids, mtp=mtp)
+
 
 class Qwen35TEModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}, model_type="qwen35_2b"):
-        clip_model = lambda **kw: Qwen35ClipModel(**kw, model_type=model_type)
+    def __init__(self, device="cpu", dtype=None, model_options={}, model_type="qwen35_2b", mtp=False):
+        clip_model = lambda **kw: Qwen35ClipModel(**kw, model_type=model_type, mtp=mtp)
         super().__init__(device=device, dtype=dtype, name=model_type, clip_model=clip_model, model_options=model_options)
 
 
@@ -800,7 +1110,7 @@ def tokenizer(model_type="qwen35_2b"):
     return Qwen35ImageTokenizer_
 
 
-def te(dtype_llama=None, llama_quantization_metadata=None, model_type="qwen35_2b"):
+def te(dtype_llama=None, llama_quantization_metadata=None, model_type="qwen35_2b", mtp=False):
     class Qwen35TEModel_(Qwen35TEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if dtype_llama is not None:
@@ -808,5 +1118,5 @@ def te(dtype_llama=None, llama_quantization_metadata=None, model_type="qwen35_2b
             if llama_quantization_metadata is not None:
                 model_options = model_options.copy()
                 model_options["quantization_metadata"] = llama_quantization_metadata
-            super().__init__(device=device, dtype=dtype, model_options=model_options, model_type=model_type)
+            super().__init__(device=device, dtype=dtype, model_options=model_options, model_type=model_type, mtp=mtp)
     return Qwen35TEModel_

comfy/text_encoders/qwen3vl.py

@@ -127,7 +127,7 @@ class Qwen3VLClipModel(sd1_clip.SDClipModel):
                          model_class=_make_qwen3vl_model(model_type), enable_attention_masks=attention_mask,
                          return_attention_masks=attention_mask, model_options=model_options)
 
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0, mtp=True):
         if isinstance(tokens, dict):
             tokens = next(iter(tokens.values()))
         tokens_only = [[t[0] for t in b] for b in tokens]

comfy_api_nodes/apis/openrouter.py

@@ -104,6 +104,7 @@ class OpenRouterChatResponse(BaseModel):
 
 class OpenRouterImageData(BaseModel):
     b64_json: str | None = Field(None)
+    url: str | None = Field(None)
     media_type: str | None = Field(None)
 
 

comfy_api_nodes/nodes_bytedance.py

@@ -3362,6 +3362,7 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
             cls,
             ApiEndpoint(path="/proxy/byteplus/api/v3/tts/create", method="POST"),
             response_model=SeedAudioResponse,
+            asset_urls=True,
             data=SeedAudioRequest(
                 model=model,
                 text_prompt=text_prompt,
@@ -3374,11 +3375,15 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
                 ),
             ),
         )
-        if not response.audio:
+        if response.audio:
+            audio_bytes = base64.b64decode(response.audio)
+        elif response.url:
+            audio_bytes = (await download_url_as_bytesio(response.url, cls=cls)).getvalue()
+        else:
             raise Exception(
                 f"Seed Audio returned no audio (code={response.code}): {response.message}"
             )
-        return IO.NodeOutput(audio_bytes_to_audio_input(base64.b64decode(response.audio)))
+        return IO.NodeOutput(audio_bytes_to_audio_input(audio_bytes))
 
 
 _VCUBE_ENHANCE_VIDEO_ENDPOINT = ApiEndpoint(path="/proxy/byteplusmediakit/api/v1/tools/enhance-video", method="POST")

comfy_api_nodes/nodes_elevenlabs.py

@@ -404,6 +404,7 @@ class ElevenLabsTextToSpeech(IO.ComfyNode):
             ),
             data=request,
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 
@@ -449,6 +450,7 @@ class ElevenLabsAudioIsolation(IO.ComfyNode):
             files={"audio": ("audio.mp4", audio_bytes_io, "audio/mp4")},
             content_type="multipart/form-data",
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 
@@ -545,6 +547,7 @@ class ElevenLabsTextToSoundEffects(IO.ComfyNode):
                 loop=model.get("loop", None),
             ),
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 
@@ -762,6 +765,7 @@ class ElevenLabsSpeechToSpeech(IO.ComfyNode):
             files={"audio": ("audio.mp4", audio_bytes_io.getvalue(), "audio/mp4")},
             content_type="multipart/form-data",
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 
@@ -901,6 +905,7 @@ class ElevenLabsTextToDialogue(IO.ComfyNode):
             ),
             data=request,
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 

comfy_api_nodes/nodes_fishaudio.py

@@ -293,6 +293,7 @@ class FishAudioTextToSpeech(IO.ComfyNode):
             ),
             data=request,
             as_binary=True,
+            asset_urls=True,
         )
         return IO.NodeOutput(audio_bytes_to_audio_input(response))
 

comfy_api_nodes/nodes_gemini.py

@@ -8,6 +8,7 @@ import os
 from fnmatch import fnmatch
 from io import BytesIO
 from typing import Any, Literal
+from urllib.parse import urlparse
 
 import torch
 from typing_extensions import override
@@ -277,8 +278,12 @@ async def get_video_from_interaction(
     )
 
 
+GEMINI_FILES_HOST = "generativelanguage.googleapis.com"
+
+
 async def download_interaction_video(uri: str, cls: type[IO.ComfyNode] | None = None) -> InputImpl.VideoFromFile:
-    if "/files/" not in uri:
+    parsed = urlparse(uri)
+    if parsed.netloc != GEMINI_FILES_HOST or "/files/" not in parsed.path:
         return await download_url_to_video_output(uri, cls=cls)
     name = uri.split("?", 1)[0].rsplit("/files/", 1)[-1].split(":", 1)[0]
     await poll_op(
@@ -1921,6 +1926,7 @@ class GeminiVideoOmni(IO.ComfyNode):
                 ),
             ),
             response_model=GeminiInteraction,
+            asset_urls=True,
         )
         if interaction.status != "completed":
             model_message = get_text_from_interaction(interaction).strip()
@@ -2181,6 +2187,7 @@ class GeminiVideoOmniV2(IO.ComfyNode):
                 response_format=response_format,
             ),
             response_model=GeminiInteraction,
+            asset_urls=True,
         )
         if interaction.status != "completed":
             model_message = get_text_from_interaction(interaction).strip()

comfy_api_nodes/nodes_heygen.py

@@ -1,5 +1,3 @@
-import uuid
-
 import torch
 from typing_extensions import override
 
@@ -70,7 +68,7 @@ async def _create_and_poll_video(cls: type[IO.ComfyNode], payload: dict) -> dict
     """POST a /v3/videos payload, poll until terminal, and return the final video data."""
     created = await sync_op_raw(
         cls,
-        ApiEndpoint(path=_VIDEOS_PATH, method="POST", headers={"Idempotency-Key": uuid.uuid4().hex}),
+        ApiEndpoint(path=_VIDEOS_PATH, method="POST"),
         data=payload,
     )
     video_id = (created.get("data") or {}).get("video_id")

comfy_api_nodes/nodes_ideogram.py

@@ -356,7 +356,6 @@ class IdeogramV3(IO.ComfyNode):
                 data=edit_request,
                 files=files,
                 content_type="multipart/form-data",
-                max_retries=1,
             )
 
         elif image is not None or mask is not None:
@@ -400,7 +399,6 @@ class IdeogramV3(IO.ComfyNode):
                 data=gen_request,
                 files=files if files else None,
                 content_type="multipart/form-data",
-                max_retries=1,
             )
 
         if not response.data or len(response.data) == 0:
@@ -514,7 +512,6 @@ class IdeogramV4(IO.ComfyNode):
                 resolution=resolution.split(" ")[0] if resolution != "Auto" else None,
                 rendering_speed=rendering_speed,
             ),
-            max_retries=1,
         )
 
         if not response.data or len(response.data) == 0:
@@ -649,7 +646,6 @@ class IdeogramPImage(IO.ComfyNode):
             ApiEndpoint(path="/proxy/ideogram/text-to-image/p-image-ideogram", method="POST"),
             response_model=IdeogramGenerateResponse,
             data=request,
-            max_retries=1,
         )
         if not response.data:
             raise Exception("No images were generated in the response")

comfy_api_nodes/nodes_krea.py

@@ -35,7 +35,6 @@ async def _upload_image_to_krea_assets(cls: type[IO.ComfyNode], image: Input.Ima
         response_model=KreaAssetResponse,
         files=[("file", (img_io.name, img_io, "image/png"))],
         content_type="multipart/form-data",
-        max_retries=1,
         wait_label="Uploading reference",
     )
     return response.image_url

comfy_api_nodes/nodes_ltxv.py

@@ -59,7 +59,6 @@ async def _v25_submit_and_poll(cls: type[IO.ComfyNode], route: str, data: BaseMo
         ApiEndpoint(f"/proxy/ltx/v2/{route}", "POST"),
         response_model=Ltx25SubmitResponse,
         data=data,
-        max_retries=1,
     )
     job = await poll_op(
         cls,

comfy_api_nodes/nodes_meta.py

@@ -16,6 +16,8 @@ from comfy_api_nodes.apis.meta import (
 from comfy_api_nodes.util import (
     ApiEndpoint,
     bytesio_to_image_tensor,
+    download_url_to_image_tensor,
+    pad_images_to_common_channels,
     sync_op,
     upload_images_to_comfyapi,
     validate_string,
@@ -57,15 +59,16 @@ def _size(aspect_ratio: str) -> str | None:
     return None if aspect_ratio == "auto" else aspect_ratio.replace(":", "x")
 
 
-def _decode_images(response: MuseImageResponse) -> torch.Tensor:
-    images = [
-        bytesio_to_image_tensor(BytesIO(base64.b64decode(item.b64_json)))
-        for item in response.data
-        if item.b64_json
-    ]
+async def _decode_images(cls: type[IO.ComfyNode], response: MuseImageResponse) -> torch.Tensor:
+    images = []
+    for item in response.data:
+        if item.b64_json:
+            images.append(bytesio_to_image_tensor(BytesIO(base64.b64decode(item.b64_json))))
+        elif item.url:
+            images.append(await download_url_to_image_tensor(item.url, cls=cls))
     if not images:
         raise Exception("The response contains no images.")
-    return torch.cat(images)
+    return torch.cat(pad_images_to_common_channels(images))
 
 
 def _reasoning_strength_input() -> IO.Combo.Input:
@@ -220,6 +223,7 @@ class MetaMuseImageTextToImageApi(IO.ComfyNode):
             cls,
             ApiEndpoint(path=GENERATIONS_PATH, method="POST"),
             response_model=MuseImageResponse,
+            asset_urls=True,
             data=MuseImageRequest(
                 model=model["model"],
                 prompt=model["prompt"],
@@ -228,7 +232,7 @@ class MetaMuseImageTextToImageApi(IO.ComfyNode):
                 tool_enablement=_tool_enablement(model),
             ),
         )
-        return IO.NodeOutput(_decode_images(response))
+        return IO.NodeOutput(await _decode_images(cls, response))
 
 
 class MetaMuseImageEditApi(IO.ComfyNode):
@@ -280,6 +284,7 @@ class MetaMuseImageEditApi(IO.ComfyNode):
             cls,
             ApiEndpoint(path=EDITS_PATH, method="POST"),
             response_model=MuseImageResponse,
+            asset_urls=True,
             data=MuseImageEditRequest(
                 model=model["model"],
                 prompt=prompt,
@@ -289,7 +294,7 @@ class MetaMuseImageEditApi(IO.ComfyNode):
                 images=[MuseImageInput(image_url=url) for url in urls],
             ),
         )
-        return IO.NodeOutput(_decode_images(response))
+        return IO.NodeOutput(await _decode_images(cls, response))
 
 
 class MetaApiExtension(ComfyExtension):

comfy_api_nodes/nodes_openai.py

@@ -84,12 +84,15 @@ SUPPORTED_REASONING_EFFORTS: dict[str, tuple[str, ...]] = {
 }
 
 
-async def validate_and_cast_response(response, timeout: int = None) -> torch.Tensor:
+async def validate_and_cast_response(
+    response, timeout: int = None, cls: type[IO.ComfyNode] = None
+) -> torch.Tensor:
     """Validates and casts a response to a torch.Tensor.
 
     Args:
         response: The response to validate and cast.
         timeout: Request timeout in seconds. Defaults to None (no timeout).
+        cls: The calling node class; required for relative `/proxy/` URLs so they can be expanded and authenticated.
 
     Returns:
         A torch.Tensor of shape (N, H, W, C) with all returned images; images whose
@@ -112,7 +115,7 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten
             img_io = BytesIO(base64.b64decode(img_data.b64_json))
         elif img_data.url:
             img_io = BytesIO()
-            await download_url_to_bytesio(img_data.url, img_io, timeout=timeout)
+            await download_url_to_bytesio(img_data.url, img_io, timeout=timeout, cls=cls)
         else:
             raise ValueError("Invalid image payload – neither URL nor base64 data present.")
 
@@ -367,6 +370,7 @@ class OpenAIGPTImage1(IO.ComfyNode):
                 cls,
                 ApiEndpoint(path="/proxy/openai/images/edits", method="POST"),
                 response_model=OpenAIImageGenerationResponse,
+                asset_urls=True,
                 data=OpenAIImageEditRequest(
                     model=model,
                     prompt=prompt,
@@ -385,6 +389,7 @@ class OpenAIGPTImage1(IO.ComfyNode):
                 cls,
                 ApiEndpoint(path="/proxy/openai/images/generations", method="POST"),
                 response_model=OpenAIImageGenerationResponse,
+                asset_urls=True,
                 data=OpenAIImageGenerationRequest(
                     model=model,
                     prompt=prompt,
@@ -396,7 +401,7 @@ class OpenAIGPTImage1(IO.ComfyNode):
                     moderation="low",
                 ),
             )
-        return IO.NodeOutput(await validate_and_cast_response(response))
+        return IO.NodeOutput(await validate_and_cast_response(response, cls=cls))
 
 
 GPT_IMAGE_QUALITIES = ("low", "medium", "high")
@@ -732,6 +737,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 cls,
                 ApiEndpoint(path="/proxy/openai/images/edits", method="POST"),
                 response_model=OpenAIImageGenerationResponse,
+                asset_urls=True,
                 data=OpenAIImageEditRequest(
                     model=model_id,
                     prompt=prompt,
@@ -749,6 +755,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 cls,
                 ApiEndpoint(path="/proxy/openai/images/generations", method="POST"),
                 response_model=OpenAIImageGenerationResponse,
+                asset_urls=True,
                 data=OpenAIImageGenerationRequest(
                     model=model_id,
                     prompt=prompt,
@@ -759,7 +766,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                     moderation="low",
                 ),
             )
-        return IO.NodeOutput(await validate_and_cast_response(response))
+        return IO.NodeOutput(await validate_and_cast_response(response, cls=cls))
 
 
 class OpenAIChatNode(IO.ComfyNode):

comfy_api_nodes/nodes_openrouter.py

@@ -29,6 +29,7 @@ from comfy_api_nodes.apis.openrouter import (
 from comfy_api_nodes.util import (
     ApiEndpoint,
     bytesio_to_image_tensor,
+    download_url_to_image_tensor,
     get_number_of_images,
     pad_images_to_common_channels,
     sync_op,
@@ -279,16 +280,20 @@ def _extract_text(response: OpenRouterChatResponse) -> str:
     return message.content or ""
 
 
-def _image_data_to_tensor(item: OpenRouterImageData) -> torch.Tensor:
+async def _image_data_to_tensor(cls: type[IO.ComfyNode], item: OpenRouterImageData) -> torch.Tensor:
+    if item.b64_json:
         try:
             return bytesio_to_image_tensor(BytesIO(base64.b64decode(item.b64_json)))
         except Exception as e:
             raise ValueError(f"OpenRouter returned an image that could not be decoded: {e}") from e
+    if item.url:
+        return await download_url_to_image_tensor(item.url, cls=cls)
+    raise ValueError("OpenRouter returned an image with neither inline data nor a URL.")
 
 
-def _extract_images(response: OpenRouterImageResponse) -> torch.Tensor:
+async def _extract_images(cls: type[IO.ComfyNode], response: OpenRouterImageResponse) -> torch.Tensor:
     _raise_on_error(response.error)
-    tensors = [_image_data_to_tensor(item) for item in response.data or [] if item.b64_json]
+    tensors = [await _image_data_to_tensor(cls, item) for item in response.data or [] if item.b64_json or item.url]
     if not tensors:
         raise ValueError("OpenRouter returned no image.")
     return torch.cat(pad_images_to_common_channels(tensors))
@@ -617,6 +622,7 @@ class OpenRouterImageNode(IO.ComfyNode):
             cls,
             ApiEndpoint(path=OPENROUTER_IMAGES_ENDPOINT, method="POST"),
             response_model=OpenRouterImageResponse,
+            asset_urls=True,
             data=OpenRouterImageRequest(
                 model=slug,
                 prompt=prompt,
@@ -625,7 +631,7 @@ class OpenRouterImageNode(IO.ComfyNode):
                 input_references=input_references,
             ),
         )
-        return IO.NodeOutput(_extract_images(response))
+        return IO.NodeOutput(await _extract_images(cls, response))
 
 
 class OpenRouterExtension(ComfyExtension):

comfy_api_nodes/nodes_recraft.py

@@ -66,7 +66,6 @@ async def handle_recraft_file_request(
         files=files,
         content_type="multipart/form-data",
         multipart_parser=recraft_multipart_parser,
-        max_retries=1,
     )
     all_bytesio = []
     if response.image is not None:
@@ -450,7 +449,6 @@ class RecraftCreateStyleNode(IO.ComfyNode):
             files=files,
             data=RecraftCreateStyleRequest(style=style),
             content_type="multipart/form-data",
-            max_retries=1,
         )
 
         return IO.NodeOutput(response.id)
@@ -525,7 +523,6 @@ class RecraftV4CreateStyleNode(IO.ComfyNode):
                 model=model,
             ),
             content_type="multipart/form-data",
-            max_retries=1,
         )
         return IO.NodeOutput(response.id)
 
@@ -629,7 +626,6 @@ class RecraftTextToImageNode(IO.ComfyNode):
                 style_id=recraft_style.style_id,
                 controls=controls_api,
             ),
-            max_retries=1,
         )
         images = []
         for data in response.data:
@@ -953,7 +949,6 @@ class RecraftTextToVectorNode(IO.ComfyNode):
                 substyle=recraft_style.substyle,
                 controls=controls_api,
             ),
-            max_retries=1,
         )
         svg_data = []
         for data in response.data:
@@ -1440,7 +1435,6 @@ class RecraftV4TextToImageNode(IO.ComfyNode):
                 style_reference_urls=style_reference_urls,
                 controls=recraft_controls.create_api_model() if recraft_controls else None,
             ),
-            max_retries=1,
         )
         images = []
         for data in response.data:
@@ -1682,7 +1676,6 @@ class RecraftV4TextToVectorNode(IO.ComfyNode):
                 style_reference_urls=style_reference_urls,
                 controls=recraft_controls.create_api_model() if recraft_controls else None,
             ),
-            max_retries=1,
         )
         svg_data = []
         for data in response.data:

comfy_api_nodes/util/_helpers.py

@@ -7,7 +7,10 @@ from collections.abc import Callable
 from datetime import datetime, timezone
 from email.utils import parsedate_to_datetime
 from io import BytesIO
+from urllib.parse import urlparse
 
+import aiohttp
+from aiohttp.client_exceptions import ClientError
 from yarl import URL
 
 from comfy.cli_args import args
@@ -75,6 +78,48 @@ def default_base_url() -> str:
     return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
 
 
+async def diagnose_connectivity() -> dict[str, bool]:
+    """Best-effort connectivity diagnostics to distinguish local vs. server issues."""
+    results = {
+        "internet_accessible": False,
+        "api_accessible": False,
+    }
+    timeout = aiohttp.ClientTimeout(total=5.0)
+
+    # Probe Google and Baidu in parallel: Google is blocked by the GFW in mainland China, so a Baidu probe is required
+    # to correctly detect that Chinese users with working internet do have working internet.
+    internet_probe_urls = ("https://www.google.com", "https://www.baidu.com")
+
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+        async def _probe(url: str) -> bool:
+            try:
+                async with session.get(url) as resp:
+                    return resp.status < 500
+            except (ClientError, OSError, asyncio.TimeoutError):
+                return False
+
+        probe_tasks = [asyncio.create_task(_probe(u)) for u in internet_probe_urls]
+        try:
+            for fut in asyncio.as_completed(probe_tasks):
+                if await fut:
+                    results["internet_accessible"] = True
+                    break
+        finally:
+            for t in probe_tasks:
+                if not t.done():
+                    t.cancel()
+            await asyncio.gather(*probe_tasks, return_exceptions=True)
+        if not results["internet_accessible"]:
+            return results
+
+        parsed = urlparse(default_base_url())
+        health_url = f"{parsed.scheme}://{parsed.netloc}/health"
+        with contextlib.suppress(ClientError, OSError):
+            async with session.get(health_url) as resp:
+                results["api_accessible"] = resp.status < 500
+    return results
+
+
 async def sleep_with_interrupt(
     seconds: float,
     node_cls: type[IO.ComfyNode] | None,

comfy_api_nodes/util/client.py

@@ -30,12 +30,14 @@ from . import request_logger
 from ._helpers import (
     _retry_after_wait,
     default_base_url,
+    diagnose_connectivity,
     get_comfy_api_headers,
     get_node_id,
     is_processing_interrupted,
     sleep_with_interrupt,
 )
 from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInterrupted
+from .download_helpers import download_url_to_bytesio
 
 M = TypeVar("M", bound=BaseModel)
 
@@ -75,6 +77,8 @@ class _RequestConfig:
     price_extractor: Callable[[dict[str, Any]], float | None] | None = None
     is_rate_limited: Callable[[int, Any], bool] | None = None
     response_header_validator: Callable[[dict[str, str]], None] | None = None
+    idempotency_key: str | None = None
+    asset_urls: bool = False
 
 
 @dataclass
@@ -93,6 +97,16 @@ class _PollUIState:
 _RETRY_STATUS = {408, 500, 502, 503, 504}  # status 429 is handled separately
 _MAX_RETRY_AFTER_WAIT = 150.0  # Cap a server Retry-After at this many seconds so a large hint can't block execution
 
+IDEMPOTENCY_KEY_HEADER = "Idempotency-Key"
+ASSET_FORMAT_HEADER = "Comfy-Asset-Format"
+ASSET_FORMAT_URL = "url"
+_IDEMPOTENT_REPLAYED_HEADER = "Idempotent-Replayed"
+_ERROR_TYPE_HEADER = "X-Comfy-Error-Type"
+_IDEMPOTENCY_IN_FLIGHT = "idempotency_in_flight"
+_IDEMPOTENCY_TERMINAL = frozenset({"idempotency_consumed", "idempotency_mismatch"})
+_IDEMPOTENCY_IN_FLIGHT_WAIT = 2.0
+_IDEMPOTENCY_IN_FLIGHT_MAX_WAIT = 30.0
+
 PRICE_CREDITS_HEADER = "X-Comfy-Credits-Used"
 """Proxy response header with the actual cost in Comfy credits. When present on any successful proxied response,
 it takes precedence over ``price_extractor``."""
@@ -183,6 +197,7 @@ async def sync_op(
     monitor_progress: bool = True,
     max_retries_on_rate_limit: int = 16,
     is_rate_limited: Callable[[int, Any], bool] | None = None,
+    asset_urls: bool = False,
 ) -> M:
     raw = await sync_op_raw(
         cls,
@@ -203,6 +218,7 @@ async def sync_op(
         monitor_progress=monitor_progress,
         max_retries_on_rate_limit=max_retries_on_rate_limit,
         is_rate_limited=is_rate_limited,
+        asset_urls=asset_urls,
     )
     if not isinstance(raw, dict):
         raise Exception("Expected JSON response to validate into a Pydantic model, got non-JSON (binary or text).")
@@ -279,12 +295,16 @@ async def sync_op_raw(
     max_retries_on_rate_limit: int = 16,
     is_rate_limited: Callable[[int, Any], bool] | None = None,
     response_header_validator: Callable[[dict[str, str]], None] | None = None,
+    idempotent: bool = True,
+    asset_urls: bool = False,
 ) -> dict[str, Any] | bytes:
     """
     Make a single network request.
       - If as_binary=False (default): returns JSON dict (or {'_raw': '<text>'} if non-JSON).
       - If as_binary=True: returns bytes.
       - response_header_validator: optional callback receiving response headers dict
+      - asset_urls=True: asks the proxy for Comfy-hosted URLs in place of inline media; a JSON {"url": ...}
+        answer to an as_binary request is downloaded and returned as the bytes.
     """
     if isinstance(data, BaseModel):
         data = data.model_dump(exclude_none=True)
@@ -310,6 +330,8 @@ async def sync_op_raw(
         max_retries_on_rate_limit=max_retries_on_rate_limit,
         is_rate_limited=is_rate_limited,
         response_header_validator=response_header_validator,
+        idempotency_key=uuid.uuid4().hex if idempotent and endpoint.method != "GET" else None,
+        asset_urls=asset_urls,
     )
     return await _request_base(cfg, expect_binary=as_binary)
 
@@ -418,6 +440,7 @@ async def poll_op_raw(
                     as_binary=False,
                     final_label_on_success=None,
                     monitor_progress=False,
+                    idempotent=False,
                 )
                 if not isinstance(resp_json, dict):
                     raise Exception("Polling endpoint returned non-JSON response.")
@@ -433,6 +456,7 @@ async def poll_op_raw(
                             as_binary=False,
                             final_label_on_success=None,
                             monitor_progress=False,
+                            idempotent=False,
                         )
                 raise
 
@@ -522,6 +546,7 @@ async def poll_op_raw(
                             as_binary=False,
                             final_label_on_success=None,
                             monitor_progress=False,
+                            idempotent=False,
                         )
                 raise
             if not is_queued:
@@ -609,46 +634,50 @@ def _estimate_progress_pct(elapsed_seconds: float, p50_seconds: int | None, p90_
     return 90 + int(5.0 * (elapsed_seconds - p50_seconds) / (horizon - p50_seconds))
 
 
-async def _diagnose_connectivity() -> dict[str, bool]:
-    """Best-effort connectivity diagnostics to distinguish local vs. server issues."""
-    results = {
-        "internet_accessible": False,
-        "api_accessible": False,
-    }
-    timeout = aiohttp.ClientTimeout(total=5.0)
+def _normalize_files(files: dict[str, Any] | list[tuple[str, Any]]) -> list[tuple[str, str, Any, str]]:
+    """Flatten `files` into (field_name, filename, value, content_type) once per request.
 
-    # Probe Google and Baidu in parallel: Google is blocked by the GFW in mainland China, so a Baidu probe is required
-    # to correctly detect that Chinese users with working internet do have working internet.
-    internet_probe_urls = ("https://www.google.com", "https://www.baidu.com")
+    File-like values are read into bytes here because aiohttp closes IOBase payloads
+    after sending, which would break re-sending the same body on retry.
+    """
+    out = []
+    file_iter = files if isinstance(files, list) else files.items()
+    for field_name, file_obj in file_iter:
+        if file_obj is None:
+            continue
+        if isinstance(file_obj, tuple):
+            filename, file_value, content_type = _unpack_tuple(file_obj)
+        else:
+            filename = getattr(file_obj, "name", field_name)
+            file_value = file_obj
+            content_type = "application/octet-stream"
+        if hasattr(file_value, "read"):
+            with contextlib.suppress(Exception):
+                file_value.seek(0)
+            data = file_value.read()
+            if not isinstance(file_value, BytesIO):
+                with contextlib.suppress(Exception):
+                    file_value.close()
+            file_value = data
+        out.append((field_name, filename, file_value, content_type))
+    return out
 
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        async def _probe(url: str) -> bool:
-            try:
-                async with session.get(url) as resp:
-                    return resp.status < 500
-            except (ClientError, OSError, asyncio.TimeoutError):
-                return False
 
-        probe_tasks = [asyncio.create_task(_probe(u)) for u in internet_probe_urls]
-        try:
-            for fut in asyncio.as_completed(probe_tasks):
-                if await fut:
-                    results["internet_accessible"] = True
-                    break
-        finally:
-            for t in probe_tasks:
-                if not t.done():
-                    t.cancel()
-            await asyncio.gather(*probe_tasks, return_exceptions=True)
-        if not results["internet_accessible"]:
-            return results
-
-        parsed = urlparse(default_base_url())
-        health_url = f"{parsed.scheme}://{parsed.netloc}/health"
-        with contextlib.suppress(ClientError, OSError):
-            async with session.get(health_url) as resp:
-                results["api_accessible"] = resp.status < 500
-    return results
+def _build_multipart_form(cfg: _RequestConfig, files: list[tuple[str, str, Any, str]]) -> aiohttp.FormData:
+    if cfg.multipart_parser and cfg.data:
+        form = cfg.multipart_parser(cfg.data)
+        if not isinstance(form, aiohttp.FormData):
+            raise ValueError("multipart_parser must return aiohttp.FormData")
+    else:
+        form = aiohttp.FormData(default_to_multipart=True)
+        if cfg.data:
+            for k, v in cfg.data.items():
+                if v is None:
+                    continue
+                form.add_field(k, str(v) if not isinstance(v, (bytes, bytearray)) else v)
+    for field_name, filename, file_value, content_type in files:
+        form.add_field(field_name, file_value, filename=filename, content_type=content_type)
+    return form
 
 
 def _unpack_tuple(t: tuple) -> tuple[str, Any, str]:
@@ -675,7 +704,18 @@ _TERMINAL_SERVICE_REFUSALS = frozenset({"comfy_cloud_provider_disabled"})
 
 
 def _is_terminal_service_refusal(body: Any) -> bool:
-    return isinstance(body, dict) and body.get("error") in _TERMINAL_SERVICE_REFUSALS
+    if not isinstance(body, dict):
+        return False
+    err = body.get("error")
+    return isinstance(err, str) and err in _TERMINAL_SERVICE_REFUSALS
+
+
+def _response_detail(body: Any) -> str:
+    if isinstance(body, dict):
+        for key in ("detail", "message"):
+            if isinstance(body.get(key), str) and body[key]:
+                return body[key]
+    return ""
 
 
 def _provider_error_detail(metadata: Any) -> str | None:
@@ -707,6 +747,8 @@ def _friendly_http_message(status: int, body: Any) -> str:
         return "There is a problem with your account. Please contact support@comfy.org."
     if status == 429:
         return "Rate Limit Exceeded: The server returned 429 after all retry attempts. Please wait and try again."
+    if isinstance(body, dict) and body.get("error_type") == "service_unavailable":
+        return "The API server could not verify this request right now. Please try again in a moment."
     try:
         if isinstance(body, dict):
             err = body.get("error")
@@ -729,6 +771,9 @@ def _friendly_http_message(status: int, body: Any) -> str:
                     return f"API Error: {msg} (Type: {typ})"
                 if msg:
                     return f"API Error: {msg}"
+            detail = _response_detail(body)
+            if detail:
+                return f"API Error: {detail}"
             return f"API Error: {json.dumps(body)}"
         else:
             txt = str(body)
@@ -739,6 +784,17 @@ def _friendly_http_message(status: int, body: Any) -> str:
         return f"HTTP {status}: Unknown error"
 
 
+def _idempotency_error_message(error_type: str, body: Any) -> str:
+    if error_type == "idempotency_consumed":
+        return (
+            "The server could not return the result of the previous attempt of this request. "
+            f"Run the node again. ({error_type})"
+        )
+    detail = _response_detail(body)
+    suffix = f"{error_type}: {detail}" if detail else error_type
+    return f"The server rejected the retry of this request as different from the original attempt. ({suffix})"
+
+
 def _generate_operation_id(method: str, path: str, attempt: int) -> str:
     slug = path.strip("/").replace("/", "_") or "op"
     return f"{method}_{slug}_try{attempt}_{uuid.uuid4().hex[:8]}"
@@ -771,6 +827,43 @@ def _snapshot_request_body_for_logging(
     return data or {}
 
 
+async def _read_binary_body(resp: aiohttp.ClientResponse, cfg: _RequestConfig, start_time: float) -> bytes:
+    buff = bytearray()
+    last_tick = time.monotonic()
+    async for chunk in resp.content.iter_chunked(64 * 1024):
+        buff.extend(chunk)
+        now = time.monotonic()
+        if now - last_tick >= 1.0:
+            last_tick = now
+            if is_processing_interrupted():
+                raise ProcessingInterrupted("Task cancelled")
+            if cfg.monitor_progress:
+                _display_time_progress(cfg.node_cls, cfg.wait_label, int(now - start_time))
+    return bytes(buff)
+
+
+def _asset_url_from_envelope(resp: aiohttp.ClientResponse, body: bytes) -> str | None:
+    if resp.content_type != "application/json":
+        return None
+    try:
+        payload = json.loads(body)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return None
+    asset_url = payload.get("url") if isinstance(payload, dict) else None
+    return asset_url if isinstance(asset_url, str) and asset_url else None
+
+
+async def _download_asset(asset_url: str, cfg: _RequestConfig) -> bytes:
+    buf = BytesIO()
+    try:
+        await download_url_to_bytesio(asset_url, buf, timeout=None, cls=cfg.node_cls)
+    except (ProcessingInterrupted, LocalNetworkError, ApiServerError):
+        raise
+    except Exception as e:
+        raise Exception(f"The request completed, but its result could not be downloaded: {e}") from e
+    return buf.getvalue()
+
+
 async def _request_base(cfg: _RequestConfig, expect_binary: bool):
     """Core request with retries, per-second interruption monitoring, true cancellation, and friendly errors."""
     url = cfg.endpoint.path
@@ -781,6 +874,8 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
 
     method = cfg.endpoint.method
     params = _merge_params(cfg.endpoint.query_params, method, cfg.data if method == "GET" else None)
+    keyed = is_comfy_api_request and bool(cfg.idempotency_key)
+    multipart_files = _normalize_files(cfg.files) if cfg.content_type == "multipart/form-data" and method != "GET" and cfg.files else []
 
     async def _monitor(stop_evt: asyncio.Event, start_ts: float):
         """Every second: update elapsed time and signal interruption."""
@@ -795,15 +890,19 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
             return  # normal shutdown
 
     start_time = cfg.progress_origin_ts if cfg.progress_origin_ts is not None else time.monotonic()
+    first_attempt_ts = time.monotonic()
     attempt = 0
+    retries_used = 0
     delay = cfg.retry_delay
     rate_limit_attempts = 0
     rate_limit_delay = cfg.retry_delay
+    in_flight_waits = 0
     operation_succeeded: bool = False
     final_elapsed_seconds: int | None = None
     extracted_price: float | None = None
     while True:
         attempt += 1
+        attempt_ts = time.monotonic()
         stop_event = asyncio.Event()
         monitor_task: asyncio.Task | None = None
         sess: aiohttp.ClientSession | None = None
@@ -816,11 +915,19 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
             payload_headers.update(get_comfy_api_headers(cfg.node_cls))
         if cfg.endpoint.headers:
             payload_headers.update(cfg.endpoint.headers)
+        if keyed:
+            payload_headers[IDEMPOTENCY_KEY_HEADER] = cfg.idempotency_key
+        if cfg.asset_urls and is_comfy_api_request:
+            payload_headers[ASSET_FORMAT_HEADER] = ASSET_FORMAT_URL
 
         payload_kw: dict[str, Any] = {"headers": payload_headers}
         if method == "GET":
             payload_headers.pop("Content-Type", None)
-        request_body_log = _snapshot_request_body_for_logging(cfg.content_type, method, cfg.data, cfg.files)
+        request_body_log = (
+            _snapshot_request_body_for_logging(cfg.content_type, method, cfg.data, cfg.files)
+            if in_flight_waits == 0
+            else None
+        )
         try:
             if cfg.monitor_progress:
                 monitor_task = asyncio.create_task(_monitor(stop_event, start_time))
@@ -829,36 +936,8 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
             sess = aiohttp.ClientSession(timeout=timeout)
 
             if cfg.content_type == "multipart/form-data" and method != "GET":
-                # aiohttp will set Content-Type boundary; remove any fixed Content-Type
                 payload_headers.pop("Content-Type", None)
-                if cfg.multipart_parser and cfg.data:
-                    form = cfg.multipart_parser(cfg.data)
-                    if not isinstance(form, aiohttp.FormData):
-                        raise ValueError("multipart_parser must return aiohttp.FormData")
-                else:
-                    form = aiohttp.FormData(default_to_multipart=True)
-                    if cfg.data:
-                        for k, v in cfg.data.items():
-                            if v is None:
-                                continue
-                            form.add_field(k, str(v) if not isinstance(v, (bytes, bytearray)) else v)
-                if cfg.files:
-                    file_iter = cfg.files if isinstance(cfg.files, list) else cfg.files.items()
-                    for field_name, file_obj in file_iter:
-                        if file_obj is None:
-                            continue
-                        if isinstance(file_obj, tuple):
-                            filename, file_value, content_type = _unpack_tuple(file_obj)
-                        else:
-                            filename = getattr(file_obj, "name", field_name)
-                            file_value = file_obj
-                            content_type = "application/octet-stream"
-                        # Attempt to rewind BytesIO for retries
-                        if isinstance(file_value, BytesIO):
-                            with contextlib.suppress(Exception):
-                                file_value.seek(0)
-                        form.add_field(field_name, file_value, filename=filename, content_type=content_type)
-                payload_kw["data"] = form
+                payload_kw["data"] = _build_multipart_form(cfg, multipart_files)
             elif cfg.content_type == "application/x-www-form-urlencoded" and method != "GET":
                 payload_headers["Content-Type"] = "application/x-www-form-urlencoded"
                 payload_kw["data"] = cfg.data or {}
@@ -893,18 +972,63 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
             # Otherwise, request finished
             resp = await req_task
             async with resp:
+                if keyed and resp.headers.get(_IDEMPOTENT_REPLAYED_HEADER):
+                    logging.info("Server replayed the response of a previous attempt for %s %s", method, url)
                 if resp.status >= 400:
                     try:
                         body = await resp.json()
                     except (ContentTypeError, json.JSONDecodeError):
                         body = await resp.text()
+                    error_type = resp.headers.get(_ERROR_TYPE_HEADER) or (body.get("error_type") if isinstance(body, dict) else "")
+                    error_type = error_type.strip().lower() if isinstance(error_type, str) else ""
+                    if keyed and resp.status == 409 and error_type in _IDEMPOTENCY_TERMINAL:
+                        msg = _idempotency_error_message(error_type, body)
+                        request_logger.log_request_response(
+                            operation_id=operation_id,
+                            request_method=method,
+                            request_url=url,
+                            response_status_code=resp.status,
+                            response_headers=dict(resp.headers),
+                            response_content=body,
+                            error_message=msg,
+                        )
+                        raise Exception(msg)
                     should_retry = False
+                    in_flight = False
                     wait_time = 0.0
+                    remaining = 0.0
                     retry_label = ""
                     is_rl = resp.status == 429 or (
                         cfg.is_rate_limited is not None and cfg.is_rate_limited(resp.status, body)
                     )
-                    if is_rl and rate_limit_attempts < cfg.max_retries_on_rate_limit:
+                    if keyed and resp.status == 409 and error_type == _IDEMPOTENCY_IN_FLIGHT:
+                        remaining = cfg.timeout - (time.monotonic() - first_attempt_ts)
+                        if remaining <= 0:
+                            msg = (
+                                "The server is still processing the previous attempt of this request "
+                                "and did not finish within the node's timeout."
+                            )
+                            request_logger.log_request_response(
+                                operation_id=operation_id,
+                                request_method=method,
+                                request_url=url,
+                                response_status_code=resp.status,
+                                response_headers=dict(resp.headers),
+                                response_content=body,
+                                error_message=msg,
+                            )
+                            raise Exception(msg)
+                        in_flight_waits += 1
+                        in_flight = True
+                        retries_used = 0
+                        delay = cfg.retry_delay
+                        wait_time = min(
+                            _IDEMPOTENCY_IN_FLIGHT_MAX_WAIT,
+                            _IDEMPOTENCY_IN_FLIGHT_WAIT * 1.5 ** (in_flight_waits - 1),
+                        )
+                        retry_label = f"previous attempt still in progress, check {in_flight_waits}"
+                        should_retry = True
+                    elif is_rl and rate_limit_attempts < cfg.max_retries_on_rate_limit:
                         rate_limit_attempts += 1
                         wait_time = min(rate_limit_delay, 30.0)
                         rate_limit_delay *= cfg.retry_backoff
@@ -913,16 +1037,23 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                     elif (
                         resp.status in _RETRY_STATUS
                         and not _is_terminal_service_refusal(body)
-                        and (attempt - rate_limit_attempts) <= cfg.max_retries
+                        and retries_used < cfg.max_retries
                     ):
+                        retries_used += 1
                         wait_time = delay
                         delay *= cfg.retry_backoff
-                        retry_label = f"retry {attempt - rate_limit_attempts} of {cfg.max_retries}"
+                        retry_label = f"retry {retries_used} of {cfg.max_retries}"
                         should_retry = True
 
                     if should_retry:
-                        wait_time = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
-                        logging.warning(
+                        retry_after = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
+                        if in_flight:
+                            wait_time = min(max(wait_time, retry_after), remaining)
+                        else:
+                            wait_time = retry_after
+                            if keyed and resp.headers.get(_IDEMPOTENT_REPLAYED_HEADER):
+                                cfg.idempotency_key = uuid.uuid4().hex
+                        (logging.info if in_flight else logging.warning)(
                             "HTTP %s %s -> %s. Waiting %.2fs (%s).",
                             method,
                             url,
@@ -960,19 +1091,9 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                     raise Exception(msg)
 
                 if expect_binary:
-                    buff = bytearray()
-                    last_tick = time.monotonic()
-                    async for chunk in resp.content.iter_chunked(64 * 1024):
-                        buff.extend(chunk)
-                        now = time.monotonic()
-                        if now - last_tick >= 1.0:
-                            last_tick = now
-                            if is_processing_interrupted():
-                                raise ProcessingInterrupted("Task cancelled")
-                            if cfg.monitor_progress:
-                                _display_time_progress(cfg.node_cls, cfg.wait_label, int(now - start_time))
-                    bytes_payload = bytes(buff)
+                    bytes_payload = await _read_binary_body(resp, cfg, start_time)
                     resp_headers = {k.lower(): v for k, v in resp.headers.items()}
+                    asset_url = _asset_url_from_envelope(resp, bytes_payload) if cfg.asset_urls else None
                     if is_comfy_api_request:
                         _maybe_remember_credits_used(cfg.node_cls, resp.headers.get(PRICE_CREDITS_HEADER))
                         _maybe_remember_server_estimate(cfg.node_cls, resp.headers)
@@ -981,8 +1102,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                             extracted_price = cfg.price_extractor(resp_headers)
                     if cfg.response_header_validator:
                         cfg.response_header_validator(resp_headers)
-                    operation_succeeded = True
-                    final_elapsed_seconds = int(time.monotonic() - start_time)
                     request_logger.log_request_response(
                         operation_id=operation_id,
                         request_method=method,
@@ -991,6 +1110,10 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                         response_headers=resp_headers,
                         response_content=bytes_payload,
                     )
+                    if asset_url:
+                        bytes_payload = await _download_asset(asset_url, cfg)
+                    operation_succeeded = True
+                    final_elapsed_seconds = int(time.monotonic() - start_time)
                     return bytes_payload
                 else:
                     try:
@@ -1023,15 +1146,18 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
         except ProcessingInterrupted:
             logging.debug("Polling was interrupted by user")
             raise
-        except (ClientError, OSError) as e:
-            if (attempt - rate_limit_attempts) <= cfg.max_retries:
+        except (ClientError, OSError, asyncio.TimeoutError) as e:
+            if retries_used < cfg.max_retries:
+                retries_used += 1
                 logging.warning(
-                    "Connection error calling %s %s. Retrying in %.2fs (%d/%d): %s",
+                    "Connection error calling %s %s after %.1fs. Retrying in %.2fs (%d/%d): %s: %s",
                     method,
                     url,
+                    time.monotonic() - attempt_ts,
                     delay,
-                    attempt - rate_limit_attempts,
+                    retries_used,
                     cfg.max_retries,
+                    type(e).__name__,
                     str(e),
                 )
                 request_logger.log_request_response(
@@ -1052,7 +1178,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                 )
                 delay *= cfg.retry_backoff
                 continue
-            diag = await _diagnose_connectivity()
+            diag = await diagnose_connectivity()
             if not diag["internet_accessible"]:
                 request_logger.log_request_response(
                     operation_id=operation_id,

comfy_api_nodes/util/download_helpers.py

@@ -17,12 +17,12 @@ from folder_paths import get_output_directory
 from . import request_logger
 from ._helpers import (
     default_base_url,
+    diagnose_connectivity,
     get_comfy_api_headers,
     is_processing_interrupted,
     sleep_with_interrupt,
     to_aiohttp_url,
 )
-from .client import _diagnose_connectivity
 from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInterrupted
 from .conversions import bytesio_to_image_tensor
 
@@ -158,24 +158,29 @@ async def download_url_to_bytesio(
                     sink = dest  # BytesIO or file-like
 
                 written = 0
-                while True:
+                read_task: asyncio.Task | None = None
                 try:
-                        chunk = await asyncio.wait_for(resp.content.read(1024 * 1024), timeout=1.0)
-                    except asyncio.TimeoutError:
-                        chunk = b""
-                    except asyncio.CancelledError:
-                        raise ProcessingInterrupted("Task cancelled") from None
-
+                    while True:
+                        if read_task is None:
+                            read_task = asyncio.create_task(resp.content.read(1024 * 1024))
+                        done, _ = await asyncio.wait({read_task}, timeout=1.0)
                         if is_processing_interrupted():
                             raise ProcessingInterrupted("Task cancelled")
-
+                        if not done:
+                            continue
+                        chunk = read_task.result()
+                        read_task = None
                         if not chunk:
                             if resp.content.at_eof():
                                 break
                             continue
-
                         sink.write(chunk)
                         written += len(chunk)
+                finally:
+                    if read_task is not None:
+                        read_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError, Exception):
+                            await read_task
 
                 if isinstance(dest, BytesIO):
                     with contextlib.suppress(Exception):
@@ -192,7 +197,7 @@ async def download_url_to_bytesio(
                 return
         except asyncio.CancelledError:
             raise ProcessingInterrupted("Task cancelled") from None
-        except (ClientError, OSError) as e:
+        except (ClientError, OSError, asyncio.TimeoutError) as e:
             if attempt <= max_retries:
                 request_logger.log_request_response(
                     operation_id=op_id,
@@ -204,7 +209,7 @@ async def download_url_to_bytesio(
                 delay *= retry_backoff
                 continue
 
-            diag = await _diagnose_connectivity()
+            diag = await diagnose_connectivity()
             if not diag["internet_accessible"]:
                 raise LocalNetworkError(
                     "Unable to connect to the network. Please check your internet connection and try again."

comfy_api_nodes/util/upload_helpers.py

@@ -13,10 +13,9 @@ from pydantic import BaseModel, Field
 from comfy_api.latest import IO, Input, Types
 
 from . import request_logger
-from ._helpers import is_processing_interrupted, sleep_with_interrupt
+from ._helpers import diagnose_connectivity, is_processing_interrupted, sleep_with_interrupt
 from .client import (
     ApiEndpoint,
-    _diagnose_connectivity,
     _display_time_progress,
     sync_op,
 )
@@ -366,7 +365,7 @@ async def upload_file(
                 delay *= retry_backoff
                 continue
 
-            diag = await _diagnose_connectivity()
+            diag = await diagnose_connectivity()
             if not diag["internet_accessible"]:
                 raise LocalNetworkError(
                     "Unable to connect to the network. Please check your internet connection and try again."

comfy_extras/nodes_cosmos.py

@@ -13,6 +13,7 @@ class EmptyCosmosLatentVideo(io.ComfyNode):
     def define_schema(cls) -> io.Schema:
         return io.Schema(
             node_id="EmptyCosmosLatentVideo",
+            display_name="Empty Cosmos Latent Video",
             category="model/latent/cosmos",
             inputs=[
                 io.Int.Input("width", default=1280, min=16, max=nodes.MAX_RESOLUTION, step=16),

comfy_extras/nodes_hunyuan.py

@@ -233,8 +233,8 @@ class HunyuanVideo15LatentUpscaleWithModel(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="HunyuanVideo15LatentUpscaleWithModel",
-            display_name="Hunyuan Video 15 Latent Upscale With Model",
-            category="model/latent/hunyhuan video",
+            display_name="Hunyuan Video 1.5 Latent Upscale With Model",
+            category="model/latent/hunyuan video",
             inputs=[
                 io.LatentUpscaleModel.Input("model"),
                 io.Latent.Input("samples"),
@@ -366,6 +366,7 @@ class EmptyHunyuanImageLatent(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="EmptyHunyuanImageLatent",
+            display_name="Empty Hunyuan Image Latent",
             category="model/latent/hunyuan image",
             inputs=[
                 io.Int.Input("width", default=2048, min=64, max=nodes.MAX_RESOLUTION, step=32),

comfy_extras/nodes_hunyuan3d.py

@@ -12,6 +12,7 @@ class EmptyLatentHunyuan3Dv2(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="EmptyLatentHunyuan3Dv2",
+            display_name="Empty Hunyuan 3D v2 Latent",
             category="model/latent/hunyuan 3d",
             inputs=[
                 IO.Int.Input("resolution", default=3072, min=1, max=8192),
@@ -97,6 +98,7 @@ class VAEDecodeHunyuan3D(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VAEDecodeHunyuan3D",
+            display_name="Hunyuan 3D VAE Decode",
             category="model/latent/hunyuan 3d",
             inputs=[
                 IO.Latent.Input("samples"),

comfy_extras/nodes_lt.py

@@ -65,6 +65,7 @@ class EmptyLTXVLatentVideo(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="EmptyLTXVLatentVideo",
+            display_name="Empty LTXV Latent Video",
             category="model/latent/ltxv",
             inputs=[
                 io.Int.Input("width", default=768, min=64, max=nodes.MAX_RESOLUTION, step=32),

comfy_extras/nodes_lt_audio.py

@@ -95,7 +95,7 @@ class LTXVEmptyLatentAudio(io.ComfyNode):
     def define_schema(cls) -> io.Schema:
         return io.Schema(
             node_id="LTXVEmptyLatentAudio",
-            display_name="LTXV Empty Latent Audio",
+            display_name="Empty LTXV Latent Audio",
             category="model/latent/ltxv",
             inputs=[
                 io.Int.Input(

comfy_extras/nodes_qwen.py

@@ -12,6 +12,7 @@ class TextEncodeQwenImageEdit(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="TextEncodeQwenImageEdit",
+            display_name="Text Encode Qwen Image Edit",
             category="model/conditioning/qwen image",
             inputs=[
                 io.Clip.Input("clip"),
@@ -55,6 +56,7 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="TextEncodeQwenImageEditPlus",
+            display_name="Text Encode Qwen Image Edit Plus",
             category="model/conditioning/qwen image",
             inputs=[
                 io.Clip.Input("clip"),

comfy_extras/nodes_seedvr.py

@@ -426,7 +426,7 @@ class SeedVR2TemporalChunk(io.ComfyNode):
         return io.Schema(
             node_id="SeedVR2TemporalChunk",
             display_name="Split SeedVR2 Latent",
-            category="model/latent/batch",
+            category="model/latent/seedvr",
             description="Split a SeedVR2 video latent into overlapping temporal chunks small enough to sample one at a time within VRAM, wiring latents outputs to both Apply SeedVR2 Conditioning and the sampler latent input before recombining with Merge SeedVR2 Latents.",
             search_aliases=["seedvr2", "split", "chunk", "temporal", "video upscale", "rebatch"],
             inputs=[
@@ -520,7 +520,7 @@ class SeedVR2TemporalMerge(io.ComfyNode):
         return io.Schema(
             node_id="SeedVR2TemporalMerge",
             display_name="Merge SeedVR2 Latents",
-            category="model/latent/batch",
+            category="model/latent/seedvr",
             is_input_list=True,
             description="Recombine sampled SeedVR2 latent temporal chunks into one latent, crossfading each overlap with a Hann window sized by the temporal_overlap wired from Split SeedVR2 Latent.",
             search_aliases=["seedvr2", "merge", "temporal", "hann", "crossfade"],

comfy_extras/nodes_stable_cascade.py

@@ -29,6 +29,7 @@ class StableCascade_EmptyLatentImage(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="StableCascade_EmptyLatentImage",
+            display_name="Empty Stable Cascade Latent Image",
             category="model/latent/stable cascade",
             inputs=[
                 io.Int.Input("width", default=1024, min=256, max=nodes.MAX_RESOLUTION, step=8),
@@ -58,6 +59,7 @@ class StableCascade_StageC_VAEEncode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="StableCascade_StageC_VAEEncode",
+            display_name="Stable Cascade Stage C VAE Encode",
             category="model/latent/stable cascade",
             inputs=[
                 io.Image.Input("image"),

comfy_extras/nodes_textgen.py

@@ -40,6 +40,7 @@ class TextGenerate(io.ComfyNode):
                 io.DynamicCombo.Input("sampling_mode", options=sampling_options, display_name="Sampling Mode"),
                 io.Boolean.Input("thinking", optional=True, default=False, tooltip="Operate in thinking mode if the model supports it."),
                 io.Boolean.Input("use_default_template", optional=True, default=True, tooltip="Use the built in system prompt/template if the model has one.", advanced=True),
+                io.Combo.Input("mtp", options=["auto", "off", "2", "3", "4", "5"], default="auto", optional=True, tooltip="Speculative decoding with the checkpoint's multi-token-prediction head. No effect without MTP weights. auto adapts the draft depth; 2-5 pins it. Sampled output stays correctly distributed but differs from non-MTP output for the same seed."),
             ],
             outputs=[
                 io.String.Output(display_name="generated_text"),
@@ -47,7 +48,9 @@ class TextGenerate(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
+    def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None, mtp="auto") -> io.NodeOutput:
+
+        mtp = False if mtp == "off" else (True if mtp == "auto" else int(mtp))
 
         tokens = clip.tokenize(prompt, image=image, skip_template=not use_default_template, min_length=1, thinking=thinking, video=video, audio=audio)
 
@@ -71,7 +74,8 @@ class TextGenerate(io.ComfyNode):
             min_p=min_p,
             repetition_penalty=repetition_penalty,
             presence_penalty=presence_penalty,
-            seed=seed
+            seed=seed,
+            mtp=mtp
         )
 
         generated_text = clip.decode(generated_ids)
@@ -225,7 +229,7 @@ class TextGenerateLTX2Prompt(TextGenerate):
         )
 
     @classmethod
-    def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
+    def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None, mtp="auto") -> io.NodeOutput:
         # Gemma 3 and Gemma 4 use different chat-turn markers and image tokens.
         # The Gemma 4 text encoder is the LTX 2.4 path; Gemma 3 is LTX 2.0.
         is_gemma4 = "gemma4" in getattr(clip.tokenizer, "clip_name", "")
@@ -254,7 +258,7 @@ class TextGenerateLTX2Prompt(TextGenerate):
                 f"<start_of_turn>model\n"
             )
 
-        out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
+        out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio, mtp=mtp)
 
         # 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.

comfy_extras/nodes_trellis2.py

@@ -118,6 +118,7 @@ class VaeDecodeShapeTrellis(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VaeDecodeShapeTrellis",
+            display_name="Trellis2 VAE Decode Shape",
             category="model/latent/trellis",
             inputs=[
                 IO.Latent.Input("samples"),
@@ -196,6 +197,7 @@ class VaeDecodeTextureTrellis(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VaeDecodeTextureTrellis",
+            display_name="Trellis2 VAE Decode Texture",
             category="model/latent/trellis",
             inputs=[
                 IO.Latent.Input("samples"),
@@ -277,6 +279,7 @@ class VaeDecodeStructureTrellis2(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VaeDecodeStructureTrellis2",
+            display_name="Trellis2 VAE Decode Structure",
             category="model/latent/trellis",
             inputs=[
                 IO.Latent.Input("samples"),
@@ -318,7 +321,7 @@ class Trellis2UpsampleStage(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="Trellis2UpsampleStage",
-            category="model/conditioning/trellis2",
+            category="model/conditioning/trellis",
             display_name="Trellis2 Upsample Stage",
             inputs=[
                 IO.Conditioning.Input("positive"),
@@ -440,7 +443,8 @@ class Trellis2Conditioning(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="Trellis2Conditioning",
-            category="model/conditioning/trellis2",
+            display_name="Trellis2 Conditioning",
+            category="model/conditioning/trellis",
             inputs=[
                 IO.ClipVision.Input("clip_vision_model"),
                 IO.Image.Input("image", tooltip="Preprocessed image from ImageCropToMask (pad_factor=1.0 for TRELLIS.2)."),
@@ -501,7 +505,8 @@ class Trellis2ShapeStage(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="Trellis2ShapeStage",
-            category="model/conditioning/trellis2",
+            display_name="Trellis2 Shape Stage",
+            category="model/conditioning/trellis",
             inputs=[
                 IO.Conditioning.Input("positive"),
                 IO.Conditioning.Input("negative"),
@@ -567,7 +572,8 @@ class Trellis2TextureStage(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="Trellis2TextureStage",
-            category="model/conditioning/trellis2",
+            display_name="Trellis2 Texture Stage",
+            category="model/conditioning/trellis",
             inputs=[
                 IO.Conditioning.Input("positive"),
                 IO.Conditioning.Input("negative"),
@@ -623,6 +629,7 @@ class EmptyTrellis2LatentStructure(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="EmptyTrellis2LatentStructure",
+            display_name="Empty Trellis2 Latent Structure",
             category="model/latent/trellis",
             inputs=[
                 IO.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch."),
@@ -764,7 +771,8 @@ class Pixal3DConditioning(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="Pixal3DConditioning",
-            category="model/conditioning/trellis2",
+            display_name="Pixal3D Conditioning",
+            category="model/conditioning/trellis",
             inputs=[
                 IO.ClipVision.Input("clip_vision_model", tooltip="DINOv3 ViT-L/16 ClipVision."),
                 IO.Image.Input("image", tooltip="Preprocessed image from ImageCropToMask (pad_factor=1.1 for Pixal3D)."),
@@ -824,7 +832,7 @@ class Pixal3DMultiViewConditioning(IO.ComfyNode):
         return IO.Schema(
             node_id="Pixal3DMultiViewConditioning",
             display_name="Pixal3D Multi-View Conditioning",
-            category="model/conditioning/trellis2",
+            category="model/conditioning/trellis",
             inputs=[IO.ClipVision.Input("clip_vision_model", tooltip="DINOv3 ViT-L/16 ClipVision with bundled NAF weights."),
                     IO.Float.Input("fov", default=20.0, min=1.0, max=170.0, step=0.01, round=False,
                                    tooltip="Horizontal FOV in degrees of the views as framed: 20 for rig renders and most "

nodes.py

@@ -1248,8 +1248,8 @@ class EmptyLatentImage:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
-                "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
+                "width": ("INT", {"default": 1024, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
+                "height": ("INT", {"default": 1024, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
                 "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
             }
         }

requirements.txt

@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=17.0.0
-comfy-kitchen==0.2.34
+comfy-kitchen==0.2.35
 comfy-aimdo==0.5.5
 requests
 simpleeval>=1.0.0