ComfyUI v0.22.1

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

comfy/latent_formats.py

@@ -152,6 +152,11 @@ class StableAudio1(LatentFormat):
     latent_dimensions = 1
     temporal_downscale_ratio = 2048
 
+class StableAudio3(LatentFormat):
+    latent_channels = 256
+    latent_dimensions = 1
+    temporal_downscale_ratio = 4096
+
 class Flux(SD3):
     latent_channels = 16
     def __init__(self):

comfy/ldm/audio/dit.py

@@ -10,6 +10,17 @@ from torch import nn
 from torch.nn import functional as F
 import math
 import comfy.ops
+from .embedders import ExpoFourierFeatures
+
+
+def _left_pad_to_match(emb, target_len):
+    emb_len = emb.shape[-2]
+    if emb_len < target_len:
+        return F.pad(emb, (0, 0, target_len - emb_len, 0), value=0.)
+    elif emb_len > target_len:
+        return emb[:, -target_len:, :]
+    return emb
+
 
 class FourierFeatures(nn.Module):
     def __init__(self, in_features, out_features, std=1., dtype=None, device=None):
@@ -22,6 +33,7 @@ class FourierFeatures(nn.Module):
         f = 2 * math.pi * input @ comfy.ops.cast_to_input(self.weight.T, input)
         return torch.cat([f.cos(), f.sin()], dim=-1)
 
+
 # norms
 class LayerNorm(nn.Module):
     def __init__(self, dim, bias=False, fix_scale=False, dtype=None, device=None):
@@ -43,6 +55,16 @@ class LayerNorm(nn.Module):
             beta = comfy.ops.cast_to_input(beta, x)
         return F.layer_norm(x, x.shape[-1:], weight=comfy.ops.cast_to_input(self.gamma, x), bias=beta)
 
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, dtype=None, device=None):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+
+    def forward(self, x):
+        return F.rms_norm(x, x.shape[-1:], weight=comfy.ops.cast_to_input(self.gamma, x))
+
+
 class GLU(nn.Module):
     def __init__(
         self,
@@ -236,13 +258,6 @@ class FeedForward(nn.Module):
 
         linear_out = operations.Linear(inner_dim, dim_out, bias = not no_bias, dtype=dtype, device=device) if not use_conv else operations.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias, dtype=dtype, device=device)
 
-        # # init last linear layer to 0
-        # if zero_init_output:
-        #     nn.init.zeros_(linear_out.weight)
-        #     if not no_bias:
-        #         nn.init.zeros_(linear_out.bias)
-
-
         self.ff = nn.Sequential(
             linear_in,
             rearrange('b d n -> b n d') if use_conv else nn.Identity(),
@@ -261,8 +276,10 @@ class Attention(nn.Module):
         dim_context = None,
         causal = False,
         zero_init_output=True,
-        qk_norm = False,
+        qk_norm = "none",
+        differential = False,
         natten_kernel_size = None,
+        feat_scale = False,
         dtype=None,
         device=None,
         operations=None,
@@ -271,6 +288,7 @@ class Attention(nn.Module):
         self.dim = dim
         self.dim_heads = dim_heads
         self.causal = causal
+        self.differential = differential
 
         dim_kv = dim_context if dim_context is not None else dim
 
@@ -278,18 +296,37 @@ class Attention(nn.Module):
         self.kv_heads = dim_kv // dim_heads
 
         if dim_context is not None:
-            self.to_q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-            self.to_kv = operations.Linear(dim_kv, dim_kv * 2, bias=False, dtype=dtype, device=device)
+            if differential:
+                self.to_q = operations.Linear(dim, dim * 2, bias=False, dtype=dtype, device=device)
+                self.to_kv = operations.Linear(dim_kv, dim_kv * 3, bias=False, dtype=dtype, device=device)
+            else:
+                self.to_q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+                self.to_kv = operations.Linear(dim_kv, dim_kv * 2, bias=False, dtype=dtype, device=device)
         else:
-            self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device)
+            if differential:
+                self.to_qkv = operations.Linear(dim, dim * 5, bias=False, dtype=dtype, device=device)
+            else:
+                self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device)
 
         self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
 
-        # if zero_init_output:
-        #     nn.init.zeros_(self.to_out.weight)
-
+        # Accept bool for backward compat
+        if isinstance(qk_norm, bool):
+            qk_norm = "l2" if qk_norm else "none"
         self.qk_norm = qk_norm
 
+        if self.qk_norm == "ln":
+            self.q_norm = operations.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+            self.k_norm = operations.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+        elif self.qk_norm == "rms":
+            self.q_norm = RMSNorm(dim_heads, dtype=dtype, device=device)
+            self.k_norm = RMSNorm(dim_heads, dtype=dtype, device=device)
+
+        self.feat_scale = feat_scale
+
+        if self.feat_scale:
+            self.lambda_dc = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+            self.lambda_hf = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
 
     def forward(
         self,
@@ -306,22 +343,51 @@ class Attention(nn.Module):
         kv_input = context if has_context else x
 
         if hasattr(self, 'to_q'):
-            # Use separate linear projections for q and k/v
-            q = self.to_q(x)
-            q = rearrange(q, 'b n (h d) -> b h n d', h = h)
+            if self.differential:
+                # cross-attention differential: to_q → (q, q_diff), to_kv → (k, k_diff, v)
+                q, q_diff = self.to_q(x).chunk(2, dim=-1)
+                q      = rearrange(q,      'b n (h d) -> b h n d', h=h)
+                q_diff = rearrange(q_diff, 'b n (h d) -> b h n d', h=h)
+                q = torch.stack([q, q_diff], dim=1)  # (B, 2, H, N, D)
+                k, k_diff, v = self.to_kv(kv_input).chunk(3, dim=-1)
+                k      = rearrange(k,      'b n (h d) -> b h n d', h=kv_h)
+                k_diff = rearrange(k_diff, 'b n (h d) -> b h n d', h=kv_h)
+                v      = rearrange(v,      'b n (h d) -> b h n d', h=kv_h)
+                k = torch.stack([k, k_diff], dim=1)  # (B, 2, H, M, D)
+            else:
+                # Use separate linear projections for q and k/v
+                q = self.to_q(x)
+                q = rearrange(q, 'b n (h d) -> b h n d', h = h)
 
-            k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+                k, v = self.to_kv(kv_input).chunk(2, dim=-1)
 
-            k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
+                k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
         else:
-            # Use fused linear projection
-            q, k, v = self.to_qkv(x).chunk(3, dim=-1)
-            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+            if self.differential:
+                # self-attention differential: to_qkv → (q, k, v, q_diff, k_diff)
+                q, k, v, q_diff, k_diff = self.to_qkv(x).chunk(5, dim=-1)
+                q, k, v, q_diff, k_diff = map(
+                    lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h),
+                    (q, k, v, q_diff, k_diff)
+                )
+                q = torch.stack([q, q_diff], dim=1)  # (B, 2, H, N, D)
+                k = torch.stack([k, k_diff], dim=1)
+            else:
+                # Use fused linear projection
+                q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+                q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
 
         # Normalize q and k for cosine sim attention
-        if self.qk_norm:
+        if self.qk_norm == "l2":
             q = F.normalize(q, dim=-1)
             k = F.normalize(k, dim=-1)
+        elif self.qk_norm == "rms":
+            q_type, k_type = q.dtype, k.dtype
+            q = self.q_norm(q).to(q_type)
+            k = self.k_norm(k).to(k_type)
+        elif self.qk_norm != 'none':
+            q = self.q_norm(q)
+            k = self.k_norm(k)
 
         if rotary_pos_emb is not None and not has_context:
             freqs, _ = rotary_pos_emb
@@ -364,9 +430,24 @@ class Attention(nn.Module):
             heads_per_kv_head = h // kv_h
             k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
 
-        out = optimized_attention(q, k, v, h, skip_reshape=True, transformer_options=transformer_options)
+        if self.differential:
+            q, q_diff = q.unbind(dim=1)
+            k, k_diff = k.unbind(dim=1)
+            out      = optimized_attention(q,      k,      v, h, skip_reshape=True, transformer_options=transformer_options)
+            out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, transformer_options=transformer_options)
+            out = out - out_diff
+        else:
+            out = optimized_attention(q, k, v, h, skip_reshape=True, transformer_options=transformer_options)
+
         out = self.to_out(out)
 
+        if self.feat_scale:
+            out_dc = out.mean(dim=-2, keepdim=True)
+            out_hf = out - out_dc
+
+            # Selectively modulate DC and high frequency components
+            out = out + comfy.ops.cast_to_input(self.lambda_dc, out) * out_dc + comfy.ops.cast_to_input(self.lambda_hf, out) * out_hf
+
         if mask is not None:
             mask = rearrange(mask, 'b n -> b n 1')
             out = out.masked_fill(~mask, 0.)
@@ -417,11 +498,14 @@ class TransformerBlock(nn.Module):
             cross_attend = False,
             dim_context = None,
             global_cond_dim = None,
+            global_cond_shared_embed = False,
+            local_add_cond_dim = None,
             causal = False,
             zero_init_branch_outputs = True,
             conformer = False,
             layer_ix = -1,
             remove_norms = False,
+            norm_type = "layer_norm",
             attn_kwargs = {},
             ff_kwargs = {},
             norm_kwargs = {},
@@ -436,8 +520,20 @@ class TransformerBlock(nn.Module):
         self.cross_attend = cross_attend
         self.dim_context = dim_context
         self.causal = causal
+        self.global_cond_shared_embed = global_cond_shared_embed
 
-        self.pre_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
+        norm_layer_map = {
+            "layer_norm": LayerNorm,
+            "rms_norm": RMSNorm,
+        }
+        norm_cls = norm_layer_map.get(norm_type, LayerNorm)
+
+        def make_norm():
+            if remove_norms:
+                return nn.Identity()
+            return norm_cls(dim, dtype=dtype, device=device, **norm_kwargs)
+
+        self.pre_norm = make_norm()
 
         self.self_attn = Attention(
             dim,
@@ -451,7 +547,7 @@ class TransformerBlock(nn.Module):
         )
 
         if cross_attend:
-            self.cross_attend_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
+            self.cross_attend_norm = make_norm()
             self.cross_attn = Attention(
                 dim,
                 dim_heads = dim_heads,
@@ -464,37 +560,56 @@ class TransformerBlock(nn.Module):
                 **attn_kwargs
             )
 
-        self.ff_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
-        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, dtype=dtype, device=device, operations=operations,**ff_kwargs)
+        self.ff_norm = make_norm()
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, dtype=dtype, device=device, operations=operations, **ff_kwargs)
 
         self.layer_ix = layer_ix
 
         self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs) if conformer else None
 
-        self.global_cond_dim = global_cond_dim
+        # Global conditioning
+        self.has_global_cond = (global_cond_dim is not None) or global_cond_shared_embed
 
-        if global_cond_dim is not None:
+        if global_cond_shared_embed:
+            # SA3 style: learnable per-block additive bias; global_cond is pre-projected to (B, dim*6)
+            self.to_scale_shift_gate = nn.Parameter(torch.empty(dim * 6, device=device, dtype=dtype))
+        elif global_cond_dim is not None:
+            # SA1 style: per-block MLP projects global_cond → (B, dim*6)
             self.to_scale_shift_gate = nn.Sequential(
                 nn.SiLU(),
-                nn.Linear(global_cond_dim, dim * 6, bias=False)
+                operations.Linear(global_cond_dim, dim * 6, bias=False, device=device, dtype=dtype)
             )
 
-            nn.init.zeros_(self.to_scale_shift_gate[1].weight)
-            #nn.init.zeros_(self.to_scale_shift_gate_self[1].bias)
+        # Local additive conditioning (e.g. inpaint mask + masked latent)
+        self.local_add_cond_dim = local_add_cond_dim
+        if local_add_cond_dim is not None:
+            self.to_local_embed = nn.Sequential(
+                operations.Linear(local_add_cond_dim, dim, bias=True, dtype=dtype, device=device),
+                nn.SiLU(),
+                operations.Linear(dim, dim, bias=True, dtype=dtype, device=device),
+            )
+        else:
+            self.to_local_embed = None
 
     def forward(
         self,
         x,
         context = None,
         global_cond=None,
+        local_add_cond=None,
         mask = None,
         context_mask = None,
         rotary_pos_emb = None,
         transformer_options={}
     ):
-        if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None:
+        if self.has_global_cond and global_cond is not None:
+            if self.global_cond_shared_embed:
+                # global_cond already has shape (B, dim*6)
+                ssg = (comfy.ops.cast_to_input(self.to_scale_shift_gate, global_cond) + global_cond).unsqueeze(1)
+            else:
+                ssg = self.to_scale_shift_gate(global_cond).unsqueeze(1)
 
-            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = self.to_scale_shift_gate(global_cond).unsqueeze(1).chunk(6, dim = -1)
+            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = ssg.chunk(6, dim = -1)
 
             # self-attention with adaLN
             residual = x
@@ -510,6 +625,9 @@ class TransformerBlock(nn.Module):
             if self.conformer is not None:
                 x = x + self.conformer(x)
 
+            if local_add_cond is not None and self.to_local_embed is not None:
+                x = x + _left_pad_to_match(self.to_local_embed(local_add_cond), x.shape[-2])
+
             # feedforward with adaLN
             residual = x
             x = self.ff_norm(x)
@@ -527,6 +645,9 @@ class TransformerBlock(nn.Module):
             if self.conformer is not None:
                 x = x + self.conformer(x)
 
+            if local_add_cond is not None and self.to_local_embed is not None:
+                x = x + _left_pad_to_match(self.to_local_embed(local_add_cond), x.shape[-2])
+
             x = x + self.ff(self.ff_norm(x))
 
         return x
@@ -543,6 +664,8 @@ class ContinuousTransformer(nn.Module):
         cross_attend=False,
         cond_token_dim=None,
         global_cond_dim=None,
+        global_cond_shared_embed=False,
+        local_add_cond_dim=None,
         causal=False,
         rotary_pos_emb=True,
         zero_init_branch_outputs=True,
@@ -550,6 +673,7 @@ class ContinuousTransformer(nn.Module):
         use_sinusoidal_emb=False,
         use_abs_pos_emb=False,
         abs_pos_emb_max_length=10000,
+        num_memory_tokens=0,
         dtype=None,
         device=None,
         operations=None,
@@ -562,6 +686,8 @@ class ContinuousTransformer(nn.Module):
         self.depth = depth
         self.causal = causal
         self.layers = nn.ModuleList([])
+        self.num_memory_tokens = num_memory_tokens
+        self.global_cond_shared_embed = global_cond_shared_embed
 
         self.project_in = operations.Linear(dim_in, dim, bias=False, dtype=dtype, device=device) if dim_in is not None else nn.Identity()
         self.project_out = operations.Linear(dim, dim_out, bias=False, dtype=dtype, device=device) if dim_out is not None else nn.Identity()
@@ -577,7 +703,22 @@ class ContinuousTransformer(nn.Module):
 
         self.use_abs_pos_emb = use_abs_pos_emb
         if use_abs_pos_emb:
-            self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length)
+            self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length + num_memory_tokens)
+
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.empty(num_memory_tokens, dim, device=device, dtype=dtype))
+
+        # Shared global-cond embedder (SA3 style): projects (B, global_cond_dim) → (B, dim*6)
+        self.global_cond_embedder = None
+        if global_cond_shared_embed and global_cond_dim is not None:
+            self.global_cond_embedder = nn.Sequential(
+                operations.Linear(global_cond_dim, dim, bias=True, dtype=dtype, device=device),
+                nn.SiLU(),
+                operations.Linear(dim, dim * 6, bias=True, dtype=dtype, device=device),
+            )
+
+        # When using shared embed, TransformerBlocks use per-block Parameter (not per-block MLP)
+        block_global_cond_dim = None if global_cond_shared_embed else global_cond_dim
 
         for i in range(depth):
             self.layers.append(
@@ -586,7 +727,9 @@ class ContinuousTransformer(nn.Module):
                     dim_heads = dim_heads,
                     cross_attend = cross_attend,
                     dim_context = cond_token_dim,
-                    global_cond_dim = global_cond_dim,
+                    global_cond_dim = block_global_cond_dim,
+                    global_cond_shared_embed = global_cond_shared_embed,
+                    local_add_cond_dim = local_add_cond_dim,
                     causal = causal,
                     zero_init_branch_outputs = zero_init_branch_outputs,
                     conformer=conformer,
@@ -605,6 +748,7 @@ class ContinuousTransformer(nn.Module):
         prepend_embeds = None,
         prepend_mask = None,
         global_cond = None,
+        local_add_cond = None,
         return_info = False,
         **kwargs
     ):
@@ -632,7 +776,9 @@ class ContinuousTransformer(nn.Module):
 
                 mask = torch.cat((prepend_mask, mask), dim = -1)
 
-        # Attention layers
+        if self.num_memory_tokens > 0:
+            memory_tokens = comfy.ops.cast_to_input(self.memory_tokens, x).expand(batch, -1, -1)
+            x = torch.cat((memory_tokens, x), dim=1)
 
         if self.rotary_pos_emb is not None:
             rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=torch.float, device=x.device)
@@ -642,6 +788,10 @@ class ContinuousTransformer(nn.Module):
         if self.use_sinusoidal_emb or self.use_abs_pos_emb:
             x = x + self.pos_emb(x)
 
+        # Project global_cond once (SA3 shared-embed path)
+        if global_cond is not None and self.global_cond_embedder is not None:
+            global_cond = self.global_cond_embedder(global_cond)
+
         blocks_replace = patches_replace.get("dit", {})
         # Iterate over the transformer layers
         for i, layer in enumerate(self.layers):
@@ -654,12 +804,17 @@ class ContinuousTransformer(nn.Module):
                 out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": global_cond, "pe": rotary_pos_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
                 x = out["img"]
             else:
-                x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, context=context, transformer_options=transformer_options)
-            # x = checkpoint(layer, x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, **kwargs)
+                x = layer(x, rotary_pos_emb=rotary_pos_emb, global_cond=global_cond,
+                          local_add_cond=local_add_cond, context=context,
+                          transformer_options=transformer_options)
 
             if return_info:
                 info["hidden_states"].append(x)
 
+        # Strip memory tokens before projecting out
+        if self.num_memory_tokens > 0:
+            x = x[:, self.num_memory_tokens:, :]
+
         x = self.project_out(x)
 
         if return_info:
@@ -682,6 +837,7 @@ class AudioDiffusionTransformer(nn.Module):
         num_heads=24,
         transformer_type: tp.Literal["continuous_transformer"] = "continuous_transformer",
         global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
+        timestep_features_type: str = "learned",
         audio_model="",
         dtype=None,
         device=None,
@@ -696,7 +852,10 @@ class AudioDiffusionTransformer(nn.Module):
         # Timestep embeddings
         timestep_features_dim = 256
 
-        self.timestep_features = FourierFeatures(1, timestep_features_dim, dtype=dtype, device=device)
+        if timestep_features_type == "expo":
+            self.timestep_features = ExpoFourierFeatures(timestep_features_dim, 0.5, 10000.0)
+        else:
+            self.timestep_features = FourierFeatures(1, timestep_features_dim, dtype=dtype, device=device)
 
         self.to_timestep_embed = nn.Sequential(
             operations.Linear(timestep_features_dim, embed_dim, bias=True, dtype=dtype, device=device),
@@ -781,6 +940,7 @@ class AudioDiffusionTransformer(nn.Module):
         cross_attn_cond=None,
         cross_attn_cond_mask=None,
         input_concat_cond=None,
+        local_add_cond=None,
         global_embed=None,
         prepend_cond=None,
         prepend_cond_mask=None,
@@ -802,9 +962,13 @@ class AudioDiffusionTransformer(nn.Module):
             prepend_cond = self.to_prepend_embed(prepend_cond)
 
             prepend_inputs = prepend_cond
+            prepend_length = prepend_cond.shape[1]
             if prepend_cond_mask is not None:
                 prepend_mask = prepend_cond_mask
 
+        if local_add_cond is not None and local_add_cond.dim() == 3:
+            local_add_cond = local_add_cond.permute(0, 2, 1)
+
         if input_concat_cond is not None:
 
             # Interpolate input_concat_cond to the same length as x
@@ -850,7 +1014,7 @@ class AudioDiffusionTransformer(nn.Module):
         if self.transformer_type == "x-transformers":
             output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, **extra_args, **kwargs)
         elif self.transformer_type == "continuous_transformer":
-            output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, return_info=return_info, **extra_args, **kwargs)
+            output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, return_info=return_info, local_add_cond=local_add_cond, **extra_args, **kwargs)
 
             if return_info:
                 output, info = output
@@ -876,6 +1040,7 @@ class AudioDiffusionTransformer(nn.Module):
         context=None,
         context_mask=None,
         input_concat_cond=None,
+        local_add_cond=None,
         global_embed=None,
         negative_global_embed=None,
         prepend_cond=None,
@@ -890,6 +1055,7 @@ class AudioDiffusionTransformer(nn.Module):
                 cross_attn_cond=context,
                 cross_attn_cond_mask=context_mask,
                 input_concat_cond=input_concat_cond,
+                local_add_cond=local_add_cond,
                 global_embed=global_embed,
                 prepend_cond=prepend_cond,
                 prepend_cond_mask=prepend_cond_mask,

comfy/ldm/audio/embedders.py

@@ -31,15 +31,39 @@ def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
     )
 
 
+class ExpoFourierFeatures(nn.Module):
+    """Exponentially-spaced Fourier features (no learnable parameters)."""
+    def __init__(self, dim, min_freq=0.5, max_freq=10000.0):
+        super().__init__()
+        self.dim = dim
+        self.min_freq = min_freq
+        self.max_freq = max_freq
+
+    def forward(self, t):
+        in_dtype = t.dtype
+        t = t.float()
+        if t.dim() == 1:
+            t = t.unsqueeze(-1)
+        half_dim = self.dim // 2
+        ramp = torch.linspace(0, 1, half_dim, device=t.device, dtype=torch.float32)
+        freqs = torch.exp(ramp * (math.log(self.max_freq) - math.log(self.min_freq)) + math.log(self.min_freq))
+        args = t * freqs * 2 * math.pi
+        return torch.cat([args.cos(), args.sin()], dim=-1).to(in_dtype)
+
+
 class NumberEmbedder(nn.Module):
     def __init__(
         self,
         features: int,
         dim: int = 256,
+        fourier_features_type="learned",
     ):
         super().__init__()
         self.features = features
-        self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
+        if fourier_features_type == "expo":
+            self.embedding = nn.Sequential(ExpoFourierFeatures(dim=dim), comfy.ops.manual_cast.Linear(in_features=dim, out_features=features))
+        else:
+            self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
 
     def forward(self, x: Union[List[float], Tensor]) -> Tensor:
         if not torch.is_tensor(x):
@@ -77,14 +101,15 @@ class NumberConditioner(Conditioner):
     def __init__(self,
                 output_dim: int,
                 min_val: float=0,
-                max_val: float=1
+                max_val: float=1,
+                fourier_features_type: str = "learned",
                 ):
         super().__init__(output_dim, output_dim)
 
         self.min_val = min_val
         self.max_val = max_val
 
-        self.embedder = NumberEmbedder(features=output_dim)
+        self.embedder = NumberEmbedder(features=output_dim, fourier_features_type=fourier_features_type)
 
     def forward(self, floats, device=None):
             # Cast the inputs to floats

comfy/ldm/audio/vae_sa3.py

@@ -0,0 +1,533 @@
+import torch
+import torch.nn as nn
+
+import comfy.ops
+import comfy.model_management
+from comfy.ldm.modules.attention import optimized_attention
+from comfy.ldm.audio.autoencoder import WNConv1d
+
+ops = comfy.ops.disable_weight_init
+
+class Transpose(nn.Module):
+    def forward(self, x, **kwargs):
+        return x.transpose(-2, -1)
+
+
+def _zero_pad_modulo_sequence(x, size, dim=-2):
+    input_len = x.shape[dim]
+    pad_len = (size - input_len % size) % size
+    if pad_len > 0:
+        pad_shape = list(x.shape)
+        pad_shape[dim] = pad_len
+        x = torch.cat([x, torch.zeros(pad_shape, device=x.device, dtype=x.dtype)], dim=dim)
+    return x
+
+
+def _sliding_window_mask(seq_len, window, device, dtype):
+    """Additive attention mask enforcing a ±window local window (matches flash_attn window_size)."""
+    i = torch.arange(seq_len, device=device).unsqueeze(1)
+    j = torch.arange(seq_len, device=device).unsqueeze(0)
+    out_of_window = (j - i).abs() > window
+    return torch.where(
+        out_of_window,
+        torch.full((1,), torch.finfo(dtype).min / 4, device=device, dtype=dtype),
+        torch.zeros(1, device=device, dtype=dtype),
+    )
+
+
+class DynamicTanh(nn.Module):
+    def __init__(self, dim, init_alpha=4.0, dtype=None, device=None, **kwargs):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.empty(1, dtype=dtype, device=device))
+        self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+        self.beta = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+
+    def forward(self, x):
+        alpha = comfy.ops.cast_to_input(self.alpha, x)
+        gamma = comfy.ops.cast_to_input(self.gamma, x)
+        beta = comfy.ops.cast_to_input(self.beta, x)
+        return gamma * torch.tanh(alpha * x) + beta
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, base=10000, base_rescale_factor=1., dtype=None, device=None):
+        super().__init__()
+        base = base * base_rescale_factor ** (dim / (dim - 2))
+        self.register_buffer("inv_freq", torch.empty(dim // 2, dtype=dtype, device=device))
+
+    def forward_from_seq_len(self, seq_len, device, dtype=None):
+        t = torch.arange(seq_len, device=device, dtype=torch.float32)
+        return self.forward(t)
+
+    def forward(self, t):
+        freqs = torch.outer(t.float(), comfy.model_management.cast_to(self.inv_freq, dtype=torch.float32, device=t.device))
+        freqs = torch.cat((freqs, freqs), dim=-1)
+        return freqs, 1.
+
+
+def _rotate_half(x):
+    d = x.shape[-1] // 2
+    return torch.cat((-x[..., d:], x[..., :d]), dim=-1)
+
+
+def _apply_rotary_pos_emb(t, freqs):
+    out_dtype = t.dtype
+    rot_dim = freqs.shape[-1]
+    seq_len = t.shape[-2]
+    freqs = freqs[-seq_len:]
+    t_rot, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+    t_rot = t_rot * freqs.cos() + _rotate_half(t_rot) * freqs.sin()
+    return torch.cat((t_rot.to(out_dtype), t_pass.to(out_dtype)), dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, dim_heads=64, qk_norm="none", qk_norm_eps=1e-6,
+                 differential=False, zero_init_output=True,
+                 dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        self.num_heads = dim // dim_heads
+        self.differential = differential
+        self.qk_norm = qk_norm
+
+        self.to_qkv = operations.Linear(
+            dim, dim * (5 if differential else 3), bias=False, dtype=dtype, device=device)
+        self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+
+        if qk_norm == "dyt":
+            self.q_norm = DynamicTanh(dim_heads, dtype=dtype, device=device)
+            self.k_norm = DynamicTanh(dim_heads, dtype=dtype, device=device)
+        elif qk_norm == "rms":
+            self.q_norm = operations.RMSNorm(dim_heads, eps=qk_norm_eps, dtype=dtype, device=device)
+            self.k_norm = operations.RMSNorm(dim_heads, eps=qk_norm_eps, dtype=dtype, device=device)
+
+    def forward(self, x, rotary_pos_emb=None, mask=None, **kwargs):
+        B, N, _ = x.shape
+        h = self.num_heads
+
+        qkv = self.to_qkv(x)
+        if self.differential:
+            q, k, v, q_diff, k_diff = qkv.chunk(5, dim=-1)
+            del qkv
+            q = q.view(B, N, h, -1).transpose(1, 2)
+            k = k.view(B, N, h, -1).transpose(1, 2)
+            v = v.view(B, N, h, -1).transpose(1, 2)
+            q_diff = q_diff.view(B, N, h, -1).transpose(1, 2)
+            k_diff = k_diff.view(B, N, h, -1).transpose(1, 2)
+        else:
+            q, k, v = qkv.chunk(3, dim=-1)
+            del qkv
+            q = q.view(B, N, h, -1).transpose(1, 2)
+            k = k.view(B, N, h, -1).transpose(1, 2)
+            v = v.view(B, N, h, -1).transpose(1, 2)
+
+        if self.qk_norm != "none":
+            q_dtype, k_dtype = q.dtype, k.dtype
+            q = self.q_norm(q).to(q_dtype)
+            k = self.k_norm(k).to(k_dtype)
+            if self.differential:
+                q_diff = self.q_norm(q_diff).to(q_dtype)
+                k_diff = self.k_norm(k_diff).to(k_dtype)
+
+        if rotary_pos_emb is not None:
+            freqs, _ = rotary_pos_emb
+            q_dtype, k_dtype = q.dtype, k.dtype
+            q = _apply_rotary_pos_emb(q.float(), freqs).to(q_dtype)
+            k = _apply_rotary_pos_emb(k.float(), freqs).to(k_dtype)
+            if self.differential:
+                q_diff = _apply_rotary_pos_emb(q_diff.float(), freqs).to(q_dtype)
+                k_diff = _apply_rotary_pos_emb(k_diff.float(), freqs).to(k_dtype)
+
+        if self.differential:
+            out = (optimized_attention(q, k, v, h, mask=mask, skip_reshape=True)
+                   - optimized_attention(q_diff, k_diff, v, h, mask=mask, skip_reshape=True))
+            del q, k, v, q_diff, k_diff
+        else:
+            out = optimized_attention(q, k, v, h, mask=mask, skip_reshape=True)
+            del q, k, v
+
+        return self.to_out(out)
+
+
+class _Sin(nn.Module):
+    def forward(self, x):
+        return torch.sin(3.14159265359 * x)
+
+
+class _GLU(nn.Module):
+    def __init__(self, dim_in, dim_out, activation, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.act = activation
+        self.proj = operations.Linear(dim_in, dim_out * 2, dtype=dtype, device=device)
+
+    def forward(self, x):
+        x = self.proj(x)
+        x, gate = x.chunk(2, dim=-1)
+        return x * self.act(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4, no_bias=False, zero_init_output=True,
+                 sinusoidal=False, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        act = _Sin() if sinusoidal else nn.SiLU()
+        self.ff = nn.Sequential(
+            _GLU(dim, inner_dim, act, dtype=dtype, device=device, operations=operations),
+            nn.Identity(),
+            operations.Linear(inner_dim, dim, bias=not no_bias, dtype=dtype, device=device),
+            nn.Identity(),
+        )
+
+    def forward(self, x, **kwargs):
+        return self.ff(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, dim_heads=64, causal=False, zero_init_branch_outputs=True,
+                 norm_type="dyt", add_rope=False, attn_kwargs=None, ff_kwargs=None,
+                 norm_kwargs=None, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        if attn_kwargs is None:
+            attn_kwargs = {}
+        if ff_kwargs is None:
+            ff_kwargs = {}
+        if norm_kwargs is None:
+            norm_kwargs = {}
+        dim_heads = min(dim_heads, dim)
+
+        Norm = DynamicTanh if norm_type == "dyt" else operations.RMSNorm
+        norm_kw = {**norm_kwargs, "dtype": dtype, "device": device}
+
+        self.pre_norm = Norm(dim, **norm_kw)
+        self.self_attn = Attention(dim, dim_heads=dim_heads,
+                                   zero_init_output=zero_init_branch_outputs,
+                                   dtype=dtype, device=device, operations=operations,
+                                   **attn_kwargs)
+        self.ff_norm = Norm(dim, **norm_kw)
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs,
+                              dtype=dtype, device=device, operations=operations, **ff_kwargs)
+        self.rope = RotaryEmbedding(dim_heads // 2, dtype=dtype, device=device) if add_rope else None
+
+    def forward(self, x, mask=None, **kwargs):
+        rope = self.rope.forward_from_seq_len(x.shape[-2], device=x.device) \
+               if self.rope is not None else None
+        x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rope, mask=mask)
+        x = x + self.ff(self.ff_norm(x))
+        return x
+
+
+class TransformerResamplingBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, type="encoder",
+                 transformer_depth=3, dim_heads=128, differential=True,
+                 sliding_window=None, chunk_size=128, chunk_midpoint_shift=False,
+                 dyt=True, ff_mult=3, mapping_bias=True, variable_stride=False,
+                 sinusoidal_blocks=0, conv_mapping=False, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        if type not in ("encoder", "decoder"):
+            raise ValueError(f"type must be 'encoder' or 'decoder', got {type!r}")
+
+        self.type = type
+        self.stride = stride
+        self.chunk_size = chunk_size
+        self.chunk_midpoint_shift = chunk_midpoint_shift
+        self.variable_stride = variable_stride
+        self.transformer_depth = transformer_depth
+
+        transformer_dim = out_channels if type == "encoder" else in_channels
+
+        self.mapping = (WNConv1d(in_channels, out_channels, 3 if conv_mapping else 1, padding="same", bias=mapping_bias)
+                        if in_channels != out_channels else nn.Identity())
+
+        self.sliding_window_latents = sliding_window
+        self.sliding_window_seq = self._get_sliding_window_size(sliding_window, stride)
+        self.input_seg_size, self.output_seg_size, self.sub_chunk_size = self._get_seg_sizes(stride)
+
+        token_seq = 1 if variable_stride else self.output_seg_size
+        self.new_tokens = nn.Parameter(torch.empty(1, token_seq, transformer_dim, dtype=dtype, device=device))
+
+        norm_type = "dyt" if dyt else "rms_norm"
+        attn_kwargs = {"qk_norm": "dyt" if dyt else "rms", "qk_norm_eps": 1e-3,
+                       "differential": differential}
+        norm_kwargs = {"eps": 1e-3}
+        transformers = []
+        for i in range(transformer_depth):
+            sinusoidal = (transformer_depth - i) < sinusoidal_blocks
+            transformers.append(TransformerBlock(
+                transformer_dim,
+                dim_heads=dim_heads,
+                causal=False,
+                zero_init_branch_outputs=True,
+                norm_type=norm_type,
+                add_rope=True,
+                attn_kwargs=attn_kwargs,
+                ff_kwargs={"mult": ff_mult, "no_bias": False, "sinusoidal": sinusoidal},
+                norm_kwargs=norm_kwargs,
+                dtype=dtype, device=device, operations=operations,
+            ))
+        self.transformers = nn.ModuleList(transformers)
+
+    def _get_sliding_window_size(self, window, stride, prepend_cond_length=0):
+        if window is None:
+            return None
+        return [w * (stride + 1 + prepend_cond_length) for w in window]
+
+    def _get_seg_sizes(self, stride, prepend_cond_length=0):
+        sub_chunk_size = stride + 1 + prepend_cond_length
+        input_seg_size = stride if self.type == "encoder" else 1
+        output_seg_size = 1 if self.type == "encoder" else stride
+        return input_seg_size, output_seg_size, sub_chunk_size
+
+    def forward(self, x, stride=None, **kwargs):
+        B = x.shape[0]
+
+        if stride is None:
+            input_seg = self.input_seg_size
+            output_seg = self.output_seg_size
+            sub_chunk = self.sub_chunk_size
+            sliding_window = self.sliding_window_seq
+        else:
+            input_seg, output_seg, sub_chunk = self._get_seg_sizes(stride)
+            sliding_window = self._get_sliding_window_size(self.sliding_window_latents, stride)
+
+        if self.type == "encoder":
+            if self.transformer_depth > 0:
+                pad_mod = self.chunk_size if sliding_window is None else input_seg
+                x = _zero_pad_modulo_sequence(x, pad_mod, dim=-1)
+            x = self.mapping(x)
+
+        if self.transformer_depth > 0:
+            x = x.permute(0, 2, 1)
+
+            if self.type != "encoder":
+                pad_mod = 1 if sliding_window is not None else (
+                    self.chunk_size // (stride if stride is not None else self.stride))
+                x = _zero_pad_modulo_sequence(x, pad_mod)
+
+            C = x.shape[2]
+            x = x.reshape(-1, input_seg, C)
+
+            new_tokens = self.new_tokens.expand(x.shape[0], output_seg, -1)
+            x = torch.cat([x, comfy.ops.cast_to_input(new_tokens, x)], dim=-2)
+            del new_tokens
+
+            x = x.reshape(B, -1, C)
+
+            if sliding_window is None:
+                eff_chunk = self.chunk_size + self.chunk_size // (stride if stride is not None else self.stride)
+
+            if sliding_window is None and self.chunk_midpoint_shift:
+                split = self.transformer_depth // 2
+                shift = eff_chunk // 2
+
+                x = x.reshape(-1, eff_chunk, C)
+                for layer in self.transformers[:split]:
+                    x = layer(x)
+                x = x.reshape(B, -1, C)
+
+                shifted = torch.cat([x[:, :shift, :], x, x[:, -shift:, :]], dim=1)
+                del x
+                x = shifted.reshape(-1, eff_chunk, C)
+                del shifted
+                for layer in self.transformers[split:]:
+                    x = layer(x)
+                x = x.reshape(B, -1, C)
+                x = x[:, shift:-shift, :]
+            elif sliding_window is None:
+                x = x.reshape(-1, eff_chunk, C)
+                for layer in self.transformers:
+                    x = layer(x)
+                x = x.reshape(B, -1, C)
+            else:
+                attn_mask = _sliding_window_mask(x.shape[1], sliding_window[0], x.device, x.dtype)
+                for layer in self.transformers:
+                    x = layer(x, mask=attn_mask)
+
+            x = x.reshape(-1, sub_chunk, C)
+            x = x[:, -output_seg:, :]
+            x = x.reshape(B, -1, C).transpose(1, 2)
+
+        if self.type == "decoder":
+            x = self.mapping(x)
+
+        return x
+
+
+class SAMEEncoder(nn.Module):
+    def __init__(self, in_channels=2, channels=128, latent_dim=32,
+                 c_mults=(1, 2, 4, 8), strides=(2, 4, 8, 8),
+                 transformer_depths=(3, 3, 3, 3),
+                 dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        channel_dims = [in_channels] + [channels * c for c in c_mults]
+        layers = []
+        for i in range(len(c_mults)):
+            layers.append(TransformerResamplingBlock(
+                in_channels=channel_dims[i], out_channels=channel_dims[i + 1],
+                stride=strides[i], type="encoder",
+                transformer_depth=transformer_depths[i],
+                dtype=dtype, device=device, operations=operations, **kwargs))
+        layers += [
+            Transpose(),
+            operations.Linear(channel_dims[-1], latent_dim, dtype=dtype, device=device),
+            Transpose(),
+        ]
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, x, **kwargs):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class SAMEDecoder(nn.Module):
+    def __init__(self, out_channels=2, channels=128, latent_dim=32,
+                 c_mults=(1, 2, 4, 8), strides=(2, 4, 8, 8),
+                 transformer_depths=(3, 3, 3, 3), sinusoidal_blocks=None,
+                 dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        if sinusoidal_blocks is None:
+            sinusoidal_blocks = [0] * len(c_mults)
+        channel_dims = [out_channels] + [channels * c for c in c_mults]
+        layers = [
+            Transpose(),
+            operations.Linear(latent_dim, channel_dims[-1], dtype=dtype, device=device),
+            Transpose(),
+        ]
+        for i in range(len(c_mults) - 1, -1, -1):
+            layers.append(TransformerResamplingBlock(
+                in_channels=channel_dims[i + 1], out_channels=channel_dims[i],
+                stride=strides[i], type="decoder",
+                transformer_depth=transformer_depths[i],
+                sinusoidal_blocks=sinusoidal_blocks[i],
+                dtype=dtype, device=device, operations=operations, **kwargs))
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, x, **kwargs):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class SoftNormBottleneck(nn.Module):
+    def __init__(self, dim=32, noise_augment_dim=0, noise_regularize=False,
+                 auto_scale=False, freeze=False, dtype=None, device=None, **kwargs):
+        super().__init__()
+        self.noise_augment_dim = noise_augment_dim
+        self.noise_regularize = noise_regularize
+        self.scaling_factor = nn.Parameter(torch.empty(1, dim, 1, dtype=dtype, device=device))
+        self.bias = nn.Parameter(torch.empty(1, dim, 1, dtype=dtype, device=device))
+        self.noise_scaling_factor = nn.Parameter(torch.empty(1, noise_augment_dim, 1, dtype=dtype, device=device))
+        if auto_scale:
+            self.register_parameter("running_std", nn.Parameter(
+                torch.empty(1, dtype=dtype, device=device), requires_grad=False))
+        if freeze:
+            for p in self.parameters():
+                p.requires_grad = False
+
+    def encode(self, x, return_info=False, **kwargs):
+        x = x * comfy.ops.cast_to_input(self.scaling_factor, x) \
+              + comfy.ops.cast_to_input(self.bias, x)
+        if hasattr(self, "running_std"):
+            x = x / comfy.ops.cast_to_input(self.running_std, x)
+        if return_info:
+            return x, {}
+        return x
+
+    def decode(self, x, **kwargs):
+        if hasattr(self, "running_std"):
+            x = x * comfy.ops.cast_to_input(self.running_std, x)
+        if self.noise_regularize:
+            scaling = self.running_std if hasattr(self, "running_std") \
+                      else x.std(dim=-1, keepdim=True)
+            noise = torch.randn_like(x) * comfy.ops.cast_to_input(scaling, x) * 1e-3
+            x = x + noise
+        if self.noise_augment_dim > 0:
+            noise = comfy.ops.cast_to_input(self.noise_scaling_factor, x) * torch.randn(
+                x.shape[0], self.noise_augment_dim, x.shape[-1], device=x.device, dtype=x.dtype)
+            x = torch.cat([x, noise], dim=1)
+        return x
+
+
+class PatchedPretransform(nn.Module):
+    def __init__(self, channels, patch_size, **kwargs):
+        super().__init__()
+        self.channels = channels
+        self.patch_size = patch_size
+        self.enable_grad = False
+
+    def _pad(self, x):
+        pad_len = (self.patch_size - x.shape[-1] % self.patch_size) % self.patch_size
+        if pad_len > 0:
+            x = torch.cat([x, torch.zeros_like(x[:, :, :pad_len])], dim=-1)
+        return x
+
+    def encode(self, x):
+        x = self._pad(x)
+        B, C, T = x.shape
+        h = self.patch_size
+        L = T // h
+        # b c (l h) -> b (c h) l
+        return x.reshape(B, C, L, h).permute(0, 1, 3, 2).reshape(B, C * h, L)
+
+    def decode(self, x):
+        B, Ch, L = x.shape
+        h = self.patch_size
+        C = Ch // h
+        # b (c h) l -> b c (l h)
+        return x.reshape(B, C, h, L).permute(0, 1, 3, 2).reshape(B, C, L * h)
+
+
+class SA3AudioVAE(nn.Module):
+    """SA3 VAE. State dict keys match checkpoint after stripping 'pretransform.model.'"""
+
+    def __init__(self, channels=256, transformer_depths=12, sinusoidal_blocks=8,
+                 sliding_window=None, decoder_conv_mapping=False,
+                 chunk_size=128, chunk_midpoint_shift=False,
+                 dtype=None, device=None, operations=None):
+        super().__init__()
+        if operations is None:
+            operations = ops
+
+        self.pretransform = PatchedPretransform(channels=2, patch_size=256)
+
+        common_kwargs = dict(
+            differential=True, dyt=True, dim_heads=64,
+            sliding_window=sliding_window, variable_stride=True,
+            chunk_size=chunk_size, chunk_midpoint_shift=chunk_midpoint_shift,
+            dtype=dtype, device=device, operations=operations,
+        )
+        self.encoder = SAMEEncoder(
+            in_channels=512, channels=channels, c_mults=[6], strides=[16],
+            latent_dim=256, transformer_depths=[transformer_depths],
+            conv_mapping=False, **common_kwargs,
+        )
+        self.decoder = SAMEDecoder(
+            out_channels=512, channels=channels, c_mults=[6], strides=[16],
+            latent_dim=256, transformer_depths=[transformer_depths], sinusoidal_blocks=[sinusoidal_blocks],
+            conv_mapping=decoder_conv_mapping, **common_kwargs,
+        )
+        self.bottleneck = SoftNormBottleneck(
+            dim=256, noise_augment_dim=0, noise_regularize=True,
+            auto_scale=True, freeze=True,
+            dtype=dtype, device=device,
+        )
+
+    @torch.no_grad()
+    def _pretransform_encode(self, x):
+        return self.pretransform.encode(x)
+
+    @torch.no_grad()
+    def _pretransform_decode(self, x):
+        return self.pretransform.decode(x)
+
+    def encode(self, x):
+        x = self._pretransform_encode(x)
+        x = self.encoder(x)
+        x = self.bottleneck.encode(x)
+        return x
+
+    def decode(self, x):
+        x = self.bottleneck.decode(x)
+        x = self.decoder(x)
+        x = self._pretransform_decode(x)
+        return x

comfy/ldm/hunyuan3dv2_1/hunyuandit.py

@@ -328,7 +328,7 @@ class CrossAttention(nn.Module):
         kv = torch.cat((k, v), dim=-1)
         split_size = kv.shape[-1] // self.num_heads // 2
 
-        kv = kv.view(1, -1, self.num_heads, split_size * 2)
+        kv = kv.view(b, -1, self.num_heads, split_size * 2)
         k, v = torch.split(kv, split_size, dim=-1)
 
         q = q.view(b, s1, self.num_heads, self.head_dim)
@@ -398,7 +398,7 @@ class Attention(nn.Module):
         qkv_combined = torch.cat((query, key, value), dim=-1)
         split_size = qkv_combined.shape[-1] // self.num_heads // 3
 
-        qkv = qkv_combined.view(1, -1, self.num_heads, split_size * 3)
+        qkv = qkv_combined.view(B, -1, self.num_heads, split_size * 3)
         query, key, value = torch.split(qkv, split_size, dim=-1)
 
         query = query.reshape(B, N, self.num_heads, self.head_dim)
@@ -607,9 +607,9 @@ class HunYuanDiTPlain(nn.Module):
     def forward(self, x, t, context, transformer_options = {}, **kwargs):
 
         x = x.movedim(-1, -2)
-        uncond_emb, cond_emb = context.chunk(2, dim = 0)
-
-        context = torch.cat([cond_emb, uncond_emb], dim = 0)
+        if context.shape[0] >= 2:
+            uncond_emb, cond_emb = context.chunk(2, dim = 0)
+            context = torch.cat([cond_emb, uncond_emb], dim = 0)
         main_condition = context
 
         t = 1.0 - t
@@ -657,5 +657,8 @@ class HunYuanDiTPlain(nn.Module):
         output = self.final_layer(combined)
         output =  output.movedim(-2, -1) * (-1.0)
 
-        cond_emb, uncond_emb = output.chunk(2, dim = 0)
-        return torch.cat([uncond_emb, cond_emb])
+        if output.shape[0] >= 2:
+            cond_emb, uncond_emb = output.chunk(2, dim = 0)
+            return torch.cat([uncond_emb, cond_emb])
+        else:
+            return output

comfy/model_base.py

@@ -813,6 +813,85 @@ class StableAudio1(BaseModel):
                 sd["{}{}".format(k, l)] = s[l]
         return sd
 
+class StableAudio3(BaseModel):
+    def __init__(self, model_config, seconds_total_embedder_weights, padding_embedding=None, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.audio.dit.AudioDiffusionTransformer)
+        self.seconds_total_embedder = comfy.ldm.audio.embedders.NumberConditioner(768, min_val=0, max_val=384, fourier_features_type=model_config.unet_config["timestep_features_type"])
+        self.seconds_total_embedder.load_state_dict(seconds_total_embedder_weights)
+        if padding_embedding is not None:
+            self.padding_embedding = torch.nn.Parameter(padding_embedding, requires_grad=False)
+        else:
+            self.padding_embedding = None
+
+    def concat_cond(self, **kwargs):
+        noise = kwargs.get("noise", None)
+        image = kwargs.get("concat_latent_image", None)
+
+        if image is None:
+            shape_image = list(noise.shape)
+            image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device)
+        else:
+            image = self.process_latent_in(image)
+            # TODO: scale if not match
+            image = utils.resize_to_batch_size(image, noise.shape[0])
+
+        mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
+        if mask is None:
+            mask = torch.zeros_like(noise)[:, :1]
+        else:
+            if mask.shape[1] != 1:
+                mask = torch.mean(mask, dim=1, keepdim=True)
+            mask = 1.0 - mask
+            # TODO: scale if not match
+            mask = utils.resize_to_batch_size(mask, noise.shape[0])
+
+        return torch.cat((mask, image), dim=1)
+
+    def extra_conds(self, **kwargs):
+        out = {}
+
+        concat_cond = self.concat_cond(**kwargs)
+        if concat_cond is not None:
+            out['local_add_cond'] = comfy.conds.CONDNoiseShape(concat_cond)
+
+        noise = kwargs.get("noise", None)
+        device = kwargs["device"]
+
+        seconds_total = kwargs.get("seconds_total", int(noise.shape[-1] / 10.7666))
+        seconds_total_embed = self.seconds_total_embedder([seconds_total])[0].to(device)
+
+        global_embed = seconds_total_embed.reshape((1, -1))
+        out['global_embed'] = comfy.conds.CONDRegular(global_embed)
+
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            cross_attn = cross_attn.to(device)
+            if self.padding_embedding is not None:
+                pe = self.padding_embedding.to(device=device, dtype=cross_attn.dtype)
+                max_text_tokens = self.model_config.unet_config.get("max_text_tokens", 256)
+                n_text = cross_attn.shape[1]
+                if n_text < max_text_tokens:
+                    pad = pe.view(1, 1, -1).expand(cross_attn.shape[0], max_text_tokens - n_text, -1)
+                    cross_attn = torch.cat([cross_attn, pad], dim=1)
+            cross_attn = torch.cat([cross_attn, seconds_total_embed.repeat((cross_attn.shape[0], 1, 1))], dim=1)
+            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+
+        return out
+
+    def state_dict_for_saving(self, unet_state_dict, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
+        sd = super().state_dict_for_saving(unet_state_dict, clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
+
+        d = {"conditioner.conditioners.seconds_total.": self.seconds_total_embedder.state_dict()}
+
+        for k in d:
+            s = d[k]
+            for l in s:
+                sd["{}{}".format(k, l)] = s[l]
+
+        if self.padding_embedding is not None:
+            sd["conditioner.conditioners.prompt.padding_embedding"] = self.padding_embedding.data
+        return sd
+
 
 class HunyuanDiT(BaseModel):
     def __init__(self, model_config, model_type=ModelType.V_PREDICTION, device=None):

comfy/model_detection.py

@@ -116,6 +116,45 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
     if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit
         unet_config = {}
         unet_config["audio_model"] = "dit1.0"
+        unet_config["global_cond_dim"] = state_dict['{}to_global_embed.0.weight'.format(key_prefix)].shape[1]
+        cond_embed = state_dict['{}to_cond_embed.0.weight'.format(key_prefix)]
+        unet_config["project_cond_tokens"] = cond_embed.shape[0] != cond_embed.shape[1]
+        unet_config["embed_dim"] = state_dict['{}to_timestep_embed.0.weight'.format(key_prefix)].shape[0]
+        mem_tokens = state_dict.get('{}transformer.memory_tokens'.format(key_prefix), None)
+        to_qkv = state_dict.get('{}transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix), None)
+        differential = False
+        if to_qkv is not None:
+            if to_qkv.shape[0] == to_qkv.shape[1] * 5:
+                differential = True
+        if mem_tokens is not None:
+            unet_config["num_memory_tokens"] = mem_tokens.shape[0]
+        if '{}transformer.layers.0.self_attn.q_norm.weight'.format(key_prefix) in state_dict:
+            unet_config["attn_kwargs"] = {"qk_norm": "ln", "feat_scale": True}
+        rms_norm = state_dict.get('{}transformer.layers.0.self_attn.q_norm.gamma'.format(key_prefix), None)
+        if rms_norm is not None:
+            unet_config["attn_kwargs"] = {"qk_norm": "rms", "differential": differential}
+            unet_config["norm_type"] = "rms_norm"
+            unet_config["num_heads"] = unet_config["embed_dim"] // rms_norm.shape[0]
+
+        if '{}timestep_features.weight'.format(key_prefix) in state_dict:
+            unet_config["timestep_features_type"] = "learned"
+        else:
+            unet_config["timestep_features_type"] = "expo"
+
+        io_channels = state_dict['{}postprocess_conv.weight'.format(key_prefix)].shape[0]
+        unet_config["io_channels"] = io_channels
+        unet_config["input_concat_dim"] = state_dict['{}transformer.project_in.weight'.format(key_prefix)].shape[1] - io_channels
+
+        local_add_cond = state_dict.get('{}transformer.layers.0.to_local_embed.0.weight'.format(key_prefix), None)
+        if local_add_cond is not None:
+            unet_config["local_add_cond_dim"] = local_add_cond.shape[1]
+
+        global_cond_embed = state_dict.get('{}transformer.global_cond_embedder.0.weight'.format(key_prefix), None)
+        if global_cond_embed is not None:
+            unet_config["global_cond_shared_embed"] = True
+            unet_config["global_cond_type"] = "adaLN"
+
+        unet_config["depth"] = count_blocks(state_dict_keys, '{}transformer.layers.'.format(key_prefix) + '{}.')
         return unet_config
 
     if '{}double_layers.0.attn.w1q.weight'.format(key_prefix) in state_dict_keys: #aura flow dit

comfy/ops.py

@@ -260,7 +260,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
 
 
 def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
-    # NOTE: offloadable=False is a a legacy and if you are a custom node author reading this please pass
+    # NOTE: offloadable=False is a legacy mode and if you are a custom node author reading this please pass
     # offloadable=True and call uncast_bias_weight() after your last usage of the weight/bias. This
     # will add async-offload support to your cast and improve performance.
     if input is not None:

comfy/sd.py

@@ -21,6 +21,7 @@ import comfy.ldm.ace.vae.music_dcae_pipeline
 import comfy.ldm.cogvideo.vae
 import comfy.ldm.hunyuan_video.vae
 import comfy.ldm.mmaudio.vae.autoencoder
+import comfy.ldm.audio.vae_sa3
 import comfy.pixel_space_convert
 import comfy.weight_adapter
 import yaml
@@ -67,6 +68,7 @@ import comfy.text_encoders.qwen35
 import comfy.text_encoders.ernie
 import comfy.text_encoders.gemma4
 import comfy.text_encoders.cogvideo
+import comfy.text_encoders.sa3
 
 import comfy.model_patcher
 import comfy.lora
@@ -854,6 +856,34 @@ class VAE:
                 self.working_dtypes = [torch.float32]
                 self.disable_offload = True
                 self.extra_1d_channel = 16
+            elif "decoder.layers.3.transformers.0.pre_norm.alpha" in sd:  # Stable Audio 3 VAE
+                if "decoder.layers.3.transformers.11.self_attn.to_out.weight" in sd:
+                    config = {"channels": 256, "transformer_depths": 12, "sinusoidal_blocks": 8,
+                              "sliding_window": [1, 1], "decoder_conv_mapping": False,
+                              "chunk_size": 128, "chunk_midpoint_shift": False}
+                    self.memory_used_encode = lambda shape, dtype: (1500 * shape[2]) * model_management.dtype_size(dtype)
+                    self.memory_used_decode = lambda shape, dtype: (1500 * shape[2] * 4096) * model_management.dtype_size(dtype)
+                else:
+                    config = {"channels": 128, "transformer_depths": 6, "sinusoidal_blocks": 0,
+                              "sliding_window": None, "decoder_conv_mapping": True,
+                              "chunk_size": 32, "chunk_midpoint_shift": True}
+                    self.memory_used_encode = lambda shape, dtype: (72 * shape[2]) * model_management.dtype_size(dtype)
+                    self.memory_used_decode = lambda shape, dtype: (72 * shape[2] * 4096) * model_management.dtype_size(dtype)
+
+                self.first_stage_model = comfy.ldm.audio.vae_sa3.SA3AudioVAE(**config)
+                self.latent_channels = 256
+                self.output_channels = 2
+                self.upscale_ratio = 4096
+                self.downscale_ratio = 4096
+                self.latent_dim = 1
+                self.audio_sample_rate = 44100
+                self.process_output = lambda audio: audio
+                self.process_input = lambda audio: audio
+                self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+                #This VAE has Parameters and Buffers the non-dynamic caster cannot handle
+                #Force cast it for --disable-dynamic-vram users until there is a true core fix.
+                if not comfy.memory_management.aimdo_enabled:
+                    self.disable_offload = True
             else:
                 logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
                 self.first_stage_model = None
@@ -1290,6 +1320,7 @@ class TEModel(Enum):
     GEMMA_4_E4B = 29
     GEMMA_4_E2B = 30
     GEMMA_4_31B = 31
+    T5_GEMMA = 32
 
 
 def detect_te_model(sd):
@@ -1314,6 +1345,8 @@ def detect_te_model(sd):
         if weight.shape[0] == 384:
             return TEModel.BYT5_SMALL_GLYPH
         return TEModel.T5_BASE
+    if "model.encoder.layers.0.pre_self_attn_layernorm.weight" in sd:
+        return TEModel.T5_GEMMA
     if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
         if 'model.layers.59.self_attn.q_norm.weight' in sd:
             return TEModel.GEMMA_4_31B
@@ -1463,6 +1496,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             else:
                 clip_target.clip = comfy.text_encoders.sa_t5.SAT5Model
                 clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer
+        elif te_model == TEModel.T5_GEMMA:
+            clip_target.clip = comfy.text_encoders.sa3.SAT5GemmaModel
+            clip_target.tokenizer = comfy.text_encoders.sa3.SAT5GemmaTokenizer
+            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
         elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B):
             variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
                        TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,

comfy/supported_models.py

@@ -7,6 +7,7 @@ from . import sdxl_clip
 import comfy.text_encoders.sd2_clip
 import comfy.text_encoders.sd3_clip
 import comfy.text_encoders.sa_t5
+import comfy.text_encoders.sa3
 import comfy.text_encoders.aura_t5
 import comfy.text_encoders.pixart_t5
 import comfy.text_encoders.hydit
@@ -603,6 +604,29 @@ class StableAudio(supported_models_base.BASE):
     def clip_target(self, state_dict={}):
         return supported_models_base.ClipTarget(comfy.text_encoders.sa_t5.SAT5Tokenizer, comfy.text_encoders.sa_t5.SAT5Model)
 
+class StableAudio3(StableAudio):
+    unet_config = {
+        "audio_model": "dit1.0",
+        "global_cond_shared_embed": True,
+    }
+
+    sampling_settings = {
+        "multiplier": 1.0,
+        "shift": 2.0,
+    }
+
+    latent_format = latent_formats.StableAudio3
+
+    memory_usage_factor = 7
+
+    def get_model(self, state_dict, prefix="", device=None):
+        seconds_total_sd = utils.state_dict_prefix_replace(state_dict, {"conditioner.conditioners.seconds_total.": ""}, filter_keys=True)
+        padding_embedding = state_dict.get("conditioner.conditioners.prompt.padding_embedding", None)
+        return model_base.StableAudio3(self, seconds_total_embedder_weights=seconds_total_sd, padding_embedding=padding_embedding, device=device)
+
+    def clip_target(self, state_dict={}):
+        return supported_models_base.ClipTarget(comfy.text_encoders.sa3.SAT5GemmaTokenizer, comfy.text_encoders.sa3.SAT5GemmaModel)
+
 class AuraFlow(supported_models_base.BASE):
     unet_config = {
         "cond_seq_dim": 2048,
@@ -2018,6 +2042,7 @@ models = [
     SV3D_u,
     SV3D_p,
     SD3,
+    StableAudio3,
     StableAudio,
     AuraFlow,
     PixArtAlpha,

comfy/text_encoders/sa3.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+from comfy import sd1_clip
+from comfy.text_encoders.llama import Attention as LlamaAttention, RMSNorm, MLP, precompute_freqs_cis, apply_rope, _make_scaled_embedding
+from comfy.text_encoders.spiece_tokenizer import SPieceTokenizer
+
+
+class T5GemmaEncoderConfig:
+    def __init__(self):
+        self.vocab_size = 256000
+        self.hidden_size = 768
+        self.intermediate_size = 2048
+        self.num_hidden_layers = 12
+        self.num_attention_heads = 12
+        self.num_key_value_heads = 12
+        self.head_dim = 64
+        self.rms_norm_eps = 1e-6
+        self.rms_norm_add = False
+        self.rope_theta = 10000.0
+        self.attn_logit_softcapping = 50.0
+        self.query_pre_attn_scalar = 64
+        self.sliding_window = 4096
+        self.mlp_activation = "gelu_pytorch_tanh"
+        self.layer_types = ["sliding_attention", "full_attention"] * 6
+        self.qkv_bias = False
+        self.q_norm = None
+        self.k_norm = None
+        self.rms_norm_add = True
+
+
+class T5GemmaAttention(LlamaAttention):
+    """Reuses LlamaAttention projection setup; overrides forward for softcap attention.
+
+    T5Gemma applies tanh(QK^T * scale / cap) * cap between the matmul and softmax.
+    This nonlinearity is incompatible with fused SDPA kernels, so attention is
+    computed manually. Everything else (projections, RoPE, GQA expansion) is identical
+    to LlamaAttention so __init__ is inherited unchanged.
+    """
+
+    def __init__(self, config, device=None, dtype=None, ops=None):
+        super().__init__(config, device=device, dtype=dtype, ops=ops)
+        self.scale = config.query_pre_attn_scalar ** -0.5
+        self.softcap = config.attn_logit_softcapping
+
+    def forward(self, hidden_states, attention_mask=None, freqs_cis=None, **kwargs):
+        B, S, _ = hidden_states.shape
+        xq = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        xk = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        xv = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        xq, xk = apply_rope(xq, xk, freqs_cis)
+        xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
+        xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
+        attn = torch.matmul(xq * self.scale, xk.transpose(-2, -1))
+        attn = torch.tanh(attn / self.softcap) * self.softcap
+        if attention_mask is not None:
+            attn = attn + attention_mask
+        attn = torch.nn.functional.softmax(attn.float(), dim=-1).to(xq.dtype)
+        out = torch.matmul(attn, xv).transpose(1, 2).reshape(B, S, self.inner_size)
+        return self.o_proj(out), None
+
+
+class T5GemmaBlock(nn.Module):
+    def __init__(self, config, layer_type, device=None, dtype=None, ops=None):
+        super().__init__()
+        self.self_attn = T5GemmaAttention(config, device=device, dtype=dtype, ops=ops)
+        self.mlp = MLP(config, device=device, dtype=dtype, ops=ops)
+        # Names match checkpoint keys: model.encoder.layers.X.<name>.weight
+        self.pre_self_attn_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=True, device=device, dtype=dtype)
+        self.post_self_attn_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=True, device=device, dtype=dtype)
+        self.pre_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=True, device=device, dtype=dtype)
+        self.post_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=True, device=device, dtype=dtype)
+        self.is_sliding = (layer_type == "sliding_attention")
+        self.sliding_window = config.sliding_window
+
+    def forward(self, x, attention_mask=None, freqs_cis=None):
+        attn_mask = attention_mask
+        if self.is_sliding and x.shape[1] > self.sliding_window:
+            S = x.shape[1]
+            pos = torch.arange(S, device=x.device)
+            dist = (pos.unsqueeze(0) - pos.unsqueeze(1)).abs()
+            sw_mask = torch.zeros(S, S, dtype=x.dtype, device=x.device)
+            sw_mask.masked_fill_(dist > self.sliding_window, -torch.finfo(x.dtype).max)
+            sw_mask = sw_mask.unsqueeze(0).unsqueeze(0)
+            attn_mask = (attention_mask + sw_mask) if attention_mask is not None else sw_mask
+        residual = x
+        x = self.pre_self_attn_layernorm(x)
+        x, _ = self.self_attn(x, attention_mask=attn_mask, freqs_cis=freqs_cis)
+        x = self.post_self_attn_layernorm(x)
+        x = residual + x
+        residual = x
+        x = self.pre_feedforward_layernorm(x)
+        x = self.mlp(x)
+        x = self.post_feedforward_layernorm(x)
+        x = residual + x
+        return x
+
+
+class T5GemmaEncoder(nn.Module):
+    """Encoder stack: embed_tokens, layers, norm.
+    Keys: embed_tokens.*, layers.X.*, norm.*"""
+
+    def __init__(self, config, device, dtype, ops):
+        super().__init__()
+        self.config = config
+        # Gemma-style scaled embedding: output *= sqrt(hidden_size)
+        self.embed_tokens = _make_scaled_embedding(
+            ops, config.vocab_size, config.hidden_size, config.hidden_size ** 0.5, device, dtype)
+        self.layers = nn.ModuleList([
+            T5GemmaBlock(config, config.layer_types[i], device=device, dtype=dtype, ops=ops)
+            for i in range(config.num_hidden_layers)
+        ])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=True, device=device, dtype=dtype)
+
+    def forward(self, input_ids, attention_mask=None, embeds=None, intermediate_output=None,
+                final_layer_norm_intermediate=True, dtype=None, num_layers=None):
+        x = embeds if embeds is not None else self.embed_tokens(input_ids, out_dtype=dtype or torch.float32)
+        seq_len = x.shape[1]
+        position_ids = torch.arange(seq_len, device=x.device).unsqueeze(0)
+        freqs_cis = precompute_freqs_cis(self.config.head_dim, position_ids, self.config.rope_theta, device=x.device)
+        mask = None
+        if attention_mask is not None:
+            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).max)
+        intermediate = None
+        for i, layer in enumerate(self.layers):
+            x = layer(x, attention_mask=mask, freqs_cis=freqs_cis)
+            if i == intermediate_output:
+                intermediate = x.clone()
+        x = self.norm(x)
+        if intermediate is not None and final_layer_norm_intermediate:
+            intermediate = self.norm(intermediate)
+        return x, intermediate
+
+
+class T5GemmaBody(nn.Module):
+    """Provides the 'encoder' sub-module.
+    Keys: encoder.*"""
+
+    def __init__(self, config, device, dtype, ops):
+        super().__init__()
+        self.encoder = T5GemmaEncoder(config, device, dtype, ops)
+
+
+class T5GemmaModel(nn.Module):
+    """Top-level model class passed to SDClipModel as model_class.
+    Module layout: self.model.encoder.* → matches checkpoint keys model.encoder.*"""
+
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        config = T5GemmaEncoderConfig()
+        self.num_layers = config.num_hidden_layers
+        self.dtype = dtype
+        self.model = T5GemmaBody(config, device, dtype, operations)
+
+    def get_input_embeddings(self):
+        return self.model.encoder.embed_tokens
+
+    def set_input_embeddings(self, embeddings):
+        self.model.encoder.embed_tokens = embeddings
+
+    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None,
+                intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, **kwargs):
+        if intermediate_output is not None and intermediate_output < 0:
+            intermediate_output = self.num_layers + intermediate_output
+        return self.model.encoder(
+            input_ids, attention_mask=attention_mask, embeds=embeds,
+            intermediate_output=intermediate_output,
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
+            dtype=dtype, num_layers=self.num_layers)
+
+
+class T5GemmaSDClipModel(sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}):
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx,
+                         textmodel_json_config={}, dtype=dtype,
+                         special_tokens={"pad": 0},
+                         model_class=T5GemmaModel,
+                         enable_attention_masks=True, zero_out_masked=True,
+                         model_options=model_options)
+
+
+class T5GemmaSDTokenizer(sd1_clip.SDTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        tokenizer_model = tokenizer_data.get("spiece_model", None)
+        super().__init__(tokenizer_model, pad_with_end=False, embedding_size=768,
+                         embedding_key="t5gemma", tokenizer_class=SPieceTokenizer,
+                         has_start_token=False, has_end_token=False, pad_to_max_length=False,
+                         max_length=99999999, min_length=1, pad_token=0,
+                         tokenizer_data=tokenizer_data,
+                         tokenizer_args={"add_bos": False, "add_eos": False})
+
+    def state_dict(self):
+        return {"spiece_model": self.tokenizer.serialize_model()}
+
+
+class SAT5GemmaTokenizer(sd1_clip.SD1Tokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        super().__init__(embedding_directory=embedding_directory,
+                         tokenizer_data=tokenizer_data, clip_name="t5gemma", tokenizer=T5GemmaSDTokenizer)
+
+
+class SAT5GemmaModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs):
+        super().__init__(device=device, dtype=dtype, model_options=model_options,
+                         name="t5gemma", clip_model=T5GemmaSDClipModel, **kwargs)

comfy_api_nodes/apis/anthropic.py

@@ -35,6 +35,19 @@ class AnthropicMessage(BaseModel):
     content: list[AnthropicTextContent | AnthropicImageContent] = Field(...)
 
 
+class AnthropicThinkingConfig(BaseModel):
+    type: Literal["enabled", "disabled", "adaptive"] = Field(...)
+    budget_tokens: int | None = Field(
+        None, ge=1024,
+        description="Reasoning budget in tokens. Used when type is 'enabled'. Must be less than max_tokens.",
+    )
+
+
+class AnthropicOutputConfig(BaseModel):
+    """Used with `thinking.type='adaptive'` on models like Opus 4.7."""
+    effort: Literal["low", "medium", "high"] | None = Field(None)
+
+
 class AnthropicMessagesRequest(BaseModel):
     model: str = Field(...)
     messages: list[AnthropicMessage] = Field(...)
@@ -44,6 +57,8 @@ class AnthropicMessagesRequest(BaseModel):
     top_p: float | None = Field(None, ge=0.0, le=1.0)
     top_k: int | None = Field(None, ge=0)
     stop_sequences: list[str] | None = Field(None)
+    thinking: AnthropicThinkingConfig | None = Field(None)
+    output_config: AnthropicOutputConfig | None = Field(None)
 
 
 class AnthropicResponseTextBlock(BaseModel):
@@ -51,6 +66,14 @@ class AnthropicResponseTextBlock(BaseModel):
     text: str = Field(...)
 
 
+class AnthropicResponseThinkingBlock(BaseModel):
+    type: Literal["thinking"] = "thinking"
+    thinking: str = Field(...)
+
+
+AnthropicResponseBlock = AnthropicResponseTextBlock | AnthropicResponseThinkingBlock
+
+
 class AnthropicCacheCreationUsage(BaseModel):
     ephemeral_5m_input_tokens: int | None = Field(None)
     ephemeral_1h_input_tokens: int | None = Field(None)
@@ -69,7 +92,7 @@ class AnthropicMessagesResponse(BaseModel):
     type: str | None = Field(None)
     role: str | None = Field(None)
     model: str | None = Field(None)
-    content: list[AnthropicResponseTextBlock] | None = Field(None)
+    content: list[AnthropicResponseBlock] | None = Field(None)
     stop_reason: str | None = Field(None)
     stop_sequence: str | None = Field(None)
     usage: AnthropicMessagesUsage | None = Field(None)

comfy_api_nodes/apis/openrouter.py

@@ -0,0 +1,93 @@
+"""Pydantic models for the OpenRouter chat completions API.
+
+See: https://openrouter.ai/docs/api/api-reference/chat/send-chat-completion-request
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class OpenRouterTextContent(BaseModel):
+    type: Literal["text"] = "text"
+    text: str = Field(...)
+
+
+class OpenRouterImageUrl(BaseModel):
+    url: str = Field(...)
+
+
+class OpenRouterImageContent(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    image_url: OpenRouterImageUrl = Field(...)
+
+
+class OpenRouterVideoUrl(BaseModel):
+    url: str = Field(...)
+
+
+class OpenRouterVideoContent(BaseModel):
+    type: Literal["video_url"] = "video_url"
+    video_url: OpenRouterVideoUrl = Field(...)
+
+
+OpenRouterContentBlock = OpenRouterTextContent | OpenRouterImageContent | OpenRouterVideoContent
+
+
+class OpenRouterMessage(BaseModel):
+    role: Literal["system", "user", "assistant"] = Field(...)
+    content: str | list[OpenRouterContentBlock] = Field(...)
+
+
+class OpenRouterReasoningConfig(BaseModel):
+    effort: str | None = Field(None)
+    exclude: bool | None = Field(None, description="If true, model reasons but reasoning is excluded from response.")
+
+
+class OpenRouterWebSearchOptions(BaseModel):
+    search_context_size: str | None = Field(None)
+
+
+class OpenRouterChatRequest(BaseModel):
+    model: str = Field(...)
+    messages: list[OpenRouterMessage] = Field(...)
+    seed: int | None = Field(None)
+    reasoning: OpenRouterReasoningConfig | None = Field(None)
+    web_search_options: OpenRouterWebSearchOptions | None = Field(None)
+    stream: bool = Field(False)
+
+
+class OpenRouterUsage(BaseModel):
+    prompt_tokens: int | None = Field(None)
+    completion_tokens: int | None = Field(None)
+    total_tokens: int | None = Field(None)
+    cost: float | None = Field(None, description="Server-side authoritative USD cost of the call.")
+
+
+class OpenRouterResponseMessage(BaseModel):
+    role: str | None = Field(None)
+    content: str | None = Field(None)
+    reasoning: str | None = Field(None)
+    refusal: str | None = Field(None)
+
+
+class OpenRouterChoice(BaseModel):
+    index: int | None = Field(None)
+    message: OpenRouterResponseMessage | None = Field(None)
+    finish_reason: str | None = Field(None)
+
+
+class OpenRouterError(BaseModel):
+    code: int | str | None = Field(None)
+    message: str | None = Field(None)
+    metadata: dict | None = Field(None)
+
+
+class OpenRouterChatResponse(BaseModel):
+    id: str | None = Field(None)
+    model: str | None = Field(None)
+    object: str | None = Field(None)
+    provider: str | None = Field(None)
+    choices: list[OpenRouterChoice] | None = Field(None)
+    usage: OpenRouterUsage | None = Field(None)
+    error: OpenRouterError | None = Field(None)

comfy_api_nodes/nodes_anthropic.py

@@ -9,8 +9,11 @@ from comfy_api_nodes.apis.anthropic import (
     AnthropicMessage,
     AnthropicMessagesRequest,
     AnthropicMessagesResponse,
+    AnthropicOutputConfig,
+    AnthropicResponseTextBlock,
     AnthropicRole,
     AnthropicTextContent,
+    AnthropicThinkingConfig,
 )
 from comfy_api_nodes.util import (
     ApiEndpoint,
@@ -32,15 +35,29 @@ CLAUDE_MODELS: dict[str, str] = {
     "Haiku 4.5": "claude-haiku-4-5-20251001",
 }
 
+_THINKING_UNSUPPORTED = {"Haiku 4.5"}
+# Models that use the newer "adaptive" thinking mode (Opus 4.7 requires it; older models keep the explicit budget API).
+# Anthropic decides the actual budget when adaptive is used, based on the `output_config.effort` hint.
+_ADAPTIVE_THINKING_MODELS = {"Opus 4.7", "Opus 4.6", "Sonnet 4.6"}
 
-def _claude_model_inputs():
-    return [
+# Budget mode (Sonnet 4.5): effort -> reasoning budget in tokens. Must be < max_tokens.
+# Sized so even the "high" budget fits comfortably under the default max_tokens=32768.
+_REASONING_BUDGET: dict[str, int] = {
+    "low": 2048,
+    "medium": 8192,
+    "high": 16384,
+}
+_REASONING_EFFORTS = ["off", "low", "medium", "high"]
+
+
+def _claude_model_inputs(model_label: str):
+    inputs: list = [
         IO.Int.Input(
             "max_tokens",
-            default=16000,
-            min=32,
-            max=32000,
-            tooltip="Maximum number of tokens to generate before stopping.",
+            default=32768,
+            min=4096,
+            max=64000,
+            tooltip="Maximum number of tokens to generate (includes reasoning tokens when enabled).",
             advanced=True,
         ),
         IO.Float.Input(
@@ -49,10 +66,24 @@ def _claude_model_inputs():
             min=0.0,
             max=1.0,
             step=0.01,
-            tooltip="Controls randomness. 0.0 is deterministic, 1.0 is most random. Ignored for Opus 4.7.",
+            tooltip=(
+                "Controls randomness. 0.0 is deterministic, 1.0 is most random. "
+                "Ignored for Opus 4.7 and any model when reasoning_effort is set."
+            ),
             advanced=True,
         ),
     ]
+    if model_label not in _THINKING_UNSUPPORTED:
+        inputs.append(
+            IO.Combo.Input(
+                "reasoning_effort",
+                options=_REASONING_EFFORTS,
+                default="off",
+                tooltip="Extended thinking effort. 'off' disables reasoning.",
+                advanced=True,
+            )
+        )
+    return inputs
 
 
 def _model_price_per_million(model: str) -> tuple[float, float] | None:
@@ -95,7 +126,11 @@ def calculate_tokens_price(response: AnthropicMessagesResponse) -> float | None:
 def _get_text_from_response(response: AnthropicMessagesResponse) -> str:
     if not response.content:
         return ""
-    return "\n".join(block.text for block in response.content if block.text)
+    # Thinking blocks are silently dropped — we never want reasoning in the output.
+    return "\n".join(
+        block.text for block in response.content
+        if isinstance(block, AnthropicResponseTextBlock) and block.text
+    )
 
 
 async def _build_image_content_blocks(
@@ -133,7 +168,10 @@ class ClaudeNode(IO.ComfyNode):
                 ),
                 IO.DynamicCombo.Input(
                     "model",
-                    options=[IO.DynamicCombo.Option(label, _claude_model_inputs()) for label in CLAUDE_MODELS],
+                    options=[
+                        IO.DynamicCombo.Option(label, _claude_model_inputs(label))
+                        for label in CLAUDE_MODELS
+                    ],
                     tooltip="The Claude model used to generate the response.",
                 ),
                 IO.Int.Input(
@@ -207,8 +245,29 @@ class ClaudeNode(IO.ComfyNode):
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         model_label = model["model"]
-        max_tokens = model["max_tokens"]
-        temperature = None if model_label == "Opus 4.7" else model["temperature"]
+        max_tokens = model.get("max_tokens", 32768)
+        reasoning_effort = model.get("reasoning_effort", "off")
+        thinking_enabled = reasoning_effort not in ("off", None) and model_label not in _THINKING_UNSUPPORTED
+
+        # Anthropic requires temperature to be unset (defaults to 1.0) when thinking is enabled.
+        # Opus 4.7 also rejects user-supplied temperature.
+        if thinking_enabled or model_label == "Opus 4.7":
+            temperature = None
+        else:
+            temperature = model.get("temperature", 1.0)
+
+        thinking_cfg: AnthropicThinkingConfig | None = None
+        output_cfg: AnthropicOutputConfig | None = None
+        if thinking_enabled:
+            if model_label in _ADAPTIVE_THINKING_MODELS:
+                # Adaptive mode - Anthropic chooses the budget based on effort hint
+                thinking_cfg = AnthropicThinkingConfig(type="adaptive")
+                output_cfg = AnthropicOutputConfig(effort=reasoning_effort)
+            else:
+                # Budget mode (Sonnet 4.5). Leave at least 1024 tokens for the actual response
+                budget = _REASONING_BUDGET[reasoning_effort]
+                budget = min(budget, max(1024, max_tokens - 1024))
+                thinking_cfg = AnthropicThinkingConfig(type="enabled", budget_tokens=budget)
 
         image_tensors: list[Input.Image] = [t for t in (images or {}).values() if t is not None]
         if sum(get_number_of_images(t) for t in image_tensors) > CLAUDE_MAX_IMAGES:
@@ -229,6 +288,8 @@ class ClaudeNode(IO.ComfyNode):
                 messages=[AnthropicMessage(role=AnthropicRole.user, content=content)],
                 system=system_prompt or None,
                 temperature=temperature,
+                thinking=thinking_cfg,
+                output_config=output_cfg,
             ),
             price_extractor=calculate_tokens_price,
         )

comfy_api_nodes/nodes_bytedance.py

@@ -43,15 +43,16 @@ from comfy_api_nodes.util import (
     ApiEndpoint,
     download_url_to_image_tensor,
     download_url_to_video_output,
+    downscale_video_to_max_pixels,
     get_number_of_images,
     image_tensor_pair_to_batch,
     poll_op,
-    resize_video_to_pixel_budget,
     sync_op,
     upload_audio_to_comfyapi,
     upload_image_to_comfyapi,
     upload_images_to_comfyapi,
     upload_video_to_comfyapi,
+    upscale_video_to_min_pixels,
     validate_image_aspect_ratio,
     validate_image_dimensions,
     validate_string,
@@ -110,12 +111,13 @@ def _validate_ref_video_pixels(video: Input.Video, model_id: str, resolution: st
     max_px = limits.get("max")
     if min_px and pixels < min_px:
         raise ValueError(
-            f"Reference video {index} is too small: {w}x{h} = {pixels:,}px. " f"Minimum is {min_px:,}px for this model."
+            f"Reference video {index} is too small: {w}x{h} = {pixels:,} total pixels. "
+            f"Minimum for this model is {min_px:,} total pixels."
         )
     if max_px and pixels > max_px:
         raise ValueError(
-            f"Reference video {index} is too large: {w}x{h} = {pixels:,}px. "
-            f"Maximum is {max_px:,}px for this model. Try downscaling the video."
+            f"Reference video {index} is too large: {w}x{h} = {pixels:,} total pixels. "
+            f"Maximum for this model is {max_px:,} total pixels. Try downscaling the video."
         )
 
 
@@ -1676,14 +1678,14 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
                     "first_frame_asset_id",
                     default="",
                     tooltip="Seedance asset_id to use as the first frame. "
-                            "Mutually exclusive with the first_frame image input.",
+                    "Mutually exclusive with the first_frame image input.",
                     optional=True,
                 ),
                 IO.String.Input(
                     "last_frame_asset_id",
                     default="",
                     tooltip="Seedance asset_id to use as the last frame. "
-                            "Mutually exclusive with the last_frame image input.",
+                    "Mutually exclusive with the last_frame image input.",
                     optional=True,
                 ),
                 IO.Int.Input(
@@ -1865,11 +1867,20 @@ def _seedance2_reference_inputs(resolutions: list[str], default_ratio: str = "16
         IO.Boolean.Input(
             "auto_downscale",
             default=False,
-            advanced=True,
             optional=True,
             tooltip="Automatically downscale reference videos that exceed the model's pixel budget "
             "for the selected resolution. Aspect ratio is preserved; videos already within limits are untouched.",
         ),
+        IO.Boolean.Input(
+            "auto_upscale",
+            default=False,
+            advanced=True,
+            optional=True,
+            tooltip="Automatically upscale reference videos that are below the model's minimum pixel count "
+            "for the selected resolution. Aspect ratio is preserved; videos already meeting the minimum are "
+            "untouched. Note: upscaling a low-resolution source does not add real detail and may produce "
+            "lower-quality generations.",
+        ),
         IO.Autogrow.Input(
             "reference_assets",
             template=IO.Autogrow.TemplateNames(
@@ -2030,7 +2041,13 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
             max_px = SEEDANCE2_REF_VIDEO_PIXEL_LIMITS.get(model_id, {}).get(model["resolution"], {}).get("max")
             if max_px:
                 for key in reference_videos:
-                    reference_videos[key] = resize_video_to_pixel_budget(reference_videos[key], max_px)
+                    reference_videos[key] = downscale_video_to_max_pixels(reference_videos[key], max_px)
+
+        if model.get("auto_upscale") and reference_videos:
+            min_px = SEEDANCE2_REF_VIDEO_PIXEL_LIMITS.get(model_id, {}).get(model["resolution"], {}).get("min")
+            if min_px:
+                for key in reference_videos:
+                    reference_videos[key] = upscale_video_to_min_pixels(reference_videos[key], min_px)
 
         total_video_duration = 0.0
         for i, key in enumerate(reference_videos, 1):

comfy_api_nodes/nodes_openrouter.py

@@ -0,0 +1,374 @@
+"""API Nodes for OpenRouter LLM chat completions."""
+
+from dataclasses import dataclass
+from typing import Literal
+
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension, Input
+from comfy_api_nodes.apis.openrouter import (
+    OpenRouterChatRequest,
+    OpenRouterChatResponse,
+    OpenRouterContentBlock,
+    OpenRouterImageContent,
+    OpenRouterImageUrl,
+    OpenRouterMessage,
+    OpenRouterReasoningConfig,
+    OpenRouterTextContent,
+    OpenRouterVideoContent,
+    OpenRouterVideoUrl,
+    OpenRouterWebSearchOptions,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    get_number_of_images,
+    sync_op,
+    upload_images_to_comfyapi,
+    upload_video_to_comfyapi,
+    validate_string,
+)
+
+OPENROUTER_CHAT_ENDPOINT = "/proxy/openrouter/api/v1/chat/completions"
+
+
+Profile = Literal["standard", "reasoning", "frontier_reasoning", "perplexity", "perplexity_reasoning"]
+
+
+@dataclass(frozen=True)
+class _ModelSpec:
+    slug: str  # exact OpenRouter model id
+    profile: Profile
+    price_in: float  # USD per token (prompt)
+    price_out: float  # USD per token (completion)
+    max_images: int = 0  # 0 = no image input; otherwise max URL-passed images supported
+    max_videos: int = 0  # 0 = no video input; otherwise max URL-passed videos supported
+
+
+MODELS: list[_ModelSpec] = [
+    _ModelSpec("anthropic/claude-opus-4.7", "frontier_reasoning", 0.000005, 0.000025, max_images=20),
+    _ModelSpec("openai/gpt-5.5-pro", "frontier_reasoning", 0.00003, 0.00018, max_images=20),
+    _ModelSpec("openai/gpt-5.5", "frontier_reasoning", 0.000005, 0.00003, max_images=20),
+    _ModelSpec("google/gemini-3.5-flash", "reasoning", 0.0000015, 0.000009, max_images=20, max_videos=4),
+    _ModelSpec("x-ai/grok-4.20", "reasoning", 0.00000125, 0.0000025, max_images=20),
+    _ModelSpec("x-ai/grok-4.3", "reasoning", 0.00000125, 0.0000025, max_images=20),
+    _ModelSpec("deepseek/deepseek-v4-pro", "reasoning", 0.000000435, 0.00000087),
+    _ModelSpec("deepseek/deepseek-v4-flash", "reasoning", 0.000000112, 0.000000224),
+    _ModelSpec("deepseek/deepseek-v3.2", "reasoning", 0.000000252, 0.000000378),
+    _ModelSpec("qwen/qwen3.6-max-preview", "reasoning", 0.00000104, 0.00000624),
+    _ModelSpec("qwen/qwen3.6-plus", "reasoning", 0.000000325, 0.00000195, max_images=10, max_videos=4),
+    _ModelSpec("qwen/qwen3.6-flash", "reasoning", 0.0000001875, 0.000001125, max_images=10, max_videos=4),
+    _ModelSpec("mistralai/mistral-large-2512", "standard", 0.0000005, 0.0000015, max_images=8),
+    _ModelSpec("mistralai/mistral-medium-3-5", "reasoning", 0.0000015, 0.0000075, max_images=8),
+    _ModelSpec("z-ai/glm-4.6", "reasoning", 0.00000043, 0.00000174),
+    _ModelSpec("z-ai/glm-5", "reasoning", 0.0000006, 0.00000192),
+    _ModelSpec("moonshotai/kimi-k2.6", "reasoning", 0.00000073, 0.00000349, max_images=10),
+    _ModelSpec("moonshotai/kimi-k2-thinking", "reasoning", 0.0000006, 0.0000025),
+    _ModelSpec("perplexity/sonar-pro", "perplexity", 0.000003, 0.000015),
+    _ModelSpec("perplexity/sonar-reasoning-pro", "perplexity_reasoning", 0.000002, 0.000008),
+    _ModelSpec("perplexity/sonar-deep-research", "perplexity_reasoning", 0.000002, 0.000008),
+]
+
+_MODELS_BY_SLUG: dict[str, _ModelSpec] = {m.slug: m for m in MODELS}
+_REASONING_EFFORTS = ["off", "low", "medium", "high"]
+_SEARCH_CONTEXT_SIZES = ["low", "medium", "high"]
+
+
+def _reasoning_extra_inputs() -> list:
+    return [
+        IO.Combo.Input(
+            "reasoning_effort",
+            options=_REASONING_EFFORTS,
+            default="off",
+            tooltip="Reasoning effort. 'off' disables reasoning entirely.",
+            advanced=True,
+        ),
+    ]
+
+
+def _perplexity_extra_inputs() -> list:
+    return [
+        IO.Combo.Input(
+            "search_context_size",
+            options=_SEARCH_CONTEXT_SIZES,
+            default="medium",
+            tooltip="How much web search context to retrieve. Larger = more grounded but slower/pricier.",
+            advanced=True,
+        ),
+    ]
+
+
+def _profile_inputs(profile: Profile) -> list:
+    if profile == "standard":
+        return []
+    if profile in ("reasoning", "frontier_reasoning"):
+        return _reasoning_extra_inputs()
+    if profile == "perplexity":
+        return _perplexity_extra_inputs()
+    if profile == "perplexity_reasoning":
+        return _perplexity_extra_inputs() + _reasoning_extra_inputs()
+    raise ValueError(f"Unknown profile: {profile}")
+
+
+def _media_inputs(spec: _ModelSpec) -> list:
+    extras: list = []
+    if spec.max_images > 0:
+        extras.append(
+            IO.Autogrow.Input(
+                "images",
+                template=IO.Autogrow.TemplateNames(
+                    IO.Image.Input("image"),
+                    names=[f"image_{i}" for i in range(1, spec.max_images + 1)],
+                    min=0,
+                ),
+                tooltip=f"Optional reference image(s) — up to {spec.max_images}. Sent as URLs.",
+            )
+        )
+    if spec.max_videos > 0:
+        extras.append(
+            IO.Autogrow.Input(
+                "videos",
+                template=IO.Autogrow.TemplateNames(
+                    IO.Video.Input("video"),
+                    names=[f"video_{i}" for i in range(1, spec.max_videos + 1)],
+                    min=0,
+                ),
+                tooltip=f"Optional reference video(s) — up to {spec.max_videos}. Sent as URLs.",
+            )
+        )
+    return extras
+
+
+def _inputs_for_model(spec: _ModelSpec) -> list:
+    return _profile_inputs(spec.profile) + _media_inputs(spec)
+
+
+def _build_model_options() -> list[IO.DynamicCombo.Option]:
+    return [IO.DynamicCombo.Option(spec.slug, _inputs_for_model(spec)) for spec in MODELS]
+
+
+def _calculate_price(response: OpenRouterChatResponse) -> float | None:
+    if response.usage and response.usage.cost is not None:
+        return float(response.usage.cost)
+    return None
+
+
+def _price_badge_jsonata() -> str:
+    rates_pairs = []
+    for spec in MODELS:
+        prompt_per_1k = spec.price_in * 1000
+        completion_per_1k = spec.price_out * 1000
+        rates_pairs.append(f'  "{spec.slug}": [{prompt_per_1k:.8g}, {completion_per_1k:.8g}]')
+    rates_block = ",\n".join(rates_pairs)
+    return (
+        "(\n"
+        "  $rates := {\n"
+        f"{rates_block}\n"
+        "  };\n"
+        "  $r := $lookup($rates, widgets.model);\n"
+        "  $r ? {\n"
+        '    "type": "list_usd",\n'
+        '    "usd": $r,\n'
+        '    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }\n'
+        '  } : {"type": "text", "text": "Token-based"}\n'
+        ")"
+    )
+
+
+async def _build_image_blocks(
+    cls: type[IO.ComfyNode], spec: _ModelSpec, images: list[Input.Image]
+) -> list[OpenRouterImageContent]:
+    urls = await upload_images_to_comfyapi(
+        cls,
+        images,
+        max_images=spec.max_images,
+        total_pixels=2048 * 2048,
+        mime_type="image/png",
+        wait_label="Uploading reference images",
+    )
+    return [OpenRouterImageContent(image_url=OpenRouterImageUrl(url=url)) for url in urls]
+
+
+async def _build_video_blocks(cls: type[IO.ComfyNode], videos: list[Input.Video]) -> list[OpenRouterVideoContent]:
+    blocks: list[OpenRouterVideoContent] = []
+    total = len(videos)
+    for idx, video in enumerate(videos):
+        label = "Uploading reference video"
+        if total > 1:
+            label = f"{label} ({idx + 1}/{total})"
+        url = await upload_video_to_comfyapi(cls, video, wait_label=label)
+        blocks.append(OpenRouterVideoContent(video_url=OpenRouterVideoUrl(url=url)))
+    return blocks
+
+
+def _user_message(prompt: str, media_blocks: list[OpenRouterContentBlock]) -> OpenRouterMessage:
+    if not media_blocks:
+        return OpenRouterMessage(role="user", content=prompt)
+    blocks: list[OpenRouterContentBlock] = list(media_blocks)
+    blocks.append(OpenRouterTextContent(text=prompt))
+    return OpenRouterMessage(role="user", content=blocks)
+
+
+def _build_messages(
+    system_prompt: str, prompt: str, media_blocks: list[OpenRouterContentBlock]
+) -> list[OpenRouterMessage]:
+    messages: list[OpenRouterMessage] = []
+    if system_prompt:
+        messages.append(OpenRouterMessage(role="system", content=system_prompt))
+    messages.append(_user_message(prompt, media_blocks))
+    return messages
+
+
+def _build_request(
+    slug: str,
+    system_prompt: str,
+    prompt: str,
+    media_blocks: list[OpenRouterContentBlock],
+    *,
+    seed: int,
+    reasoning_effort: str | None,
+    search_context_size: str | None,
+) -> OpenRouterChatRequest:
+    reasoning_cfg: OpenRouterReasoningConfig | None = None
+    if reasoning_effort and reasoning_effort != "off":
+        # exclude=True asks providers to reason internally but not return the trace
+        reasoning_cfg = OpenRouterReasoningConfig(effort=reasoning_effort, exclude=True)
+    web_search_cfg: OpenRouterWebSearchOptions | None = None
+    if search_context_size:
+        web_search_cfg = OpenRouterWebSearchOptions(search_context_size=search_context_size)
+    return OpenRouterChatRequest(
+        model=slug,
+        messages=_build_messages(system_prompt, prompt, media_blocks),
+        seed=seed if seed > 0 else None,
+        reasoning=reasoning_cfg,
+        web_search_options=web_search_cfg,
+    )
+
+
+def _extract_text(response: OpenRouterChatResponse) -> str:
+    if response.error:
+        code = response.error.code if response.error.code is not None else "unknown"
+        raise ValueError(f"OpenRouter error ({code}): {response.error.message or 'no message'}")
+    if not response.choices:
+        raise ValueError("Empty response from OpenRouter (no choices).")
+    message = response.choices[0].message
+    if not message:
+        raise ValueError("Empty response from OpenRouter (no message).")
+    if message.refusal:
+        raise ValueError(f"Model refused to respond: {message.refusal}")
+    return message.content or ""
+
+
+class OpenRouterLLMNode(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="OpenRouterLLMNode",
+            display_name="OpenRouter LLM",
+            category="api node/text/OpenRouter",
+            essentials_category="Text Generation",
+            description=(
+                "Generate text responses through OpenRouter. Routes to a curated set of popular "
+                "models from xAI, DeepSeek, Qwen, Mistral, Z.AI (GLM), Moonshot (Kimi), and "
+                "Perplexity Sonar."
+            ),
+            inputs=[
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Text input to the model.",
+                ),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=_build_model_options(),
+                    tooltip="The OpenRouter model used to generate the response.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    tooltip="Seed for sampling. Set to 0 to omit. Most models treat this as a hint only.",
+                ),
+                IO.String.Input(
+                    "system_prompt",
+                    multiline=True,
+                    default="",
+                    optional=True,
+                    advanced=True,
+                    tooltip="Foundational instructions that dictate the model's behavior.",
+                ),
+            ],
+            outputs=[IO.String.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(widgets=["model"]),
+                expr=_price_badge_jsonata(),
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        prompt: str,
+        model: dict,
+        seed: int,
+        system_prompt: str = "",
+    ) -> IO.NodeOutput:
+        validate_string(prompt, strip_whitespace=True, min_length=1)
+        slug: str = model["model"]
+        spec = _MODELS_BY_SLUG.get(slug)
+        if spec is None:
+            raise ValueError(f"Unknown OpenRouter model: {slug}")
+
+        reasoning_effort: str | None = model.get("reasoning_effort")
+        search_context_size: str | None = model.get("search_context_size")
+
+        image_tensors: list[Input.Image] = [t for t in (model.get("images") or {}).values() if t is not None]
+        if image_tensors and sum(get_number_of_images(t) for t in image_tensors) > spec.max_images:
+            raise ValueError(f"Up to {spec.max_images} images are supported for {slug}.")
+        video_inputs: list[Input.Video] = [v for v in (model.get("videos") or {}).values() if v is not None]
+        if video_inputs and len(video_inputs) > spec.max_videos:
+            raise ValueError(f"Up to {spec.max_videos} videos are supported for {slug}.")
+
+        media_blocks: list[OpenRouterContentBlock] = []
+        if image_tensors:
+            media_blocks.extend(await _build_image_blocks(cls, spec, image_tensors))
+        if video_inputs:
+            media_blocks.extend(await _build_video_blocks(cls, video_inputs))
+
+        request = _build_request(
+            slug,
+            system_prompt,
+            prompt,
+            media_blocks,
+            seed=seed,
+            reasoning_effort=reasoning_effort,
+            search_context_size=search_context_size,
+        )
+
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=OPENROUTER_CHAT_ENDPOINT, method="POST"),
+            response_model=OpenRouterChatResponse,
+            data=request,
+            price_extractor=_calculate_price,
+        )
+        return IO.NodeOutput(_extract_text(response))
+
+
+class OpenRouterExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [OpenRouterLLMNode]
+
+
+async def comfy_entrypoint() -> OpenRouterExtension:
+    return OpenRouterExtension()

comfy_api_nodes/util/__init__.py

@@ -16,16 +16,17 @@ from .conversions import (
     convert_mask_to_image,
     downscale_image_tensor,
     downscale_image_tensor_by_max_side,
+    downscale_video_to_max_pixels,
     image_tensor_pair_to_batch,
     pil_to_bytesio,
     resize_mask_to_image,
-    resize_video_to_pixel_budget,
     tensor_to_base64_string,
     tensor_to_bytesio,
     tensor_to_pil,
     text_filepath_to_base64_string,
     text_filepath_to_data_uri,
     trim_video,
+    upscale_video_to_min_pixels,
     video_to_base64_string,
 )
 from .download_helpers import (
@@ -88,16 +89,17 @@ __all__ = [
     "convert_mask_to_image",
     "downscale_image_tensor",
     "downscale_image_tensor_by_max_side",
+    "downscale_video_to_max_pixels",
     "image_tensor_pair_to_batch",
     "pil_to_bytesio",
     "resize_mask_to_image",
-    "resize_video_to_pixel_budget",
     "tensor_to_base64_string",
     "tensor_to_bytesio",
     "tensor_to_pil",
     "text_filepath_to_base64_string",
     "text_filepath_to_data_uri",
     "trim_video",
+    "upscale_video_to_min_pixels",
     "video_to_base64_string",
     # Validation utilities
     "get_image_dimensions",

comfy_api_nodes/util/conversions.py

@@ -415,14 +415,48 @@ def trim_video(video: Input.Video, duration_sec: float) -> Input.Video:
         raise RuntimeError(f"Failed to trim video: {str(e)}") from e
 
 
-def resize_video_to_pixel_budget(video: Input.Video, total_pixels: int) -> Input.Video:
-    """Downscale a video to fit within ``total_pixels`` (w * h), preserving aspect ratio.
+def downscale_video_to_max_pixels(video: Input.Video, max_pixels: int) -> Input.Video:
+    """Downscale a video to fit within ``max_pixels`` (w * h), preserving aspect ratio.
 
     Returns the original video object untouched when it already fits. Preserves frame rate, duration, and audio.
     Aspect ratio is preserved up to a fraction of a percent (even-dim rounding).
     """
     src_w, src_h = video.get_dimensions()
-    scale_dims = _compute_downscale_dims(src_w, src_h, total_pixels)
+    scale_dims = _compute_downscale_dims(src_w, src_h, max_pixels)
+    if scale_dims is None:
+        return video
+    return _apply_video_scale(video, scale_dims)
+
+
+def _compute_upscale_dims(src_w: int, src_h: int, total_pixels: int) -> tuple[int, int] | None:
+    """Return upscaled (w, h) with even dims meeting at least ``total_pixels``, or None if already large enough.
+
+    Source aspect ratio is preserved; output may drift by a fraction of a percent because both dimensions
+    are rounded up to even values (many codecs require divisible-by-2). The result is guaranteed to be at
+    least ``total_pixels``.
+    """
+    pixels = src_w * src_h
+    if pixels >= total_pixels:
+        return None
+    scale = math.sqrt(total_pixels / pixels)
+    new_w = math.ceil(src_w * scale)
+    new_h = math.ceil(src_h * scale)
+    if new_w % 2:
+        new_w += 1
+    if new_h % 2:
+        new_h += 1
+    return new_w, new_h
+
+
+def upscale_video_to_min_pixels(video: Input.Video, min_pixels: int) -> Input.Video:
+    """Upscale a video to meet at least ``min_pixels`` (w * h), preserving aspect ratio.
+
+    Returns the original video object untouched when it already meets the minimum. Preserves frame rate,
+    duration, and audio. Aspect ratio is preserved up to a fraction of a percent (even-dim rounding).
+    Note: upscaling a low-resolution source does not add real detail; downstream model quality may suffer.
+    """
+    src_w, src_h = video.get_dimensions()
+    scale_dims = _compute_upscale_dims(src_w, src_h, min_pixels)
     if scale_dims is None:
         return video
     return _apply_video_scale(video, scale_dims)

comfy_extras/nodes_lt.py

@@ -77,7 +77,7 @@ class EmptyLTXVLatentVideo(io.ComfyNode):
     @classmethod
     def execute(cls, width, height, length, batch_size=1) -> io.NodeOutput:
         latent = torch.zeros([batch_size, 128, ((length - 1) // 8) + 1, height // 32, width // 32], device=comfy.model_management.intermediate_device())
-        return io.NodeOutput({"samples": latent})
+        return io.NodeOutput({"samples": latent, "downscale_ratio_spacial": 32})
 
     generate = execute  # TODO: remove
 

comfy_extras/nodes_string.py

@@ -1,10 +1,41 @@
 import re
 import json
+import string
 from typing_extensions import override
 
 from comfy_api.latest import ComfyExtension, io
 
 
+class StringFormat(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        autogrow = io.Autogrow.TemplateNames(
+            input=io.AnyType.Input("value"),
+            names=list(string.ascii_lowercase),
+            min=0,
+        )
+        return io.Schema(
+            node_id="StringFormat",
+            display_name="Format Text",
+            category="text",
+            search_aliases=["string", "format"],
+            description="Same as Python's string format method. Supports all of Python's format options and features.",
+            inputs=[
+                io.Autogrow.Input("values", template=autogrow),
+                io.String.Input("f_string", default="{a}", multiline=True),
+            ],
+            outputs=[
+                io.String.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(
+        cls, values: io.Autogrow.Type, f_string: str
+    ) -> io.NodeOutput:
+        return io.NodeOutput(f_string.format(**values))
+
+
 class StringConcatenate(io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -413,6 +444,7 @@ class StringExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
         return [
+            StringFormat,
             StringConcatenate,
             StringSubstring,
             StringLength,

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.21.1"
+__version__ = "0.22.1"

openapi.yaml

@@ -4160,6 +4160,10 @@ paths:
                 name:
                   type: string
                   description: Display name for the API key
+                description:
+                  type: string
+                  description: User-provided description of the key's purpose
+                  maxLength: 5000
       responses:
         "201":
           description: API key created
@@ -6351,14 +6355,6 @@ components:
           type: integer
           format: int64
           description: Size of the asset in bytes
-        width:
-          type: integer
-          nullable: true
-          description: "Original image width in pixels. Null for non-image assets or assets ingested before dimension extraction."
-        height:
-          type: integer
-          nullable: true
-          description: "Original image height in pixels. Null for non-image assets or assets ingested before dimension extraction."
         mime_type:
           type: string
           description: MIME type of the asset
@@ -7685,11 +7681,16 @@ components:
       required:
         - id
         - name
+        - description
       properties:
         id:
           type: string
         name:
           type: string
+        description:
+          type: string
+          maxLength: 5000
+          description: User-provided description of the key's purpose. Always present in responses; empty string when no description was supplied on create.
         prefix:
           type: string
           description: First few characters of the key for identification
@@ -7710,12 +7711,17 @@ components:
       required:
         - id
         - name
+        - description
         - key
       properties:
         id:
           type: string
         name:
           type: string
+        description:
+          type: string
+          maxLength: 5000
+          description: User-provided description of the key's purpose. Always present in responses; empty string when no description was supplied on create.
         key:
           type: string
           description: Full API key value (only returned on creation)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.21.1"
+version = "0.22.1"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.43.18
-comfyui-workflow-templates==0.9.77
+comfyui-workflow-templates==0.9.82
 comfyui-embedded-docs==0.5.0
 torch
 torchsde