From 78b5dec6b6beefb9fb40f917d33d2f10a40d9e53 Mon Sep 17 00:00:00 2001
From: Cezarijus Kivylius <ltcezaris@gmail.com>
Date: Wed, 20 May 2026 12:58:49 +0100
Subject: [PATCH 01/14] fix: Hunyuan3D 2.1 batch size crashes in attention and
 forward pass (#13699)

---
 comfy/ldm/hunyuan3dv2_1/hunyuandit.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py b/comfy/ldm/hunyuan3dv2_1/hunyuandit.py
index f67ba84e9..bc36b8998 100644
--- a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py
+++ b/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

From f9c84c94b43855ec46e26afc944c06fb121cb9a8 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Wed, 20 May 2026 08:34:22 -0700
Subject: [PATCH 02/14] Support Stable Audio 3 model. (#14010)

---
 comfy/latent_formats.py      |   5 +
 comfy/ldm/audio/dit.py       | 250 +++++++++++++---
 comfy/ldm/audio/embedders.py |  31 +-
 comfy/ldm/audio/vae_sa3.py   | 533 +++++++++++++++++++++++++++++++++++
 comfy/model_base.py          |  79 ++++++
 comfy/model_detection.py     |  39 +++
 comfy/sd.py                  |  37 +++
 comfy/supported_models.py    |  25 ++
 comfy/text_encoders/sa3.py   | 207 ++++++++++++++
 9 files changed, 1161 insertions(+), 45 deletions(-)
 create mode 100644 comfy/ldm/audio/vae_sa3.py
 create mode 100644 comfy/text_encoders/sa3.py

diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py
index 6e37080bb..75d459b59 100644
--- a/comfy/latent_formats.py
+++ b/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):
diff --git a/comfy/ldm/audio/dit.py b/comfy/ldm/audio/dit.py
index ca865189e..a6258b755 100644
--- a/comfy/ldm/audio/dit.py
+++ b/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,
diff --git a/comfy/ldm/audio/embedders.py b/comfy/ldm/audio/embedders.py
index 20edb365a..ba9a62837 100644
--- a/comfy/ldm/audio/embedders.py
+++ b/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
diff --git a/comfy/ldm/audio/vae_sa3.py b/comfy/ldm/audio/vae_sa3.py
new file mode 100644
index 000000000..276846444
--- /dev/null
+++ b/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
diff --git a/comfy/model_base.py b/comfy/model_base.py
index c22705655..d81f13c69 100644
--- a/comfy/model_base.py
+++ b/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):
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index bc0b933bc..70b4df8b3 100644
--- a/comfy/model_detection.py
+++ b/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
diff --git a/comfy/sd.py b/comfy/sd.py
index 2443353a4..7bd07ed3a 100644
--- a/comfy/sd.py
+++ b/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,
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 1e4434fd5..617db4f28 100644
--- a/comfy/supported_models.py
+++ b/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,
diff --git a/comfy/text_encoders/sa3.py b/comfy/text_encoders/sa3.py
new file mode 100644
index 000000000..0a1c73ec1
--- /dev/null
+++ b/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)

From 4efe1ddb5c7933e9db28d5ecb0be4fa77e0edcde Mon Sep 17 00:00:00 2001
From: "Daxiong (Lin)" <contact@comfyui-wiki.com>
Date: Wed, 20 May 2026 23:46:20 +0800
Subject: [PATCH 03/14] chore: update workflow templates to v0.9.79 (#14011)

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index f499a10ae..1c87690da 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.43.18
-comfyui-workflow-templates==0.9.77
+comfyui-workflow-templates==0.9.79
 comfyui-embedded-docs==0.5.0
 torch
 torchsde

From a8d2519058ea766ca3b14916bcc01ecef5efd235 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 20 May 2026 13:49:36 -0400
Subject: [PATCH 04/14] ComfyUI v0.22.0

---
 comfyui_version.py | 2 +-
 pyproject.toml     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/comfyui_version.py b/comfyui_version.py
index 4c6f5eb2a..0bb0f780c 100644
--- a/comfyui_version.py
+++ b/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.0"
diff --git a/pyproject.toml b/pyproject.toml
index 0a1554428..1e449b4a3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.21.1"
+version = "0.22.0"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"

From 4d6a058bf1dd18fb6d4594081c3f9a7575c97256 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?=
 <40791699+kijai@users.noreply.github.com>
Date: Thu, 21 May 2026 02:07:48 +0300
Subject: [PATCH 05/14] feat: MediaPipe face detection (CORE-235) (#14009)

* Initial mediapipe face detection support

* Update face_geometry.py

* Account for diff sized batch input

* Model folder placeholder
---
 comfy_extras/mediapipe/face_geometry.py    | 111 ++++
 comfy_extras/mediapipe/face_landmarker.py  | 682 +++++++++++++++++++++
 comfy_extras/nodes_mediapipe.py            | 502 +++++++++++++++
 folder_paths.py                            |   2 +
 models/mediapipe/put_mediapipe_models_here |   0
 nodes.py                                   |   1 +
 6 files changed, 1298 insertions(+)
 create mode 100644 comfy_extras/mediapipe/face_geometry.py
 create mode 100644 comfy_extras/mediapipe/face_landmarker.py
 create mode 100644 comfy_extras/nodes_mediapipe.py
 create mode 100644 models/mediapipe/put_mediapipe_models_here

diff --git a/comfy_extras/mediapipe/face_geometry.py b/comfy_extras/mediapipe/face_geometry.py
new file mode 100644
index 000000000..04b2b0557
--- /dev/null
+++ b/comfy_extras/mediapipe/face_geometry.py
@@ -0,0 +1,111 @@
+"""Pure-numpy port of MediaPipe's face_geometry (FACE_LANDMARK_PIPELINE mode)
++ weighted Procrustes solver. Computes the 4x4 facial transformation matrix.
+"""
+
+from __future__ import annotations
+
+import math
+import numpy as np
+
+
+def _solve_weighted_orthogonal_problem(src: np.ndarray, tgt: np.ndarray, weights: np.ndarray) -> np.ndarray:
+    """Weighted orthogonal Procrustes (similarity). Returns 4x4 M with
+    `target ≈ M @ homogeneous(source)` in the weighted LS sense. fp64 for
+    SVD stability. Port of procrustes_solver.cc."""
+    sqrt_w = np.sqrt(weights.astype(np.float64))
+    w_total = float((sqrt_w ** 2).sum())
+    ws = src.astype(np.float64) * sqrt_w
+    wt = tgt.astype(np.float64) * sqrt_w
+
+    c_w = (ws @ sqrt_w) / w_total
+    centered = ws - np.outer(c_w, sqrt_w)
+    U, _S, Vt = np.linalg.svd(wt @ centered.T, full_matrices=True)
+    # Disallow reflection: flip the least-significant axis when det(U)·det(V)<0.
+    post, pre = U.copy(), Vt.T.copy()
+    if np.linalg.det(post) * np.linalg.det(pre) < 0:
+        post[:, 2] *= -1.0
+    R = post @ pre.T
+
+    denom = float((centered * ws).sum())
+    if denom < 1e-12:
+        raise ValueError("Procrustes denominator collapsed (degenerate source).")
+    scale = float((R @ centered * wt).sum()) / denom
+    translation = ((wt - scale * (R @ ws)) @ sqrt_w) / w_total
+
+    M = np.eye(4, dtype=np.float64)
+    M[:3, :3] = scale * R
+    M[:3, 3] = translation
+    return M
+
+
+def _estimate_scale(canonical: np.ndarray, runtime: np.ndarray, weights: np.ndarray) -> float:
+    """scale = ‖first column of M[:3]‖ per geometry_pipeline.cc::EstimateScale."""
+    return float(np.linalg.norm(_solve_weighted_orthogonal_problem(canonical, runtime, weights)[:3, 0]))
+
+
+def solve_facial_transformation_matrix(
+    landmarks_normalized: np.ndarray,
+    canonical_vertices: np.ndarray,
+    procrustes_indices: np.ndarray,
+    procrustes_weights: np.ndarray,
+    image_width: int,
+    image_height: int,
+    # face_geometry_calculator_options.pbtxt defaults
+    vertical_fov_degrees: float = 63.0,
+    near: float = 1.0,
+) -> np.ndarray:
+    """4x4 facial transformation matrix via two-pass scale recovery
+    `landmarks_normalized` is (N, 3) in MediaPipe normalized convention: x, y
+    in [0,1] with TOP-LEFT origin, z in width-scaled units.
+    """
+
+    h_near = 2.0 * near * math.tan(0.5 * math.radians(vertical_fov_degrees))
+    w_near = image_width * h_near / image_height
+
+    sub = procrustes_indices.astype(np.int64)
+    screen = landmarks_normalized[sub].T.astype(np.float64).copy()
+    canon = canonical_vertices[sub].T.astype(np.float64).copy()
+    weights = procrustes_weights.astype(np.float64)
+
+    # ProjectXY (TOP_LEFT y-flip, then scale all 3 axes; z uses x-scale).
+    screen[1] = 1.0 - screen[1]
+    screen[0] = screen[0] * w_near - 0.5 * w_near
+    screen[1] = screen[1] * h_near - 0.5 * h_near
+    screen[2] = screen[2] * w_near
+    depth_offset = float(screen[2].mean())
+
+    def _unproject(s: np.ndarray, scale: float) -> np.ndarray:
+        s = s.copy()
+        s[2] = (s[2] - depth_offset + near) / scale
+        s[0] *= s[2] / near
+        s[1] *= s[2] / near
+        s[2] *= -1.0
+        return s
+
+    first = screen.copy()
+    first[2] *= -1.0
+    s1 = _estimate_scale(canon, first, weights) # 1st pass: Procrustes on projected XY
+    s2 = _estimate_scale(canon, _unproject(screen, s1), weights) # 2nd pass: rescale z by s1, un-project XY
+    return _solve_weighted_orthogonal_problem(canon, _unproject(screen, s1 * s2), weights).astype(np.float32)
+
+
+def transformation_matrix_from_detection(face_dict: dict, image_width: int, image_height: int, canonical_data: dict) -> np.ndarray:
+    """Adapt a FaceLandmarker face dict to MP's normalized convention and solve.
+    FaceMesh emits (x, y, z) in 192-canonical units; MP's geometry expects
+    z_norm = z_canonical * scale_x / image_width"""
+
+    lmks_xy, lmks_3d = face_dict["landmarks_xy"], face_dict["landmarks_3d"]
+    aug = np.concatenate([lmks_3d[:, :2].astype(np.float64), np.ones((lmks_xy.shape[0], 1))], axis=1)
+    M, *_ = np.linalg.lstsq(aug, lmks_xy.astype(np.float64), rcond=None)
+    scale_x = float(np.linalg.norm(M[0]))
+    z_scale = scale_x / image_width if scale_x > 1e-6 else 1.0 / image_width
+
+    normalized = np.empty((lmks_xy.shape[0], 3), dtype=np.float32)
+    normalized[:, 0] = lmks_xy[:, 0] / image_width
+    normalized[:, 1] = lmks_xy[:, 1] / image_height
+    normalized[:, 2] = lmks_3d[:, 2] * z_scale
+    return solve_facial_transformation_matrix(
+        normalized, canonical_data["canonical_vertices"],
+        canonical_data["procrustes_indices"], canonical_data["procrustes_weights"],
+        image_width=image_width, image_height=image_height,
+    )
diff --git a/comfy_extras/mediapipe/face_landmarker.py b/comfy_extras/mediapipe/face_landmarker.py
new file mode 100644
index 000000000..a792b6046
--- /dev/null
+++ b/comfy_extras/mediapipe/face_landmarker.py
@@ -0,0 +1,682 @@
+"""Pure-PyTorch port of MediaPipe's face_landmarker_v2_with_blendshapes.task:
+BlazeFace detector → FaceMesh v2 → ARKit-52 blendshapes."""
+
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.special import expit
+from torch import Tensor, nn
+
+
+# Values below must stay verbatim with the published face_landmarker_v2 graph
+
+# face_blendshapes_graph.cc::kLandmarksSubsetIdxs
+_BS_INPUT_INDICES: Tuple[int, ...] = (
+    0, 1, 4, 5, 6, 7, 8, 10, 13, 14, 17, 21, 33, 37, 39, 40, 46, 52, 53, 54,
+    55, 58, 61, 63, 65, 66, 67, 70, 78, 80, 81, 82, 84, 87, 88, 91, 93, 95,
+    103, 105, 107, 109, 127, 132, 133, 136, 144, 145, 146, 148, 149, 150, 152,
+    153, 154, 155, 157, 158, 159, 160, 161, 162, 163, 168, 172, 173, 176, 178,
+    181, 185, 191, 195, 197, 234, 246, 249, 251, 263, 267, 269, 270, 276, 282,
+    283, 284, 285, 288, 291, 293, 295, 296, 297, 300, 308, 310, 311, 312, 314,
+    317, 318, 321, 323, 324, 332, 334, 336, 338, 356, 361, 362, 365, 373, 374,
+    375, 377, 378, 379, 380, 381, 382, 384, 385, 386, 387, 388, 389, 390, 397,
+    398, 400, 402, 405, 409, 415, 454, 466, 468, 469, 470, 471, 472, 473, 474,
+    475, 476, 477,
+)
+
+# face_blendshapes_graph.cc::kCategoryNames
+BLENDSHAPE_NAMES: Tuple[str, ...] = (
+    "_neutral", "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft",
+    "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight",
+    "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight",
+    "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight",
+    "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight",
+    "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen",
+    "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight",
+    "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft",
+    "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft",
+    "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower",
+    "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft",
+    "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight",
+    "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", "noseSneerRight",
+)
+
+# face_detection.pbtxt — short-range BlazeFace.
+_BF_NUM_LAYERS = 4
+_BF_INPUT_SIZE = 128
+_BF_STRIDES = (8, 16, 16, 16)
+_BF_ANCHOR_OFFSET_X = 0.5
+_BF_ANCHOR_OFFSET_Y = 0.5
+_BF_ASPECT_RATIOS = (1.0,)
+_BF_INTERP_SCALE_AR = 1.0
+_BF_BOX_SCALE = 128.0
+_BF_KP_OFFSET = 4
+_BF_SCORE_CLIP = 100.0
+_BF_MIN_SCORE = 0.5
+
+# face_detection_full_range.pbtxt — 48x48 grid at stride 4, 1 anchor/cell.
+_BF_FR_INPUT_SIZE = 192
+_BF_FR_GRID = 48
+_BF_FR_NUM_ANCHORS = _BF_FR_GRID * _BF_FR_GRID
+_BF_FR_BOX_SCALE = 192.0
+_BF_FR_SCORE_CLIP = 100.0
+
+_FM_INPUT_SIZE = 192
+
+# Face ROI: 1.5xbbox rect warped anisotropically into 192x192.
+_FACE_LEFT_EYE_KP = 0
+_FACE_RIGHT_EYE_KP = 1
+_FACE_ROI_SCALE_X = 1.5
+_FACE_ROI_SCALE_Y = 1.5
+_FACE_ROI_TARGET_ANGLE = 0.0
+
+
+def _tf_same_pad(x: Tensor, kernel: int, stride: int) -> Tensor:
+    """TF SAME pad (asymmetric on stride-2; PyTorch's symmetric pad undershoots by 1 px)."""
+    H, W = x.shape[-2], x.shape[-1]
+    pad_h = max(((H + stride - 1) // stride - 1) * stride + kernel - H, 0)
+    pad_w = max(((W + stride - 1) // stride - 1) * stride + kernel - W, 0)
+    if pad_h == 0 and pad_w == 0:
+        return x
+    return F.pad(x, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2))
+
+
+# BlazeFace short-range: stem 5x5/s2 → 16 BlazeBlocks → parallel heads at
+# 16²x88 (2 anchors/cell) and 8²x96 (6/cell) = 896 anchors. (in, out, stride):
+_BLAZEFACE_BLOCKS = [
+    (24, 24, 1), (24, 28, 1), (28, 32, 2), (32, 36, 1),
+    (36, 42, 1), (42, 48, 2), (48, 56, 1), (56, 64, 1),
+    (64, 72, 1), (72, 80, 1), (80, 88, 1), (88, 96, 2),
+    (96, 96, 1), (96, 96, 1), (96, 96, 1), (96, 96, 1),
+]
+
+
+class BlazeFaceBlock(nn.Module):
+    """DW 3x3 + PW + residual. Residual max-pools on stride>1, channel-pads on out_ch>in_ch."""
+
+    def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride
+        self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, device=device, dtype=dtype)
+        self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, device=device, dtype=dtype)
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x
+        if self.out_ch > self.in_ch:
+            residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))
+        x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))
+        return F.relu(self.pointwise(self.depthwise(x)) + residual)
+
+
+class BlazeFace(nn.Module):
+    """Short-range BlazeFace: (B, 3, 128, 128) in [-1, 1] → 896 anchors x 17."""
+
+    def __init__(self, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        self.stem = ops.Conv2d(3, 24, 5, stride=2, padding=0, bias=True, **kw)
+        self.blocks = nn.ModuleList(BlazeFaceBlock(i, o, s, device=device, dtype=dtype, operations=operations)
+                                    for (i, o, s) in _BLAZEFACE_BLOCKS)
+        # 16²x2 + 8²x6 = 512 + 384 = 896 anchors.
+        self.cls_16 = ops.Conv2d(88, 2, 1, padding=0, bias=True, **kw)
+        self.cls_8 = ops.Conv2d(96, 6, 1, padding=0, bias=True, **kw)
+        self.reg_16 = ops.Conv2d(88, 32, 1, padding=0, bias=True, **kw)
+        self.reg_8 = ops.Conv2d(96, 96, 1, padding=0, bias=True, **kw)
+
+    def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
+        x = F.relu(self.stem(_tf_same_pad(image_chw_normalized, 5, 2)))
+        # 16x16 tap is block-10 output (before the 88→96 stride-2 in block 11).
+        for i in range(11):
+            x = self.blocks[i](x)
+        feat_16 = x
+        for i in range(11, 16):
+            x = self.blocks[i](x)
+        feat_8 = x
+
+        def flat(t, a, k):  # NHWC flatten → (B, H*W*A, K)
+            B, _, H, W = t.shape
+            return t.permute(0, 2, 3, 1).reshape(B, H * W * a, k)
+
+        cls = torch.cat([flat(self.cls_16(feat_16), 2, 1), flat(self.cls_8(feat_8), 6, 1)], dim=1)
+        reg = torch.cat([flat(self.reg_16(feat_16), 2, 16), flat(self.reg_8(feat_8), 6, 16)], dim=1)
+        return reg, cls
+
+
+# BlazeFace full-range (face_detection_full_range_sparse.tflite): MobileNetV2-ish
+# backbone + top-down FPN, 192² input → 2304 anchors at the 48x48 grid.
+class FRBlock(nn.Module):
+    """Double inverted residual: DW → PW(mid) → DW → PW(out) [+ residual].
+
+    Per source tflite: dw* have no fused activation, pw1 is always ReLU, pw2
+    is ReLU only when no residual (else ReLU fuses into the ADD).
+    """
+
+    def __init__(self, in_ch: int, mid_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        self.has_residual = (in_ch == out_ch and stride == 1)
+        self.dw1 = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)
+        self.pw1 = ops.Conv2d(in_ch, mid_ch, 1, padding=0, bias=True, **kw)
+        self.dw2 = ops.Conv2d(mid_ch, mid_ch, 3, stride=1, padding=0, groups=mid_ch, bias=True, **kw)
+        self.pw2 = ops.Conv2d(mid_ch, out_ch, 1, padding=0, bias=True, **kw)
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = x if self.has_residual else None
+        x = F.relu(self.pw1(self.dw1(F.pad(x, (1, 1, 1, 1)))))
+        x = self.pw2(self.dw2(F.pad(x, (1, 1, 1, 1))))
+        return F.relu(x + residual) if residual is not None else F.relu(x)
+
+
+# (in_ch, mid_ch, out_ch, stride). Stages downsample 96²x32 → 48²x64 → 24²x128
+# → 12²x192 → 6²x384. Lateral taps at indices 4, 7, 10 (see _FR_LATERAL_*).
+_FR_BACKBONE_BLOCKS = [
+    (32, 8, 32, 1),    (32, 8, 32, 1),                                            # 96²x32
+    (32, 16, 64, 2),   (64, 16, 64, 1),   (64, 16, 64, 1),                        # 48²x64 — tap[0]
+    (64, 32, 128, 2),  (128, 32, 128, 1), (128, 32, 128, 1),                      # 24²x128 — tap[1]
+    (128, 48, 192, 2), (192, 48, 192, 1), (192, 48, 192, 1),                      # 12²x192 — tap[2]
+    (192, 96, 384, 2), (384, 96, 384, 1), (384, 96, 384, 1), (384, 96, 384, 1),   # 6²x384
+]
+_FR_LATERAL_TAP_INDICES = (4, 7, 10)
+_FR_LATERAL_CHANNELS = ((64, 48), (128, 64), (192, 96))  # (in, out) per side-conv
+
+# Decoder blocks per FPN level (after upsample-and-merge with the lateral).
+_FR_DECODER_BLOCKS = [
+    [(96, 48, 96, 1), (96, 48, 96, 1)],  # 12²x96
+    [(64, 32, 64, 1), (64, 32, 64, 1)],  # 24²x64
+    [(48, 24, 48, 1)],                   # 48²x48 — feeds the heads
+]
+
+
+def _dcr_depth_to_space(t: Tensor, r: int, c_out: int) -> Tensor:
+    """TF DEPTH_TO_SPACE in DCR layout (input channels = (i, j, c_out)).
+    pixel_shuffle uses CRD which permutes output channels for c_out > 1."""
+    B_, _, H_, W_ = t.shape
+    t = t.reshape(B_, r, r, c_out, H_, W_)
+    t = t.permute(0, 3, 4, 1, 5, 2).contiguous()
+    return t.reshape(B_, c_out, H_ * r, W_ * r)
+
+
+class BlazeFaceFullRange(nn.Module):
+    """Full-range face detector: (B, 3, 192, 192) in [-1, 1] → 2304 anchors x 17 values."""
+
+    def __init__(self, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        mk_block = lambda i, m, o, s: FRBlock(i, m, o, s, device=device, dtype=dtype, operations=operations)
+        self.stem = ops.Conv2d(3, 32, 3, stride=2, padding=0, bias=True, **kw)
+        self.backbone = nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in _FR_BACKBONE_BLOCKS)
+        self.lateral_convs = nn.ModuleList(ops.Conv2d(i, o, 1, padding=0, bias=True, **kw) for (i, o) in _FR_LATERAL_CHANNELS)
+        self.top_conv = ops.Conv2d(384, 96, 1, padding=0, bias=True, **kw)
+        self.decoder_levels = nn.ModuleList(
+            nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in lvl) for lvl in _FR_DECODER_BLOCKS
+        )
+        # 96→64 before 12→24, 64→48 before 24→48.
+        self.decoder_reduce_convs = nn.ModuleList([
+            ops.Conv2d(96, 64, 1, padding=0, bias=True, **kw),
+            ops.Conv2d(64, 48, 1, padding=0, bias=True, **kw),
+        ])
+        # Heads mix 2x2-cell info via DW-stride-2 + depth_to_space block_size=2.
+        self.cls_conv = ops.Conv2d(48, 4, 1, padding=0, bias=True, **kw)
+        self.cls_dw = ops.Conv2d(4, 4, 3, stride=2, padding=0, groups=4, bias=True, **kw)
+        self.reg_conv = ops.Conv2d(48, 64, 1, padding=0, bias=True, **kw)
+        self.reg_dw = ops.Conv2d(64, 64, 3, stride=2, padding=0, groups=64, bias=True, **kw)
+
+    def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
+        # Symmetric pad-1 throughout (full-range tflite uses explicit TF PAD, not SAME).
+        x = F.relu(self.stem(F.pad(image_chw_normalized, (1, 1, 1, 1))))
+        tap_set = set(_FR_LATERAL_TAP_INDICES)
+        laterals: list[Tensor] = []
+        for i, blk in enumerate(self.backbone):
+            x = blk(x)
+            if i in tap_set:
+                laterals.append(x)
+
+        # top_conv / lateral_convs / decoder_reduce_convs all have fused ReLU in the tflite.
+        p = F.relu(self.top_conv(x))
+        laterals_rev = list(reversed(laterals))
+        lateral_convs_rev = list(reversed(self.lateral_convs))
+        for level in range(len(self.decoder_levels)):
+            lateral = laterals_rev[level]
+            p = F.interpolate(p, size=lateral.shape[-2:], mode="bilinear", align_corners=False)
+            p = p + F.relu(lateral_convs_rev[level](lateral))
+            for blk in self.decoder_levels[level]:
+                p = blk(p)
+            if level < len(self.decoder_reduce_convs):
+                p = F.relu(self.decoder_reduce_convs[level](p))
+
+        c = self.cls_dw(F.pad(self.cls_conv(p), (1, 1, 1, 1)))
+        c = _dcr_depth_to_space(c, r=2, c_out=1)
+        r = self.reg_dw(F.pad(self.reg_conv(p), (1, 1, 1, 1)))
+        r = _dcr_depth_to_space(r, r=2, c_out=16)
+        B = c.shape[0]
+        cls_out = c.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 1)
+        reg_out = r.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 16)
+        return reg_out, cls_out
+
+
+@lru_cache(maxsize=1)
+def _blazeface_full_range_anchors() -> np.ndarray:
+    """2304 anchors over 48x48; anchor_w=anchor_h=1 (fixed_anchor_size)."""
+    feat = _BF_FR_GRID
+    yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")
+    cx, cy, ones = (xx + 0.5) / feat, (yy + 0.5) / feat, np.ones_like(xx)
+    return np.stack([cx, cy, ones, ones], axis=-1).reshape(_BF_FR_NUM_ANCHORS, 4)
+
+
+def _decode_blazeface_full_range(regressors: np.ndarray, classificators: np.ndarray,
+                                 score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:
+    """Same decode as short-range with 2304-anchor grid and box_scale=192."""
+    scores = expit(np.clip(classificators[:, 0], -_BF_FR_SCORE_CLIP, _BF_FR_SCORE_CLIP))
+    keep = scores >= score_thresh
+    if not keep.any():
+        return np.empty((0, 17), dtype=np.float32)
+    r = regressors[keep] / _BF_FR_BOX_SCALE
+    a = _blazeface_full_range_anchors()[keep]
+    cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
+    xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys
+    w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs
+    out = np.empty((r.shape[0], 17), dtype=np.float32)
+    out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
+    out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs
+    out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys
+    out[:, 16] = scores[keep]
+    return out
+
+
+# FaceMesh (face_landmarks_detector.tflite): PReLU variant of BlazeBlock,
+# 17 blocks, heads for 478x3 landmarks + presence.
+_FACEMESH_BLOCKS = [  # (in_ch, out_ch, stride)
+    (16, 16, 1),  (16, 16, 1),  (16, 32, 2),  (32, 32, 1), (32, 32, 1), (32, 64, 2),
+    (64, 64, 1),  (64, 64, 1),  (64, 128, 2), (128, 128, 1), (128, 128, 1), (128, 128, 2),
+    (128, 128, 1), (128, 128, 1), (128, 128, 2), (128, 128, 1), (128, 128, 1),
+]
+
+
+class FaceMeshBlock(nn.Module):
+    """PReLU BlazeBlock: PReLU between DW and PW, and after the residual add."""
+
+    def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride
+        self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)
+        self.prelu_dwise = nn.PReLU(num_parameters=in_ch, **kw)
+        self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, **kw)
+        self.prelu_out = nn.PReLU(num_parameters=out_ch, **kw)
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x
+        if self.out_ch > self.in_ch:
+            residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))
+        x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))
+        return self.prelu_out(self.pointwise(self.prelu_dwise(self.depthwise(x))) + residual)
+
+
+class FaceMesh(nn.Module):
+    NUM_LANDMARKS = 478
+
+    def __init__(self, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        self.stem = ops.Conv2d(3, 16, 3, stride=2, padding=0, bias=True, **kw)
+        self.prelu_stem = nn.PReLU(num_parameters=16, **kw)
+        self.blocks = nn.ModuleList(FaceMeshBlock(i, o, s, device=device, dtype=dtype, operations=operations)
+                                    for (i, o, s) in _FACEMESH_BLOCKS)
+        self.head_reduce = ops.Conv2d(128, 8, 1, padding=0, bias=True, **kw)
+        self.prelu_head_reduce = nn.PReLU(num_parameters=8, **kw)
+        self.head_block = FaceMeshBlock(8, 8, 1, device=device, dtype=dtype, operations=operations)
+        self.head_presence = ops.Conv2d(8, 1, 3, padding=0, bias=True, **kw)
+        self.head_landmarks = ops.Conv2d(8, self.NUM_LANDMARKS * 3, 3, padding=0, bias=True, **kw)
+
+    def forward(self, face_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
+        """(B, 3, 192, 192) in [0, 1] → ((B, 478, 3) landmarks in 192-canonical, (B,) presence)."""
+        x = self.prelu_stem(self.stem(_tf_same_pad(face_chw_normalized, 3, 2)))
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.prelu_head_reduce(self.head_reduce(x))
+        x = self.head_block(x)
+        B = x.shape[0]
+        presence = self.head_presence(x).reshape(B)
+        lmks = self.head_landmarks(x).reshape(B, self.NUM_LANDMARKS, 3)
+        return lmks, presence
+
+
+# FaceBlendshapes (MLP-Mixer "GhumMarkerPoserMlpMixerGeneral"):
+# 146x2 → token-reduce 146→96 → embed 2→64 → +cls token → 4x mixer → cls→52.
+_BS_NUM_INPUT_LANDMARKS = 146
+_BS_NUM_TOKENS_REDUCED = 96
+_BS_NUM_TOKENS = 97  # +1 cls
+_BS_TOKEN_DIM = 64
+_BS_TOKEN_MIX_HIDDEN = 384
+_BS_CHANNEL_MIX_HIDDEN = 256
+_BS_NUM_BLENDSHAPES = 52
+_BS_LN_EPS = 1e-6
+
+
+class MlpMixerBlock(nn.Module):
+    """MLP-Mixer block: token-mixing MLP (over tokens) → channel-mixing MLP (over dim).
+    Both pre-LN, both residual. LN has no beta (bias=False) to match MP."""
+
+    def __init__(self, num_tokens: int, token_dim: int, token_hidden: int, channel_hidden: int,
+                 device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        # bias=False → no LN beta (matches MP).
+        self.ln1 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)
+        self.ln2 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)
+        self.token_mlp1 = ops.Linear(num_tokens, token_hidden, bias=True, **kw)
+        self.token_mlp2 = ops.Linear(token_hidden, num_tokens, bias=True, **kw)
+        self.channel_mlp1 = ops.Linear(token_dim, channel_hidden, bias=True, **kw)
+        self.channel_mlp2 = ops.Linear(channel_hidden, token_dim, bias=True, **kw)
+
+    def forward(self, x: Tensor) -> Tensor:
+        y = self.ln1(x).transpose(1, 2)
+        x = x + self.token_mlp2(F.relu(self.token_mlp1(y))).transpose(1, 2)
+        return x + self.channel_mlp2(F.relu(self.channel_mlp1(self.ln2(x))))
+
+
+class FaceBlendshapes(nn.Module):
+    def __init__(self, device=None, dtype=None, operations=None):
+        super().__init__()
+        ops = operations if operations is not None else nn
+        kw = dict(device=device, dtype=dtype)
+        self.token_reduce = ops.Linear(_BS_NUM_INPUT_LANDMARKS, _BS_NUM_TOKENS_REDUCED, bias=True, **kw)
+        self.token_embed = ops.Linear(2, _BS_TOKEN_DIM, bias=True, **kw)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, _BS_TOKEN_DIM, **kw))
+        self.blocks = nn.ModuleList(
+            MlpMixerBlock(_BS_NUM_TOKENS, _BS_TOKEN_DIM, _BS_TOKEN_MIX_HIDDEN, _BS_CHANNEL_MIX_HIDDEN,
+                          device=device, dtype=dtype, operations=operations) for _ in range(4)
+        )
+        self.head = ops.Linear(_BS_TOKEN_DIM, _BS_NUM_BLENDSHAPES, bias=True, **kw)
+
+    @staticmethod
+    def _input_normalize(landmarks_2d: Tensor) -> Tensor:
+        # Centroid-subtract → L2 scale → x0.5. The 0.5 is baked into training.
+        centroid = landmarks_2d.mean(dim=1, keepdim=True)
+        x = landmarks_2d - centroid
+        mag = torch.sqrt((x * x).sum(dim=-1, keepdim=True))
+        scale = mag.mean(dim=1, keepdim=True)
+        return (x / scale.clamp(min=1e-12)) * 0.5
+
+    def forward(self, landmarks_2d: Tensor) -> Tensor:
+        """(B, 146, 2) → (B, 52) in [0, 1]. Input units don't matter (centroid + L2 normalize)."""
+        x = self._input_normalize(landmarks_2d)
+        x = self.token_reduce(x.transpose(1, 2)).transpose(1, 2)
+        x = self.token_embed(x)
+        cls = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat([cls, x], dim=1)
+        for blk in self.blocks:
+            x = blk(x)
+        return torch.sigmoid(self.head(x[:, 0]))
+
+
+@lru_cache(maxsize=1)
+def _blazeface_anchors() -> np.ndarray:
+    """896 anchors per SsdAnchorsCalculator (fixed_anchor_size → anchor_w=anchor_h=1)."""
+    per_ar = len(_BF_ASPECT_RATIOS) + (1 if _BF_INTERP_SCALE_AR > 0 else 0)
+    layer_anchors: List[np.ndarray] = []
+    layer = 0
+    while layer < _BF_NUM_LAYERS:
+        stride = _BF_STRIDES[layer]
+        last = layer
+        while last < _BF_NUM_LAYERS and _BF_STRIDES[last] == stride:
+            last += 1
+        per_cell = per_ar * (last - layer)
+        feat = (_BF_INPUT_SIZE + stride - 1) // stride
+        yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")
+        cx, cy, ones = (xx + _BF_ANCHOR_OFFSET_X) / feat, (yy + _BF_ANCHOR_OFFSET_Y) / feat, np.ones_like(xx)
+        cell = np.stack([cx, cy, ones, ones], axis=-1).reshape(-1, 4)
+        layer_anchors.append(np.repeat(cell, per_cell, axis=0))
+        layer = last
+    out = np.concatenate(layer_anchors, axis=0)
+    assert out.shape == (896, 4), out.shape
+    return out
+
+
+def _decode_blazeface(regressors: np.ndarray, classificators: np.ndarray,
+                      score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:
+    """Decode (regs (896,16), cls (896,1)) → (N, 17) = [xyxy, kp0x..kp5y, score] in [0, 1]."""
+    scores = expit(np.clip(classificators[:, 0], -_BF_SCORE_CLIP, _BF_SCORE_CLIP))
+    keep = scores >= score_thresh
+    if not keep.any():
+        return np.empty((0, 17), dtype=np.float32)
+    r = regressors[keep] / _BF_BOX_SCALE
+    a = _blazeface_anchors()[keep]  # (N, 4) cx, cy, 1, 1
+    cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
+    xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys
+    w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs
+    out = np.empty((r.shape[0], 17), dtype=np.float32)
+    out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
+    out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs
+    out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys
+    out[:, 16] = scores[keep]
+    return out
+
+
+def _weighted_nms(detections: np.ndarray, iou_thresh: float = 0.5) -> np.ndarray:
+    """MP weighted NMS — kept boxes are score-weighted averages of overlapping detections."""
+    if detections.shape[0] == 0:
+        return detections
+    dets = detections[np.argsort(-detections[:, 16])]
+    N = dets.shape[0]
+    areas = np.clip(dets[:, 2] - dets[:, 0], 0, None) * np.clip(dets[:, 3] - dets[:, 1], 0, None)
+    kept: List[np.ndarray] = []
+    used = np.zeros(N, dtype=bool)
+    for i in range(N):
+        if used[i]:
+            continue
+        ax1, ay1, ax2, ay2 = dets[i, 0:4]
+        merge_idx = [i]
+        for j in range(i + 1, N):
+            if used[j]:
+                continue
+            bx1, by1, bx2, by2 = dets[j, 0:4]
+            iw = max(0.0, min(ax2, bx2) - max(ax1, bx1))
+            ih = max(0.0, min(ay2, by2) - max(ay1, by1))
+            inter = iw * ih
+            union = areas[i] + areas[j] - inter
+            if union > 0 and inter / union > iou_thresh:  # strict > matches MP
+                merge_idx.append(j)
+                used[j] = True
+        used[i] = True
+        cluster = dets[merge_idx]
+        ws = cluster[:, 16:17]
+        ws_sum = ws.sum()
+        merged = np.copy(cluster[0])
+        if ws_sum > 0:
+            merged[:16] = (cluster[:, :16] * ws).sum(axis=0) / ws_sum
+        kept.append(merged)
+    return np.stack(kept, axis=0) if kept else np.empty((0, 17), dtype=np.float32)
+
+
+def _detection_to_face_rect(detection: np.ndarray, image_w: int, image_h: int) -> Tuple[float, float, float, float, float]:
+    """Detection (normalized) → rotated 1.5xbbox ROI in image pixels (anisotropic)."""
+    xmin, ymin, xmax, ymax = detection[0:4]
+    lx = detection[4 + _FACE_LEFT_EYE_KP * 2 + 0] * image_w
+    ly = detection[4 + _FACE_LEFT_EYE_KP * 2 + 1] * image_h
+    rx = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 0] * image_w
+    ry = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 1] * image_h
+    # Image-y-down convention: angle = target - atan2(-dy, dx).
+    angle = _FACE_ROI_TARGET_ANGLE - math.atan2(ly - ry, rx - lx)
+    return (float((xmin + xmax) * 0.5 * image_w),
+            float((ymin + ymax) * 0.5 * image_h),
+            float((xmax - xmin) * image_w * _FACE_ROI_SCALE_X),
+            float((ymax - ymin) * image_h * _FACE_ROI_SCALE_Y),
+            float(angle))
+
+
+def _sample_warp(image_chw: Tensor, src_x: Tensor, src_y: Tensor, padding_mode: str) -> Tensor:
+    """Bilinear-sample image_chw at corner-aligned (src_x, src_y)."""
+    H, W = int(image_chw.shape[-2]), int(image_chw.shape[-1])
+    grid = torch.stack([(2.0 * src_x + 1.0) / W - 1.0,
+                        (2.0 * src_y + 1.0) / H - 1.0], dim=-1).unsqueeze(0)
+    return F.grid_sample(image_chw.unsqueeze(0), grid, mode="bilinear",
+                         align_corners=False, padding_mode=padding_mode).squeeze(0)
+
+
+def _warp_face_crop(image_chw: Tensor, cx: float, cy: float, width: float, height: float,
+                    angle: float, output_size: int = _FM_INPUT_SIZE) -> Tensor:
+    """Rotated rect → output_size² with BORDER_REPLICATE. image_chw must be in [0, 1]."""
+    s_x, s_y = width / output_size, height / output_size
+    cos_a, sin_a = math.cos(angle), math.sin(angle)
+    arange = torch.arange(output_size, dtype=image_chw.dtype, device=image_chw.device) - output_size * 0.5
+    v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")
+    src_x = cx + u_grid * s_x * cos_a - v_grid * s_y * sin_a
+    src_y = cy + u_grid * s_x * sin_a + v_grid * s_y * cos_a
+    return _sample_warp(image_chw, src_x, src_y, "border")
+
+
+def _blazeface_input_warp(image_chw_raw: Tensor, target: int = _BF_INPUT_SIZE) -> Tuple[Tensor, float, float, float]:
+    """Centered max(W,H) square → target² with BORDER_ZERO + [-1, 1] norm.
+
+    Sub-pixel grid_sample matters; integer-pad-then-resize drifts the bbox ~5%.
+    Returns (warped, sub_rect_cx, sub_rect_cy, sub_rect_size) — the triplet maps
+    tensor-normalized [0,1] detections back to image pixels.
+    """
+    H, W = int(image_chw_raw.shape[1]), int(image_chw_raw.shape[2])
+    sub_rect_size = float(max(W, H))
+    sub_rect_cx, sub_rect_cy = W * 0.5, H * 0.5
+    s = sub_rect_size / target
+    arange = torch.arange(target, dtype=image_chw_raw.dtype, device=image_chw_raw.device) - target * 0.5
+    v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")
+    out = _sample_warp(image_chw_raw, sub_rect_cx + u_grid * s, sub_rect_cy + v_grid * s, "zeros")
+    return (out / 127.5) - 1.0, sub_rect_cx, sub_rect_cy, sub_rect_size
+
+
+class FaceLandmarker(nn.Module):
+    """BlazeFace → FaceMesh v2 → blendshapes. `detector_variant` selects 'short'
+    (128², ≤2m) or 'full' (192² FPN, ≤5m). State dict uses inner-module prefixes
+    `detector.*` / `mesh.*` / `blendshapes.*`; the outer FaceLandmarkerModel
+    wrapper rewrites `detector_{variant}.*` keys to `detector.*` before loading.
+    """
+
+    def __init__(self, device=None, dtype=None, operations=None, detector_variant: str = "short"):
+        super().__init__()
+        det_cls = {"short": BlazeFace, "full": BlazeFaceFullRange}.get(detector_variant)
+
+        self.detector_variant = detector_variant
+        self.detector = det_cls(device=device, dtype=dtype, operations=operations)
+        self.mesh = FaceMesh(device=device, dtype=dtype, operations=operations)
+        self.blendshapes = FaceBlendshapes(device=device, dtype=dtype, operations=operations)
+        self.register_buffer("_bs_idx", torch.tensor(_BS_INPUT_INDICES, dtype=torch.long), persistent=False)
+
+    def run_detector_batch(self, images_rgb_uint8: List[np.ndarray],
+                           score_thresh: float = _BF_MIN_SCORE,
+                           iou_thresh: float = 0.5):
+        """Batched detector pass. Returns (img_raws, sub_rects, sizes, per_frame_decoded)
+        where per_frame_decoded[b] is (N, 17) in tensor-normalized [0,1] coords."""
+        if not images_rgb_uint8:
+            return [], [], [], []
+        device, dtype = self.detector.stem.weight.device, self.detector.stem.weight.dtype
+        det_input_size, decode_fn = ((_BF_FR_INPUT_SIZE, _decode_blazeface_full_range)
+                                     if self.detector_variant == "full"
+                                     else (_BF_INPUT_SIZE, _decode_blazeface))
+
+        # Same-size frames: stack once and transfer once. Variable size falls back
+        # to per-image (only triggers for SAM3DBody's head crops).
+        sizes = [tuple(img.shape[:2]) for img in images_rgb_uint8]
+        if len(set(sizes)) == 1:
+            batch_chw = torch.from_numpy(np.stack(images_rgb_uint8, axis=0)).to(device, dtype).movedim(-1, -3).contiguous()
+            img_raws = [batch_chw[bi] for bi in range(batch_chw.shape[0])]
+        else:
+            img_raws = [torch.from_numpy(img).to(device, dtype).movedim(-1, -3).contiguous() for img in images_rgb_uint8]
+
+        warps = [_blazeface_input_warp(img_raw, det_input_size) for img_raw in img_raws]
+        det_crops = [w[0] for w in warps]
+        sub_rects = [(w[1], w[2], w[3]) for w in warps]
+
+        regs_b, cls_b = self.detector(torch.stack(det_crops, dim=0))
+        regs_np, cls_np = regs_b.float().cpu().numpy(), cls_b.float().cpu().numpy()
+        per_frame = []
+        for b in range(len(images_rgb_uint8)):
+            decoded = decode_fn(regs_np[b], cls_np[b], score_thresh=score_thresh)
+            per_frame.append(_weighted_nms(decoded, iou_thresh=iou_thresh) if decoded.shape[0] > 0 else decoded)
+        return img_raws, sub_rects, sizes, per_frame
+
+    def detect_batch(self, images_rgb_uint8: List[np.ndarray], num_faces: int = 1,
+                     score_thresh: float = _BF_MIN_SCORE) -> List[List[dict]]:
+        """Full pipeline batched across `images_rgb_uint8`. Returns one face-dict
+        list per image (empty if nothing detected). Face dict:
+            bbox_xyxy (4,) image pixels, blendshapes {52} ∈ [0,1],
+            landmarks_xy (478, 2) image pixels, landmarks_3d (478, 3) in
+            192-canonical (pre-transformation) units, presence float (raw logit).
+        """
+        img_raws, sub_rects, sizes, per_frame_dets = self.run_detector_batch(
+            images_rgb_uint8, score_thresh=score_thresh,
+        )
+        # tensor-normalized → image-normalized [0,1] for _detection_to_face_rect.
+        for b, decoded in enumerate(per_frame_dets):
+            if decoded.shape[0] == 0:
+                continue
+            cx, cy, size = sub_rects[b]
+            H, W = sizes[b]
+            sx0, sy0 = cx - size * 0.5, cy - size * 0.5
+            decoded[:, 0:16:2] = (sx0 + size * decoded[:, 0:16:2]) / W
+            decoded[:, 1:16:2] = (sy0 + size * decoded[:, 1:16:2]) / H
+            if num_faces > 0:
+                per_frame_dets[b] = decoded[: int(num_faces)]
+
+        # Collect every detected face across all frames into one mesh input.
+        face_params: List[Tuple[int, float, float, float, float, float, float]] = []
+        mesh_crops: List[Tensor] = []
+        for b, dets in enumerate(per_frame_dets):
+            if dets.shape[0] == 0:
+                continue
+            H, W = sizes[b]
+            img_for_mesh = img_raws[b] / 255.0
+            for det in dets:
+                cx, cy, w, h, angle = _detection_to_face_rect(det, W, H)
+                mesh_crops.append(_warp_face_crop(img_for_mesh, cx, cy, w, h, angle, _FM_INPUT_SIZE))
+                face_params.append((b, float(det[16]), cx, cy, w, h, angle))
+
+        results: List[List[dict]] = [[] for _ in range(len(images_rgb_uint8))]
+        if not mesh_crops:
+            return results
+
+        lmks_canon_b, presence_b = self.mesh(torch.stack(mesh_crops, dim=0))
+        bs_out_b = self.blendshapes(lmks_canon_b[:, self._bs_idx, :2])
+
+        # Batched canonical→image affine
+        params_t = torch.tensor(
+            [(cx, cy, w, h, math.cos(a), math.sin(a)) for (_b, _s, cx, cy, w, h, a) in face_params],
+            device=lmks_canon_b.device, dtype=lmks_canon_b.dtype,
+        )
+        cxs, cys, ws, hs, cos_a, sin_a = params_t.unbind(dim=1)
+        inv = 1.0 / _FM_INPUT_SIZE
+        u = lmks_canon_b[..., 0] - _FM_INPUT_SIZE * 0.5
+        v = lmks_canon_b[..., 1] - _FM_INPUT_SIZE * 0.5
+        lmks_xy_t = torch.stack([
+            cxs[:, None] + u * (ws * inv * cos_a)[:, None] - v * (hs * inv * sin_a)[:, None],
+            cys[:, None] + u * (ws * inv * sin_a)[:, None] + v * (hs * inv * cos_a)[:, None],
+        ], dim=-1)
+
+        lmks_xy_np = lmks_xy_t.float().cpu().numpy()
+        lmks_canon_np = lmks_canon_b.float().cpu().numpy()
+        presence_np = presence_b.float().cpu().numpy()
+        bs_np = bs_out_b.float().cpu().numpy()
+
+        for i, (b, score, *_) in enumerate(face_params):
+            lmks_xy = lmks_xy_np[i]
+            mn, mx = lmks_xy.min(0), lmks_xy.max(0)
+            results[b].append({
+                "bbox_xyxy": np.array([mn[0], mn[1], mx[0], mx[1]], dtype=np.float32),
+                "blendshapes": dict(zip(BLENDSHAPE_NAMES, bs_np[i].tolist())),
+                "landmarks_xy": lmks_xy,
+                "landmarks_3d": lmks_canon_np[i],
+                "presence": float(presence_np[i]),
+                "score": score,
+            })
+        return results
diff --git a/comfy_extras/nodes_mediapipe.py b/comfy_extras/nodes_mediapipe.py
new file mode 100644
index 000000000..2e67ae83f
--- /dev/null
+++ b/comfy_extras/nodes_mediapipe.py
@@ -0,0 +1,502 @@
+"""ComfyUI nodes for the pure-PyTorch MediaPipe Face Landmarker port.
+
+Custom IO types:
+  FACE_LANDMARKER  — FaceLandmarkerModel wrapper (ModelPatcher inside)
+  FACE_LANDMARKS   — {"frames": List[List[face_dict]], "image_size": (H, W),
+                      "connection_sets": dict[str, frozenset[(int, int)]]}
+                     face_dict: bbox_xyxy, blendshapes, landmarks_xy,
+                                landmarks_3d, presence, score, transformation_matrix
+
+MediaPipeFaceLandmarker also emits the core BOUNDING_BOX type — pair with DrawBBoxes.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from PIL import Image, ImageColor, ImageDraw
+from tqdm.auto import tqdm
+from typing_extensions import override
+
+import comfy.model_management
+import comfy.model_patcher
+import comfy.utils
+import folder_paths
+from comfy_api.latest import ComfyExtension, io
+
+from comfy_extras.mediapipe.face_landmarker import FaceLandmarker
+from comfy_extras.mediapipe.face_geometry import transformation_matrix_from_detection
+
+
+FaceLandmarkerType = io.Custom("FACE_LANDMARKER")
+FaceLandmarksType = io.Custom("FACE_LANDMARKS")
+
+_CANONICAL_KEYS = ("canonical_vertices", "procrustes_indices", "procrustes_weights")
+_CONTOUR_PARTS = ("face_oval", "left_eye", "right_eye", "left_eyebrow", "right_eyebrow", "lips")
+
+
+class FaceLandmarkerModel:
+    """Loaded FaceLandmarker variants + ModelPatcher per variant.
+
+    Safetensors layout: `detector_short.*` / `detector_full.*` plus shared
+    `mesh.*`, `blendshapes.*`, `canonical_*`, and `topology.*`.
+    PReLU forces plain-nn / fp32 (manual_cast strands buffers across devices).
+    """
+
+    def __init__(self, state_dict: dict):
+        self.load_device = comfy.model_management.text_encoder_device()
+        offload_device = comfy.model_management.text_encoder_offload_device()
+        self.dtype = torch.float32
+
+        # FACEMESH_* connection sets, embedded as int32 (N, 2) under topology.*.
+        base: dict[str, frozenset] = {}
+        for k in [k for k in state_dict if k.startswith("topology.")]:
+            base[k[len("topology."):]] = frozenset(map(tuple, state_dict.pop(k).tolist()))
+        base["contours"] = frozenset().union(*(base[p] for p in _CONTOUR_PARTS))
+        base["all"] = base["contours"] | base["irises"] | base["nose"]
+
+        self.connection_sets: dict[str, frozenset] = base
+        self.canonical_data: dict[str, np.ndarray] = {k: state_dict.pop(k).numpy() for k in _CANONICAL_KEYS}
+
+        shared = {k: v for k, v in state_dict.items() if k.startswith(("mesh.", "blendshapes."))}
+
+        self.models: dict[str, FaceLandmarker] = {}
+        self.patchers: dict[str, comfy.model_patcher.ModelPatcher] = {}
+        for variant in ("short", "full"):
+            prefix = f"detector_{variant}."
+            sub = dict(shared)
+            sub.update({f"detector.{k[len(prefix):]}": v for k, v in state_dict.items() if k.startswith(prefix)})
+            fl = FaceLandmarker(device=offload_device, dtype=self.dtype, operations=None, detector_variant=variant).eval()
+            fl.load_state_dict(sub, strict=False)
+
+            self.models[variant] = fl
+            self.patchers[variant] = comfy.model_patcher.CoreModelPatcher(
+                fl, load_device=self.load_device, offload_device=offload_device,
+                size=comfy.model_management.module_size(fl),
+            )
+
+    def detect_batch(self, images, num_faces: int, score_thresh: float, variant: str):
+        comfy.model_management.load_model_gpu(self.patchers[variant])
+        return self.models[variant].detect_batch(images, num_faces=num_faces, score_thresh=score_thresh)
+
+
+def _image_to_uint8(image: torch.Tensor) -> np.ndarray:
+    return image[..., :3].mul(255.0).add_(0.5).clamp_(0, 255).to(torch.uint8).cpu().numpy()
+
+
+def _parse_color(color: str) -> tuple[int, int, int]:
+    try:
+        return ImageColor.getrgb(color)[:3]
+    except ValueError:
+        return (0, 255, 0)
+
+
+def _copy_face(face: dict) -> dict:
+    """Shallow copy of a face_dict with array-fields cloned so callers can mutate."""
+    return {
+        "bbox_xyxy":    face["bbox_xyxy"].copy(),
+        "blendshapes":  dict(face["blendshapes"]),
+        "landmarks_xy": face["landmarks_xy"].copy(),
+        "landmarks_3d": face["landmarks_3d"].copy(),
+        "presence":     face["presence"],
+        "score":        face["score"],
+    }
+
+
+def _lerp_face(a: dict, b: dict, t: float) -> dict:
+    return {
+        "bbox_xyxy":    (1 - t) * a["bbox_xyxy"]    + t * b["bbox_xyxy"],
+        "blendshapes":  {k: (1 - t) * a["blendshapes"][k] + t * b["blendshapes"][k] for k in a["blendshapes"]},
+        "landmarks_xy": (1 - t) * a["landmarks_xy"] + t * b["landmarks_xy"],
+        "landmarks_3d": (1 - t) * a["landmarks_3d"] + t * b["landmarks_3d"],
+        "presence":     (1 - t) * a["presence"] + t * b["presence"],
+        "score":        (1 - t) * a["score"]    + t * b["score"],
+    }
+
+
+def _match_faces(a: list[dict], b: list[dict]) -> list[tuple[int, int]]:
+    """Greedy nearest-neighbour pairing of faces between two frames by bbox
+    centre distance. Unmatched (when counts differ) are dropped."""
+    if not a or not b:
+        return []
+    centers_a = np.array([(0.5 * (f["bbox_xyxy"][0] + f["bbox_xyxy"][2]),
+                           0.5 * (f["bbox_xyxy"][1] + f["bbox_xyxy"][3])) for f in a])
+    centers_b = np.array([(0.5 * (f["bbox_xyxy"][0] + f["bbox_xyxy"][2]),
+                           0.5 * (f["bbox_xyxy"][1] + f["bbox_xyxy"][3])) for f in b])
+    dists = np.linalg.norm(centers_a[:, None] - centers_b[None], axis=-1)
+    pairs: list[tuple[int, int]] = []
+    used_a: set[int] = set()
+    used_b: set[int] = set()
+    candidates = sorted((dists[ia, ib], ia, ib) for ia in range(len(a)) for ib in range(len(b)))
+    for _, ia, ib in candidates:
+        if ia in used_a or ib in used_b:
+            continue
+        pairs.append((ia, ib))
+        used_a.add(ia)
+        used_b.add(ib)
+    return pairs
+
+
+def _fill_missing_frames(frames: list[list[dict]], mode: str) -> None:
+    """In-place fill empty frame slots from neighbouring detections. Multi-face
+    aware: pairs faces across bracketing frames by greedy bbox-centre NN.
+    When counts differ, unmatched faces are dropped from the synthesised frame."""
+    if mode == "empty":
+        return
+    valid = [i for i, fr in enumerate(frames) if fr]
+    if not valid:
+        return  # nothing to fill from
+    if mode == "previous":
+        last: list[dict] = []
+        for i, fr in enumerate(frames):
+            if fr:
+                last = fr
+            elif last:
+                frames[i] = [_copy_face(f) for f in last]
+        return
+    # interpolate: lerp between bracketing valid frames; clamp at ends.
+    for i in range(len(frames)):
+        if frames[i]:
+            continue
+        prev_i = max((v for v in valid if v < i), default=None)
+        next_i = min((v for v in valid if v > i), default=None)
+        if prev_i is None:
+            frames[i] = [_copy_face(f) for f in frames[next_i]]
+        elif next_i is None:
+            frames[i] = [_copy_face(f) for f in frames[prev_i]]
+        else:
+            t = (i - prev_i) / (next_i - prev_i)
+            pairs = _match_faces(frames[prev_i], frames[next_i])
+            frames[i] = [_lerp_face(frames[prev_i][a], frames[next_i][b], t) for a, b in pairs]
+
+
+def _ordered_rings(edges: frozenset[tuple[int, int]]) -> list[list[int]]:
+    """Walk an unordered edge set into one or more closed-loop vertex rings
+    (handles multi-loop sets like FACEMESH_LIPS: outer + inner)."""
+    adj: dict[int, set[int]] = {}
+    for a, b in edges:
+        adj.setdefault(a, set()).add(b)
+        adj.setdefault(b, set()).add(a)
+    visited: set[int] = set()
+    rings: list[list[int]] = []
+    for start in adj:
+        if start in visited:
+            continue
+        ring = [start]
+        visited.add(start)
+        prev, cur = -1, start
+        while True:
+            nxt = next((v for v in adj[cur] if v != prev), None)
+            if nxt is None or nxt == start:
+                break
+            ring.append(nxt)
+            visited.add(nxt)
+            prev, cur = cur, nxt
+        rings.append(ring)
+    return rings
+
+
+class LoadMediaPipeFaceLandmarker(io.ComfyNode):
+    """Load MediaPipe Face Landmarker v2 weights. Contains both detector variants
+    (short / full), shared mesh, blendshapes, and canonical geometry."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoadMediaPipeFaceLandmarker",
+            display_name="Load MediaPipe Face Landmarker",
+            category="loaders",
+            inputs=[
+                io.Combo.Input("model_name", options=folder_paths.get_filename_list("mediapipe"),
+                               tooltip="Face Landmarker safetensors from models/mediapipe/."),
+            ],
+            outputs=[FaceLandmarkerType.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model_name) -> io.NodeOutput:
+        sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("mediapipe", model_name), safe_load=True)
+        wrapper = FaceLandmarkerModel(sd)
+        return io.NodeOutput(wrapper)
+
+
+# Per-frame fallback modes for detection failures in a batch.
+_FALLBACK_MODES = ("empty", "previous", "interpolate")
+
+
+class MediaPipeFaceLandmarker(io.ComfyNode):
+    """BlazeFace → FaceMesh v2 → ARKit-52 blendshapes, batched across the
+    input. Also emits a BOUNDING_BOX list (landmark-extent bbox per face) —
+    pair with DrawBBoxes for detector-only viz or MediaPipeFaceMeshVisualize
+    for the mesh overlay."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MediaPipeFaceLandmarker",
+            display_name="MediaPipe Face Landmarker",
+            category="image/detection",
+            inputs=[
+                FaceLandmarkerType.Input("face_landmarker"),
+                io.Image.Input("image"),
+                io.Combo.Input("detector_variant", options=["short", "full", "both"], default="short",
+                               tooltip="Face detector range. 'short' is tuned for close-up faces "
+                                       "(within ~2 m of the camera); 'full' covers farther / smaller "
+                                       "faces (up to ~5 m) but is slower. 'both' runs both detectors and "
+                                       "keeps whichever found more faces per frame (~2× detection cost)."),
+                io.Int.Input("num_faces", default=1, min=0, max=16, step=1,
+                             tooltip="Maximum faces to return per frame. 0 = no cap (return all detected)."),
+                io.Float.Input("min_confidence", default=0.5, min=0.0, max=1.0, step=0.01, advanced=True,
+                               tooltip="BlazeFace score threshold. Lower to catch small/occluded faces."),
+                io.Combo.Input("missing_frame_fallback", options=list(_FALLBACK_MODES), default="empty", advanced=True,
+                               tooltip="Per-frame behaviour when detection fails in a batch. "
+                                       "'empty' leaves the frame faceless. 'previous' copies the most recent successful "
+                                       "detection. 'interpolate' lerps landmarks/bbox/blendshapes between bracketing "
+                                       "successful frames. Multi-face: pairs faces across frames by greedy bbox-centre NN."),
+            ],
+            outputs=[
+                FaceLandmarksType.Output(display_name="face_landmarks"),
+                io.BoundingBox.Output("bboxes"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, face_landmarker, image, detector_variant, num_faces, min_confidence,
+                missing_frame_fallback) -> io.NodeOutput:
+        canonical = face_landmarker.canonical_data
+        img_np = _image_to_uint8(image)
+        B, H, W = img_np.shape[:3]
+        chunk = 16
+        is_both = detector_variant == "both"
+        total_work = 2 * B if is_both else B
+        pbar = comfy.utils.ProgressBar(total_work)
+
+        def _run(variant: str) -> list[list[dict]]:
+            res: list[list[dict]] = []
+            with tqdm(total=B, desc=f"MediaPipe Face Landmarker ({variant})") as tq:
+                for i in range(0, B, chunk):
+                    end = min(i + chunk, B)
+                    res.extend(face_landmarker.detect_batch(
+                        [img_np[bi] for bi in range(i, end)],
+                        num_faces=int(num_faces),
+                        score_thresh=float(min_confidence),
+                        variant=variant,
+                    ))
+                    pbar.update_absolute(min(pbar.current + (end - i), total_work))
+                    tq.update(end - i)
+            return res
+
+        if is_both:
+            short_res = _run("short")
+            full_res = _run("full")
+            # Per-frame keep whichever found more faces (tie → short).
+            frames: list[list[dict]] = [
+                short_res[bi] if len(short_res[bi]) >= len(full_res[bi]) else full_res[bi]
+                for bi in range(B)
+            ]
+        else:
+            frames = _run(detector_variant)
+        _fill_missing_frames(frames, missing_frame_fallback)
+        bboxes = []
+        for per_frame in frames:
+            per_bb = []
+            for f in per_frame:
+                f["transformation_matrix"] = transformation_matrix_from_detection(f, W, H, canonical)
+                x1, y1, x2, y2 = (float(v) for v in f["bbox_xyxy"])
+                per_bb.append({"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1, "label": "face", "score": float(f["score"])})
+            bboxes.append(per_bb)
+        return io.NodeOutput({"frames": frames, "image_size": (H, W),
+                              "connection_sets": face_landmarker.connection_sets}, bboxes)
+
+
+# Topology keys unioned by the 'all' connections preset (contour parts + irises + nose).
+_ALL_CONNECTION_PARTS: tuple[str, ...] = (*_CONTOUR_PARTS, "irises", "nose")
+_CUSTOM_FEATURES: tuple[tuple[str, bool], ...] = (
+    ("face_oval",     True),
+    ("lips",          True),
+    ("left_eye",      True),
+    ("right_eye",     True),
+    ("left_eyebrow",  True),
+    ("right_eyebrow", True),
+    ("irises",        True),
+    ("nose",          True),
+    ("tesselation",   False),
+)
+
+
+class MediaPipeFaceMeshVisualize(io.ComfyNode):
+    """Draw a FACEMESH_* subset over an image. Topology travels with the
+    FACE_LANDMARKS payload (set at detection time)."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MediaPipeFaceMeshVisualize",
+            display_name="MediaPipe Face Mesh Visualize",
+            category="image/detection",
+            inputs=[
+                FaceLandmarksType.Input("face_landmarks"),
+                io.Image.Input("image", optional=True, tooltip="If not connected, a black canvas will be used."),
+                io.DynamicCombo.Input(
+                    "connections",
+                    tooltip="'all' = oval+eyes+brows+lips+irises+nose. 'fill' = solid face_oval polygon (silhouette mask). 'custom' = toggle each feature individually (including 'tesselation', the full 2547-edge wireframe).",
+                    options=[
+                        io.DynamicCombo.Option("all", []),
+                        io.DynamicCombo.Option("fill", []),
+                        io.DynamicCombo.Option("custom", [
+                            io.Boolean.Input(feat, default=default,
+                                             tooltip=f"Draw the '{feat}' connection set.")
+                            for feat, default in _CUSTOM_FEATURES
+                        ]),
+                    ],
+                ),
+                io.Color.Input("color", default="#00ff00"),
+                io.Int.Input("thickness", default=1, min=0, max=8, step=1,
+                             tooltip="Edge line thickness in pixels. 0 disables edge drawing."),
+                io.Int.Input("point_size", default=2, min=0, max=16, step=1,
+                             tooltip="Landmark dot radius in pixels. 0 disables point drawing."),
+            ],
+            outputs=[io.Image.Output()],
+        )
+
+    @classmethod
+    def execute(cls, face_landmarks, connections, color, thickness, point_size, image=None) -> io.NodeOutput:
+        sets = face_landmarks["connection_sets"]
+        sel = connections["connections"]
+        fill_rings: list[list[int]] | None = None
+        if sel == "fill":
+            fill_rings = _ordered_rings(sets["face_oval"])
+            edges = frozenset()
+        elif sel == "custom":
+            parts = [feat for feat, _ in _CUSTOM_FEATURES if connections.get(feat, False)]
+            edges = frozenset().union(*(sets[p] for p in parts))
+        else:  # "all"
+            edges = frozenset().union(*(sets[p] for p in _ALL_CONNECTION_PARTS))
+        rgb, thick, psize = _parse_color(color), int(thickness), int(point_size)
+        frames = face_landmarks["frames"]
+        if image is None:
+            H, W = face_landmarks["image_size"]
+            img_np = np.zeros((len(frames), H, W, 3), dtype=np.uint8)
+        else:
+            img_np = _image_to_uint8(image)
+        B = img_np.shape[0]
+        n_frames = len(frames)
+        pbar = comfy.utils.ProgressBar(B)
+        out = np.empty_like(img_np)
+        for bi in range(B):
+            faces = frames[bi] if bi < n_frames else []
+            out[bi] = _draw_mesh(img_np[bi], faces, edges, rgb, thick, psize, fill_rings)
+            pbar.update_absolute(bi + 1)
+        return io.NodeOutput(torch.from_numpy(out).to(
+            device=comfy.model_management.intermediate_device(),
+            dtype=comfy.model_management.intermediate_dtype(),
+        ).div_(255.0))
+
+
+def _draw_mesh(image_rgb: np.ndarray, faces: list, edges,
+               rgb: tuple[int, int, int], thickness: int,
+               point_size: int, fill_rings: list[list[int]] | None = None) -> np.ndarray:
+    draw_edges = thickness > 0 and edges
+    if not faces or (fill_rings is None and not draw_edges and point_size <= 0):
+        return image_rgb.copy()
+    pil = Image.fromarray(image_rgb)
+    draw = ImageDraw.Draw(pil)
+    r = point_size * 0.5
+    if fill_rings is not None:
+        for f in faces:
+            lmks = f["landmarks_xy"]
+            for ring in fill_rings:
+                draw.polygon([(float(lmks[i, 0]), float(lmks[i, 1])) for i in ring], fill=rgb)
+        return np.asarray(pil)
+    for f in faces:
+        lmks = f["landmarks_xy"]
+        n = lmks.shape[0]
+        if draw_edges:
+            for a, b in edges:
+                if a < n and b < n:
+                    draw.line([(float(lmks[a, 0]), float(lmks[a, 1])),
+                               (float(lmks[b, 0]), float(lmks[b, 1]))], fill=rgb, width=thickness)
+        if point_size == 1:
+            draw.point(lmks.flatten().tolist(), fill=rgb)
+        elif point_size > 1:
+            for x, y in lmks:
+                draw.ellipse((float(x) - r, float(y) - r, float(x) + r, float(y) + r), fill=rgb)
+    return np.asarray(pil)
+
+
+# Mask region presets — closed-loop topologies only.
+_MASK_REGIONS: tuple[str, ...] = ("face_oval", "lips", "left_eye", "right_eye", "irises")
+_MASK_CUSTOM_FEATURES: tuple[tuple[str, bool], ...] = (
+    ("face_oval",  True),
+    ("lips",       False),
+    ("left_eye",   False),
+    ("right_eye",  False),
+    ("irises",     False),
+)
+
+
+class MediaPipeFaceMask(io.ComfyNode):
+    """Binary mask from face landmarks, filled polygon per face. One mask per
+    frame in the batch; faces in the same frame composite (union)."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MediaPipeFaceMask",
+            display_name="MediaPipe Face Mask",
+            category="image/detection",
+            inputs=[
+                FaceLandmarksType.Input("face_landmarks"),
+                io.DynamicCombo.Input(
+                    "regions",
+                    tooltip="'all' = union of face_oval+lips+eyes+irises (which collapses to face_oval since it encloses the rest). 'custom' = toggle each region individually for combos like lips+eyes.",
+                    options=[
+                        io.DynamicCombo.Option("all", []),
+                        io.DynamicCombo.Option("custom", [
+                            io.Boolean.Input(reg, default=default,
+                                             tooltip=f"Include the '{reg}' region in the mask.")
+                            for reg, default in _MASK_CUSTOM_FEATURES
+                        ]),
+                    ],
+                ),
+            ],
+            outputs=[io.Mask.Output()],
+        )
+
+    @classmethod
+    def execute(cls, face_landmarks, regions) -> io.NodeOutput:
+        sets = face_landmarks["connection_sets"]
+        sel = regions["regions"]
+        if sel == "custom":
+            picked = [reg for reg, _ in _MASK_CUSTOM_FEATURES if regions.get(reg, False)]
+        else:
+            picked = list(_MASK_REGIONS)
+        rings = [r for reg in picked for r in _ordered_rings(sets[reg])]
+        frames = face_landmarks["frames"]
+        H, W = face_landmarks["image_size"]
+        masks = np.zeros((len(frames), H, W), dtype=np.uint8)
+        pbar = comfy.utils.ProgressBar(len(frames))
+        for bi, per_frame in enumerate(frames):
+            if per_frame:
+                pil = Image.new("L", (W, H), 0)
+                draw = ImageDraw.Draw(pil)
+                for f in per_frame:
+                    lmks = f["landmarks_xy"]
+                    for ring in rings:
+                        draw.polygon([(float(lmks[i, 0]), float(lmks[i, 1])) for i in ring], fill=255)
+                masks[bi] = np.asarray(pil)
+            pbar.update_absolute(bi + 1)
+        return io.NodeOutput(torch.from_numpy(masks).to(
+            device=comfy.model_management.intermediate_device(),
+            dtype=comfy.model_management.intermediate_dtype(),
+        ).div_(255.0))
+
+
+class MediaPipeFaceExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [LoadMediaPipeFaceLandmarker, MediaPipeFaceLandmarker, MediaPipeFaceMeshVisualize, MediaPipeFaceMask]
+
+
+async def comfy_entrypoint() -> MediaPipeFaceExtension:
+    return MediaPipeFaceExtension()
diff --git a/folder_paths.py b/folder_paths.py
index ad7f0f4fc..ce152eb37 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -60,6 +60,8 @@ folder_names_and_paths["geometry_estimation"] = ([os.path.join(models_dir, "geom
 
 folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flow")], supported_pt_extensions)
 
+folder_names_and_paths["mediapipe"] = ([os.path.join(models_dir, "mediapipe")], supported_pt_extensions)
+
 output_directory = os.path.join(base_path, "output")
 temp_directory = os.path.join(base_path, "temp")
 input_directory = os.path.join(base_path, "input")
diff --git a/models/mediapipe/put_mediapipe_models_here b/models/mediapipe/put_mediapipe_models_here
new file mode 100644
index 000000000..e69de29bb
diff --git a/nodes.py b/nodes.py
index fdd6eeb5f..13e46ac8a 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2444,6 +2444,7 @@ async def init_builtin_extra_nodes():
         "nodes_hidream_o1.py",
         "nodes_save_3d.py",
         "nodes_moge.py",
+        "nodes_mediapipe.py",
     ]
 
     import_failed = []

From 5aa5ccc9e02aec94cf43e0f71d4b2f62b204b5b6 Mon Sep 17 00:00:00 2001
From: rattus <46076784+rattus128@users.noreply.github.com>
Date: Thu, 21 May 2026 10:03:58 +1000
Subject: [PATCH 06/14] Multi-threaded load of models from disk (big load time
 speedups & Offload to disk) (CORE-43,CORE-152,CORE-164,CORE-165,CORE-117)
 (#13802)

* model_management: disable non-dynamic smart memory

Disable smart memory outright for non dynamic models.

This is a minor step towards deprecation of --disable-dynamic-vram
and the legacy ModelPatcher.

This is needed for estimate-free model development, where new models
can opt-out of supplying a memory estimate and not have to worry
about hard VRAM allocations due to legacy non-dynamic model patchers

This is also a general stability increase for a lot of stray use cases
where estimates may still be off and going forward we are not going
to accurately maintain such estimates.

* pinned_memory: implement with aimdo growable buffer

Use a single growable buffer so we can do threaded pre-warming on
pinned memory.

* mm: use aimdo to do transfer from disk to pin

Aimdo implements a faster threaded loader.

* Add stream host pin buffer for AIMDO casts

Introduce per-offload-stream HostBuffer reuse for pinned staging,
include it in cast buffer reset synchronization.

Defer actual casts that go via this pin path to a separate pass
such that the buffer can be allocated monolithically (to avoid
cudaHostRegister thrash).

* remove old pin path

* Implement JIT pinned memory pressure

Replace the predictive pin pressure mechanism with JIT PIN memory
pressure.

* LowVRAMPatch: change to two-phase visit

* lora: re-implement as inplace swiss-army-knife operation

* prepare for multiple pin sets

* implement pinned loras

* requirements: comfy-aimdo 0.4.0

* ops: remove unused arg

This was defeatured in aimdo iteration

* ops: sync the CPU with only the offload stream activity

This was syncing with the offload stream which itself is synced with the
compute stream, so this was syncing CPU with compute transitively. Define
the event to sync it more gently.

* pins: implement freeing intermediate for pinned memory

Pinning is more important than inactive intermediates and the stream
pin buffer is more important than even active intermediates.

* execution: implement pin eviction on RAM presure

Add back proper pin freeing on RAM pressure

* implement pin registration swaps

Uncap the windows pins from 50% by extending the pool and have a pressure
mechanism to move the pin reservations om demand.

This unfortunately implies a GPU sync to do the freeing so significant
hysterisis needs to be added to consolidate these pressure events.

* cli_args/execution: Implement lower background cache-ram threshold

Limit the amount of RAM background intermediates can use, so that
switching workflows doesn't degrade performance too much.

* make default

* bump aimdo

* model-patcher: force-cast tiny weights

Flux 2 gets crazy stalls due to a mix of tiny and giant weights
creating lopsided steam buffer rotations which creates stalls.

* ops: refactor in prep for chunking

* mm: delegate pin-on-the-way to aimdo

Aimdo is able to chunk and slice this on the way for better CPU->GPU
overlap. The main advantage is the ability to shorten the bus contention
window between previous weight transfer and the next weights vbar
fault.

* bump aimdo

* pinning updates

* specify hostbuf max allocation size

There a signs of virtual memory exhaustion on some linux systems when
throwing 128GB for every little piece. Pass the actual to save aimdo
from over-estimates

* tests: update execution tests for caching

The default caching changed to ram-cache so update these tests
accordingly.

Remove the LRU 0 test as this also falls through to RAM cache.
---
 comfy/cli_args.py                   |   7 +-
 comfy/lora.py                       |  19 ++-
 comfy/memory_management.py          |  24 +++-
 comfy/model_management.py           | 189 +++++++++++++++++-----------
 comfy/model_patcher.py              | 138 +++++++++++++++-----
 comfy/ops.py                        |  88 +++++++++++--
 comfy/pinned_memory.py              |  68 ++++++----
 comfy/utils.py                      |   2 -
 comfy/windows.py                    |  52 --------
 execution.py                        |  12 +-
 main.py                             |  20 +--
 requirements.txt                    |   2 +-
 tests/execution/test_async_nodes.py |   3 +-
 tests/execution/test_execution.py   |   3 +-
 14 files changed, 408 insertions(+), 219 deletions(-)
 delete mode 100644 comfy/windows.py

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 76faed3ad..9d88c8517 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -110,13 +110,11 @@ parser.add_argument("--preview-method", type=LatentPreviewMethod, default=Latent
 
 parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
 
-CACHE_RAM_AUTO_GB = -1.0
-
 cache_group = parser.add_mutually_exclusive_group()
+cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 25%% of system RAM (min 4GB, max 32GB), inactive 75%% of system RAM (min 12GB, max 96GB).")
 cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
 cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
 cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
-cache_group.add_argument("--cache-ram", nargs='?', const=CACHE_RAM_AUTO_GB, type=float, default=0, help="Use RAM pressure caching with the specified headroom threshold. If available RAM drops below the threshold the cache removes large items to free RAM. Default (when no value is provided): 25%% of system RAM (min 4GB, max 32GB).")
 
 attn_group = parser.add_mutually_exclusive_group()
 attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
@@ -245,6 +243,9 @@ if comfy.options.args_parsing:
 else:
     args = parser.parse_args([])
 
+if args.cache_ram is not None and len(args.cache_ram) > 2:
+    parser.error("--cache-ram accepts at most two values: active GB and inactive GB")
+
 if args.windows_standalone_build:
     args.auto_launch = True
 
diff --git a/comfy/lora.py b/comfy/lora.py
index f11e26ec9..c0e8b865c 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -484,16 +484,23 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori
 
     return weight
 
-def prefetch_prepared_value(value, allocate_buffer, stream):
+def prefetch_prepared_value(value, counter, destination, stream, copy):
     if isinstance(value, torch.Tensor):
-        dest = allocate_buffer(comfy.memory_management.vram_aligned_size(value))
-        comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
+        size = comfy.memory_management.vram_aligned_size(value)
+        offset = counter[0]
+        counter[0] += size
+        if destination is None:
+            return value
+
+        dest = destination[offset:offset + size]
+        if copy:
+            comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
         return comfy.memory_management.interpret_gathered_like([value], dest)[0]
     elif isinstance(value, weight_adapter.WeightAdapterBase):
-        return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, allocate_buffer, stream))
+        return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, counter, destination, stream, copy))
     elif isinstance(value, tuple):
-        return tuple(prefetch_prepared_value(item, allocate_buffer, stream) for item in value)
+        return tuple(prefetch_prepared_value(item, counter, destination, stream, copy) for item in value)
     elif isinstance(value, list):
-        return [prefetch_prepared_value(item, allocate_buffer, stream) for item in value]
+        return [prefetch_prepared_value(item, counter, destination, stream, copy) for item in value]
 
     return value
diff --git a/comfy/memory_management.py b/comfy/memory_management.py
index 48e3c11da..c43f0c4a2 100644
--- a/comfy/memory_management.py
+++ b/comfy/memory_management.py
@@ -15,7 +15,7 @@ class TensorFileSlice(NamedTuple):
     size: int
 
 
-def read_tensor_file_slice_into(tensor, destination):
+def read_tensor_file_slice_into(tensor, destination, stream=None, destination2=None):
 
     if isinstance(tensor, QuantizedTensor):
         if not isinstance(destination, QuantizedTensor):
@@ -23,12 +23,17 @@ def read_tensor_file_slice_into(tensor, destination):
         if tensor._layout_cls != destination._layout_cls:
             return False
 
-        if not read_tensor_file_slice_into(tensor._qdata, destination._qdata):
+        if not read_tensor_file_slice_into(tensor._qdata, destination._qdata, stream=stream,
+                                           destination2=(destination2._qdata if destination2 is not None else None)):
             return False
 
         dst_orig_dtype = destination._params.orig_dtype
         destination._params.copy_from(tensor._params, non_blocking=False)
         destination._params = dataclasses.replace(destination._params, orig_dtype=dst_orig_dtype)
+        if destination2 is not None:
+            dst_orig_dtype = destination2._params.orig_dtype
+            destination2._params.copy_from(destination._params, non_blocking=True)
+            destination2._params = dataclasses.replace(destination2._params, orig_dtype=dst_orig_dtype)
         return True
 
     info = getattr(tensor.untyped_storage(), "_comfy_tensor_file_slice", None)
@@ -48,6 +53,17 @@ def read_tensor_file_slice_into(tensor, destination):
     if info.size == 0:
         return True
 
+    hostbuf = getattr(destination.untyped_storage(), "_comfy_hostbuf", None)
+    if hostbuf is not None:
+        stream_ptr = getattr(stream, "cuda_stream", 0) if stream is not None else 0
+        device_ptr = destination2.data_ptr() if destination2 is not None else 0
+        hostbuf.read_file_slice(file_obj, info.offset, info.size,
+                                offset=destination.data_ptr() - hostbuf.get_raw_address(),
+                                stream=stream_ptr,
+                                device_ptr=device_ptr,
+                                device=None if destination2 is None else destination2.device.index)
+        return True
+
     buf_type = ctypes.c_ubyte * info.size
     view = memoryview(buf_type.from_address(destination.data_ptr()))
 
@@ -151,7 +167,7 @@ def set_ram_cache_release_state(callback, headroom):
     extra_ram_release_callback = callback
     RAM_CACHE_HEADROOM = max(0, int(headroom))
 
-def extra_ram_release(target):
+def extra_ram_release(target, free_active=False):
     if extra_ram_release_callback is None:
         return 0
-    return extra_ram_release_callback(target)
+    return extra_ram_release_callback(target, free_active=free_active)
diff --git a/comfy/model_management.py b/comfy/model_management.py
index 21738a4c7..3894dfa9c 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -31,6 +31,7 @@ from contextlib import nullcontext
 import comfy.memory_management
 import comfy.utils
 import comfy.quant_ops
+import comfy_aimdo.host_buffer
 import comfy_aimdo.vram_buffer
 
 class VRAMState(Enum):
@@ -495,6 +496,14 @@ except:
 
 current_loaded_models = []
 
+DIRTY_MMAPS = set()
+
+PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
+
+#Freeing registerables on pressure does imply a GPU sync, so go big on
+#the hysteresis so each expensive sync gives us back a good chunk.
+REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
+
 def module_size(module):
     module_mem = 0
     sd = module.state_dict()
@@ -503,27 +512,46 @@ def module_size(module):
         module_mem += t.nbytes
     return module_mem
 
-def module_mmap_residency(module, free=False):
-    mmap_touched_mem = 0
-    module_mem = 0
-    bounced_mmaps = set()
-    sd = module.state_dict()
-    for k in sd:
-        t = sd[k]
-        module_mem += t.nbytes
-        storage = t._qdata.untyped_storage() if isinstance(t, comfy.quant_ops.QuantizedTensor) else t.untyped_storage()
-        if not getattr(storage, "_comfy_tensor_mmap_touched", False):
-            continue
-        mmap_touched_mem += t.nbytes
-        if not free:
-            continue
-        storage._comfy_tensor_mmap_touched = False
-        mmap_obj = storage._comfy_tensor_mmap_refs[0]
-        if mmap_obj in bounced_mmaps:
-            continue
-        mmap_obj.bounce()
-        bounced_mmaps.add(mmap_obj)
-    return mmap_touched_mem, module_mem
+def mark_mmap_dirty(storage):
+    mmap_refs = getattr(storage, "_comfy_tensor_mmap_refs", None)
+    if mmap_refs is not None:
+        DIRTY_MMAPS.add(mmap_refs[0])
+
+def free_pins(size, evict_active=False):
+    freed_total = 0
+    for loaded_model in reversed(current_loaded_models):
+        if size <= 0:
+            return freed_total
+        model = loaded_model.model
+        if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
+            freed = model.partially_unload_ram(size)
+            freed_total += freed
+            size -= freed
+    return freed_total
+
+def ensure_pin_budget(size, evict_active=False):
+    shortfall = size + comfy.memory_management.RAM_CACHE_HEADROOM / 2 - psutil.virtual_memory().available
+    if shortfall <= 0:
+        return True
+
+    to_free = shortfall + PIN_PRESSURE_HYSTERESIS
+    return free_pins(to_free, evict_active=evict_active) >= shortfall
+
+def ensure_pin_registerable(size, evict_active=False):
+    shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
+    if MAX_PINNED_MEMORY <= 0:
+        return False
+    if shortfall <= 0:
+        return True
+
+    shortfall += REGISTERABLE_PIN_HYSTERESIS
+    for loaded_model in reversed(current_loaded_models):
+        model = loaded_model.model
+        if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
+            shortfall -= model.unregister_inactive_pins(shortfall)
+            if shortfall <= 0:
+                return True
+    return shortfall <= REGISTERABLE_PIN_HYSTERESIS
 
 class LoadedModel:
     def __init__(self, model):
@@ -553,9 +581,6 @@ class LoadedModel:
     def model_memory(self):
         return self.model.model_size()
 
-    def model_mmap_residency(self, free=False):
-        return self.model.model_mmap_residency(free=free)
-
     def model_loaded_memory(self):
         return self.model.loaded_size()
 
@@ -635,15 +660,9 @@ WINDOWS = any(platform.win32_ver())
 
 EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
 if WINDOWS:
-    import comfy.windows
     EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue
     if total_vram > (15 * 1024):  # more extra reserved vram on 16GB+ cards
         EXTRA_RESERVED_VRAM += 100 * 1024 * 1024
-    def get_free_ram():
-        return comfy.windows.get_free_ram()
-else:
-    def get_free_ram():
-        return psutil.virtual_memory().available
 
 if args.reserve_vram is not None:
     EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024
@@ -657,7 +676,6 @@ def minimum_inference_memory():
 
 def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins_required=0, ram_required=0):
     cleanup_models_gc()
-    comfy.memory_management.extra_ram_release(max(pins_required, ram_required))
     unloaded_model = []
     can_unload = []
     unloaded_models = []
@@ -673,11 +691,9 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
     for x in can_unload_sorted:
         i = x[-1]
         memory_to_free = 1e32
-        pins_to_free = 1e32
-        if not DISABLE_SMART_MEMORY or device is None:
+        if current_loaded_models[i].model.is_dynamic() and (not DISABLE_SMART_MEMORY or device is None):
             memory_to_free = 0 if device is None else memory_required - get_free_memory(device)
-            pins_to_free = pins_required - get_free_ram()
-            if current_loaded_models[i].model.is_dynamic() and for_dynamic:
+            if for_dynamic:
                 #don't actually unload dynamic models for the sake of other dynamic models
                 #as that works on-demand.
                 memory_required -= current_loaded_models[i].model.loaded_size()
@@ -685,18 +701,6 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
         if memory_to_free > 0 and current_loaded_models[i].model_unload(memory_to_free):
             logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
             unloaded_model.append(i)
-        if pins_to_free > 0:
-            logging.debug(f"PIN Unloading {current_loaded_models[i].model.model.__class__.__name__}")
-            current_loaded_models[i].model.partially_unload_ram(pins_to_free)
-
-    for x in can_unload_sorted:
-        i = x[-1]
-        ram_to_free = ram_required - psutil.virtual_memory().available
-        if ram_to_free <= 0 and i not in unloaded_model:
-            continue
-        resident_memory, _ = current_loaded_models[i].model_mmap_residency(free=True)
-        if resident_memory > 0:
-            logging.debug(f"RAM Unloading {current_loaded_models[i].model.model.__class__.__name__}")
 
     for i in sorted(unloaded_model, reverse=True):
         unloaded_models.append(current_loaded_models.pop(i))
@@ -762,29 +766,16 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
             model_to_unload.model.detach(unpatch_all=False)
             model_to_unload.model_finalizer.detach()
 
-
     total_memory_required = {}
-    total_pins_required = {}
-    total_ram_required = {}
     for loaded_model in models_to_load:
         device = loaded_model.device
         total_memory_required[device] = total_memory_required.get(device, 0) + loaded_model.model_memory_required(device)
-        resident_memory, model_memory = loaded_model.model_mmap_residency()
-        pinned_memory = loaded_model.model.pinned_memory_size()
-        #FIXME: This can over-free the pins as it budgets to pin the entire model. We should
-        #make this JIT to keep as much pinned as possible.
-        pins_required = model_memory - pinned_memory
-        ram_required = model_memory - resident_memory
-        total_pins_required[device] = total_pins_required.get(device, 0) + pins_required
-        total_ram_required[device] = total_ram_required.get(device, 0) + ram_required
 
     for device in total_memory_required:
         if device != torch.device("cpu"):
             free_memory(total_memory_required[device] * 1.1 + extra_mem,
                         device,
-                        for_dynamic=free_for_dynamic,
-                        pins_required=total_pins_required[device],
-                        ram_required=total_ram_required[device])
+                        for_dynamic=free_for_dynamic)
 
     for device in total_memory_required:
         if device != torch.device("cpu"):
@@ -1180,6 +1171,7 @@ STREAM_CAST_BUFFERS = {}
 LARGEST_CASTED_WEIGHT = (None, 0)
 STREAM_AIMDO_CAST_BUFFERS = {}
 LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
+STREAM_PIN_BUFFERS = {}
 
 DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
 
@@ -1220,21 +1212,66 @@ def get_aimdo_cast_buffer(offload_stream, device):
     if cast_buffer is None:
         cast_buffer = comfy_aimdo.vram_buffer.VRAMBuffer(DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE, device.index)
         STREAM_AIMDO_CAST_BUFFERS[offload_stream] = cast_buffer
-
     return cast_buffer
+
+def get_pin_buffer(offload_stream):
+    pin_buffer = STREAM_PIN_BUFFERS.get(offload_stream, None)
+    if pin_buffer is None:
+        pin_buffer = comfy_aimdo.host_buffer.HostBuffer(0, 0, pinned_hostbuf_size(8 * 1024**3))
+        STREAM_PIN_BUFFERS[offload_stream] = pin_buffer
+    elif offload_stream is not None:
+        event = getattr(pin_buffer, "_comfy_event", None)
+        if event is not None:
+            event.synchronize()
+            delattr(pin_buffer, "_comfy_event")
+    return pin_buffer
+
+def resize_pin_buffer(pin_buffer, size):
+    global TOTAL_PINNED_MEMORY
+    old_size = pin_buffer.size
+    if size <= old_size:
+        return True
+    growth = size - old_size
+    comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
+    ensure_pin_budget(growth, evict_active=True)
+    ensure_pin_registerable(growth, evict_active=True)
+    try:
+        pin_buffer.extend(size=size, reallocate=True)
+    except RuntimeError:
+        return False
+    TOTAL_PINNED_MEMORY += pin_buffer.size - old_size
+    return True
+
 def reset_cast_buffers():
+    global TOTAL_PINNED_MEMORY
     global LARGEST_CASTED_WEIGHT
     global LARGEST_AIMDO_CASTED_WEIGHT
 
     LARGEST_CASTED_WEIGHT = (None, 0)
     LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
-    for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS):
+    for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS) | set(STREAM_PIN_BUFFERS):
         if offload_stream is not None:
             offload_stream.synchronize()
     synchronize()
 
+    for mmap_obj in DIRTY_MMAPS:
+        mmap_obj.bounce()
+    DIRTY_MMAPS.clear()
+
+    for pin_buffer in STREAM_PIN_BUFFERS.values():
+        TOTAL_PINNED_MEMORY -= pin_buffer.size
+    TOTAL_PINNED_MEMORY = max(0, TOTAL_PINNED_MEMORY)
+
+    for loaded_model in current_loaded_models:
+        model = loaded_model.model
+        if model is not None and model.is_dynamic():
+            model.model.dynamic_pins[model.load_device]["active"] = False
+            model.partially_unload_ram(1e30, subsets=[ "patches" ])
+            model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0])
+
     STREAM_CAST_BUFFERS.clear()
     STREAM_AIMDO_CAST_BUFFERS.clear()
+    STREAM_PIN_BUFFERS.clear()
     soft_empty_cache()
 
 def get_offload_stream(device):
@@ -1280,7 +1317,7 @@ def sync_stream(device, stream):
     current_stream(device).wait_stream(stream)
 
 
-def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
+def cast_to_gathered(tensors, r, non_blocking=False, stream=None, r2=None):
     wf_context = nullcontext()
     if stream is not None:
        wf_context = stream
@@ -1288,17 +1325,20 @@ def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
            wf_context = wf_context.as_context(stream)
 
     dest_views = comfy.memory_management.interpret_gathered_like(tensors, r)
+    dest2_views = comfy.memory_management.interpret_gathered_like(tensors, r2) if r2 is not None else None
     with wf_context:
         for tensor in tensors:
             dest_view = dest_views.pop(0)
+            dest2_view = dest2_views.pop(0) if dest2_views is not None else None
             if tensor is None:
                 continue
-            if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view):
+            if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view, stream=stream, destination2=dest2_view):
                 continue
             storage = tensor._qdata.untyped_storage() if isinstance(tensor, comfy.quant_ops.QuantizedTensor) else tensor.untyped_storage()
-            if hasattr(storage, "_comfy_tensor_mmap_touched"):
-                storage._comfy_tensor_mmap_touched = True
+            mark_mmap_dirty(storage)
             dest_view.copy_(tensor, non_blocking=non_blocking)
+            if dest2_view is not None:
+                dest2_view.copy_(dest_view, non_blocking=non_blocking)
 
 
 def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None, r=None):
@@ -1339,14 +1379,18 @@ TOTAL_PINNED_MEMORY = 0
 MAX_PINNED_MEMORY = -1
 if not args.disable_pinned_memory:
     if is_nvidia() or is_amd():
+        ram = get_total_memory(torch.device("cpu"))
         if WINDOWS:
-            MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.40  # Windows limit is apparently 50%
+            MAX_PINNED_MEMORY = ram * 0.40  # Windows limit is apparently 50%
         else:
-            MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.90
+            MAX_PINNED_MEMORY = ram * 0.90
         logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
 
 PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
 
+def pinned_hostbuf_size(size):
+    return max(0, int(min(size, MAX_PINNED_MEMORY) * 2))
+
 def discard_cuda_async_error():
     try:
         a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
@@ -1378,8 +1422,8 @@ def pin_memory(tensor):
         return False
 
     size = tensor.nbytes
-    if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
-        return False
+    comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
+    ensure_pin_registerable(size)
 
     ptr = tensor.data_ptr()
     if ptr == 0:
@@ -1416,7 +1460,8 @@ def unpin_memory(tensor):
         return False
 
     if torch.cuda.cudart().cudaHostUnregister(ptr) == 0:
-        TOTAL_PINNED_MEMORY -= PINNED_MEMORY.pop(ptr)
+        size = PINNED_MEMORY.pop(ptr)
+        TOTAL_PINNED_MEMORY -= size
         return True
     else:
         logging.warning("Unpin error.")
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index 4f9d8403e..c8ed02e70 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -35,6 +35,7 @@ import comfy.model_management
 import comfy.ops
 import comfy.patcher_extension
 import comfy.utils
+import comfy_aimdo.host_buffer
 from comfy.comfy_types import UnetWrapperFunction
 from comfy.quant_ops import QuantizedTensor
 from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
@@ -117,6 +118,8 @@ def string_to_seed(data):
     return comfy.utils.string_to_seed(data)
 
 class LowVramPatch:
+    is_lowvram_patch = True
+
     def __init__(self, key, patches, convert_func=None, set_func=None):
         self.key = key
         self.patches = patches
@@ -124,11 +127,21 @@ class LowVramPatch:
         self.set_func = set_func
         self.prepared_patches = None
 
-    def prepare(self, allocate_buffer, stream):
-        self.prepared_patches = [
-            (patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4])
+    def memory_required(self):
+        counter = [0]
+        for patch in self.patches[self.key]:
+            comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False)
+        return counter[0]
+
+    def prepare(self, destination, stream, copy=True, commit=True):
+        counter = [0]
+        prepared_patches = [
+            (patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4])
             for patch in self.patches[self.key]
         ]
+        if commit:
+            self.prepared_patches = prepared_patches
+        return prepared_patches
 
     def clear_prepared(self):
         self.prepared_patches = None
@@ -341,9 +354,6 @@ class ModelPatcher:
         self.size = comfy.model_management.module_size(self.model)
         return self.size
 
-    def model_mmap_residency(self, free=False):
-        return comfy.model_management.module_mmap_residency(self.model, free=free)
-
     def loaded_size(self):
         return self.model.model_loaded_weight_memory
 
@@ -1118,8 +1128,12 @@ class ModelPatcher:
         # Pinned memory pressure tracking is only implemented for DynamicVram loading
         return 0
 
+    def loaded_ram_size(self):
+        # Loaded RAM pressure tracking is only implemented for DynamicVram loading
+        return 0
+
     def partially_unload_ram(self, ram_to_unload):
-        pass
+        return 0
 
     def detach(self, unpatch_all=True):
         self.eject_model()
@@ -1550,6 +1564,16 @@ class ModelPatcherDynamic(ModelPatcher):
         super().__init__(model, load_device, offload_device, size, weight_inplace_update)
         if not hasattr(self.model, "dynamic_vbars"):
             self.model.dynamic_vbars = {}
+        if not hasattr(self.model, "dynamic_pins"):
+            self.model.dynamic_pins = {}
+        if self.load_device not in self.model.dynamic_pins:
+            self.model.dynamic_pins[self.load_device] = {
+                "weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
+                "patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
+                "hostbufs_initialized": False,
+                "failed": False,
+                "active": False,
+            }
         self.non_dynamic_delegate_model = None
         assert load_device is not None
 
@@ -1611,6 +1635,14 @@ class ModelPatcherDynamic(ModelPatcher):
             self.unpatch_hooks()
 
             vbar = self._vbar_get(create=True)
+            pin_state = self.model.dynamic_pins[self.load_device]
+            if not pin_state["hostbufs_initialized"]:
+                hostbuf_size = comfy.model_management.pinned_hostbuf_size(self.model_size())
+                pin_state["weights"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0])
+                pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0])
+                pin_state["hostbufs_initialized"] = True
+            pin_state["failed"] = False
+            pin_state["active"] = True
             if vbar is not None:
                 vbar.prioritize()
 
@@ -1636,7 +1668,9 @@ class ModelPatcherDynamic(ModelPatcher):
                     if key in self.patches:
                         if comfy.lora.calculate_shape(self.patches[key], weight, key) != weight.shape:
                             return (True, 0)
-                        setattr(m, param_key + "_lowvram_function", LowVramPatch(key, self.patches))
+                        lowvram_patch = LowVramPatch(key, self.patches)
+                        lowvram_patch._pin_state = pin_state
+                        setattr(m, param_key + "_lowvram_function", lowvram_patch)
                         num_patches += 1
                     else:
                         setattr(m, param_key + "_lowvram_function", None)
@@ -1653,6 +1687,9 @@ class ModelPatcherDynamic(ModelPatcher):
 
                 def force_load_param(self, param_key, device_to):
                     key = key_param_name_to_key(n, param_key)
+                    weight, _, _ = get_key_weight(self.model, key)
+                    if weight is None:
+                        return
                     if key in self.backup:
                         comfy.utils.set_attr_param(self.model, key, self.backup[key].weight)
                     self.patch_weight_to_device(key, device_to=device_to, force_cast=True)
@@ -1662,17 +1699,23 @@ class ModelPatcherDynamic(ModelPatcher):
 
                 if hasattr(m, "comfy_cast_weights"):
                     m.comfy_cast_weights = True
-                    m.pin_failed = False
                     m.seed_key = n
+                    m._pin_state = pin_state
                     set_dirty(m, dirty)
 
-                    force_load, v_weight_size = setup_param(self, m, n, "weight")
-                    force_load_bias, v_weight_bias = setup_param(self, m, n, "bias")
-                    force_load = force_load or force_load_bias
-                    v_weight_size += v_weight_bias
+                    #Models that mix tiny and giant weights can causing lopsided stream buffer
+                    #rotations and stall. force the tinys over.
+                    if module_mem > 16 * 1024:
+                        force_load, v_weight_size = setup_param(self, m, n, "weight")
+                        force_load_bias, v_weight_bias = setup_param(self, m, n, "bias")
+                        force_load = force_load or force_load_bias
+                        v_weight_size += v_weight_bias
+                        if force_load:
+                            logging.info(f"Module {n} has resizing Lora - force loading")
+                    else:
+                        force_load=True
 
                     if force_load:
-                        logging.info(f"Module {n} has resizing Lora - force loading")
                         force_load_param(self, "weight", device_to)
                         force_load_param(self, "bias", device_to)
                     else:
@@ -1740,23 +1783,58 @@ class ModelPatcherDynamic(ModelPatcher):
 
         return freed
 
-    def pinned_memory_size(self):
-        total = 0
-        loading = self._load_list(for_dynamic=True)
-        for x in loading:
-            _, _, _, _, m, _ = x
-            pin = comfy.pinned_memory.get_pin(m)
-            if pin is not None:
-                total += pin.numel() * pin.element_size()
-        return total
+    def loaded_ram_size(self):
+        return (self.model.dynamic_pins[self.load_device]["weights"][0].size +
+                self.model.dynamic_pins[self.load_device]["patches"][0].size)
 
-    def partially_unload_ram(self, ram_to_unload):
-        loading = self._load_list(for_dynamic=True, default_device=self.offload_device)
-        for x in loading:
-            *_, m, _ = x
-            ram_to_unload -= comfy.pinned_memory.unpin_memory(m)
-            if ram_to_unload <= 0:
-                return
+    def pinned_memory_size(self):
+        return (self.model.dynamic_pins[self.load_device]["weights"][3][0] +
+                self.model.dynamic_pins[self.load_device]["patches"][3][0])
+
+    def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights", "patches" ]):
+        freed = 0
+        pin_state = self.model.dynamic_pins[self.load_device]
+        for subset in subsets:
+            hostbuf, stack, stack_split, pinned_size = pin_state[subset]
+            split = stack_split[0]
+            while split >= 0:
+                module, offset = stack[split]
+                split -= 1
+                stack_split[0] = split
+                if not module._pin_registered:
+                    continue
+                size = module._pin.numel() * module._pin.element_size()
+                if torch.cuda.cudart().cudaHostUnregister(module._pin.data_ptr()) != 0:
+                    comfy.model_management.discard_cuda_async_error()
+                    continue
+                module._pin_registered = False
+                comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
+                pinned_size[0] = max(0, pinned_size[0] - size)
+                freed += size
+                ram_to_unload -= size
+                if ram_to_unload <= 0:
+                    return freed
+        return freed
+
+    def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
+        freed = 0
+        pin_state = self.model.dynamic_pins[self.load_device]
+        for subset in subsets:
+            hostbuf, stack, stack_split, pinned_size = pin_state[subset]
+            while len(stack) > 0:
+                module, offset = stack.pop()
+                size = module._pin.numel() * module._pin.element_size()
+                del module._pin
+                hostbuf.truncate(offset, do_unregister=module._pin_registered)
+                stack_split[0] = min(stack_split[0], len(stack) - 1)
+                if module._pin_registered:
+                    comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
+                    pinned_size[0] = max(0, pinned_size[0] - size)
+                freed += size
+                ram_to_unload -= size
+                if ram_to_unload <= 0:
+                    return freed
+        return freed
 
     def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):
         #This isn't used by the core at all and can only be to load a model out of
diff --git a/comfy/ops.py b/comfy/ops.py
index eae3bd873..9bcd6c900 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -75,6 +75,8 @@ except:
 
 cast_to = comfy.model_management.cast_to #TODO: remove once no more references
 
+STREAM_PIN_BUFFER_HEADROOM = 8 * 1024 * 1024
+
 def cast_to_input(weight, input, non_blocking=False, copy=True):
     return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
 
@@ -91,6 +93,9 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
     offload_stream = None
     cast_buffer = None
     cast_buffer_offset = 0
+    stream_pin_hostbuf = None
+    stream_pin_offset = 0
+    stream_pin_queue = []
 
     def ensure_offload_stream(module, required_size, check_largest):
         nonlocal offload_stream
@@ -124,6 +129,22 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
         cast_buffer_offset += buffer_size
         return buffer
 
+    def get_stream_pin_buffer_offset(buffer_size):
+        nonlocal stream_pin_hostbuf
+        nonlocal stream_pin_offset
+
+        if buffer_size == 0 or offload_stream is None:
+            return None
+
+        if stream_pin_hostbuf is None:
+            stream_pin_hostbuf = comfy.model_management.get_pin_buffer(offload_stream)
+            if stream_pin_hostbuf is None:
+                return None
+
+        offset = stream_pin_offset
+        stream_pin_offset += buffer_size
+        return offset
+
     for s in comfy_modules:
         signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
         resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
@@ -162,23 +183,47 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
         if xfer_dest is None:
             xfer_dest = get_cast_buffer(dest_size)
 
-        if signature is None and pin is None:
-            comfy.pinned_memory.pin_memory(s)
-            pin = comfy.pinned_memory.get_pin(s)
-        else:
-            pin = None
+        def cast_maybe_lowvram_patch(xfer_source, xfer_dest, stream):
+            if xfer_source is not None:
+                if getattr(xfer_source, "is_lowvram_patch", False):
+                    xfer_source.prepare(xfer_dest, stream, copy=True, commit=False)
+                else:
+                    comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream)
 
-        if pin is not None:
-            comfy.model_management.cast_to_gathered(xfer_source, pin)
-            xfer_source = [ pin ]
-        #send it over
-        comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=offload_stream)
+        def handle_pin(m, pin, source, dest, subset="weights", size=None):
+            if pin is not None:
+                cast_maybe_lowvram_patch([pin], dest, offload_stream)
+                return
+            if signature is None:
+                comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
+                pin = comfy.pinned_memory.get_pin(m, subset=subset)
+                if pin is not None:
+                    if isinstance(source, list):
+                        comfy.model_management.cast_to_gathered(source, pin, non_blocking=non_blocking, stream=offload_stream, r2=dest)
+                    else:
+                        cast_maybe_lowvram_patch(source, pin, None)
+                        cast_maybe_lowvram_patch([ pin ], dest, offload_stream)
+                    return
+            if pin is None:
+                pin_offset = get_stream_pin_buffer_offset(size)
+                if pin_offset is not None:
+                    stream_pin_queue.append((source, pin_offset, size, dest))
+                    return
+            cast_maybe_lowvram_patch(source, dest, offload_stream)
+
+        handle_pin(s, pin, xfer_source, xfer_dest, size=dest_size)
 
         for param_key in ("weight", "bias"):
-            lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
-            if lowvram_fn is not None:
+            lowvram_source = getattr(s, param_key + "_lowvram_function", None)
+            if lowvram_source is not None:
                 ensure_offload_stream(s, cast_buffer_offset, False)
-                lowvram_fn.prepare(lambda size: get_cast_buffer(size), offload_stream)
+                lowvram_size = lowvram_source.memory_required()
+                lowvram_dest = get_cast_buffer(lowvram_size)
+                lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
+
+                pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches")
+                handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset="patches", size=lowvram_size)
+
 
         prefetch["xfer_dest"] = xfer_dest
         prefetch["cast_dest"] = cast_dest
@@ -186,6 +231,23 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
         prefetch["needs_cast"] = needs_cast
         s._prefetch = prefetch
 
+    if stream_pin_offset > 0:
+        if stream_pin_hostbuf.size < stream_pin_offset:
+            if not comfy.model_management.resize_pin_buffer(stream_pin_hostbuf, stream_pin_offset + STREAM_PIN_BUFFER_HEADROOM):
+                for xfer_source, _, _, xfer_dest in stream_pin_queue:
+                    cast_maybe_lowvram_patch(xfer_source, xfer_dest, offload_stream)
+                return offload_stream
+        stream_pin_tensor = comfy_aimdo.torch.hostbuf_to_tensor(stream_pin_hostbuf)
+        stream_pin_tensor.untyped_storage()._comfy_hostbuf = stream_pin_hostbuf
+        for xfer_source, pin_offset, pin_size, xfer_dest in stream_pin_queue:
+            pin = stream_pin_tensor[pin_offset:pin_offset + pin_size]
+            if isinstance(xfer_source, list):
+                comfy.model_management.cast_to_gathered(xfer_source, pin, non_blocking=non_blocking, stream=offload_stream, r2=xfer_dest)
+            else:
+                cast_maybe_lowvram_patch(xfer_source, pin, None)
+                comfy.model_management.cast_to_gathered([ pin ], xfer_dest, non_blocking=non_blocking, stream=offload_stream)
+        stream_pin_hostbuf._comfy_event = offload_stream.record_event()
+
     return offload_stream
 
 
diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py
index 6d3ba367a..0e8f573ba 100644
--- a/comfy/pinned_memory.py
+++ b/comfy/pinned_memory.py
@@ -2,42 +2,62 @@ import comfy.model_management
 import comfy.memory_management
 import comfy_aimdo.host_buffer
 import comfy_aimdo.torch
+import torch
 
 from comfy.cli_args import args
 
-def get_pin(module):
-    return getattr(module, "_pin", None)
+def get_pin(module, subset="weights"):
+    pin = getattr(module, "_pin", None)
+    if pin is None or module._pin_registered or args.disable_pinned_memory:
+        return pin
 
-def pin_memory(module):
-    if module.pin_failed or args.disable_pinned_memory or get_pin(module) is not None:
+    _, _, stack_split, pinned_size = module._pin_state[subset]
+    size = pin.nbytes
+    comfy.model_management.ensure_pin_registerable(size)
+
+    if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
+        comfy.model_management.discard_cuda_async_error()
+        return pin
+
+    module._pin_registered = True
+    stack_split[0] = max(stack_split[0], module._pin_stack_index)
+    comfy.model_management.TOTAL_PINNED_MEMORY += size
+    pinned_size[0] += size
+    return pin
+
+def pin_memory(module, subset="weights", size=None):
+    pin_state = module._pin_state
+    if args.disable_pinned_memory:
         return
 
-    size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
+    pin = get_pin(module, subset)
+    if pin is not None or pin_state["failed"]:
+        return
 
-    if comfy.model_management.MAX_PINNED_MEMORY <= 0 or (comfy.model_management.TOTAL_PINNED_MEMORY + size) > comfy.model_management.MAX_PINNED_MEMORY:
-        module.pin_failed = True
+    hostbuf, stack, stack_split, pinned_size = pin_state[subset]
+    if size is None:
+        size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
+    offset = hostbuf.size
+    registerable_size = size + max(0, hostbuf.size - pinned_size[0])
+
+    comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
+    if (not comfy.model_management.ensure_pin_budget(size) or
+        not comfy.model_management.ensure_pin_registerable(registerable_size)):
+        pin_state["failed"] = True
         return False
 
     try:
-        hostbuf = comfy_aimdo.host_buffer.HostBuffer(size)
+        hostbuf.extend(size=size)
     except RuntimeError:
-        module.pin_failed = True
+        pin_state["failed"] = True
         return False
 
-    module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)
-    module._pin_hostbuf = hostbuf
+    module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size]
+    module._pin.untyped_storage()._comfy_hostbuf = hostbuf
+    stack.append((module, offset))
+    module._pin_registered = True
+    module._pin_stack_index = len(stack) - 1
+    stack_split[0] = max(stack_split[0], module._pin_stack_index)
     comfy.model_management.TOTAL_PINNED_MEMORY += size
+    pinned_size[0] += size
     return True
-
-def unpin_memory(module):
-    if get_pin(module) is None:
-        return 0
-    size = module._pin.numel() * module._pin.element_size()
-
-    comfy.model_management.TOTAL_PINNED_MEMORY -= size
-    if comfy.model_management.TOTAL_PINNED_MEMORY < 0:
-        comfy.model_management.TOTAL_PINNED_MEMORY = 0
-
-    del module._pin
-    del module._pin_hostbuf
-    return size
diff --git a/comfy/utils.py b/comfy/utils.py
index 66682690a..00e382fac 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -113,7 +113,6 @@ def load_safetensors(ckpt):
                         "_comfy_tensor_file_slice",
                         comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start))
                 setattr(storage, "_comfy_tensor_mmap_refs", (model_mmap, mv))
-                setattr(storage, "_comfy_tensor_mmap_touched", False)
                 sd[name] = tensor
 
     return sd, header.get("__metadata__", {}),
@@ -1451,4 +1450,3 @@ def deepcopy_list_dict(obj, memo=None):
 
     memo[obj_id] = res
     return res
-
diff --git a/comfy/windows.py b/comfy/windows.py
deleted file mode 100644
index 213dc481d..000000000
--- a/comfy/windows.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import ctypes
-import logging
-import psutil
-from ctypes import wintypes
-
-import comfy_aimdo.control
-
-psapi = ctypes.WinDLL("psapi")
-kernel32 = ctypes.WinDLL("kernel32")
-
-class PERFORMANCE_INFORMATION(ctypes.Structure):
-    _fields_ = [
-        ("cb", wintypes.DWORD),
-        ("CommitTotal", ctypes.c_size_t),
-        ("CommitLimit", ctypes.c_size_t),
-        ("CommitPeak", ctypes.c_size_t),
-        ("PhysicalTotal", ctypes.c_size_t),
-        ("PhysicalAvailable", ctypes.c_size_t),
-        ("SystemCache", ctypes.c_size_t),
-        ("KernelTotal", ctypes.c_size_t),
-        ("KernelPaged", ctypes.c_size_t),
-        ("KernelNonpaged", ctypes.c_size_t),
-        ("PageSize", ctypes.c_size_t),
-        ("HandleCount", wintypes.DWORD),
-        ("ProcessCount", wintypes.DWORD),
-        ("ThreadCount", wintypes.DWORD),
-    ]
-
-def get_free_ram():
-    #Windows is way too conservative and chalks recently used uncommitted model RAM
-    #as "in-use". So, calculate free RAM for the sake of general use as the greater of:
-    #
-    #1: What psutil says
-    #2: Total Memory - (Committed Memory - VRAM in use)
-    #
-    #We have to subtract VRAM in use from the comitted memory as WDDM creates a naked
-    #commit charge for all VRAM used just incase it wants to page it all out. This just
-    #isn't realistic so "overcommit" on our calculations by just subtracting it off.
-
-    pi = PERFORMANCE_INFORMATION()
-    pi.cb = ctypes.sizeof(pi)
-
-    if not psapi.GetPerformanceInfo(ctypes.byref(pi), pi.cb):
-        logging.warning("WARNING: Failed to query windows performance info. RAM usage may be sub optimal")
-        return psutil.virtual_memory().available
-
-    committed = pi.CommitTotal * pi.PageSize
-    total = pi.PhysicalTotal * pi.PageSize
-
-    return max(psutil.virtual_memory().available,
-               total - (committed - comfy_aimdo.control.get_total_vram_usage()))
-
diff --git a/execution.py b/execution.py
index 4c7de2e84..5246d651c 100644
--- a/execution.py
+++ b/execution.py
@@ -2,6 +2,7 @@ import copy
 import heapq
 import inspect
 import logging
+import psutil
 import sys
 import threading
 import time
@@ -727,6 +728,7 @@ class PromptExecutor:
 
         self._notify_prompt_lifecycle("start", prompt_id)
         ram_headroom = int(self.cache_args["ram"] * (1024 ** 3))
+        ram_inactive_headroom = int(self.cache_args["ram_inactive"] * (1024 ** 3))
         ram_release_callback = self.caches.outputs.ram_release if self.cache_type == CacheType.RAM_PRESSURE else None
         comfy.memory_management.set_ram_cache_release_state(ram_release_callback, ram_headroom)
 
@@ -780,8 +782,14 @@ class PromptExecutor:
                         execution_list.complete_node_execution()
 
                     if self.cache_type == CacheType.RAM_PRESSURE:
-                        comfy.model_management.free_memory(0, None, pins_required=ram_headroom, ram_required=ram_headroom)
-                        ram_release_callback(ram_headroom, free_active=True)
+                        ram_release_callback(ram_inactive_headroom)
+                        ram_shortfall = ram_headroom - psutil.virtual_memory().available
+                        freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
+                        if freed < ram_shortfall:
+                            if freed > 64 * (1024 ** 2):
+                                # AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
+                                time.sleep(0.05)
+                            ram_release_callback(ram_headroom, free_active=True)
                 else:
                     # Only execute when the while-loop ends without break
                     # Send cached UI for intermediate output nodes that weren't executed
diff --git a/main.py b/main.py
index a6fdaf43c..1e47cab84 100644
--- a/main.py
+++ b/main.py
@@ -283,19 +283,25 @@ def _collect_output_absolute_paths(history_result: dict) -> list[str]:
 
 def prompt_worker(q, server_instance):
     current_time: float = 0.0
-    cache_ram = args.cache_ram
-    if cache_ram < 0:
+    cache_ram = 0
+    cache_ram_inactive = 0
+    if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
         cache_ram = min(32.0, max(4.0, comfy.model_management.total_ram * 0.25 / 1024.0))
+        cache_ram_inactive = min(96.0, max(12.0, comfy.model_management.total_ram * 0.75 / 1024.0))
+        if len(args.cache_ram) > 0:
+            cache_ram = args.cache_ram[0]
+        if len(args.cache_ram) > 1:
+            cache_ram_inactive = args.cache_ram[1]
 
-    cache_type = execution.CacheType.CLASSIC
-    if args.cache_lru > 0:
+    cache_type = execution.CacheType.RAM_PRESSURE
+    if args.cache_classic:
+        cache_type = execution.CacheType.CLASSIC
+    elif args.cache_lru > 0:
         cache_type = execution.CacheType.LRU
-    elif cache_ram > 0:
-        cache_type = execution.CacheType.RAM_PRESSURE
     elif args.cache_none:
         cache_type = execution.CacheType.NONE
 
-    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram } )
+    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive } )
     last_gc_collect = 0
     need_gc = False
     gc_collect_interval = 10.0
diff --git a/requirements.txt b/requirements.txt
index 1c87690da..d2986eda8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
 filelock
 av>=14.2.0
 comfy-kitchen>=0.2.8
-comfy-aimdo==0.3.0
+comfy-aimdo==0.4.3
 requests
 simpleeval>=1.0.0
 blake3
diff --git a/tests/execution/test_async_nodes.py b/tests/execution/test_async_nodes.py
index c771b4b36..54660c112 100644
--- a/tests/execution/test_async_nodes.py
+++ b/tests/execution/test_async_nodes.py
@@ -14,7 +14,6 @@ from tests.execution.test_execution import ComfyClient, run_warmup
 class TestAsyncNodes:
     @fixture(scope="class", autouse=True, params=[
         (False, 0),
-        (True, 0),
         (True, 100),
     ])
     def _server(self, args_pytest, request):
@@ -29,6 +28,8 @@ class TestAsyncNodes:
         use_lru, lru_size = request.param
         if use_lru:
             pargs += ['--cache-lru', str(lru_size)]
+        else:
+            pargs += ['--cache-classic']
         # Running server with args: pargs
         p = subprocess.Popen(pargs)
         yield
diff --git a/tests/execution/test_execution.py b/tests/execution/test_execution.py
index f73ca7e3c..15e2304fc 100644
--- a/tests/execution/test_execution.py
+++ b/tests/execution/test_execution.py
@@ -183,8 +183,7 @@ class TestExecution:
     # Initialize server and client
     #
     @fixture(scope="class", autouse=True, params=[
-        { "extra_args" : [], "should_cache_results" : True },
-        { "extra_args" : ["--cache-lru", 0], "should_cache_results" : True },
+        { "extra_args" : ["--cache-classic"], "should_cache_results" : True },
         { "extra_args" : ["--cache-lru", 100], "should_cache_results" : True },
         { "extra_args" : ["--cache-none"], "should_cache_results" : False },
     ])

From 95fdc6cf910f809e39edc3254470e619ffa9dbf8 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Wed, 20 May 2026 17:17:55 -0700
Subject: [PATCH 07/14] Repo security stuff. (#14019)

---
 CODEOWNERS | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/CODEOWNERS b/CODEOWNERS
index 946dbf946..043c0ec75 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -1,2 +1,5 @@
-# Admins
 * @comfyanonymous @kosinkadink @guill @alexisrolland @rattus128 @kijai
+
+/CODEOWNERS @comfyanonymous
+/.ci/ @comfyanonymous
+/.github/ @comfyanonymous

From 9f9b32ed978045262b71e6b27093e4ae80c29804 Mon Sep 17 00:00:00 2001
From: Matt Miller <mattmiller@comfy.org>
Date: Wed, 20 May 2026 21:22:12 -0700
Subject: [PATCH 08/14] feat: add OAuth 2.1 + RFC 7591 DCR endpoints to
 openapi.yaml (#14026)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add the OAuth 2.1 authorization flow and RFC 7591 Dynamic Client
Registration endpoints to the shared spec, alongside the existing
auth-tagged operations (/api/auth/session, /api/auth/token,
/.well-known/jwks.json). All tagged x-runtime: [cloud] with a
[cloud-only] description prefix, following the established
convention for cloud-runtime-only operations.

Endpoints:

- GET  /.well-known/oauth-authorization-server  (RFC 8414 metadata)
- GET  /.well-known/oauth-protected-resource    (RFC 9728 metadata)
- GET  /oauth/authorize                         (consent challenge)
- POST /oauth/authorize                         (consent submission)
- POST /oauth/token                             (RFC 6749 §3.2)
- POST /oauth/register                          (RFC 7591 §3.1 DCR)

Component schemas added:

- OAuthAuthorizationServerMetadata
- OAuthProtectedResourceMetadata
- OAuthConsentChallenge, OAuthConsentChallengeWorkspace
- OAuthAuthorizeRedirectResponse
- OAuthTokenResponse, OAuthTokenError
- OAuthRegisterRequest, OAuthRegisterResponse, OAuthRegisterError

These endpoints are implemented in the cloud runtime today and
are called by browser frontends rendering the consent UI and by
MCP-spec-compliant clients (Claude Desktop, Cursor, etc.) doing
auto-discovery + self-registration. Documenting them in the
shared spec lets the cloud frontend generate types directly from
this spec instead of maintaining a parallel definition.

Spectral lints clean (0 errors). The hint-level findings on
OAuthTokenError / OAuthRegisterError ("standard error schema")
match the same hint on CloudError — these are protocol-specific
RFC-shaped errors, not generic application errors.
---
 openapi.yaml | 608 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 608 insertions(+)

diff --git a/openapi.yaml b/openapi.yaml
index 2658b9b86..92f7eaccc 100644
--- a/openapi.yaml
+++ b/openapi.yaml
@@ -3790,6 +3790,295 @@ paths:
               schema:
                 $ref: "#/components/schemas/JwksResponse"
 
+  # ---------------------------------------------------------------------------
+  # OAuth 2.1 / RFC 7591 Dynamic Client Registration (cloud)
+  # ---------------------------------------------------------------------------
+  /.well-known/oauth-authorization-server:
+    get:
+      operationId: getOAuthAuthorizationServer
+      tags: [auth]
+      summary: "[cloud-only] OAuth 2.1 authorization-server metadata (RFC 8414)"
+      description: "[cloud-only] Public metadata document for OAuth 2.1 clients. Cached 5 minutes."
+      x-runtime: [cloud]
+      security: []
+      responses:
+        "200":
+          description: Authorization-server metadata
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthAuthorizationServerMetadata"
+        "404":
+          description: OAuth disabled
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+
+  /.well-known/oauth-protected-resource:
+    get:
+      operationId: getOAuthProtectedResource
+      tags: [auth]
+      summary: "[cloud-only] OAuth 2.1 protected-resource metadata (RFC 9728)"
+      description: "[cloud-only] Public metadata describing the currently advertised protected resource. Cached 5 minutes."
+      x-runtime: [cloud]
+      security: []
+      responses:
+        "200":
+          description: Protected-resource metadata
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthProtectedResourceMetadata"
+        "404":
+          description: OAuth disabled or no active resource configured
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+
+  /oauth/authorize:
+    get:
+      operationId: getOAuthAuthorize
+      tags: [auth]
+      summary: "[cloud-only] Begin or resume an OAuth 2.1 authorization request"
+      description: |
+        [cloud-only] Two modes:
+        - **Initial entry** (OAuth params present): validates client/redirect/resource/scopes, persists a server-side authorization-request row, and either redirects (no session / unverified email) to the configured frontend login URL carrying only the opaque `oauth_request_id`, or returns the JSON consent challenge for the frontend to render.
+        - **Resume** (`oauth_request_id` present): loads the server-side row, fails closed if expired/consumed/unknown, returns the JSON consent challenge. Browser-replayed OAuth params are intentionally ignored.
+
+        The frontend renders the consent UI from the JSON payload and POSTs the user's decision back to this endpoint.
+      x-runtime: [cloud]
+      security: []
+      parameters:
+        - { name: response_type,         in: query, required: false, schema: { type: string } }
+        - { name: client_id,             in: query, required: false, schema: { type: string } }
+        - { name: redirect_uri,          in: query, required: false, schema: { type: string } }
+        - { name: scope,                 in: query, required: false, schema: { type: string } }
+        - name: state
+          in: query
+          required: false
+          schema: { type: string }
+          description: |
+            RFC 6749 §10.12 marks `state` as RECOMMENDED. Cloud hardening makes it REQUIRED on the initial-entry path (omitted only on the resume path where `oauth_request_id` is supplied instead). This parameter is `required: false` at the spec level only because the operation is dual-mode (initial entry vs. resume); the runtime rejects empty `state` on the initial-entry path with a stable `invalid_request` 400.
+        - { name: code_challenge,        in: query, required: false, schema: { type: string } }
+        - { name: code_challenge_method, in: query, required: false, schema: { type: string } }
+        - { name: resource,              in: query, required: false, schema: { type: string } }
+        - { name: oauth_request_id,      in: query, required: false, schema: { type: string } }
+      responses:
+        "200":
+          description: Consent challenge payload (session present, email verified). Frontend renders the consent UI from this payload and POSTs back to /oauth/authorize.
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthConsentChallenge"
+        "302":
+          description: Redirect to login (no session / unverified email) or to registered redirect_uri (pre-validated client error)
+          headers:
+            Location:
+              schema:
+                type: string
+        "400":
+          description: Invalid authorize request (pre-redirect failure — unknown client, redirect mismatch, malformed params)
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+        "404":
+          description: OAuth disabled
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+    post:
+      operationId: postOAuthAuthorize
+      tags: [auth]
+      summary: "[cloud-only] Submit OAuth consent decision"
+      description: |
+        [cloud-only] JSON-only consent submission. The handler verifies the per-row CSRF token, atomically marks the authorization request consumed (single-use covers both allow and deny paths), then returns the redirect URL the browser must navigate to. The URL contains either `code` + original `state` for allow, or the RFC 6749 §5.2 error and `state` for deny.
+
+        Workspace membership is re-checked at submission time. Consent is persisted keyed by `(user_id, client_id, resource_id, workspace_id)`; broadening the previously approved scope set requires a fresh consent flow.
+      x-runtime: [cloud]
+      security: []
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              required: [oauth_request_id, csrf_token, decision, workspace_id]
+              properties:
+                oauth_request_id: { type: string, format: uuid }
+                csrf_token:       { type: string }
+                decision:         { type: string, enum: [allow, deny] }
+                workspace_id:     { type: string }
+      responses:
+        "200":
+          description: Redirect URL for the frontend to navigate to (allow → with code+state; deny → with error+state)
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthAuthorizeRedirectResponse"
+        "400":
+          description: Bad request (CSRF mismatch, expired/consumed request, inaccessible workspace)
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+        "403":
+          description: Scope broadening on consent re-grant — fresh consent flow required
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+        "404":
+          description: OAuth disabled
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+
+  /oauth/token:
+    post:
+      operationId: postOAuthToken
+      tags: [auth]
+      summary: "[cloud-only] Exchange authorization code or refresh token for a resource-bound access token"
+      description: |
+        [cloud-only] OAuth 2.1 token endpoint (RFC 6749 §3.2). Public clients only — `client_secret` is rejected.
+
+        Two grant types are supported:
+        - `authorization_code` — exchanges the code minted by `/oauth/authorize` (with PKCE verifier) for an access token + first refresh token. Single-use; reuse fails closed.
+        - `refresh_token` — rotates the refresh token. Old token immediately invalid; presenting an already-rotated token revokes the entire token family and emits a security metric.
+
+        Both grant types re-validate canonical user state, current workspace membership, and the resource's active flag at every mint. A code or refresh token bound to a deactivated resource fails closed.
+
+        Errors follow RFC 6749 §5.2. Logs never contain raw codes, refresh tokens, or minted tokens.
+
+        Per RFC 6749 §5.1, every 200 and 400 response carries `Cache-Control: no-store` and `Pragma: no-cache` so intermediaries cannot cache token-bearing or state-change-reason responses.
+      x-runtime: [cloud]
+      security: []
+      requestBody:
+        required: true
+        content:
+          application/x-www-form-urlencoded:
+            schema:
+              type: object
+              required: [grant_type, client_id]
+              properties:
+                grant_type:    { type: string, enum: [authorization_code, refresh_token] }
+                client_id:     { type: string }
+                code:          { type: string }
+                redirect_uri:  { type: string }
+                code_verifier: { type: string }
+                refresh_token: { type: string }
+                scope:         { type: string }
+                client_secret: { type: string }
+      responses:
+        "200":
+          description: New token pair
+          headers:
+            Cache-Control:
+              schema:
+                type: string
+              description: 'Always "no-store" per RFC 6749 §5.1'
+            Pragma:
+              schema:
+                type: string
+              description: 'Always "no-cache" per RFC 6749 §5.1'
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthTokenResponse"
+        "400":
+          description: RFC 6749 §5.2 error
+          headers:
+            Cache-Control:
+              schema:
+                type: string
+              description: 'Always "no-store" per RFC 6749 §5.1'
+            Pragma:
+              schema:
+                type: string
+              description: 'Always "no-cache" per RFC 6749 §5.1'
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthTokenError"
+        "404":
+          description: OAuth disabled
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+
+  /oauth/register:
+    post:
+      operationId: postOAuthRegister
+      tags: [auth]
+      summary: "[cloud-only] Dynamic Client Registration (RFC 7591)"
+      description: |
+        [cloud-only] Public, unauthenticated, insert-only RFC 7591 §3.1 client registration. Used by MCP-spec-compliant clients to self-register a public OAuth client without operator involvement.
+
+        Policy:
+
+        - Public clients only — `token_endpoint_auth_method` is forced to `none`. Confidential-client registration is out of scope this phase.
+        - Server-owned `resource_grants`. Caller-supplied `scope` or `resource_grants` is rejected as `invalid_client_metadata` (would be a privilege-escalation surface). Dynamic clients receive the same scopes the active resource publishes.
+        - Application-type-aware redirect URI policy. `application_type=native` accepts loopback (`127.0.0.1`, `::1`, `localhost`) and reverse-DNS-shaped custom schemes; `application_type=web` accepts HTTPS to hosts in an operator-controlled allowlist only. `application_type` is REQUIRED on the request — missing or empty rejects with `invalid_client_metadata`.
+        - Anti-impersonation: reserved client names are rejected from third parties via NFKC-folded compare.
+        - Generated `client_id` carries a stable prefix to distinguish dynamic from seeded clients in audit logs.
+        - Cache-Control: `no-store` on every 201 and 400 response (the response carries fresh credentials and rejection reasons).
+      x-runtime: [cloud]
+      security: []
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/OAuthRegisterRequest"
+      responses:
+        "201":
+          description: Registered. Body echoes the metadata RFC 7591 §3.2.1 requires.
+          headers:
+            Cache-Control:
+              schema:
+                type: string
+              description: 'Always "no-store"'
+            Pragma:
+              schema:
+                type: string
+              description: 'Always "no-cache"'
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthRegisterResponse"
+        "400":
+          description: RFC 7591 §3.2.2 invalid client metadata
+          headers:
+            Cache-Control:
+              schema:
+                type: string
+              description: 'Always "no-store"'
+            Pragma:
+              schema:
+                type: string
+              description: 'Always "no-cache"'
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/OAuthRegisterError"
+        "404":
+          description: OAuth disabled
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+        "503":
+          description: No active resource is configured — DCR cannot mint a usable client until an active resource row is seeded.
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/CloudError"
+
   # ---------------------------------------------------------------------------
   # Billing (cloud)
   # ---------------------------------------------------------------------------
@@ -7424,6 +7713,325 @@ components:
                 description: RSA exponent (base64url)
             additionalProperties: true
 
+    OAuthAuthorizationServerMetadata:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] OAuth 2.1 authorization-server metadata (RFC 8414)."
+      required:
+        - issuer
+        - authorization_endpoint
+        - token_endpoint
+        - jwks_uri
+        - response_types_supported
+        - grant_types_supported
+        - code_challenge_methods_supported
+        - token_endpoint_auth_methods_supported
+      properties:
+        issuer:
+          type: string
+          format: uri
+        authorization_endpoint:
+          type: string
+          format: uri
+        token_endpoint:
+          type: string
+          format: uri
+        jwks_uri:
+          type: string
+          format: uri
+        registration_endpoint:
+          type: string
+          format: uri
+          description: "[cloud-only] RFC 7591 §3.1 Dynamic Client Registration endpoint. Advertised so MCP-spec-compliant clients can auto-discover and self-register without operator involvement. Present only when DCR is enabled."
+        response_types_supported:
+          type: array
+          items:
+            type: string
+        grant_types_supported:
+          type: array
+          items:
+            type: string
+        code_challenge_methods_supported:
+          type: array
+          items:
+            type: string
+        token_endpoint_auth_methods_supported:
+          type: array
+          items:
+            type: string
+        scopes_supported:
+          type: array
+          items:
+            type: string
+
+    OAuthProtectedResourceMetadata:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] OAuth 2.1 protected-resource metadata (RFC 9728)."
+      required:
+        - resource
+        - authorization_servers
+        - scopes_supported
+      properties:
+        resource:
+          type: string
+          format: uri
+        authorization_servers:
+          type: array
+          items:
+            type: string
+            format: uri
+        scopes_supported:
+          type: array
+          items:
+            type: string
+        bearer_methods_supported:
+          type: array
+          items:
+            type: string
+
+    OAuthConsentChallenge:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] Server-side state describing the OAuth consent decision the user is being asked to make. Returned by GET /oauth/authorize when a valid session exists; the frontend renders the consent UI from this payload and POSTs the decision back. Browser never sees the original OAuth params on resume."
+      required:
+        - oauth_request_id
+        - csrf_token
+        - client_display_name
+        - resource_display_name
+        - scopes
+        - workspaces
+      properties:
+        oauth_request_id:
+          type: string
+          format: uuid
+          description: Opaque server-side identifier for the authorization-request row. Carried back unchanged in the consent submission.
+        csrf_token:
+          type: string
+          description: Per-row CSRF token bound to this authorization request (not to the session). Must be echoed back on POST.
+        client_display_name:
+          type: string
+          description: Human-readable name of the OAuth client requesting authorization.
+        resource_display_name:
+          type: string
+          description: Human-readable name of the protected resource.
+        scopes:
+          type: array
+          description: Scopes the client is requesting for this resource. The frontend should present these for the user to approve.
+          items:
+            type: string
+        workspaces:
+          type: array
+          description: Workspaces the user can select from. Membership is re-checked on POST.
+          items:
+            $ref: "#/components/schemas/OAuthConsentChallengeWorkspace"
+
+    OAuthConsentChallengeWorkspace:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] One workspace option presented in the OAuth consent challenge."
+      required: [id, name, type, role]
+      properties:
+        id:   { type: string }
+        name: { type: string }
+        type: { type: string, enum: [personal, team] }
+        role: { type: string, enum: [owner, member] }
+
+    OAuthAuthorizeRedirectResponse:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] Redirect target produced after a JSON consent submission. The frontend must navigate the browser to this URL so custom-scheme client callbacks work without relying on fetch-visible 302 headers."
+      required:
+        - redirect_url
+      properties:
+        redirect_url:
+          type: string
+          format: uri
+          description: OAuth client redirect URI with either code+state for allow, or error+state for deny.
+
+    OAuthTokenResponse:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] RFC 6749 §5.1 successful token response."
+      required: [access_token, token_type, expires_in, refresh_token, scope]
+      properties:
+        access_token:
+          type: string
+          description: Resource-bound access token (audience matches the protected resource).
+        token_type:
+          type: string
+          enum: [Bearer]
+        expires_in:
+          type: integer
+          description: Access token lifetime in seconds.
+        refresh_token:
+          type: string
+          description: Opaque refresh token. Rotates on every successful refresh; presenting an already-rotated token revokes the entire family.
+        scope:
+          type: string
+          description: Space-delimited scopes granted with this token.
+
+    OAuthTokenError:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] RFC 6749 §5.2 error response."
+      required: [error]
+      properties:
+        error:
+          type: string
+          description: 'RFC 6749 §5.2 error code: invalid_request, invalid_client, invalid_grant, unauthorized_client, unsupported_grant_type, invalid_scope.'
+        error_description:
+          type: string
+          description: Human-readable, no leak of internal storage state.
+
+    OAuthRegisterRequest:
+      type: object
+      x-runtime: [cloud]
+      additionalProperties: false
+      description: "[cloud-only] RFC 7591 §2 client metadata document. Only the fields the server honors are listed; presence of `scope` or `resource_grants` in the request is rejected (`invalid_client_metadata`) because those are server-owned for dynamic clients."
+      required:
+        - redirect_uris
+        - application_type
+      properties:
+        redirect_uris:
+          type: array
+          items:
+            type: string
+          minItems: 1
+          maxItems: 5
+          description: 1–5 redirect URIs. Validated against `application_type` policy.
+        client_name:
+          type: string
+          maxLength: 100
+          description: Human-readable name shown in the consent UI. Reserved-name list rejects impersonation of major clients.
+        application_type:
+          type: string
+          enum: [native, web]
+          description: |
+            RFC 7591 §2 application_type. **REQUIRED** — clients MUST declare intent; the server does not default this field. `native` for desktop / CLI / MCP-spec-strict clients (loopback redirects); `web` for hosted clients (HTTPS only, host must be allowlisted). A missing or explicitly empty `application_type` rejects with `invalid_client_metadata`.
+        token_endpoint_auth_method:
+          type: string
+          enum: [none]
+          description: 'Public clients only this phase — must be `none` if present. The server forces `none` regardless.'
+        grant_types:
+          type: array
+          items:
+            type: string
+            enum: [authorization_code, refresh_token]
+          description: Optional. Defaults to `["authorization_code","refresh_token"]`.
+        response_types:
+          type: array
+          items:
+            type: string
+            enum: [code]
+          description: Optional. Defaults to `["code"]`.
+        scope:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Dynamic clients do not pick scopes — the server assigns scopes from the active resource's published list. Sending `scope` in the registration body is treated as a privilege-escalation attempt and returns `invalid_client_metadata`."
+        resource_grants:
+          type: object
+          nullable: true
+          additionalProperties:
+            type: array
+            items:
+              type: string
+          description: "**REJECTED IF PRESENT.** Same reason as `scope`. The set of resources and scopes a dynamic client may request is server-policy, not request-driven."
+        client_uri:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        logo_uri:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        tos_uri:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        policy_uri:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        software_id:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        software_version:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        contacts:
+          type: array
+          nullable: true
+          items:
+            type: string
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        jwks:
+          type: object
+          nullable: true
+          additionalProperties: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+        jwks_uri:
+          type: string
+          nullable: true
+          description: "**REJECTED IF PRESENT.** Unsupported RFC 7591 metadata for this public-client phase."
+
+    OAuthRegisterResponse:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] RFC 7591 §3.2.1 successful registration response."
+      required:
+        - client_id
+        - client_id_issued_at
+        - redirect_uris
+        - grant_types
+        - response_types
+        - token_endpoint_auth_method
+        - application_type
+      properties:
+        client_id:
+          type: string
+          description: Server-generated client_id.
+        client_id_issued_at:
+          type: integer
+          format: int64
+          description: Unix timestamp (seconds) when the client was registered.
+        client_name:
+          type: string
+        redirect_uris:
+          type: array
+          items:
+            type: string
+        grant_types:
+          type: array
+          items:
+            type: string
+        response_types:
+          type: array
+          items:
+            type: string
+        token_endpoint_auth_method:
+          type: string
+          enum: [none]
+        application_type:
+          type: string
+          enum: [native, web]
+
+    OAuthRegisterError:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] RFC 7591 §3.2.2 error response."
+      required:
+        - error
+      properties:
+        error:
+          type: string
+          enum: [invalid_redirect_uri, invalid_client_metadata]
+        error_description:
+          type: string
+          nullable: true
+
     BillingBalance:
       type: object
       x-runtime: [cloud]

From ea174d3f120bf43c0219eb341e9373834036083c Mon Sep 17 00:00:00 2001
From: Matt Miller <mattmiller@comfy.org>
Date: Wed, 20 May 2026 21:28:16 -0700
Subject: [PATCH 09/14] fix(openapi): correct POST /api/assets/import to
 importPublishedAssets (#14027)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The operation at POST /api/assets/import was defined as `importAssets`
with a URL-list body shape, but no runtime actually serves that
operation at this path. The cloud runtime serves a different operation
here — `importPublishedAssets` — which imports published-workflow
assets into the caller's library by ID, not by URL.

Cloud's URL-based asset ingestion lives at separate paths
(POST /assets/download + GET /assets/remote-metadata) tracked
elsewhere; nothing in this PR affects that work.

Changes:

- Replace the operation at POST /api/assets/import with
  `importPublishedAssets`, taking ImportPublishedAssetsRequest
  (published_asset_ids + optional share_id) and returning
  ImportPublishedAssetsResponse (list of AssetInfo).
- Remove the unused AssetImportRequest component schema (no other
  references in the spec).
- Operation and schemas tagged x-runtime: [cloud] with [cloud-only]
  description prefix, matching the existing convention for
  cloud-runtime-only operations elsewhere in the spec.

Spectral lint passes (0 errors); the two hint-level findings on
the spec are pre-existing and unrelated.

No FE consumer references AssetImportRequest today; this is a pure
spec correction to match what the cloud runtime actually serves.
---
 openapi.yaml | 59 ++++++++++++++++++++++++++--------------------------
 1 file changed, 29 insertions(+), 30 deletions(-)

diff --git a/openapi.yaml b/openapi.yaml
index 92f7eaccc..0e7a9b4a7 100644
--- a/openapi.yaml
+++ b/openapi.yaml
@@ -2514,37 +2514,25 @@ paths:
 
   /api/assets/import:
     post:
-      operationId: importAssets
+      operationId: importPublishedAssets
       tags: [assets]
-      summary: Import assets from external URLs
-      description: "[cloud-only] Imports one or more assets from external URLs into the cloud asset store."
+      summary: "[cloud-only] Import published assets into the caller's library"
+      description: |
+        [cloud-only] Imports the specified published assets into the caller's asset library. New DB records reference the same storage objects; no file copying occurs. Assets the caller already owns (by hash) are deduplicated. The `id` field on each returned `AssetInfo` is the caller's newly-created private asset ID, not the published asset ID supplied in the request.
       x-runtime: [cloud]
       requestBody:
         required: true
         content:
           application/json:
             schema:
-              type: object
-              required:
-                - imports
-              properties:
-                imports:
-                  type: array
-                  items:
-                    $ref: "#/components/schemas/AssetImportRequest"
-                  description: Assets to import
+              $ref: "#/components/schemas/ImportPublishedAssetsRequest"
       responses:
         "200":
-          description: Import initiated
+          description: Successfully imported assets
           content:
             application/json:
               schema:
-                type: object
-                properties:
-                  assets:
-                    type: array
-                    items:
-                      $ref: "#/components/schemas/Asset"
+                $ref: "#/components/schemas/ImportPublishedAssetsResponse"
         "400":
           description: Bad request
           content:
@@ -7379,24 +7367,35 @@ components:
           type: string
           description: Target path on the runtime filesystem
 
-    AssetImportRequest:
+    ImportPublishedAssetsRequest:
       type: object
       x-runtime: [cloud]
-      description: "[cloud-only] A single asset to import from an external URL."
+      description: "[cloud-only] Request body for importing published assets into the caller's library."
       required:
-        - url
+        - published_asset_ids
       properties:
-        url:
-          type: string
-          format: uri
-          description: URL of the asset to import
-        name:
-          type: string
-          description: Display name for the imported asset
-        tags:
+        published_asset_ids:
           type: array
+          description: IDs of published assets (inputs and models) to import.
           items:
             type: string
+        share_id:
+          type: string
+          nullable: true
+          description: |
+            Optional. Share ID of the published workflow these assets belong to. When provided (non-null, non-empty): all `published_asset_ids` must belong to this share's workflow version; returns 400 if the share is not found or any asset does not belong to it. When omitted, null, or empty string: no share-scoped validation is performed and the assets are validated only against global rules (preserved for clients that have not yet adopted `share_id`).
+
+    ImportPublishedAssetsResponse:
+      type: object
+      x-runtime: [cloud]
+      description: "[cloud-only] Response after importing published assets. Each returned `AssetInfo.id` is the caller's newly-created private asset ID, not the published asset ID supplied in the request."
+      required:
+        - assets
+      properties:
+        assets:
+          type: array
+          items:
+            $ref: "#/components/schemas/AssetInfo"
 
     RemoteAssetMetadata:
       type: object

From 1668aaf0378db1fe8ddd2c0572e7312a9ebbdd41 Mon Sep 17 00:00:00 2001
From: Matt Miller <mattmiller@comfy.org>
Date: Wed, 20 May 2026 21:32:08 -0700
Subject: [PATCH 10/14] openapi: remove cloud-only job_ids query param from GET
 /api/assets (#14016)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The job_ids query parameter on GET /api/assets is tagged x-runtime:
[cloud] and only exists for cloud's variant of this endpoint. Cloud
removed all consumers and the cloud-side handler/codegen/tests in
Comfy-Org/cloud#3778. With cloud no longer accepting this parameter,
the [cloud-only] documentation here is wrong — drop it so the daily
sync to cloud/services/ingest/vendor/openapi.yaml propagates the
removal.
---
 openapi.yaml | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/openapi.yaml b/openapi.yaml
index 0e7a9b4a7..885231acc 100644
--- a/openapi.yaml
+++ b/openapi.yaml
@@ -1556,12 +1556,6 @@ paths:
             type: string
             enum: [asc, desc]
           description: Sort direction
-        - name: job_ids
-          in: query
-          schema:
-            type: string
-          x-runtime: [cloud]
-          description: "[cloud-only] Comma-separated UUIDs to filter assets by associated job."
         - name: include_public
           in: query
           schema:

From 7b7c5fed7ce978b05da27b13e26ef340d284b60e Mon Sep 17 00:00:00 2001
From: Alexis Rolland <alexisrolland@hotmail.com>
Date: Thu, 21 May 2026 14:39:30 +0800
Subject: [PATCH 11/14] Update MediaPipe nodes to standardize with existing
 code base (CORE-242) (#14025)

---
 comfy_extras/nodes_mediapipe.py               | 35 +++++++++++--------
 folder_paths.py                               |  2 +-
 .../put_detection_models_here}                |  0
 3 files changed, 22 insertions(+), 15 deletions(-)
 rename models/{mediapipe/put_mediapipe_models_here => detection/put_detection_models_here} (100%)

diff --git a/comfy_extras/nodes_mediapipe.py b/comfy_extras/nodes_mediapipe.py
index 2e67ae83f..6b7916aee 100644
--- a/comfy_extras/nodes_mediapipe.py
+++ b/comfy_extras/nodes_mediapipe.py
@@ -28,7 +28,7 @@ from comfy_extras.mediapipe.face_landmarker import FaceLandmarker
 from comfy_extras.mediapipe.face_geometry import transformation_matrix_from_detection
 
 
-FaceLandmarkerType = io.Custom("FACE_LANDMARKER")
+FaceDetectionType = io.Custom("FACE_DETECTION_MODEL")
 FaceLandmarksType = io.Custom("FACE_LANDMARKS")
 
 _CANONICAL_KEYS = ("canonical_vertices", "procrustes_indices", "procrustes_weights")
@@ -204,18 +204,19 @@ class LoadMediaPipeFaceLandmarker(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="LoadMediaPipeFaceLandmarker",
-            display_name="Load MediaPipe Face Landmarker",
+            search_aliases=["face", "facial", "mediapipe", "face landmark", "face mesh", "blazeface", "face detection"],
+            display_name="Load Face Detection Model (MediaPipe)",
             category="loaders",
             inputs=[
-                io.Combo.Input("model_name", options=folder_paths.get_filename_list("mediapipe"),
-                               tooltip="Face Landmarker safetensors from models/mediapipe/."),
+                io.Combo.Input("model_name", options=folder_paths.get_filename_list("detection"),
+                               tooltip="Face detection model from models/detection/."),
             ],
-            outputs=[FaceLandmarkerType.Output()],
+            outputs=[FaceDetectionType.Output()],
         )
 
     @classmethod
     def execute(cls, model_name) -> io.NodeOutput:
-        sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("mediapipe", model_name), safe_load=True)
+        sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("detection", model_name), safe_load=True)
         wrapper = FaceLandmarkerModel(sd)
         return io.NodeOutput(wrapper)
 
@@ -234,10 +235,12 @@ class MediaPipeFaceLandmarker(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MediaPipeFaceLandmarker",
-            display_name="MediaPipe Face Landmarker",
+            search_aliases=["face", "facial", "mediapipe", "face landmark", "face mesh", "blazeface", "face detection"],
+            display_name="Detect Face Landmarks (MediaPipe)",
             category="image/detection",
+            description="Detects facial landmarks using MediaPipe model.",
             inputs=[
-                FaceLandmarkerType.Input("face_landmarker"),
+                FaceDetectionType.Input("face_detection_model"),
                 io.Image.Input("image"),
                 io.Combo.Input("detector_variant", options=["short", "full", "both"], default="short",
                                tooltip="Face detector range. 'short' is tuned for close-up faces "
@@ -261,9 +264,9 @@ class MediaPipeFaceLandmarker(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, face_landmarker, image, detector_variant, num_faces, min_confidence,
+    def execute(cls, face_detection_model, image, detector_variant, num_faces, min_confidence,
                 missing_frame_fallback) -> io.NodeOutput:
-        canonical = face_landmarker.canonical_data
+        canonical = face_detection_model.canonical_data
         img_np = _image_to_uint8(image)
         B, H, W = img_np.shape[:3]
         chunk = 16
@@ -276,7 +279,7 @@ class MediaPipeFaceLandmarker(io.ComfyNode):
             with tqdm(total=B, desc=f"MediaPipe Face Landmarker ({variant})") as tq:
                 for i in range(0, B, chunk):
                     end = min(i + chunk, B)
-                    res.extend(face_landmarker.detect_batch(
+                    res.extend(face_detection_model.detect_batch(
                         [img_np[bi] for bi in range(i, end)],
                         num_faces=int(num_faces),
                         score_thresh=float(min_confidence),
@@ -306,7 +309,7 @@ class MediaPipeFaceLandmarker(io.ComfyNode):
                 per_bb.append({"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1, "label": "face", "score": float(f["score"])})
             bboxes.append(per_bb)
         return io.NodeOutput({"frames": frames, "image_size": (H, W),
-                              "connection_sets": face_landmarker.connection_sets}, bboxes)
+                              "connection_sets": face_detection_model.connection_sets}, bboxes)
 
 
 # Topology keys unioned by the 'all' connections preset (contour parts + irises + nose).
@@ -332,8 +335,10 @@ class MediaPipeFaceMeshVisualize(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MediaPipeFaceMeshVisualize",
-            display_name="MediaPipe Face Mesh Visualize",
+            search_aliases=["face", "facial", "mediapipe", "face landmark", "face mesh", "blazeface", "face detection", "visualize"],
+            display_name="Visualize Face Landmarks (MediaPipe)",
             category="image/detection",
+            description="Draws face landmarks mesh on the input image.",
             inputs=[
                 FaceLandmarksType.Input("face_landmarks"),
                 io.Image.Input("image", optional=True, tooltip="If not connected, a black canvas will be used."),
@@ -443,8 +448,10 @@ class MediaPipeFaceMask(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MediaPipeFaceMask",
-            display_name="MediaPipe Face Mask",
+            search_aliases=["face", "facial", "mediapipe", "face mask", "blazeface", "face detection", "visualize"],
+            display_name="Draw Face Mask (MediaPipe)",
             category="image/detection",
+            description="Draws a mask from face landmarks.",
             inputs=[
                 FaceLandmarksType.Input("face_landmarks"),
                 io.DynamicCombo.Input(
diff --git a/folder_paths.py b/folder_paths.py
index ce152eb37..36d61fcd0 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -60,7 +60,7 @@ folder_names_and_paths["geometry_estimation"] = ([os.path.join(models_dir, "geom
 
 folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flow")], supported_pt_extensions)
 
-folder_names_and_paths["mediapipe"] = ([os.path.join(models_dir, "mediapipe")], supported_pt_extensions)
+folder_names_and_paths["detection"] = ([os.path.join(models_dir, "detection")], supported_pt_extensions)
 
 output_directory = os.path.join(base_path, "output")
 temp_directory = os.path.join(base_path, "temp")
diff --git a/models/mediapipe/put_mediapipe_models_here b/models/detection/put_detection_models_here
similarity index 100%
rename from models/mediapipe/put_mediapipe_models_here
rename to models/detection/put_detection_models_here

From af3d9b60afddbe6f7c82e31ee688f7f5c9af39d0 Mon Sep 17 00:00:00 2001
From: Alexis Rolland <alexisrolland@hotmail.com>
Date: Thu, 21 May 2026 15:14:16 +0800
Subject: [PATCH 12/14] chore: Dataset nodes clean-up (CORE-237) (#14002)

---
 comfy_extras/nodes_audio.py     |   7 +-
 comfy_extras/nodes_dataset.py   | 188 ++++++++++++++++++++++----------
 comfy_extras/nodes_hunyuan3d.py |   9 +-
 comfy_extras/nodes_images.py    |   3 +-
 comfy_extras/nodes_lt_audio.py  |   8 +-
 5 files changed, 145 insertions(+), 70 deletions(-)

diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py
index 2d6b3c7ea..d5084497e 100644
--- a/comfy_extras/nodes_audio.py
+++ b/comfy_extras/nodes_audio.py
@@ -543,7 +543,7 @@ class AudioConcat(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioConcat",
             search_aliases=["join audio", "combine audio", "append audio"],
-            display_name="Audio Concat",
+            display_name="Concatenate Audio",
             description="Concatenates the audio1 to audio2 in the specified direction.",
             category="audio",
             inputs=[
@@ -597,7 +597,7 @@ class AudioMerge(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioMerge",
             search_aliases=["mix audio", "overlay audio", "layer audio"],
-            display_name="Audio Merge",
+            display_name="Merge Audio",
             description="Combine two audio tracks by overlaying their waveforms.",
             category="audio",
             inputs=[
@@ -667,8 +667,9 @@ class AudioAdjustVolume(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioAdjustVolume",
             search_aliases=["audio gain", "loudness", "audio level"],
-            display_name="Audio Adjust Volume",
+            display_name="Adjust Audio Volume",
             category="audio",
+            description="Adjust the volume of the audio by a specified amount in decibels (dB).",
             inputs=[
                 IO.Audio.Input("audio"),
                 IO.Int.Input(
diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py
index 98ed25d7e..22f5ff203 100644
--- a/comfy_extras/nodes_dataset.py
+++ b/comfy_extras/nodes_dataset.py
@@ -47,8 +47,10 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="LoadImageDataSetFromFolder",
-            display_name="Load Image Dataset from Folder",
-            category="dataset",
+            search_aliases=["load folder", "load from folder", "load dataset", "load images", "import dataset"],
+            display_name="Load Image (from Folder)",
+            category="image",
+            description="Load a dataset of images from a specified folder and return a list of images. Supported formats: PNG, JPG, JPEG, WEBP.",
             is_experimental=True,
             inputs=[
                 io.Combo.Input(
@@ -84,14 +86,16 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="LoadImageTextDataSetFromFolder",
-            display_name="Load Image and Text Dataset from Folder",
-            category="dataset",
+            search_aliases=["load folder", "load from folder", "load dataset", "load images", "import dataset"],
+            display_name="Load Image-Text (from Folder)",
+            category="image",
+            description="Load a dataset of pairs of images and text captions from a specified folder and return them as a list. Supported formats: PNG, JPG, JPEG, WEBP.",
             is_experimental=True,
             inputs=[
                 io.Combo.Input(
                     "folder",
                     options=folder_paths.get_input_subfolders(),
-                    tooltip="The folder to load images from.",
+                    tooltip="The folder to load images and text captions from.",
                 )
             ],
             outputs=[
@@ -206,8 +210,10 @@ class SaveImageDataSetToFolderNode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="SaveImageDataSetToFolder",
-            display_name="Save Image Dataset to Folder",
-            category="dataset",
+            search_aliases=["save folder", "save to folder", "save dataset", "save images", "export dataset"],
+            display_name="Save Image (to Folder) (DEPRECATED)",
+            category="image",
+            description="Save a dataset of images to a specified folder. Supported formats: PNG.",
             is_experimental=True,
             is_output_node=True,
             is_input_list=True,  # Receive images as list
@@ -226,6 +232,7 @@ class SaveImageDataSetToFolderNode(io.ComfyNode):
                 ),
             ],
             outputs=[],
+            is_deprecated=True,  # This node is redundant and superseded by existing Save Image nodes where the target folder can be specified in the filename_prefix
         )
 
     @classmethod
@@ -246,14 +253,20 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="SaveImageTextDataSetToFolder",
-            display_name="Save Image and Text Dataset to Folder",
-            category="dataset",
+            search_aliases=["save folder", "save to folder", "save dataset", "save images", "save text", "export dataset"],
+            display_name="Save Image-Text (to Folder)",
+            category="image",
+            description="Save a dataset of pairs of images and text captions to a specified folder. Images are saved as PNG files and captions are saved as TXT files with the same filename_prefix.",
             is_experimental=True,
             is_output_node=True,
             is_input_list=True,  # Receive both images and texts as lists
             inputs=[
                 io.Image.Input("images", tooltip="List of images to save."),
-                io.String.Input("texts", tooltip="List of text captions to save."),
+                io.String.Input("texts",
+                    optional=True,
+                    force_input=True,
+                    tooltip="List of text captions to save."
+                ),
                 io.String.Input(
                     "folder_name",
                     default="dataset",
@@ -270,7 +283,7 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, images, texts, folder_name, filename_prefix):
+    def execute(cls, images, folder_name, filename_prefix, texts=None):
         # Extract scalar values
         folder_name = folder_name[0]
         filename_prefix = filename_prefix[0]
@@ -279,11 +292,12 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode):
         saved_files = save_images_to_folder(images, output_dir, filename_prefix)
 
         # Save captions
-        for idx, (filename, caption) in enumerate(zip(saved_files, texts)):
-            caption_filename = filename.replace(".png", ".txt")
-            caption_path = os.path.join(output_dir, caption_filename)
-            with open(caption_path, "w", encoding="utf-8") as f:
-                f.write(caption)
+        if texts:
+            for idx, (filename, caption) in enumerate(zip(saved_files, texts)):
+                caption_filename = filename.replace(".png", ".txt")
+                caption_path = os.path.join(output_dir, caption_filename)
+                with open(caption_path, "w", encoding="utf-8") as f:
+                    f.write(caption)
 
         logging.info(f"Saved {len(saved_files)} images and captions to {output_dir}.")
         return io.NodeOutput()
@@ -314,11 +328,13 @@ class ImageProcessingNode(io.ComfyNode):
 
     Child classes should set:
         node_id: Unique node identifier (required)
+        search_aliases: List of search aliases (optional)
         display_name: Display name (optional, defaults to node_id)
         description: Node description (optional)
         extra_inputs: List of additional io.Input objects beyond "images" (optional)
         is_group_process: None (auto-detect), True (group), or False (individual) (optional)
         is_output_list: True (list output) or False (single output) (optional, default True)
+        is_deprecated: True if the node is deprecated (optional, default False)
 
     Child classes must implement ONE of:
         _process(cls, image, **kwargs) -> tensor  (for single-item processing)
@@ -326,12 +342,13 @@ class ImageProcessingNode(io.ComfyNode):
     """
 
     node_id = None
+    search_aliases = []
     display_name = None
     description = None
     extra_inputs = []
     is_group_process = None  # None = auto-detect, True/False = explicit
     is_output_list = None  # None = auto-detect based on processing mode
-
+    is_deprecated = False
     @classmethod
     def _detect_processing_mode(cls):
         """Detect whether this node uses group or individual processing.
@@ -402,8 +419,10 @@ class ImageProcessingNode(io.ComfyNode):
 
         return io.Schema(
             node_id=cls.node_id,
+            search_aliases=cls.search_aliases,
             display_name=cls.display_name or cls.node_id,
-            category="dataset/image",
+            category=cls.category,
+            description=cls.description,
             is_experimental=True,
             is_input_list=is_group,  # True for group, False for individual
             inputs=inputs,
@@ -472,11 +491,13 @@ class TextProcessingNode(io.ComfyNode):
 
     Child classes should set:
         node_id: Unique node identifier (required)
+        search_aliases: List of search aliases (optional)
         display_name: Display name (optional, defaults to node_id)
         description: Node description (optional)
         extra_inputs: List of additional io.Input objects beyond "texts" (optional)
         is_group_process: None (auto-detect), True (group), or False (individual) (optional)
         is_output_list: True (list output) or False (single output) (optional, default True)
+        is_deprecated: True if the node is deprecated (optional, default False)
 
     Child classes must implement ONE of:
         _process(cls, text, **kwargs) -> str  (for single-item processing)
@@ -484,12 +505,13 @@ class TextProcessingNode(io.ComfyNode):
     """
 
     node_id = None
+    search_aliases = []
     display_name = None
     description = None
     extra_inputs = []
     is_group_process = None  # None = auto-detect, True/False = explicit
     is_output_list = None  # None = auto-detect based on processing mode
-
+    is_deprecated = False
     @classmethod
     def _detect_processing_mode(cls):
         """Detect whether this node uses group or individual processing.
@@ -627,15 +649,17 @@ class TextProcessingNode(io.ComfyNode):
 
 class ResizeImagesByShorterEdgeNode(ImageProcessingNode):
     node_id = "ResizeImagesByShorterEdge"
-    display_name = "Resize Images by Shorter Edge"
-    description = "Resize images so that the shorter edge matches the specified length while preserving aspect ratio."
+    display_name = "Resize Images by Shorter Edge (DEPRECATED)"
+    category = "image/transform"
+    description = "Resize images so that the shorter edge matches the specified dimension while preserving aspect ratio."
+    is_deprecated = True  # This node is superseded by Resize Image/Mask with resize_type = scale shorter dimension
     extra_inputs = [
         io.Int.Input(
             "shorter_edge",
             default=512,
             min=1,
             max=8192,
-            tooltip="Target length for the shorter edge.",
+            tooltip="Target dimension for the shorter edge.",
         ),
     ]
 
@@ -655,15 +679,17 @@ class ResizeImagesByShorterEdgeNode(ImageProcessingNode):
 
 class ResizeImagesByLongerEdgeNode(ImageProcessingNode):
     node_id = "ResizeImagesByLongerEdge"
-    display_name = "Resize Images by Longer Edge"
-    description = "Resize images so that the longer edge matches the specified length while preserving aspect ratio."
+    display_name = "Resize Images by Longer Edge (DEPRECATED)"
+    category = "image/transform"
+    description = "Resize images so that the longer edge matches the specified dimension while preserving aspect ratio."
+    is_deprecated = True  # This node is superseded by Resize Image/Mask with resize_type = scale longer dimension
     extra_inputs = [
         io.Int.Input(
             "longer_edge",
             default=1024,
             min=1,
             max=8192,
-            tooltip="Target length for the longer edge.",
+            tooltip="Target dimension for the longer edge.",
         ),
     ]
 
@@ -686,8 +712,10 @@ class ResizeImagesByLongerEdgeNode(ImageProcessingNode):
 
 class CenterCropImagesNode(ImageProcessingNode):
     node_id = "CenterCropImages"
-    display_name = "Center Crop Images"
-    description = "Center crop all images to the specified dimensions."
+    search_aliases=["crop", "cut", "trim"]
+    display_name="Crop Image (Center)"
+    category="image/transform"
+    description = "Center crop an image to the specified dimensions."
     extra_inputs = [
         io.Int.Input("width", default=512, min=1, max=8192, tooltip="Crop width."),
         io.Int.Input("height", default=512, min=1, max=8192, tooltip="Crop height."),
@@ -706,10 +734,11 @@ class CenterCropImagesNode(ImageProcessingNode):
 
 class RandomCropImagesNode(ImageProcessingNode):
     node_id = "RandomCropImages"
-    display_name = "Random Crop Images"
-    description = (
-        "Randomly crop all images to the specified dimensions (for data augmentation)."
-    )
+    search_aliases=["crop", "cut", "trim"]
+    display_name = "Crop Image (Random)"
+    category="image/transform"
+    description = "Randomly crop an image to the specified dimensions."
+
     extra_inputs = [
         io.Int.Input("width", default=512, min=1, max=8192, tooltip="Crop width."),
         io.Int.Input("height", default=512, min=1, max=8192, tooltip="Crop height."),
@@ -734,7 +763,9 @@ class RandomCropImagesNode(ImageProcessingNode):
 
 class NormalizeImagesNode(ImageProcessingNode):
     node_id = "NormalizeImages"
-    display_name = "Normalize Images"
+    search_aliases=["normalize", "normalize colors"]
+    display_name = "Normalize Image Colors"
+    category = "image/color"
     description = "Normalize images using mean and standard deviation."
     extra_inputs = [
         io.Float.Input(
@@ -762,8 +793,10 @@ class NormalizeImagesNode(ImageProcessingNode):
 
 class AdjustBrightnessNode(ImageProcessingNode):
     node_id = "AdjustBrightness"
+    search_aliases=["brightness"]
     display_name = "Adjust Brightness"
-    description = "Adjust brightness of all images."
+    category="image/adjustments"
+    description = "Adjust the brightness of an image."
     extra_inputs = [
         io.Float.Input(
             "factor",
@@ -781,8 +814,10 @@ class AdjustBrightnessNode(ImageProcessingNode):
 
 class AdjustContrastNode(ImageProcessingNode):
     node_id = "AdjustContrast"
+    search_aliases=["contrast"]
     display_name = "Adjust Contrast"
-    description = "Adjust contrast of all images."
+    category="image/adjustments"
+    description = "Adjust the contrast of an image."
     extra_inputs = [
         io.Float.Input(
             "factor",
@@ -800,8 +835,10 @@ class AdjustContrastNode(ImageProcessingNode):
 
 class ShuffleDatasetNode(ImageProcessingNode):
     node_id = "ShuffleDataset"
-    display_name = "Shuffle Image Dataset"
-    description = "Randomly shuffle the order of images in the dataset."
+    search_aliases=["shuffle", "randomize", "mix"]
+    display_name = "Shuffle Images List"
+    category = "image/batch"
+    description = "Randomly shuffle the order of images in a list."
     is_group_process = True  # Requires full list to shuffle
     extra_inputs = [
         io.Int.Input(
@@ -823,13 +860,15 @@ class ShuffleImageTextDatasetNode(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="ShuffleImageTextDataset",
-            display_name="Shuffle Image-Text Dataset",
-            category="dataset/image",
+            search_aliases=["shuffle", "randomize", "mix"],
+            display_name = "Shuffle Pairs of Image-Text",
+            category = "image/batch",
+            description = "Randomly shuffle the order of pairs of image-text in a list.",
             is_experimental=True,
             is_input_list=True,
             inputs=[
                 io.Image.Input("images", tooltip="List of images to shuffle."),
-                io.String.Input("texts", tooltip="List of texts to shuffle."),
+                io.String.Input("texts", tooltip="List of texts to shuffle.", force_input=True),
                 io.Int.Input(
                     "seed",
                     default=0,
@@ -865,8 +904,11 @@ class ShuffleImageTextDatasetNode(io.ComfyNode):
 
 class TextToLowercaseNode(TextProcessingNode):
     node_id = "TextToLowercase"
-    display_name = "Text to Lowercase"
-    description = "Convert all texts to lowercase."
+    search_aliases=["lowercase"]
+    display_name = "Convert Text to Lowercase (DEPRECATED)"
+    category = "text"
+    description = "Convert text to lowercase."
+    is_deprecated = True  # This node is superseded by the Convert Text Case node
 
     @classmethod
     def _process(cls, text):
@@ -875,8 +917,11 @@ class TextToLowercaseNode(TextProcessingNode):
 
 class TextToUppercaseNode(TextProcessingNode):
     node_id = "TextToUppercase"
-    display_name = "Text to Uppercase"
-    description = "Convert all texts to uppercase."
+    search_aliases=["uppercase"]
+    display_name = "Convert Text to Uppercase (DEPRECATED)"
+    category = "text"
+    description = "Convert text to uppercase."
+    is_deprecated = True  # This node is superseded by the Convert Text Case node
 
     @classmethod
     def _process(cls, text):
@@ -885,8 +930,10 @@ class TextToUppercaseNode(TextProcessingNode):
 
 class TruncateTextNode(TextProcessingNode):
     node_id = "TruncateText"
+    search_aliases=["truncate", "cut", "shorten"]
     display_name = "Truncate Text"
-    description = "Truncate all texts to a maximum length."
+    category = "text"
+    description = "Truncate text to a maximum length."
     extra_inputs = [
         io.Int.Input(
             "max_length", default=77, min=1, max=10000, tooltip="Maximum text length."
@@ -900,8 +947,10 @@ class TruncateTextNode(TextProcessingNode):
 
 class AddTextPrefixNode(TextProcessingNode):
     node_id = "AddTextPrefix"
-    display_name = "Add Text Prefix"
+    display_name = "Add Text Prefix (DEPRECATED)"
+    category = "text"
     description = "Add a prefix to all texts."
+    is_deprecated = True  # This node is superseded by the Concatenate Text node
     extra_inputs = [
         io.String.Input("prefix", default="", tooltip="Prefix to add."),
     ]
@@ -913,8 +962,10 @@ class AddTextPrefixNode(TextProcessingNode):
 
 class AddTextSuffixNode(TextProcessingNode):
     node_id = "AddTextSuffix"
-    display_name = "Add Text Suffix"
+    display_name = "Add Text Suffix (DEPRECATED)"
+    category = "text"
     description = "Add a suffix to all texts."
+    is_deprecated = True  # This node is superseded by the Concatenate Text node
     extra_inputs = [
         io.String.Input("suffix", default="", tooltip="Suffix to add."),
     ]
@@ -926,8 +977,10 @@ class AddTextSuffixNode(TextProcessingNode):
 
 class ReplaceTextNode(TextProcessingNode):
     node_id = "ReplaceText"
-    display_name = "Replace Text"
+    display_name = "Replace Text (DEPRECATED)"
+    category = "text"
     description = "Replace text in all texts."
+    is_deprecated = True  # This node is superseded by the other Replace Text node
     extra_inputs = [
         io.String.Input("find", default="", tooltip="Text to find."),
         io.String.Input("replace", default="", tooltip="Text to replace with."),
@@ -940,8 +993,10 @@ class ReplaceTextNode(TextProcessingNode):
 
 class StripWhitespaceNode(TextProcessingNode):
     node_id = "StripWhitespace"
-    display_name = "Strip Whitespace"
+    display_name = "Strip Whitespace (DEPRECATED)"
+    category = "text"
     description = "Strip leading and trailing whitespace from all texts."
+    is_deprecated = True  # This node is superseded by the Trim Text node
 
     @classmethod
     def _process(cls, text):
@@ -952,11 +1007,13 @@ class StripWhitespaceNode(TextProcessingNode):
 
 
 class ImageDeduplicationNode(ImageProcessingNode):
-    """Remove duplicate or very similar images from the dataset using perceptual hashing."""
+    """Remove duplicate or very similar images from a list using perceptual hashing."""
 
     node_id = "ImageDeduplication"
-    display_name = "Image Deduplication"
-    description = "Remove duplicate or very similar images from the dataset."
+    search_aliases=["deduplicate", "remove duplicates", "similarity filter"]
+    display_name = "Deduplicate Images"
+    category = "image/batch"
+    description = "Remove duplicate or very similar images from a list."
     is_group_process = True  # Requires full list to compare images
     extra_inputs = [
         io.Float.Input(
@@ -1026,7 +1083,9 @@ class ImageGridNode(ImageProcessingNode):
     """Combine multiple images into a single grid/collage."""
 
     node_id = "ImageGrid"
-    display_name = "Image Grid"
+    search_aliases=["grid", "collage", "combine"]
+    display_name = "Make Image Grid"
+    category="image/batch"
     description = "Arrange multiple images into a grid layout."
     is_group_process = True  # Requires full list to create grid
     is_output_list = False  # Outputs single grid image
@@ -1102,9 +1161,12 @@ class MergeImageListsNode(ImageProcessingNode):
     """Merge multiple image lists into a single list."""
 
     node_id = "MergeImageLists"
-    display_name = "Merge Image Lists"
+    search_aliases=["list", "merge list", "make list"]
+    display_name = "Merge Image Lists (DEPRECATED)"
+    category = "image/batch"
     description = "Concatenate multiple image lists into one."
     is_group_process = True  # Receives images as list
+    is_deprecated = True  # This node is superseded by the Create List node
 
     @classmethod
     def _group_process(cls, images):
@@ -1119,9 +1181,11 @@ class MergeTextListsNode(TextProcessingNode):
     """Merge multiple text lists into a single list."""
 
     node_id = "MergeTextLists"
-    display_name = "Merge Text Lists"
+    display_name = "Merge Text Lists (DEPRECATED)"
+    category = "text"
     description = "Concatenate multiple text lists into one."
     is_group_process = True  # Receives texts as list
+    is_deprecated = True  # This node is superseded by the Create List node
 
     @classmethod
     def _group_process(cls, texts):
@@ -1142,8 +1206,10 @@ class ResolutionBucket(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="ResolutionBucket",
+            search_aliases=["bucket by resolution", "group by resolution", "batch by resolution"],
             display_name="Resolution Bucket",
-            category="dataset",
+            category="training",
+            description="Group latents and conditionings into buckets",
             is_experimental=True,
             is_input_list=True,
             inputs=[
@@ -1236,7 +1302,8 @@ class MakeTrainingDataset(io.ComfyNode):
             node_id="MakeTrainingDataset",
             search_aliases=["encode dataset"],
             display_name="Make Training Dataset",
-            category="dataset",
+            category="training",
+            description="Encode images with VAE and texts with CLIP to create a training dataset of latents and conditionings.",
             is_experimental=True,
             is_input_list=True,  # images and texts as lists
             inputs=[
@@ -1251,6 +1318,7 @@ class MakeTrainingDataset(io.ComfyNode):
                     "texts",
                     optional=True,
                     tooltip="List of text captions. Can be length n (matching images), 1 (repeated for all), or omitted (uses empty string).",
+                    force_input=True
                 ),
             ],
             outputs=[
@@ -1320,9 +1388,10 @@ class SaveTrainingDataset(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="SaveTrainingDataset",
-            search_aliases=["export training data"],
+            search_aliases=["export dataset", "save dataset"],
             display_name="Save Training Dataset",
-            category="dataset",
+            category="training",
+            description="Save encoded training dataset (latents + conditioning) to disk for efficient loading during training.",
             is_experimental=True,
             is_output_node=True,
             is_input_list=True,  # Receive lists
@@ -1424,7 +1493,8 @@ class LoadTrainingDataset(io.ComfyNode):
             node_id="LoadTrainingDataset",
             search_aliases=["import dataset", "training data"],
             display_name="Load Training Dataset",
-            category="dataset",
+            category="training",
+            description="Load encoded training dataset (latents + conditioning) from disk for use in training.",
             is_experimental=True,
             inputs=[
                 io.String.Input(
diff --git a/comfy_extras/nodes_hunyuan3d.py b/comfy_extras/nodes_hunyuan3d.py
index 403eb855b..bcd3f9198 100644
--- a/comfy_extras/nodes_hunyuan3d.py
+++ b/comfy_extras/nodes_hunyuan3d.py
@@ -419,15 +419,17 @@ class VoxelToMeshBasic(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VoxelToMeshBasic",
-            display_name="Voxel to Mesh (Basic)",
+            display_name="Voxel to Mesh (Basic) (DEPRECATED)",
             category="3d",
+            description="Converts a voxel grid to a mesh.",
+            is_deprecated=True, # This node is superseded by the Voxel To Mesh node
             inputs=[
                 IO.Voxel.Input("voxel"),
                 IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
             ],
             outputs=[
                 IO.Mesh.Output(),
-            ]
+            ],
         )
 
     @classmethod
@@ -453,9 +455,10 @@ class VoxelToMesh(IO.ComfyNode):
             node_id="VoxelToMesh",
             display_name="Voxel to Mesh",
             category="3d",
+            description="Converts a voxel grid to a mesh.",
             inputs=[
                 IO.Voxel.Input("voxel"),
-                IO.Combo.Input("algorithm", options=["surface net", "basic"], advanced=True),
+                IO.Combo.Input("algorithm", options=["surface net", "basic"]),
                 IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
             ],
             outputs=[
diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index 6326c5be8..4856346d7 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -55,9 +55,10 @@ class ImageCropV2(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="ImageCropV2",
-            search_aliases=["trim"],
+            search_aliases=["crop", "cut", "trim"],
             display_name="Crop Image",
             category="image/transform",
+            description = "Crop an image to the specified dimensions.",
             essentials_category="Image Tools",
             has_intermediate_output=True,
             inputs=[
diff --git a/comfy_extras/nodes_lt_audio.py b/comfy_extras/nodes_lt_audio.py
index 2c1f63afb..51ddf584a 100644
--- a/comfy_extras/nodes_lt_audio.py
+++ b/comfy_extras/nodes_lt_audio.py
@@ -11,8 +11,8 @@ class LTXVAudioVAELoader(io.ComfyNode):
     def define_schema(cls) -> io.Schema:
         return io.Schema(
             node_id="LTXVAudioVAELoader",
-            display_name="LTXV Audio VAE Loader",
-            category="audio",
+            display_name="Load LTXV Audio VAE",
+            category="loaders",
             inputs=[
                 io.Combo.Input(
                     "ckpt_name",
@@ -40,7 +40,7 @@ class LTXVAudioVAEEncode(VAEEncodeAudio):
         return io.Schema(
             node_id="LTXVAudioVAEEncode",
             display_name="LTXV Audio VAE Encode",
-            category="audio",
+            category="latent/audio",
             inputs=[
                 io.Audio.Input("audio", tooltip="The audio to be encoded."),
                 io.Vae.Input(
@@ -63,7 +63,7 @@ class LTXVAudioVAEDecode(io.ComfyNode):
         return io.Schema(
             node_id="LTXVAudioVAEDecode",
             display_name="LTXV Audio VAE Decode",
-            category="audio",
+            category="latent/audio",
             inputs=[
                 io.Latent.Input("samples", tooltip="The latent to be decoded."),
                 io.Vae.Input(

From 4259a0c7c3b805e3dd1f178e603e6d725780583a Mon Sep 17 00:00:00 2001
From: Alexis Rolland <alexisrolland@hotmail.com>
Date: Thu, 21 May 2026 16:50:09 +0800
Subject: [PATCH 13/14] Update MoGe nodes display names, search aliases and
 descriptions (#14030)

---
 comfy_extras/nodes_moge.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/comfy_extras/nodes_moge.py b/comfy_extras/nodes_moge.py
index d9a08ebc7..3508781a0 100644
--- a/comfy_extras/nodes_moge.py
+++ b/comfy_extras/nodes_moge.py
@@ -103,8 +103,10 @@ class MoGePanoramaInference(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MoGePanoramaInference",
-            display_name="MoGe Panorama Inference",
+            search_aliases=["moge", "panorama", "depth", "geometry", "depth estimation", "geometry estimation"],
+            display_name="Run MoGe Panorama Inference",
             category="image/geometry_estimation",
+            description="Run MoGe on an equirectangular panorama by splitting it into 12 perspective views, running inference on each, and merging the results into a single depth map.",
             inputs=[
                 MoGeModelType.Input("moge_model"),
                 io.Image.Input("image", tooltip="Equirectangular panorama (any aspect)."),
@@ -222,7 +224,9 @@ class MoGeInference(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MoGeInference",
-            display_name="MoGe Inference",
+            search_aliases=["moge", "depth", "geometry", "depth estimation", "geometry estimation"],
+            display_name="Run MoGe Inference",
+            description="Run MoGe on a single image to estimate depth and geometry.",
             category="image/geometry_estimation",
             inputs=[
                 MoGeModelType.Input("moge_model"),
@@ -277,7 +281,9 @@ class MoGeRender(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MoGeRender",
-            display_name="MoGe Render",
+            search_aliases=["moge", "render", "geometry", "depth", "normal"],
+            display_name="Render MoGe Geometry",
+            description="Render a depth map or normal map from geometry data",
             category="image/geometry_estimation",
             inputs=[
                 MoGeGeometry.Input("moge_geometry"),
@@ -342,7 +348,9 @@ class MoGePointMapToMesh(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="MoGePointMapToMesh",
-            display_name="MoGe Point Map to Mesh",
+            search_aliases=["moge", "mesh", "geometry", "point map"],
+            display_name="Convert MoGe Point Map to Mesh",
+            description="Convert a MoGe point map into a 3D mesh.",
             category="image/geometry_estimation",
             inputs=[
                 MoGeGeometry.Input("moge_geometry"),

From 7c685bc9c351f9631ac26ded994fab08aaf8481d Mon Sep 17 00:00:00 2001
From: "Yousef R. Gamaleldin" <81116377+yousef-rafat@users.noreply.github.com>
Date: Thu, 21 May 2026 17:01:38 +0300
Subject: [PATCH 14/14] Add Trellis 2 Support (CORE-7) (#12183)

---
 comfy/clip_vision.py                     |    4 +
 comfy/image_encoders/dino3.py            |  285 +++++
 comfy/image_encoders/dino3_large.json    |   23 +
 comfy/latent_formats.py                  |    2 +
 comfy/ldm/trellis2/attention.py          |  282 +++++
 comfy/ldm/trellis2/flexgemm.py           |  298 +++++
 comfy/ldm/trellis2/model.py              |  935 ++++++++++++++
 comfy/ldm/trellis2/vae.py                | 1446 ++++++++++++++++++++++
 comfy/model_base.py                      |   12 +-
 comfy/model_detection.py                 |   24 +
 comfy/sd.py                              |   13 +
 comfy/supported_models.py                |   28 +-
 comfy_api/latest/_util/geometry_types.py |    5 +-
 comfy_extras/nodes_mesh_postprocess.py   |  845 +++++++++++++
 comfy_extras/nodes_save_3d.py            |   34 +-
 comfy_extras/nodes_trellis2.py           |  667 ++++++++++
 nodes.py                                 |    6 +
 17 files changed, 4891 insertions(+), 18 deletions(-)
 create mode 100644 comfy/image_encoders/dino3.py
 create mode 100644 comfy/image_encoders/dino3_large.json
 create mode 100644 comfy/ldm/trellis2/attention.py
 create mode 100644 comfy/ldm/trellis2/flexgemm.py
 create mode 100644 comfy/ldm/trellis2/model.py
 create mode 100644 comfy/ldm/trellis2/vae.py
 create mode 100644 comfy_extras/nodes_mesh_postprocess.py
 create mode 100644 comfy_extras/nodes_trellis2.py

diff --git a/comfy/clip_vision.py b/comfy/clip_vision.py
index 1691fca81..71f2200b7 100644
--- a/comfy/clip_vision.py
+++ b/comfy/clip_vision.py
@@ -9,6 +9,7 @@ import comfy.model_management
 import comfy.utils
 import comfy.clip_model
 import comfy.image_encoders.dino2
+import comfy.image_encoders.dino3
 
 class Output:
     def __getitem__(self, key):
@@ -23,6 +24,7 @@ IMAGE_ENCODERS = {
     "siglip_vision_model": comfy.clip_model.CLIPVisionModelProjection,
     "siglip2_vision_model": comfy.clip_model.CLIPVisionModelProjection,
     "dinov2": comfy.image_encoders.dino2.Dinov2Model,
+    "dinov3": comfy.image_encoders.dino3.DINOv3ViTModel
 }
 
 class ClipVisionModel():
@@ -134,6 +136,8 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
     elif 'encoder.layer.23.layer_scale2.lambda1' in sd:
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
+    elif 'layer.9.attention.o_proj.bias' in sd: # dinov3
+        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino3_large.json")
     else:
         return None
 
diff --git a/comfy/image_encoders/dino3.py b/comfy/image_encoders/dino3.py
new file mode 100644
index 000000000..145bd5490
--- /dev/null
+++ b/comfy/image_encoders/dino3.py
@@ -0,0 +1,285 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import comfy.model_management
+from comfy.ldm.modules.attention import optimized_attention_for_device
+from comfy.image_encoders.dino2 import LayerScale as DINOv3ViTLayerScale
+
+class DINOv3ViTMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, mlp_bias, device, dtype, operations):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.up_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.down_proj = operations.Linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.act_fn = torch.nn.GELU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+def apply_rotary_pos_emb(q, k, cos, sin, **kwargs):
+    num_tokens = q.shape[-2]
+    num_patches = sin.shape[-2]
+    num_prefix_tokens = num_tokens - num_patches
+
+    q_prefix_tokens, q_patches = q.split((num_prefix_tokens, num_patches), dim=-2)
+    k_prefix_tokens, k_patches = k.split((num_prefix_tokens, num_patches), dim=-2)
+
+    q_patches = (q_patches * cos) + (rotate_half(q_patches) * sin)
+    k_patches = (k_patches * cos) + (rotate_half(k_patches) * sin)
+
+    q = torch.cat((q_prefix_tokens, q_patches), dim=-2)
+    k = torch.cat((k_prefix_tokens, k_patches), dim=-2)
+
+    return q, k
+
+class DINOv3ViTAttention(nn.Module):
+    def __init__(self, hidden_size, num_attention_heads, device, dtype, operations):
+        super().__init__()
+        self.embed_dim = hidden_size
+        self.num_heads = num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+
+        self.k_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=False, device=device, dtype=dtype) # key_bias = False
+        self.v_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+
+        self.q_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+        self.o_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+
+        batch_size, patches, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is not None:
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        attn = optimized_attention_for_device(query_states.device, mask=False)
+
+        attn_output = attn(
+            query_states, key_states, value_states, self.num_heads, attention_mask, skip_reshape=True, skip_output_reshape=True
+        )
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(batch_size, patches, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output
+
+class DINOv3ViTGatedMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, mlp_bias, device, dtype, operations):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.up_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.down_proj = operations.Linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.act_fn = torch.nn.GELU()
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+def get_patches_center_coordinates(
+    num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor:
+
+    coords_h = torch.arange(0.5, num_patches_h, dtype=dtype, device=device)
+    coords_w = torch.arange(0.5, num_patches_w, dtype=dtype, device=device)
+    coords_h = coords_h / num_patches_h
+    coords_w = coords_w / num_patches_w
+    coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)
+    coords = coords.flatten(0, 1)
+    coords = 2.0 * coords - 1.0
+    return coords
+
+class DINOv3ViTRopePositionEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, rope_theta, hidden_size, num_attention_heads, image_size, patch_size, device, dtype):
+        super().__init__()
+        self.base = rope_theta
+        self.head_dim = hidden_size // num_attention_heads
+        self.num_patches_h = image_size // patch_size
+        self.num_patches_w = image_size // patch_size
+        self.patch_size = patch_size
+
+        inv_freq = 1 / self.base ** torch.arange(0, 1, 4 / self.head_dim, dtype=torch.float32, device=device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        _, _, height, width = pixel_values.shape
+        num_patches_h = height // self.patch_size
+        num_patches_w = width // self.patch_size
+
+        device = pixel_values.device
+        device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"
+        with torch.amp.autocast(device_type = device_type, enabled=False):
+            patch_coords = get_patches_center_coordinates(
+                num_patches_h, num_patches_w, dtype=torch.float32, device=device
+            )
+
+            self.inv_freq = self.inv_freq.to(device)
+            angles = 2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :]
+            angles = angles.flatten(1, 2)
+            angles = angles.tile(2)
+
+            cos = torch.cos(angles)
+            sin = torch.sin(angles)
+
+        dtype = pixel_values.dtype
+        return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+class DINOv3ViTEmbeddings(nn.Module):
+    def __init__(self, hidden_size, num_register_tokens, num_channels, patch_size, dtype, device, operations):
+        super().__init__()
+        self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_size, device=device, dtype=dtype))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, hidden_size, device=device, dtype=dtype))
+        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, hidden_size, device=device, dtype=dtype))
+        self.patch_embeddings = operations.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
+        )
+
+    def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None):
+        batch_size = pixel_values.shape[0]
+        target_dtype = self.patch_embeddings.weight.dtype
+
+        patch_embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
+        patch_embeddings = patch_embeddings.flatten(2).transpose(1, 2)
+
+        if bool_masked_pos is not None:
+            mask_token = self.mask_token.to(patch_embeddings.dtype)
+            patch_embeddings = torch.where(bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings)
+
+        cls_token = self.cls_token.expand(batch_size, -1, -1)
+        register_tokens = self.register_tokens.expand(batch_size, -1, -1)
+        device = patch_embeddings.device
+        cls_token = cls_token.to(device)
+        register_tokens = register_tokens.to(device)
+        embeddings = torch.cat([cls_token, register_tokens, patch_embeddings], dim=1)
+
+        return embeddings
+
+class DINOv3ViTLayer(nn.Module):
+
+    def __init__(self, hidden_size, layer_norm_eps, use_gated_mlp, mlp_bias, intermediate_size, num_attention_heads,
+                 device, dtype, operations):
+        super().__init__()
+
+        self.norm1 = operations.LayerNorm(hidden_size, eps=layer_norm_eps, device=device, dtype=dtype)
+        self.attention = DINOv3ViTAttention(hidden_size, num_attention_heads, device=device, dtype=dtype, operations=operations)
+        self.layer_scale1 = DINOv3ViTLayerScale(hidden_size, device=device, dtype=dtype, operations=None)
+
+        self.norm2 = operations.LayerNorm(hidden_size, eps=layer_norm_eps, device=device, dtype=dtype)
+
+        if use_gated_mlp:
+            self.mlp = DINOv3ViTGatedMLP(hidden_size, intermediate_size, mlp_bias, device=device, dtype=dtype, operations=operations)
+        else:
+            self.mlp = DINOv3ViTMLP(hidden_size, intermediate_size=intermediate_size, mlp_bias=mlp_bias, device=device, dtype=dtype, operations=operations)
+        self.layer_scale2 = DINOv3ViTLayerScale(hidden_size, device=device, dtype=dtype, operations=None)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+        )
+        hidden_states = self.layer_scale1(hidden_states)
+        hidden_states = hidden_states + residual
+
+        residual = hidden_states
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.layer_scale2(hidden_states)
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class DINOv3ViTModel(nn.Module):
+    def __init__(self, config, dtype, device, operations):
+        super().__init__()
+        use_bf16 = comfy.model_management.should_use_bf16(device, prioritize_performance=True)
+        if dtype == torch.float16 and use_bf16:
+            dtype = torch.bfloat16
+        elif dtype == torch.float16 and not use_bf16:
+            dtype = torch.float32
+        num_hidden_layers = config["num_hidden_layers"]
+        hidden_size = config["hidden_size"]
+        num_attention_heads = config["num_attention_heads"]
+        num_register_tokens = config["num_register_tokens"]
+        intermediate_size = config["intermediate_size"]
+        layer_norm_eps = config["layer_norm_eps"]
+        num_channels = config["num_channels"]
+        patch_size = config["patch_size"]
+        rope_theta = config["rope_theta"]
+
+        self.embeddings = DINOv3ViTEmbeddings(
+            hidden_size, num_register_tokens, num_channels=num_channels, patch_size=patch_size, dtype=dtype, device=device, operations=operations
+        )
+        self.rope_embeddings = DINOv3ViTRopePositionEmbedding(
+            rope_theta, hidden_size, num_attention_heads, image_size=512, patch_size=patch_size, dtype=dtype, device=device
+        )
+        self.layer = nn.ModuleList(
+            [DINOv3ViTLayer(hidden_size, layer_norm_eps, use_gated_mlp=False, mlp_bias=True,
+                            intermediate_size=intermediate_size,num_attention_heads = num_attention_heads,
+                            dtype=dtype, device=device, operations=operations)
+            for _ in range(num_hidden_layers)])
+        self.norm = operations.LayerNorm(hidden_size, eps=layer_norm_eps, dtype=dtype, device=device)
+
+    def get_input_embeddings(self):
+        return self.embeddings.patch_embeddings
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        bool_masked_pos: torch.Tensor | None = None,
+        **kwargs,
+    ):
+
+        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
+        position_embeddings = self.rope_embeddings(pixel_values)
+
+        for i, layer_module in enumerate(self.layer):
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        if kwargs.get("skip_norm_elementwise", False):
+            sequence_output= F.layer_norm(hidden_states, hidden_states.shape[-1:])
+        else:
+            norm = self.norm.to(hidden_states.device)
+            sequence_output = norm(hidden_states)
+        pooled_output = sequence_output[:, 0, :]
+
+        return sequence_output, None, pooled_output, None
diff --git a/comfy/image_encoders/dino3_large.json b/comfy/image_encoders/dino3_large.json
new file mode 100644
index 000000000..b37b61dc8
--- /dev/null
+++ b/comfy/image_encoders/dino3_large.json
@@ -0,0 +1,23 @@
+{
+  "model_type": "dinov3",
+  "hidden_size": 1024,
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "key_bias": false,
+  "layer_norm_eps": 1e-05,
+  "mlp_bias": true,
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "num_register_tokens": 4,
+  "patch_size": 16,
+  "pos_embed_rescale": 2.0,
+  "proj_bias": true,
+  "query_bias": true,
+  "rope_theta": 100.0,
+  "use_gated_mlp": false,
+  "value_bias": true,
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_std": [0.229, 0.224, 0.225]
+}
diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py
index 75d459b59..3b7717915 100644
--- a/comfy/latent_formats.py
+++ b/comfy/latent_formats.py
@@ -760,6 +760,8 @@ class Hunyuan3Dv2_1(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1
 
+class Trellis2(LatentFormat): # TODO
+    latent_channels = 32
 class Hunyuan3Dv2mini(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1
diff --git a/comfy/ldm/trellis2/attention.py b/comfy/ldm/trellis2/attention.py
new file mode 100644
index 000000000..d95b071b5
--- /dev/null
+++ b/comfy/ldm/trellis2/attention.py
@@ -0,0 +1,282 @@
+import torch
+import math
+from comfy.ldm.modules.attention import optimized_attention
+from typing import Tuple, Union, List
+from comfy.ldm.trellis2.vae import VarLenTensor
+import comfy.ops
+
+
+# replica of the seedvr2 code
+def var_attn_arg(kwargs):
+    cu_seqlens_q = kwargs.get("cu_seqlens_q", None)
+    max_seqlen_q = kwargs.get("max_seqlen_q", None)
+    cu_seqlens_k = kwargs.get("cu_seqlens_kv", cu_seqlens_q)
+    max_seqlen_k = kwargs.get("max_kv_seqlen", max_seqlen_q)
+    assert cu_seqlens_q is not None, "cu_seqlens_q shouldn't be None when var_length is True"
+    return cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k
+
+def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    var_length = True
+    if var_length:
+        cu_seqlens_q, cu_seqlens_k, _, _ = var_attn_arg(kwargs)
+        if not skip_reshape:
+            # assumes 2D q, k,v [total_tokens, embed_dim]
+            total_tokens, embed_dim = q.shape
+            head_dim = embed_dim // heads
+            q = q.view(total_tokens, heads, head_dim)
+            k = k.view(k.shape[0], heads, head_dim)
+            v = v.view(v.shape[0], heads, head_dim)
+
+        b = q.size(0)
+        dim_head = q.shape[-1]
+        q = torch.nested.nested_tensor_from_jagged(q, offsets=cu_seqlens_q.long())
+        k = torch.nested.nested_tensor_from_jagged(k, offsets=cu_seqlens_k.long())
+        v = torch.nested.nested_tensor_from_jagged(v, offsets=cu_seqlens_k.long())
+
+        mask = None
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+    if mask is not None:
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+
+    out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+    if var_length:
+        return out.transpose(1, 2).values()
+    if not skip_output_reshape:
+        out = (
+            out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+        )
+    return out
+
+def scaled_dot_product_attention(*args, **kwargs):
+    num_all_args = len(args) + len(kwargs)
+
+    q = None
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs.get('qkv')
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs.get('q')
+        kv = args[1] if len(args) > 1 else kwargs.get('kv')
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs.get('q')
+        k = args[1] if len(args) > 1 else kwargs.get('k')
+        v = args[2] if len(args) > 2 else kwargs.get('v')
+
+    if q is not None:
+        heads = q.shape[2]
+    else:
+        heads = qkv.shape[3]
+
+    if num_all_args == 1:
+        q, k, v = qkv.unbind(dim=2)
+    elif num_all_args == 2:
+        k, v = kv.unbind(dim=2)
+
+    q = q.permute(0, 2, 1, 3)
+    k = k.permute(0, 2, 1, 3)
+    v = v.permute(0, 2, 1, 3)
+
+    out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True, **kwargs)
+
+    out = out.permute(0, 2, 1, 3)
+
+    return out
+
+def sparse_windowed_scaled_dot_product_self_attention(
+    qkv,
+    window_size: int,
+    shift_window: Tuple[int, int, int] = (0, 0, 0)
+):
+
+    serialization_spatial_cache_name = f'windowed_attention_{window_size}_{shift_window}'
+    serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+    if serialization_spatial_cache is None:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = calc_window_partition(qkv, window_size, shift_window)
+        qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, attn_func_args))
+    else:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = serialization_spatial_cache
+
+    qkv_feats = qkv.feats[fwd_indices]      # [M, 3, H, C]
+    heads = qkv_feats.shape[2]
+
+    if optimized_attention.__name__ == 'attention_xformers':
+        q, k, v = qkv_feats.unbind(dim=1)
+        q = q.unsqueeze(0)                                                              # [1, M, H, C]
+        k = k.unsqueeze(0)                                                              # [1, M, H, C]
+        v = v.unsqueeze(0)                                                              # [1, M, H, C]
+        #out = xops.memory_efficient_attention(q, k, v, **attn_func_args)[0]             # [M, H, C]
+        out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True)
+    elif optimized_attention.__name__ == 'attention_flash':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args)  # [M, H, C]
+    else:
+        out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True)
+
+    out = out[bwd_indices]      # [T, H, C]
+
+    return qkv.replace(out)
+
+def calc_window_partition(
+    tensor,
+    window_size: Union[int, Tuple[int, ...]],
+    shift_window: Union[int, Tuple[int, ...]] = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
+
+    DIM = tensor.coords.shape[1] - 1
+    shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
+    window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
+    shifted_coords = tensor.coords.clone().detach()
+    shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+
+    MAX_COORDS = [i + j for i, j in zip(tensor.spatial_shape, shift_window)]
+    NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
+    OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
+
+    shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+    shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
+    fwd_indices = torch.argsort(shifted_indices)
+    bwd_indices = torch.empty_like(fwd_indices)
+    bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
+    seq_lens = torch.bincount(shifted_indices)
+    mask = seq_lens != 0
+    seq_lens = seq_lens[mask]
+
+    if optimized_attention.__name__ == 'attention_xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        attn_func_args = {
+            'attn_bias': xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+        }
+    elif optimized_attention.__name__ == 'attention_flash':
+        attn_func_args = {
+            'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'max_seqlen': torch.max(seq_lens)
+        }
+
+    return fwd_indices, bwd_indices, seq_lens, attn_func_args
+
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+    q=None
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        device = qkv.device
+
+        s = qkv
+        q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+        kv_seqlen = q_seqlen
+        qkv = qkv.feats     # [T, 3, H, C]
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, C]
+        else:
+            s = None
+            N, L, H, C = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, C)   # [T_Q, H, C]
+
+        if isinstance(kv, VarLenTensor):
+            kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+            kv = kv.feats     # [T_KV, 2, H, C]
+        else:
+            N, L, _, H, C = kv.shape
+            kv_seqlen = [L] * N
+            kv = kv.reshape(N * L, 2, H, C)   # [T_KV, 2, H, C]
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, Ci]
+        else:
+            s = None
+            N, L, H, CI = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, CI)  # [T_Q, H, Ci]
+
+        if isinstance(k, VarLenTensor):
+            kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+            k = k.feats     # [T_KV, H, Ci]
+            v = v.feats     # [T_KV, H, Co]
+        else:
+            N, L, H, CI, CO = *k.shape, v.shape[-1]
+            kv_seqlen = [L] * N
+            k = k.reshape(N * L, H, CI)     # [T_KV, H, Ci]
+            v = v.reshape(N * L, H, CO)     # [T_KV, H, Co]
+
+    # TODO: change
+    if q is not None:
+        heads = q
+    else:
+        heads = qkv
+    heads = heads.shape[2]
+    if optimized_attention.__name__ == 'attention_xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        q = q.unsqueeze(0)
+        k = k.unsqueeze(0)
+        v = v.unsqueeze(0)
+        mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+        out = xops.memory_efficient_attention(q, k, v, mask)[0]
+    elif optimized_attention.__name__ == 'attention_flash':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+
+    elif optimized_attention.__name__ == "attention_pytorch":
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        else:
+            cu_seqlens_kv = cu_seqlens_q
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        out = attention_pytorch(q, k, v, heads=heads,cu_seqlens_q=cu_seqlens_q,
+                                cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=max(q_seqlen), max_kv_seqlen=max(kv_seqlen),
+                                skip_reshape=True, skip_output_reshape=True)
+
+    if s is not None:
+        return s.replace(out)
+    else:
+        return out.reshape(N, L, H, -1)
diff --git a/comfy/ldm/trellis2/flexgemm.py b/comfy/ldm/trellis2/flexgemm.py
new file mode 100644
index 000000000..416e322ab
--- /dev/null
+++ b/comfy/ldm/trellis2/flexgemm.py
@@ -0,0 +1,298 @@
+# will contain every cuda -> pytorch operation
+
+from typing import Optional, Tuple
+import torch
+
+UINT32_SENTINEL = 0xFFFFFFFF
+
+
+def compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device):
+    """Kernel spatial offsets in the same order as the CUDA/Triton kernels."""
+    offsets = []
+    for vx in range(Kw):
+        for vy in range(Kh):
+            for vz in range(Kd):
+                offsets.append((vx * Dw, vy * Dh, vz * Dd))
+    return torch.tensor(offsets, device=device, dtype=torch.int32)
+
+
+class TorchHashMap:
+    """Sorted-array hashmap backed by torch.searchsorted."""
+
+    def __init__(self, keys: torch.Tensor, values: torch.Tensor, default_value: int):
+        device = keys.device
+        self.sorted_keys, order = torch.sort(keys.to(torch.long))
+        self.sorted_vals = values.to(torch.long)[order]
+        self.default_value = torch.tensor(default_value, dtype=torch.long, device=device)
+        self._n = self.sorted_keys.numel()
+
+    def lookup_flat(self, flat_keys: torch.Tensor) -> torch.Tensor:
+        flat = flat_keys.to(torch.long)
+        if self._n == 0:
+            return torch.full((flat.shape[0],), -1, device=flat.device, dtype=torch.int32)
+        idx = torch.searchsorted(self.sorted_keys, flat)
+        idx_safe = torch.clamp(idx, max=self._n - 1)
+        found = (idx < self._n) & (self.sorted_keys[idx_safe] == flat)
+        out = torch.full((flat.shape[0],), -1, device=flat.device, dtype=torch.int32)
+        if found.any():
+            out[found] = self.sorted_vals[idx_safe[found]].to(torch.int32)
+        return out
+
+
+def build_submanifold_neighbor_map(
+    hashmap,
+    coords: torch.Tensor,
+    W, H, D,
+    Kw, Kh, Kd,
+    Dw, Dh, Dd,
+):
+    device = coords.device
+    M = coords.shape[0]
+    V = Kw * Kh * Kd
+    half_V = V // 2 + 1
+    INVALID = -1
+
+    # int32 neighbour map: 4 bytes/elem vs 8 bytes for int64
+    neighbor = torch.full((M, V), INVALID, device=device, dtype=torch.int32)
+
+    b = coords[:, 0].long()
+    x = coords[:, 1].long()
+    y = coords[:, 2].long()
+    z = coords[:, 3].long()
+
+    offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device)
+
+    ox = x - (Kw // 2) * Dw
+    oy = y - (Kh // 2) * Dh
+    oz = z - (Kd // 2) * Dd
+
+    for v in range(half_V):
+        if v == half_V - 1:
+            # Center voxel always maps to itself
+            neighbor[:, v] = torch.arange(M, device=device, dtype=torch.int32)
+            continue
+
+        dx, dy, dz = offsets[v]
+
+        kx = ox + dx
+        ky = oy + dy
+        kz = oz + dz
+
+        valid = (
+            (kx >= 0) & (kx < W) &
+            (ky >= 0) & (ky < H) &
+            (kz >= 0) & (kz < D)
+        )
+
+        flat = (
+            b[valid] * (W * H * D) +
+            kx[valid] * (H * D) +
+            ky[valid] * D +
+            kz[valid]
+        )
+
+        if flat.numel() > 0:
+            found = hashmap.lookup_flat(flat)
+            idx_in_M = torch.where(valid)[0]
+            neighbor[idx_in_M, v] = found.to(torch.int32)
+
+            # BUG FIX: old code used  found != hashmap.default_value  which
+            # compared int32 -1 against int64 4294967295 → always True.
+            # We now explicitly check for valid indices.
+            valid_found_mask = found >= 0
+            if valid_found_mask.any():
+                src_points = idx_in_M[valid_found_mask]
+                dst_points = found[valid_found_mask].long()
+                neighbor[dst_points, V - 1 - v] = src_points.to(torch.int32)
+
+    return neighbor
+
+def get_recommended_chunk_mem(
+    device=None,
+    safety_fraction: float = 0.4,
+    min_gb: float = 0.25,
+    max_gb: float = 8.0,
+):
+
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    else:
+        device = torch.device(device)
+
+    if device.type == 'cuda':
+        try:
+            idx = device.index if device.index is not None else 0
+            free_bytes, total_bytes = torch.cuda.mem_get_info(idx)
+            free_gb = free_bytes / (1024 ** 3)
+            total_gb = total_bytes / (1024 ** 3)
+
+            recommended = free_gb * safety_fraction
+            result = max(min_gb, min(recommended, max_gb))
+            return result
+
+        except Exception:
+            try:
+                idx = device.index if device.index is not None else 0
+                total_gb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 3)
+            except Exception:
+                total_gb = 16.0
+
+            if total_gb < 12:
+                result = 0.5
+            elif total_gb < 16:
+                result = 0.75
+            elif total_gb < 24:
+                result = 1.0
+            elif total_gb < 32:
+                result = 2.0
+            elif total_gb < 48:
+                result = 4.0
+            else:
+                result = 6.0
+            return result
+
+    else:
+        try:
+            import psutil
+            avail_gb = psutil.virtual_memory().available / (1024 ** 3)
+            recommended = avail_gb * safety_fraction
+            result = max(min_gb, min(recommended, max_gb))
+            return result
+        except ImportError:
+            return min_gb
+
+def sparse_submanifold_conv3d(
+    feats: torch.Tensor,
+    coords: torch.Tensor,
+    shape: tuple,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    neighbor_cache: Optional[torch.Tensor],
+    dilation: tuple,
+    max_chunk_mem_gb: float = 6.0,
+    accumulate_f32: bool = True,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+    if feats.shape[0] == 0:
+        Co = weight.shape[0]
+        return torch.empty((0, Co), device=feats.device, dtype=feats.dtype), None
+
+    if len(shape) == 5:
+        _, _, W, H, D = shape
+    else:
+        W, H, D = shape
+
+    Co, Kw, Kh, Kd, Ci = weight.shape
+    V = Kw * Kh * Kd
+    device = feats.device
+    sentinel = -1
+    max_chunk_mem_gb = get_recommended_chunk_mem(device)
+
+    if neighbor_cache is None:
+        b_stride = W * H * D
+        x_stride = H * D
+        y_stride = D
+        z_stride = 1
+
+        flat_keys = (coords[:, 0].long() * b_stride +
+                     coords[:, 1].long() * x_stride +
+                     coords[:, 2].long() * y_stride +
+                     coords[:, 3].long() * z_stride)
+        vals = torch.arange(coords.shape[0], dtype=torch.int32, device=device)
+        hashmap = TorchHashMap(flat_keys, vals, UINT32_SENTINEL)
+
+        neighbor = build_submanifold_neighbor_map(
+            hashmap, coords, W, H, D, Kw, Kh, Kd,
+            dilation[0], dilation[1], dilation[2]
+        )
+    else:
+        neighbor = neighbor_cache
+
+    N_pts = feats.shape[0]
+
+    if accumulate_f32:
+        weight_T = weight.view(Co, V * Ci).to(torch.float32).T.contiguous()
+        output = torch.zeros(N_pts, Co, device=device, dtype=torch.float32)
+    else:
+        weight_T = weight.view(Co, V * Ci).to(feats.dtype).T.contiguous()
+        output = torch.zeros(N_pts, Co, device=device, dtype=feats.dtype)
+
+    # ------------------------------------------------------------------
+    # Chunk size from memory budget
+    # ------------------------------------------------------------------
+    bytes_per_elem = 4 if accumulate_f32 else feats.element_size()
+    mem_per_row = V * Ci * bytes_per_elem
+    max_chunk_mem = max_chunk_mem_gb * (1024 ** 3)
+    chunk_size = max(1, int(max_chunk_mem / mem_per_row))
+    chunk_size = min(chunk_size, N_pts)
+
+    # ------------------------------------------------------------------
+    # Chunked forward pass
+    #   Each iteration:
+    #     1. gather   (chunk, V, Ci)     – memory bound
+    #     2. mask     zero invalids       – in-place, no extra alloc
+    #     3. reshape  (chunk, V*Ci)
+    #     4. GEMM     (chunk, V*Ci) @ (V*Ci, Co) → (chunk, Co)  – cuBLAS
+    #        written directly into output slice via out= argument
+    # ------------------------------------------------------------------
+    for start in range(0, N_pts, chunk_size):
+        end = min(start + chunk_size, N_pts)
+        actual_chunk = end - start
+
+        # (chunk, V) int32
+        chunk_neighbor = neighbor[start:end]
+        chunk_valid = chunk_neighbor != sentinel
+
+        # Clamp sentinel -1 → 0 for safe indexing.  No clone of the full map.
+        chunk_idx = chunk_neighbor.clamp(min=0).long()
+
+        # Gather: (chunk, V, Ci).  Memory-bound, single index_select.
+        gathered = feats[chunk_idx]
+
+        # Zero invalid neighbours in-place.  gathered is a fresh tensor from
+        # advanced indexing, so in-place mutation is safe.
+        gathered.mul_(chunk_valid.unsqueeze(-1))
+
+        # Reshape to (chunk, V*Ci)
+        gathered_flat = gathered.view(actual_chunk, V * Ci)
+        if accumulate_f32:
+            gathered_flat = gathered_flat.to(torch.float32)
+
+        # Single GEMM call per chunk, written directly into output.
+        # This avoids allocating a temporary (chunk, Co) tensor.
+        torch.matmul(gathered_flat, weight_T, out=output[start:end])
+
+    if accumulate_f32:
+        output = output.to(feats.dtype)
+
+    if bias is not None:
+        output = output + bias.unsqueeze(0).to(output.dtype)
+
+    return output, neighbor
+
+class Mesh:
+    def __init__(self,
+        vertices,
+        faces,
+        vertex_attrs=None
+    ):
+        self.vertices = vertices.float()
+        self.faces = faces.int()
+        self.vertex_attrs = vertex_attrs
+
+    @property
+    def device(self):
+        return self.vertices.device
+
+    def to(self, device, non_blocking=False):
+        return Mesh(
+            self.vertices.to(device, non_blocking=non_blocking),
+            self.faces.to(device, non_blocking=non_blocking),
+            self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None,
+        )
+
+    def cuda(self, non_blocking=False):
+        return self.to('cuda', non_blocking=non_blocking)
+
+    def cpu(self):
+        return self.to('cpu')
diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py
new file mode 100644
index 000000000..3c61a5d77
--- /dev/null
+++ b/comfy/ldm/trellis2/model.py
@@ -0,0 +1,935 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from comfy.ldm.trellis2.vae import SparseTensor, SparseLinear, sparse_cat, VarLenTensor
+from typing import Optional, Tuple, Literal, Union, List
+from comfy.ldm.trellis2.attention import (
+    sparse_windowed_scaled_dot_product_self_attention, sparse_scaled_dot_product_attention, scaled_dot_product_attention
+)
+from comfy.ldm.genmo.joint_model.layers import TimestepEmbedder
+from comfy.ldm.flux.math import apply_rope, apply_rope1
+
+class SparseGELU(nn.GELU):
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
+
+class SparseFeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            SparseLinear(channels, int(channels * mlp_ratio), device=device, dtype=dtype, operations=operations),
+            SparseGELU(approximate="tanh"),
+            SparseLinear(int(channels * mlp_ratio), channels, device=device, dtype=dtype, operations=operations),
+        )
+
+    def forward(self, x: VarLenTensor) -> VarLenTensor:
+        return self.mlp(x)
+
+def manual_cast(obj, dtype):
+    return obj.to(dtype=dtype)
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int, device, dtype):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim, device=device, dtype=dtype))
+
+    def forward(self, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        x_type = x.dtype
+        x = x.float()
+        if isinstance(x, VarLenTensor):
+            x = x.replace(F.normalize(x.feats, dim=-1) * self.gamma * self.scale)
+        else:
+            x = F.normalize(x, dim=-1) * self.gamma * self.scale
+        return x.to(x_type)
+
+class SparseRotaryPositionEmbedder(nn.Module):
+    def __init__(
+        self,
+        head_dim: int,
+        dim: int = 3,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        device=None
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.dim = dim
+        self.rope_freq = rope_freq
+        self.freq_dim = head_dim // 2 // dim
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32, device=device) / self.freq_dim
+        self.freqs = rope_freq[0] / (rope_freq[1] ** (self.freqs))
+
+    def _get_freqs_cis(self, coords: torch.Tensor) -> torch.Tensor:
+        phases_list = []
+        for i in range(self.dim):
+            phases_list.append(torch.outer(coords[..., i], self.freqs.to(coords.device)))
+
+        phases = torch.cat(phases_list, dim=-1)
+
+        if phases.shape[-1] < self.head_dim // 2:
+            padn = self.head_dim // 2 - phases.shape[-1]
+            phases = torch.cat([phases, torch.zeros(*phases.shape[:-1], padn, device=phases.device)], dim=-1)
+
+        cos = torch.cos(phases)
+        sin = torch.sin(phases)
+
+        f_cis_0 = torch.stack([cos, sin], dim=-1)
+        f_cis_1 = torch.stack([-sin, cos], dim=-1)
+        freqs_cis = torch.stack([f_cis_0, f_cis_1], dim=-1)
+
+        return freqs_cis
+
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+
+    def forward(self, q, k=None):
+        cache_name = f'rope_cis_{self.dim}d_f{self.rope_freq[1]}_hd{self.head_dim}'
+        freqs_cis = q.get_spatial_cache(cache_name)
+
+        if freqs_cis is None:
+            coords = q.coords[..., 1:].to(torch.float32)
+            freqs_cis = self._get_freqs_cis(coords)
+            q.register_spatial_cache(cache_name, freqs_cis)
+
+        if q.feats.ndim == 3:
+            f_cis = freqs_cis.unsqueeze(1)
+        else:
+            f_cis = freqs_cis
+
+        if k is None:
+            return q.replace(apply_rope1(q.feats, f_cis))
+
+        q_feats, k_feats = apply_rope(q.feats, k.feats, f_cis)
+        return q.replace(q_feats), k.replace(k_feats)
+
+    @staticmethod
+    def apply_rotary_embedding(x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases.unsqueeze(-2)
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+
+class RotaryPositionEmbedder(SparseRotaryPositionEmbedder):
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+        if torch.is_complex(phases):
+            phases = phases.to(torch.complex64)
+        else:
+            phases = phases.to(torch.float32)
+        if phases.shape[-1] < self.head_dim // 2:
+                padn = self.head_dim // 2 - phases.shape[-1]
+                phases = torch.cat([phases, torch.polar(
+                    torch.ones(*phases.shape[:-1], padn, device=phases.device, dtype=torch.float32),
+                    torch.zeros(*phases.shape[:-1], padn, device=phases.device, dtype=torch.float32)
+                )], dim=-1)
+        return phases
+
+class SparseMultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int] = None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed", "double_windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = operations.Linear(channels, channels * 3, bias=qkv_bias, device=device, dtype=dtype)
+        else:
+            self.to_q = operations.Linear(channels, channels, bias=qkv_bias, device=device, dtype=dtype)
+            self.to_kv = operations.Linear(self.ctx_channels, channels * 2, bias=qkv_bias, device=device, dtype=dtype)
+
+        if self.qk_rms_norm:
+            self.q_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+            self.k_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+
+        self.to_out = operations.Linear(channels, channels, device=device, dtype=dtype)
+
+        if use_rope:
+            self.rope = SparseRotaryPositionEmbedder(self.head_dim, rope_freq=rope_freq, device=device)
+
+    @staticmethod
+    def _linear(module: nn.Linear, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.replace(module(x.feats))
+        else:
+            return module(x)
+
+    @staticmethod
+    def _reshape_chs(x: Union[VarLenTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.reshape(*shape)
+        else:
+            return x.reshape(*x.shape[:2], *shape)
+
+    def _fused_pre(self, x: Union[VarLenTensor, torch.Tensor], num_fused: int) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            x_feats = x.feats.unsqueeze(0)
+        else:
+            x_feats = x
+        x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+        return x.replace(x_feats.squeeze(0)) if isinstance(x, VarLenTensor) else x_feats
+
+    def forward(self, x: SparseTensor, context: Optional[Union[VarLenTensor, torch.Tensor]] = None) -> SparseTensor:
+        if self._type == "self":
+            dtype = next(self.to_qkv.parameters()).dtype
+            x = x.to(dtype)
+            qkv = self._linear(self.to_qkv, x)
+            qkv = self._fused_pre(qkv, num_fused=3)
+            if self.qk_rms_norm or self.use_rope:
+                q, k, v = qkv.unbind(dim=-3)
+                if self.qk_rms_norm:
+                    q = self.q_rms_norm(q)
+                    k = self.k_rms_norm(k)
+                if self.use_rope:
+                    q, k = self.rope(q, k)
+                qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+            if self.attn_mode == "full":
+                h = sparse_scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                h = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv, self.window_size, shift_window=self.shift_window
+                )
+            elif self.attn_mode == "double_windowed":
+                qkv0 = qkv.replace(qkv.feats[:, :, self.num_heads//2:])
+                qkv1 = qkv.replace(qkv.feats[:, :, :self.num_heads//2])
+                h0 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv0, self.window_size, shift_window=(0, 0, 0)
+                )
+                h1 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv1, self.window_size, shift_window=tuple([self.window_size//2] * 3)
+                )
+                h = qkv.replace(torch.cat([h0.feats, h1.feats], dim=1))
+        else:
+            q = self._linear(self.to_q, x)
+            q = self._reshape_chs(q, (self.num_heads, -1))
+            dtype = next(self.to_kv.parameters()).dtype
+            context = context.to(dtype)
+            kv = self._linear(self.to_kv, context)
+            kv = self._fused_pre(kv, num_fused=2)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=-3)
+                k = self.k_rms_norm(k)
+                h = sparse_scaled_dot_product_attention(q, k, v)
+            else:
+                h = sparse_scaled_dot_product_attention(q, kv)
+        h = self._reshape_chs(h, (-1,))
+        h = self._linear(self.to_out, h)
+        return h
+
+class ModulatedSparseTransformerCrossBlock(nn.Module):
+    """
+    Sparse Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, device=device)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.self_attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.cross_attn = SparseMultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+            device=device, dtype=dtype, operations=operations
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(channels, 6 * channels, bias=True, device=device, dtype=dtype)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels, device=device, dtype=dtype) / channels ** 0.5)
+
+    def _forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = x.replace(self.norm1(x.feats))
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.self_attn(h)
+        h = h * gate_msa
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = x.replace(self.norm3(x.feats))
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        return self._forward(x, mod, context)
+
+
+class SLatFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "rope",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        dtype = None,
+        device = None,
+        operations = None,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = dtype
+
+        self.t_embedder = TimestepEmbedder(model_channels, device=device, dtype=dtype, operations=operations)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(model_channels, 6 * model_channels, bias=True, device=device, dtype=dtype)
+            )
+
+        self.input_layer = SparseLinear(in_channels, model_channels, device=device, dtype=dtype, operations=operations)
+
+        self.blocks = nn.ModuleList([
+            ModulatedSparseTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=self.share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                device=device, dtype=dtype, operations=operations
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = SparseLinear(model_channels, out_channels, device=device, dtype=dtype, operations=operations)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        x: SparseTensor,
+        t: torch.Tensor,
+        cond: Union[torch.Tensor, List[torch.Tensor]],
+        concat_cond: Optional[SparseTensor] = None,
+        **kwargs
+    ) -> SparseTensor:
+        if concat_cond is not None:
+            x = sparse_cat([x, concat_cond], dim=-1)
+        if isinstance(cond, list):
+            cond = VarLenTensor.from_tensor_list(cond)
+
+        dtype = next(self.input_layer.parameters()).dtype
+        x = x.to(dtype)
+        h = self.input_layer(x)
+        h = manual_cast(h, self.dtype)
+        t = t.to(dtype)
+        t_embedder = self.t_embedder.to(dtype)
+        t_emb = t_embedder(t, out_dtype = t.dtype)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+
+        h = manual_cast(h, x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return h
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            operations.Linear(channels, int(channels * mlp_ratio), device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(int(channels * mlp_ratio), channels, device=device, dtype=dtype),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int, device=None, dtype=None):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim, device=device, dtype=dtype))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int]=None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = operations.Linear(channels, channels * 3, bias=qkv_bias, dtype=dtype, device=device)
+        else:
+            self.to_q = operations.Linear(channels, channels, bias=qkv_bias, device=device, dtype=dtype)
+            self.to_kv = operations.Linear(self.ctx_channels, channels * 2, bias=qkv_bias, device=device, dtype=dtype)
+
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+
+        self.to_out = operations.Linear(channels, channels, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self._type == "self":
+            x = x.to(next(self.to_qkv.parameters()).dtype)
+            qkv = self.to_qkv(x)
+            qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+
+            if self.attn_mode == "full":
+                if self.qk_rms_norm or self.use_rope:
+                    q, k, v = qkv.unbind(dim=2)
+                    if self.qk_rms_norm:
+                        q = self.q_rms_norm(q)
+                        k = self.k_rms_norm(k)
+                    if self.use_rope:
+                        assert phases is not None, "Phases must be provided for RoPE"
+                        q = RotaryPositionEmbedder.apply_rotary_embedding(q, phases)
+                        k = RotaryPositionEmbedder.apply_rotary_embedding(k, phases)
+                    h = scaled_dot_product_attention(q, k, v)
+                else:
+                    h = scaled_dot_product_attention(qkv)
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x)
+            context = context.to(next(self.to_kv.parameters()).dtype)
+            kv = self.to_kv(context)
+            q = q.reshape(B, L, self.num_heads, -1)
+            kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q, k, v)
+            else:
+                h = scaled_dot_product_attention(q, kv)
+        h = h.reshape(B, L, -1)
+        h = self.to_out(h)
+        return h
+
+class ModulatedTransformerCrossBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, device=device)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.self_attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.cross_attn = MultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+            device=device, dtype=dtype, operations=operations
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(channels, 6 * channels, bias=True, dtype=dtype, device=device)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels, device=device, dtype=dtype) / channels ** 0.5)
+
+    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = self.norm1(x)
+        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        h = self.self_attn(h, phases=phases)
+        h = h * gate_msa.unsqueeze(1)
+        x = x + h
+        h = self.norm2(x)
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = self.norm3(x)
+        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        h = self.mlp(h)
+        h = h * gate_mlp.unsqueeze(1)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return self._forward(x, mod, context, phases)
+
+
+class SparseStructureFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "rope",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        operations=None,
+        device = None,
+        dtype = torch.float32,
+        **kwargs
+    ):
+        super().__init__()
+        self.device = device
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = dtype
+        self.device = device
+
+        self.t_embedder = TimestepEmbedder(model_channels, dtype=dtype, device=device, operations=operations)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(model_channels, 6 * model_channels, bias=True, device=device, dtype=dtype)
+            )
+
+        pos_embedder = RotaryPositionEmbedder(self.model_channels // self.num_heads, 3, device=device)
+        coords = torch.meshgrid(*[torch.arange(res, device=self.device, dtype=dtype) for res in [resolution] * 3], indexing='ij')
+        coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+        rope_phases = pos_embedder(coords)
+        self.register_buffer("rope_phases", rope_phases, persistent=False)
+
+        if pe_mode != "rope":
+            self.rope_phases = None
+
+        self.input_layer = operations.Linear(in_channels, model_channels, device=device, dtype=dtype)
+
+        self.blocks = nn.ModuleList([
+            ModulatedTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                device=device, dtype=dtype, operations=operations
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = operations.Linear(model_channels, out_channels, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        x = x.view(x.shape[0], self.in_channels, *[self.resolution] * 3)
+
+        h = x.view(*x.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+        h = h.to(next(self.input_layer.parameters()).dtype)
+        h = self.input_layer(h)
+        t_emb = self.t_embedder(t, out_dtype = t.dtype)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        h = manual_cast(h, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond, self.rope_phases)
+        h = manual_cast(h, x.dtype)
+        h = F.layer_norm(h, h.shape[-1:])
+        h = h.to(next(self.out_layer.parameters()).dtype)
+        h = self.out_layer(h)
+
+        h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution] * 3).contiguous()
+
+        return h
+
+def timestep_reshift(t_shifted, old_shift=3.0, new_shift=5.0):
+    t_shifted = t_shifted / 1000.0
+    t_linear = t_shifted / (old_shift - t_shifted * (old_shift - 1))
+    t_new = (new_shift * t_linear) / (1 + (new_shift - 1) * t_linear)
+    t_new *= 1000.0
+    return t_new
+
+class Trellis2(nn.Module):
+    def __init__(self, resolution,
+                 in_channels = 32,
+                 out_channels = 32,
+                 model_channels = 1536,
+                 cond_channels = 1024,
+                 num_blocks = 30,
+                 num_heads = 12,
+                 mlp_ratio = 5.3334,
+                 share_mod = True,
+                 qk_rms_norm = True,
+                 qk_rms_norm_cross = True,
+                 init_txt_model=False, # for now
+                 dtype=None, device=None, operations=None, **kwargs):
+
+        super().__init__()
+        self.dtype = dtype
+        operations = operations or nn
+        # for some reason it passes num_heads = -1
+        if num_heads == -1:
+            num_heads = 12
+        args = {
+            "out_channels":out_channels, "num_blocks":num_blocks, "cond_channels" :cond_channels,
+            "model_channels":model_channels, "num_heads":num_heads, "mlp_ratio": mlp_ratio, "share_mod": share_mod,
+            "qk_rms_norm": qk_rms_norm, "qk_rms_norm_cross": qk_rms_norm_cross, "device": device, "dtype": dtype, "operations": operations
+        }
+        txt_only = kwargs.get("txt_only", False)
+        if not txt_only:
+            self.img2shape = SLatFlowModel(resolution=resolution, in_channels=in_channels, **args)
+            self.shape2txt = None
+            if init_txt_model:
+                self.shape2txt = SLatFlowModel(resolution=resolution, in_channels=in_channels*2, **args)
+            self.img2shape_512 = SLatFlowModel(resolution=32, in_channels=in_channels, **args)
+            args.pop("out_channels")
+            self.structure_model = SparseStructureFlowModel(resolution=16, in_channels=8, out_channels=8, **args)
+        else:
+            self.shape2txt = SLatFlowModel(resolution=resolution, in_channels=in_channels*2, **args)
+        self.guidance_interval = [0.6, 1.0]
+        self.guidance_interval_txt = [0.6, 0.9]
+
+    def forward(self, x, timestep, context, **kwargs):
+        transformer_options = kwargs.get("transformer_options", {})
+        model_options = {}
+        if hasattr(self, "meta"):
+            model_options = self.meta
+        timestep = timestep.to(x.dtype)
+        embeds = kwargs.get("embeds")
+        if embeds is None:
+            raise ValueError("Trellis2.forward requires 'embeds' in kwargs")
+
+        is_1024 = True#self.img2shape.resolution == 1024
+        coords = model_options.get("coords", None)
+        coord_counts = model_options.get("coord_counts", None)
+        mode = model_options.get("generation_mode", "structure_generation")
+
+        is_512_run = False
+        if mode == "shape_generation_512":
+            is_512_run = True
+            mode = "shape_generation"
+
+        if coords is not None:
+            if x.ndim == 4:
+                x = x.squeeze(-1).transpose(1, 2)
+            not_struct_mode = True
+        else:
+            mode = "structure_generation"
+            not_struct_mode = False
+
+        if x.size(-1) == 16 and x.size(-2) == 16:
+            mode = "structure_generation"
+            not_struct_mode = False
+
+        if not not_struct_mode:
+            bsz = x.size(0)
+            x = x[:, :8]
+            x = x.view(bsz, 8, 16, 16, 16)
+
+        if is_1024 and not_struct_mode and not is_512_run:
+            context = embeds
+
+        sigmas = transformer_options.get("sigmas")[0].item()
+        if sigmas < 1.00001:
+            timestep *= 1000.0
+
+        if context.size(0) > 1:
+            cond = context.chunk(2)[1]
+        else:
+            cond = context
+
+        shape_rule = sigmas < self.guidance_interval[0] or sigmas > self.guidance_interval[1]
+        txt_rule = sigmas < self.guidance_interval_txt[0] or sigmas > self.guidance_interval_txt[1]
+
+        if not_struct_mode:
+            orig_bsz = x.shape[0]
+            rule = txt_rule if mode == "texture_generation" else shape_rule
+
+            # CFG Bypass Slicing
+            if rule and orig_bsz > 1:
+                half = orig_bsz // 2
+                x_eval = x[half:]
+                t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep
+                c_eval = cond
+            else:
+                x_eval = x
+                t_eval = timestep
+                c_eval = context
+
+            B, N, C = x_eval.shape
+
+            # Vectorized SparseTensor Construction
+            if mode in ["shape_generation", "texture_generation"]:
+                if coord_counts is not None:
+                    logical_batch = coord_counts.shape[0]
+                    # Duplicate coords if CFG is active
+                    if B > logical_batch:
+                        c_pos = coords.clone()
+                        c_pos[:, 0] += logical_batch
+                        batched_coords = torch.cat([coords, c_pos], dim=0)
+                        counts_eval = torch.cat([coord_counts, coord_counts], dim=0)
+                    else:
+                        batched_coords = coords
+                        counts_eval = coord_counts
+
+                    # Create boolean mask [B, N] to drop the padded zeros instantly
+                    mask = torch.arange(N, device=x.device).unsqueeze(0) < counts_eval.unsqueeze(1)
+                    feats_flat = x_eval[mask]
+                else:
+                    feats_flat = x_eval.reshape(-1, C)
+                    coords_list =[]
+                    for i in range(B):
+                        c = coords.clone()
+                        c[:, 0] = i
+                        coords_list.append(c)
+                    batched_coords = torch.cat(coords_list, dim=0)
+                    mask = None
+            else:
+                batched_coords = coords
+                feats_flat = x_eval
+                mask = None
+
+            x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32))
+
+        if mode == "shape_generation":
+            if is_512_run:
+                out = self.img2shape_512(x_st, t_eval, c_eval)
+            else:
+                out = self.img2shape(x_st, t_eval, c_eval)
+
+        elif mode == "texture_generation":
+            if self.shape2txt is None:
+                raise ValueError("Checkpoint for Trellis2 doesn't include texture generation!")
+            slat = model_options.get("shape_slat")
+            if slat is None:
+                raise ValueError("shape_slat can't be None")
+
+            slat_feats = slat
+            # Duplicate shape context if CFG is active
+            if coord_counts is not None and B > coord_counts.shape[0]:
+                slat_feats = torch.cat([slat_feats, slat_feats], dim=0)
+            elif coord_counts is None:
+                slat_feats = slat_feats[:N].repeat(B, 1)
+
+            x_st = x_st.replace(feats=torch.cat([x_st.feats, slat_feats.to(x_st.feats.device)], dim=-1))
+            out = self.shape2txt(x_st, t_eval, c_eval)
+
+        else: # structure
+            orig_bsz = x.shape[0]
+            if shape_rule and orig_bsz > 1:
+                half = orig_bsz // 2
+                x_eval = x[half:]
+                t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep
+                out = self.structure_model(x_eval, t_eval, cond)
+                out = out.repeat(2, 1, 1, 1, 1)
+            else:
+                out = self.structure_model(x, timestep, context)
+
+        if not_struct_mode:
+            if mask is not None:
+                # Instantly scatter the valid tokens back into a padded rectangular tensor
+                padded_out = torch.zeros((B, N, out.feats.shape[-1]), device=x.device, dtype=out.feats.dtype)
+                padded_out[mask] = out.feats
+                out_tensor = padded_out.transpose(1, 2).unsqueeze(-1)
+            else:
+                out_tensor = out.feats.view(B, N, -1).transpose(1, 2).unsqueeze(-1)
+
+            if rule and orig_bsz > 1:
+                out_tensor = out_tensor.repeat(2, 1, 1, 1)
+            return out_tensor
+        else:
+            out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, 0, 0, 24))
+
+        return out
diff --git a/comfy/ldm/trellis2/vae.py b/comfy/ldm/trellis2/vae.py
new file mode 100644
index 000000000..5336f4dc7
--- /dev/null
+++ b/comfy/ldm/trellis2/vae.py
@@ -0,0 +1,1446 @@
+import math
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from fractions import Fraction
+from dataclasses import dataclass
+from typing import List, Any, Dict, Optional, overload, Union, Tuple
+from comfy.ldm.trellis2.flexgemm import TorchHashMap, Mesh, sparse_submanifold_conv3d
+
+
+def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor:
+    B, C, H, W, D = x.shape
+    C_ = C // scale_factor**3
+    x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D)
+    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4)
+    x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor)
+    return x
+
+class SparseConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+        super(SparseConv3d, self).__init__()
+        sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride, dilation, padding, bias, indice_key)
+
+    def forward(self, x):
+        return sparse_conv3d_forward(self, x)
+
+
+def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size, ) * 3
+    self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, ) * 3
+    self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation, ) * 3
+
+    self.weight = nn.Parameter(torch.empty((out_channels, in_channels, *self.kernel_size)))
+    if bias:
+        self.bias = nn.Parameter(torch.empty(out_channels))
+    else:
+        self.register_parameter("bias", None)
+
+    if self.bias is not None:
+        fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight)
+        if fan_in != 0:
+            bound = 1 / math.sqrt(fan_in)
+            torch.nn.init.uniform_(self.bias, -bound, bound)
+
+    # Permute weight (Co, Ci, Kd, Kh, Kw) -> (Co, Kd, Kh, Kw, Ci)
+    self.weight = nn.Parameter(self.weight.permute(0, 2, 3, 4, 1).contiguous())
+
+
+def sparse_conv3d_forward(self, x):
+    # check if neighbor map is already computed
+    Co, Kd, Kh, Kw, Ci = self.weight.shape
+    neighbor_cache_key = f'SubMConv3d_neighbor_cache_{Kw}x{Kh}x{Kd}_dilation{self.dilation}'
+    neighbor_cache = x.get_spatial_cache(neighbor_cache_key)
+    x = x.to(self.weight.dtype).to(self.weight.device)
+
+    out, neighbor_cache_ = sparse_submanifold_conv3d(
+        x.feats,
+        x.coords,
+        x.spatial_shape,
+        self.weight,
+        self.bias,
+        neighbor_cache,
+        self.dilation
+    )
+
+    if neighbor_cache is None:
+        x.register_spatial_cache(neighbor_cache_key, neighbor_cache_)
+
+    out = x.replace(out)
+    return out
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = x.to(torch.float32)
+        w = self.weight.to(torch.float32) if self.weight is not None else None
+        b = self.bias.to(torch.float32) if self.bias is not None else None
+
+        o = F.layer_norm(x, self.normalized_shape, w, b, self.eps)
+        return o.to(x_dtype)
+
+class SparseConvNeXtBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        mlp_ratio: float = 4.0,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.use_checkpoint = use_checkpoint
+
+        self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.conv = SparseConv3d(channels, channels, 3)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.SiLU(),
+            nn.Linear(int(channels * mlp_ratio), channels),
+        )
+
+    def _forward(self, x):
+        h = self.conv(x)
+        h = h.replace(self.norm(h.feats))
+        h = h.replace(self.mlp(h.feats))
+        return h + x
+
+    def forward(self, x):
+        return self._forward(x)
+
+class SparseSpatial2Channel(nn.Module):
+    def __init__(self, factor: int = 2):
+        super(SparseSpatial2Channel, self).__init__()
+        self.factor = factor
+
+    def forward(self, x):
+        DIM = x.coords.shape[-1] - 1
+        cache = x.get_spatial_cache(f'spatial2channel_{self.factor}')
+        if cache is None:
+            coord = list(x.coords.unbind(dim=-1))
+            for i in range(DIM):
+                coord[i+1] = coord[i+1] // self.factor
+            subidx = x.coords[:, 1:] % self.factor
+            subidx = sum([subidx[..., i] * self.factor ** i for i in range(DIM)])
+
+            MAX = [(s + self.factor - 1) // self.factor for s in x.spatial_shape]
+            OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1]
+            code = sum([c * o for c, o in zip(coord, OFFSET)])
+            code, idx = code.unique(return_inverse=True)
+
+            new_coords = torch.stack(
+                [code // OFFSET[0]] +
+                [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)],
+                dim=-1
+            )
+        else:
+            new_coords, idx, subidx = cache
+
+        new_feats = torch.zeros(new_coords.shape[0] * self.factor ** DIM, x.feats.shape[1], device=x.feats.device, dtype=x.feats.dtype)
+        new_feats[idx * self.factor ** DIM + subidx] = x.feats
+
+        out = SparseTensor(new_feats.reshape(new_coords.shape[0], -1), new_coords, None if x._shape is None else torch.Size([x._shape[0], x._shape[1] * self.factor ** DIM]))
+        out._scale = tuple([s * self.factor for s in x._scale])
+        out._spatial_cache = x._spatial_cache
+
+        if cache is None:
+            x.register_spatial_cache(f'spatial2channel_{self.factor}', (new_coords, idx, subidx))
+            out.register_spatial_cache(f'channel2spatial_{self.factor}', (x.coords, idx, subidx))
+            out.register_spatial_cache('shape', torch.Size(MAX))
+
+        return out
+
+class SparseChannel2Spatial(nn.Module):
+    def __init__(self, factor: int = 2):
+        super(SparseChannel2Spatial, self).__init__()
+        self.factor = factor
+
+    def forward(self, x, subdivision = None):
+        DIM = x.coords.shape[-1] - 1
+
+        cache = x.get_spatial_cache(f'channel2spatial_{self.factor}')
+        if cache is None:
+            if subdivision is None:
+                raise ValueError('Cache not found. Provide subdivision tensor or pair SparseChannel2Spatial with SparseSpatial2Channel.')
+            else:
+                sub = subdivision.feats         # [N, self.factor ** DIM]
+                N_leaf = sub.sum(dim=-1)        # [N]
+                subidx = sub.nonzero()[:, -1]
+                new_coords = x.coords.clone().detach()
+                new_coords[:, 1:] *= self.factor
+                new_coords = torch.repeat_interleave(new_coords, N_leaf, dim=0, output_size=subidx.shape[0])
+                for i in range(DIM):
+                    new_coords[:, i+1] += subidx // self.factor ** i % self.factor
+                idx = torch.repeat_interleave(torch.arange(x.coords.shape[0], device=x.device), N_leaf, dim=0, output_size=subidx.shape[0])
+        else:
+            new_coords, idx, subidx = cache
+
+        x_feats = x.feats.reshape(x.feats.shape[0] * self.factor ** DIM, -1)
+        new_feats = x_feats[idx * self.factor ** DIM + subidx]
+        out = SparseTensor(new_feats, new_coords, None if x._shape is None else torch.Size([x._shape[0], x._shape[1] // self.factor ** DIM]))
+        out._scale = tuple([s / self.factor for s in x._scale])
+        if cache is not None:           # only keep cache when subdiv following it
+            out._spatial_cache = x._spatial_cache
+        return out
+
+class SparseResBlockC2S3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.pred_subdiv = pred_subdiv
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = SparseConv3d(channels, self.out_channels * 8, 3)
+        self.conv2 = SparseConv3d(self.out_channels, self.out_channels, 3)
+        self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
+        if pred_subdiv:
+            self.to_subdiv = SparseLinear(channels, 8)
+        self.updown = SparseChannel2Spatial(2)
+
+    def forward(self, x, subdiv = None):
+        if self.pred_subdiv:
+            dtype = next(self.to_subdiv.parameters()).dtype
+            x = x.to(dtype)
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
+        h = self.updown(h, subdiv_binarized)
+        x = self.updown(x, subdiv_binarized)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.pred_subdiv:
+            return h, subdiv
+        else:
+            return h
+
+@dataclass
+class config:
+    CONV = "flexgemm"
+    FLEX_GEMM_HASHMAP_RATIO = 2.0
+
+class VarLenTensor:
+
+    def __init__(self, feats: torch.Tensor, layout: List[slice]=None):
+        self.feats = feats
+        self.layout = layout if layout is not None else [slice(0, feats.shape[0])]
+        self._cache = {}
+
+    @staticmethod
+    def layout_from_seqlen(seqlen: list) -> List[slice]:
+        """
+        Create a layout from a tensor of sequence lengths.
+        """
+        layout = []
+        start = 0
+        for l in seqlen:
+            layout.append(slice(start, start + l))
+            start += l
+        return layout
+
+    @staticmethod
+    def from_tensor_list(tensor_list: List[torch.Tensor]) -> 'VarLenTensor':
+        """
+        Create a VarLenTensor from a list of tensors.
+        """
+        feats = torch.cat(tensor_list, dim=0)
+        layout = []
+        start = 0
+        for tensor in tensor_list:
+            layout.append(slice(start, start + tensor.shape[0]))
+            start += tensor.shape[0]
+        return VarLenTensor(feats, layout)
+
+    def __len__(self) -> int:
+        return len(self.layout)
+
+    @property
+    def shape(self) -> torch.Size:
+        return torch.Size([len(self.layout), *self.feats.shape[1:]])
+
+    def dim(self) -> int:
+        return len(self.shape)
+
+    @property
+    def ndim(self) -> int:
+        return self.dim()
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def device(self):
+        return self.feats.device
+
+    @property
+    def seqlen(self) -> torch.LongTensor:
+        if 'seqlen' not in self._cache:
+            self._cache['seqlen'] = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device)
+        return self._cache['seqlen']
+
+    @property
+    def cum_seqlen(self) -> torch.LongTensor:
+        if 'cum_seqlen' not in self._cache:
+            self._cache['cum_seqlen'] = torch.cat([
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                self.seqlen.cumsum(dim=0)
+            ], dim=0)
+        return self._cache['cum_seqlen']
+
+    @property
+    def batch_boardcast_map(self) -> torch.LongTensor:
+        """
+        Get the broadcast map for the varlen tensor.
+        """
+        if 'batch_boardcast_map' not in self._cache:
+            self._cache['batch_boardcast_map'] = torch.repeat_interleave(
+                torch.arange(len(self.layout), device=self.device),
+                self.seqlen,
+            )
+        return self._cache['batch_boardcast_map']
+
+    @overload
+    def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ...
+
+    @overload
+    def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ...
+
+    def to(self, *args, **kwargs) -> 'VarLenTensor':
+        device = None
+        dtype = None
+        if len(args) == 2:
+            device, dtype = args
+        elif len(args) == 1:
+            if isinstance(args[0], torch.dtype):
+                dtype = args[0]
+            else:
+                device = args[0]
+        if 'dtype' in kwargs:
+            assert dtype is None, "to() received multiple values for argument 'dtype'"
+            dtype = kwargs['dtype']
+        if 'device' in kwargs:
+            assert device is None, "to() received multiple values for argument 'device'"
+            device = kwargs['device']
+        non_blocking = kwargs.get('non_blocking', False)
+        copy = kwargs.get('copy', False)
+
+        new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy)
+        return self.replace(new_feats)
+
+    def type(self, dtype):
+        new_feats = self.feats.type(dtype)
+        return self.replace(new_feats)
+
+    def cpu(self) -> 'VarLenTensor':
+        new_feats = self.feats.cpu()
+        return self.replace(new_feats)
+
+    def cuda(self) -> 'VarLenTensor':
+        new_feats = self.feats.cuda()
+        return self.replace(new_feats)
+
+    def half(self) -> 'VarLenTensor':
+        new_feats = self.feats.half()
+        return self.replace(new_feats)
+
+    def float(self) -> 'VarLenTensor':
+        new_feats = self.feats.float()
+        return self.replace(new_feats)
+
+    def detach(self) -> 'VarLenTensor':
+        new_feats = self.feats.detach()
+        return self.replace(new_feats)
+
+    def reshape(self, *shape) -> 'VarLenTensor':
+        new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+        return self.replace(new_feats)
+
+    def unbind(self, dim: int) -> List['VarLenTensor']:
+        return varlen_unbind(self, dim)
+
+    def replace(self, feats: torch.Tensor) -> 'VarLenTensor':
+        new_tensor = VarLenTensor(
+            feats=feats,
+            layout=self.layout,
+        )
+        new_tensor._cache = self._cache
+        return new_tensor
+
+    def to_dense(self, max_length=None) -> torch.Tensor:
+        N = len(self)
+        L = max_length or self.seqlen.max().item()
+        spatial = self.feats.shape[1:]
+        idx = torch.arange(L, device=self.device).unsqueeze(0).expand(N, L)
+        mask = (idx < self.seqlen.unsqueeze(1))
+        mapping = mask.reshape(-1).cumsum(dim=0) - 1
+        dense = self.feats[mapping]
+        dense = dense.reshape(N, L, *spatial)
+        return dense, mask
+
+    def __neg__(self) -> 'VarLenTensor':
+        return self.replace(-self.feats)
+
+    def __elemwise__(self, other: Union[torch.Tensor, 'VarLenTensor'], op: callable) -> 'VarLenTensor':
+        if isinstance(other, torch.Tensor):
+            try:
+                other = torch.broadcast_to(other, self.shape)
+                other = other[self.batch_boardcast_map]
+            except:
+                pass
+        if isinstance(other, VarLenTensor):
+            other = other.feats
+        new_feats = op(self.feats, other)
+        new_tensor = self.replace(new_feats)
+        return new_tensor
+
+    def __add__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.add)
+
+    def __radd__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.add)
+
+    def __sub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.sub)
+
+    def __rsub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, lambda x, y: torch.sub(y, x))
+
+    def __mul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.mul)
+
+    def __rmul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.mul)
+
+    def __truediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.div)
+
+    def __rtruediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, lambda x, y: torch.div(y, x))
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            idx = [idx]
+        elif isinstance(idx, slice):
+            idx = range(*idx.indices(self.shape[0]))
+        elif isinstance(idx, list):
+            assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}"
+        elif isinstance(idx, torch.Tensor):
+            if idx.dtype == torch.bool:
+                assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+                idx = idx.nonzero().squeeze(1)
+            elif idx.dtype in [torch.int32, torch.int64]:
+                assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+            else:
+                raise ValueError(f"Unknown index type: {idx.dtype}")
+        else:
+            raise ValueError(f"Unknown index type: {type(idx)}")
+
+        new_feats = []
+        new_layout = []
+        start = 0
+        for new_idx, old_idx in enumerate(idx):
+            new_feats.append(self.feats[self.layout[old_idx]])
+            new_layout.append(slice(start, start + len(new_feats[-1])))
+            start += len(new_feats[-1])
+        new_feats = torch.cat(new_feats, dim=0).contiguous()
+        new_tensor = VarLenTensor(feats=new_feats, layout=new_layout)
+        return new_tensor
+
+    def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        if isinstance(dim, int):
+            dim = (dim,)
+
+        if op =='mean':
+            red = self.feats.mean(dim=dim, keepdim=keepdim)
+        elif op =='sum':
+            red = self.feats.sum(dim=dim, keepdim=keepdim)
+        elif op == 'prod':
+            red = self.feats.prod(dim=dim, keepdim=keepdim)
+        else:
+            raise ValueError(f"Unsupported reduce operation: {op}")
+
+        if dim is None or 0 in dim:
+            return red
+
+        red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
+        return red
+
+    def mean(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='mean', dim=dim, keepdim=keepdim)
+
+    def sum(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='sum', dim=dim, keepdim=keepdim)
+
+    def prod(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='prod', dim=dim, keepdim=keepdim)
+
+    def std(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        mean = self.mean(dim=dim, keepdim=True)
+        mean2 = self.replace(self.feats ** 2).mean(dim=dim, keepdim=True)
+        std = (mean2 - mean ** 2).sqrt()
+        return std
+
+    def __repr__(self) -> str:
+        return f"VarLenTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})"
+
+def varlen_unbind(input: VarLenTensor, dim: int) -> Union[List[VarLenTensor]]:
+
+    if dim == 0:
+        return [input[i] for i in range(len(input))]
+    else:
+        feats = input.feats.unbind(dim)
+        return [input.replace(f) for f in feats]
+
+
+class SparseTensor(VarLenTensor):
+
+    SparseTensorData = None
+
+    @overload
+    def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, **kwargs): ...
+
+    @overload
+    def __init__(self, data, shape: Optional[torch.Size] = None, **kwargs): ...
+
+    def __init__(self, *args, **kwargs):
+        # Lazy import of sparse tensor backend
+        if self.SparseTensorData is None:
+            import importlib
+            if config.CONV == 'torchsparse':
+                self.SparseTensorData = importlib.import_module('torchsparse').SparseTensor
+            elif config.CONV == 'spconv':
+                self.SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor
+
+        method_id = 0
+        if len(args) != 0:
+            method_id = 0 if isinstance(args[0], torch.Tensor) else 1
+        else:
+            method_id = 1 if 'data' in kwargs else 0
+
+        if method_id == 0:
+            feats, coords, shape = args + (None,) * (3 - len(args))
+            if 'feats' in kwargs:
+                feats = kwargs['feats']
+                del kwargs['feats']
+            if 'coords' in kwargs:
+                coords = kwargs['coords']
+                del kwargs['coords']
+            if 'shape' in kwargs:
+                shape = kwargs['shape']
+                del kwargs['shape']
+
+            if config.CONV == 'torchsparse':
+                self.data = self.SparseTensorData(feats, coords, **kwargs)
+            elif config.CONV == 'spconv':
+                spatial_shape = list(coords.max(0)[0] + 1)
+                self.data = self.SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape[1:], spatial_shape[0], **kwargs)
+                self.data._features = feats
+            else:
+                self.data = {
+                    'feats': feats,
+                    'coords': coords,
+                }
+        elif method_id == 1:
+            data, shape = args + (None,) * (2 - len(args))
+            if 'data' in kwargs:
+                data = kwargs['data']
+                del kwargs['data']
+            if 'shape' in kwargs:
+                shape = kwargs['shape']
+                del kwargs['shape']
+
+            self.data = data
+
+        self._shape = shape
+        self._scale = kwargs.get('scale', (Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)))
+        self._spatial_cache = kwargs.get('spatial_cache', {})
+
+    @staticmethod
+    def from_tensor_list(feats_list: List[torch.Tensor], coords_list: List[torch.Tensor]) -> 'SparseTensor':
+        """
+        Create a SparseTensor from a list of tensors.
+        """
+        feats = torch.cat(feats_list, dim=0)
+        coords = []
+        for i, coord in enumerate(coords_list):
+            coord = torch.cat([torch.full_like(coord[:, :1], i), coord[:, 1:]], dim=1)
+            coords.append(coord)
+        coords = torch.cat(coords, dim=0)
+        return SparseTensor(feats, coords)
+
+    def to_tensor_list(self) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+        """
+        Convert a SparseTensor to list of tensors.
+        """
+        feats_list = []
+        coords_list = []
+        for s in self.layout:
+            feats_list.append(self.feats[s])
+            coords_list.append(self.coords[s])
+        return feats_list, coords_list
+
+    def __len__(self) -> int:
+        return len(self.layout)
+
+    def __cal_shape(self, feats, coords):
+        shape = []
+        shape.append(coords[:, 0].max().item() + 1)
+        shape.extend([*feats.shape[1:]])
+        return torch.Size(shape)
+
+    def __cal_layout(self, coords, batch_size):
+        seq_len = torch.bincount(coords[:, 0], minlength=batch_size)
+        offset = torch.cumsum(seq_len, dim=0)
+        layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)]
+        return layout
+
+    def __cal_spatial_shape(self, coords):
+        return torch.Size((coords[:, 1:].max(0)[0] + 1).tolist())
+
+    @property
+    def shape(self) -> torch.Size:
+        if self._shape is None:
+            self._shape = self.__cal_shape(self.feats, self.coords)
+        return self._shape
+
+    @property
+    def layout(self) -> List[slice]:
+        layout = self.get_spatial_cache('layout')
+        if layout is None:
+            layout = self.__cal_layout(self.coords, self.shape[0])
+            self.register_spatial_cache('layout', layout)
+        return layout
+
+    @property
+    def spatial_shape(self) -> torch.Size:
+        spatial_shape = self.get_spatial_cache('shape')
+        if spatial_shape is None:
+            spatial_shape = self.__cal_spatial_shape(self.coords)
+            self.register_spatial_cache('shape', spatial_shape)
+        return spatial_shape
+
+    @property
+    def feats(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.F
+        elif config.CONV == 'spconv':
+            return self.data.features
+        else:
+            return self.data['feats']
+
+    @feats.setter
+    def feats(self, value: torch.Tensor):
+        if config.CONV == 'torchsparse':
+            self.data.F = value
+        elif config.CONV == 'spconv':
+            self.data.features = value
+        else:
+            self.data['feats'] = value
+
+    @property
+    def coords(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.C
+        elif config.CONV == 'spconv':
+            return self.data.indices
+        else:
+            return self.data['coords']
+
+    @coords.setter
+    def coords(self, value: torch.Tensor):
+        if config.CONV == 'torchsparse':
+            self.data.C = value
+        elif config.CONV == 'spconv':
+            self.data.indices = value
+        else:
+            self.data['coords'] = value
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def device(self):
+        return self.feats.device
+
+    @property
+    def seqlen(self) -> torch.LongTensor:
+        seqlen = self.get_spatial_cache('seqlen')
+        if seqlen is None:
+            seqlen = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device)
+            self.register_spatial_cache('seqlen', seqlen)
+        return seqlen
+
+    @property
+    def cum_seqlen(self) -> torch.LongTensor:
+        cum_seqlen = self.get_spatial_cache('cum_seqlen')
+        if cum_seqlen is None:
+            cum_seqlen = torch.cat([
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                self.seqlen.cumsum(dim=0)
+            ], dim=0)
+            self.register_spatial_cache('cum_seqlen', cum_seqlen)
+        return cum_seqlen
+
+    @property
+    def batch_boardcast_map(self) -> torch.LongTensor:
+        """
+        Get the broadcast map for the varlen tensor.
+        """
+        batch_boardcast_map = self.get_spatial_cache('batch_boardcast_map')
+        if batch_boardcast_map is None:
+            batch_boardcast_map = torch.repeat_interleave(
+                torch.arange(len(self.layout), device=self.device),
+                self.seqlen,
+            )
+            self.register_spatial_cache('batch_boardcast_map', batch_boardcast_map)
+        return batch_boardcast_map
+
+    @overload
+    def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'SparseTensor': ...
+
+    @overload
+    def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> 'SparseTensor': ...
+
+    def to(self, *args, **kwargs) -> 'SparseTensor':
+        device = None
+        dtype = None
+        if len(args) == 2:
+            device, dtype = args
+        elif len(args) == 1:
+            if isinstance(args[0], torch.dtype):
+                dtype = args[0]
+            else:
+                device = args[0]
+        if 'dtype' in kwargs:
+            assert dtype is None, "to() received multiple values for argument 'dtype'"
+            dtype = kwargs['dtype']
+        if 'device' in kwargs:
+            assert device is None, "to() received multiple values for argument 'device'"
+            device = kwargs['device']
+        non_blocking = kwargs.get('non_blocking', False)
+        copy = kwargs.get('copy', False)
+
+        new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy)
+        new_coords = self.coords.to(device=device, non_blocking=non_blocking, copy=copy)
+        return self.replace(new_feats, new_coords)
+
+    def type(self, dtype):
+        new_feats = self.feats.type(dtype)
+        return self.replace(new_feats)
+
+    def cpu(self) -> 'SparseTensor':
+        new_feats = self.feats.cpu()
+        new_coords = self.coords.cpu()
+        return self.replace(new_feats, new_coords)
+
+    def cuda(self) -> 'SparseTensor':
+        new_feats = self.feats.cuda()
+        new_coords = self.coords.cuda()
+        return self.replace(new_feats, new_coords)
+
+    def half(self) -> 'SparseTensor':
+        new_feats = self.feats.half()
+        return self.replace(new_feats)
+
+    def float(self) -> 'SparseTensor':
+        new_feats = self.feats.float()
+        return self.replace(new_feats)
+
+    def detach(self) -> 'SparseTensor':
+        new_coords = self.coords.detach()
+        new_feats = self.feats.detach()
+        return self.replace(new_feats, new_coords)
+
+    def reshape(self, *shape) -> 'SparseTensor':
+        new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+        return self.replace(new_feats)
+
+    def unbind(self, dim: int) -> List['SparseTensor']:
+        return sparse_unbind(self, dim)
+
+    def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor':
+        if config.CONV == 'torchsparse':
+            new_data = self.SparseTensorData(
+                feats=feats,
+                coords=self.data.coords if coords is None else coords,
+                stride=self.data.stride,
+                spatial_range=self.data.spatial_range,
+            )
+            new_data._caches = self.data._caches
+        elif config.CONV == 'spconv':
+            new_data = self.SparseTensorData(
+                self.data.features.reshape(self.data.features.shape[0], -1),
+                self.data.indices,
+                self.data.spatial_shape,
+                self.data.batch_size,
+                self.data.grid,
+                self.data.voxel_num,
+                self.data.indice_dict
+            )
+            new_data._features = feats
+            new_data.benchmark = self.data.benchmark
+            new_data.benchmark_record = self.data.benchmark_record
+            new_data.thrust_allocator = self.data.thrust_allocator
+            new_data._timer = self.data._timer
+            new_data.force_algo = self.data.force_algo
+            new_data.int8_scale = self.data.int8_scale
+            if coords is not None:
+                new_data.indices = coords
+        else:
+            new_data = {
+                'feats': feats,
+                'coords': self.data['coords'] if coords is None else coords,
+            }
+        new_tensor = SparseTensor(
+            new_data,
+            shape=torch.Size([self._shape[0]] + list(feats.shape[1:])) if self._shape is not None else None,
+            scale=self._scale,
+            spatial_cache=self._spatial_cache
+        )
+        return new_tensor
+
+    def to_dense(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.dense()
+        elif config.CONV == 'spconv':
+            return self.data.dense()
+        else:
+            spatial_shape = self.spatial_shape
+            ret = torch.zeros(*self.shape, *spatial_shape, dtype=self.dtype, device=self.device)
+            idx = [self.coords[:, 0], slice(None)] + self.coords[:, 1:].unbind(1)
+            ret[tuple(idx)] = self.feats
+            return ret
+
+    @staticmethod
+    def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor':
+        N, C = dim
+        x = torch.arange(aabb[0], aabb[3] + 1)
+        y = torch.arange(aabb[1], aabb[4] + 1)
+        z = torch.arange(aabb[2], aabb[5] + 1)
+        coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3)
+        coords = torch.cat([
+            torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1),
+            coords.repeat(N, 1),
+        ], dim=1).to(dtype=torch.int32, device=device)
+        feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device)
+        return SparseTensor(feats=feats, coords=coords)
+
+    def __merge_sparse_cache(self, other: 'SparseTensor') -> dict:
+        new_cache = {}
+        for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())):
+            if k in self._spatial_cache:
+                new_cache[k] = self._spatial_cache[k]
+            if k in other._spatial_cache:
+                if k not in new_cache:
+                    new_cache[k] = other._spatial_cache[k]
+                else:
+                    new_cache[k].update(other._spatial_cache[k])
+        return new_cache
+
+    def __elemwise__(self, other: Union[torch.Tensor, VarLenTensor], op: callable) -> 'SparseTensor':
+        if isinstance(other, torch.Tensor):
+            try:
+                other = torch.broadcast_to(other, self.shape)
+                other = other[self.batch_boardcast_map]
+            except:
+                pass
+        if isinstance(other, VarLenTensor):
+            other = other.feats
+        new_feats = op(self.feats, other)
+        new_tensor = self.replace(new_feats)
+        if isinstance(other, SparseTensor):
+            new_tensor._spatial_cache = self.__merge_sparse_cache(other)
+        return new_tensor
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            idx = [idx]
+        elif isinstance(idx, slice):
+            idx = range(*idx.indices(self.shape[0]))
+        elif isinstance(idx, list):
+            assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}"
+        elif isinstance(idx, torch.Tensor):
+            if idx.dtype == torch.bool:
+                assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+                idx = idx.nonzero().squeeze(1)
+            elif idx.dtype in [torch.int32, torch.int64]:
+                assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+            else:
+                raise ValueError(f"Unknown index type: {idx.dtype}")
+        else:
+            raise ValueError(f"Unknown index type: {type(idx)}")
+
+        new_coords = []
+        new_feats = []
+        new_layout = []
+        new_shape = torch.Size([len(idx)] + list(self.shape[1:]))
+        start = 0
+        for new_idx, old_idx in enumerate(idx):
+            new_coords.append(self.coords[self.layout[old_idx]].clone())
+            new_coords[-1][:, 0] = new_idx
+            new_feats.append(self.feats[self.layout[old_idx]])
+            new_layout.append(slice(start, start + len(new_coords[-1])))
+            start += len(new_coords[-1])
+        new_coords = torch.cat(new_coords, dim=0).contiguous()
+        new_feats = torch.cat(new_feats, dim=0).contiguous()
+        new_tensor = SparseTensor(feats=new_feats, coords=new_coords, shape=new_shape)
+        new_tensor.register_spatial_cache('layout', new_layout)
+        return new_tensor
+
+    def clear_spatial_cache(self) -> None:
+        """
+        Clear all spatial caches.
+        """
+        self._spatial_cache = {}
+
+    def register_spatial_cache(self, key, value) -> None:
+        """
+        Register a spatial cache.
+        The spatial cache can be any thing you want to cache.
+        The registery and retrieval of the cache is based on current scale.
+        """
+        scale_key = str(self._scale)
+        if scale_key not in self._spatial_cache:
+            self._spatial_cache[scale_key] = {}
+        self._spatial_cache[scale_key][key] = value
+
+    def get_spatial_cache(self, key=None):
+        """
+        Get a spatial cache.
+        """
+        scale_key = str(self._scale)
+        cur_scale_cache = self._spatial_cache.get(scale_key, {})
+        if key is None:
+            return cur_scale_cache
+        return cur_scale_cache.get(key, None)
+
+    def __repr__(self) -> str:
+        return f"SparseTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})"
+
+def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor:
+    if dim == 0:
+        start = 0
+        coords = []
+        for input in inputs:
+            coords.append(input.coords.clone())
+            coords[-1][:, 0] += start
+            start += input.shape[0]
+        coords = torch.cat(coords, dim=0)
+        feats = torch.cat([input.feats for input in inputs], dim=0)
+        output = SparseTensor(
+            coords=coords,
+            feats=feats,
+        )
+    else:
+        feats = torch.cat([input.feats for input in inputs], dim=dim)
+        output = inputs[0].replace(feats)
+
+    return output
+
+
+def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]:
+    if dim == 0:
+        return [input[i] for i in range(input.shape[0])]
+    else:
+        feats = input.feats.unbind(dim)
+        return [input.replace(f) for f in feats]
+
+# allow operations.Linear inheritance
+class SparseLinear:
+    def __new__(cls, in_features, out_features, bias=True, device=None, dtype=None, operations=nn, *args, **kwargs):
+        class _SparseLinear(operations.Linear):
+            def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
+                super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
+
+            def forward(self, input: VarLenTensor) -> VarLenTensor:
+                return input.replace(super().forward(input.feats))
+
+        return _SparseLinear(in_features, out_features, bias=bias, device=device, dtype=dtype, *args, **kwargs)
+
+MIX_PRECISION_MODULES = (
+    nn.Conv1d,
+    nn.Conv2d,
+    nn.Conv3d,
+    nn.ConvTranspose1d,
+    nn.ConvTranspose2d,
+    nn.ConvTranspose3d,
+    nn.Linear,
+    SparseConv3d,
+    SparseLinear,
+)
+
+
+def convert_module_to_f16(l):
+    if isinstance(l, MIX_PRECISION_MODULES):
+        for p in l.parameters():
+            p.data = p.data.half()
+
+class SparseUnetVaeDecoder(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.use_fp16 = use_fp16
+        self.pred_subdiv = pred_subdiv
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.low_vram = False
+
+        self.output_layer = SparseLinear(model_channels[-1], out_channels)
+        self.from_latent = SparseLinear(latent_channels, model_channels[0])
+
+        self.blocks = nn.ModuleList([])
+        for i in range(len(num_blocks)):
+            self.blocks.append(nn.ModuleList([]))
+            for j in range(num_blocks[i]):
+                self.blocks[-1].append(
+                    globals()[block_type[i]](
+                        model_channels[i],
+                        **block_args[i],
+                    )
+                )
+            if i < len(num_blocks) - 1:
+                self.blocks[-1].append(
+                    globals()[up_block_type[i]](
+                        model_channels[i],
+                        model_channels[i+1],
+                        pred_subdiv=pred_subdiv,
+                        **block_args[i],
+                    )
+                )
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(self, x: SparseTensor, guide_subs: Optional[List[SparseTensor]] = None, return_subs: bool = False) -> SparseTensor:
+
+        dtype = next(self.from_latent.parameters()).dtype
+        device = next(self.from_latent.parameters()).device
+        x.feats = x.feats.to(dtype).to(device)
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        subs = []
+        for i, res in enumerate(self.blocks):
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    if self.pred_subdiv:
+                        h, sub = block(h)
+                        subs.append(sub)
+                    else:
+                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
+                else:
+                    h = block(h)
+        h = h.type(x.feats.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.output_layer(h)
+        if return_subs:
+            return h, subs
+        else:
+            return h
+
+    def upsample(self, x: SparseTensor, upsample_times: int) -> torch.Tensor:
+
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        for i, res in enumerate(self.blocks):
+            if i == upsample_times:
+                return h.coords
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    h, sub = block(h)
+                else:
+                    h = block(h)
+
+class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        voxel_margin: float = 0.5,
+        use_fp16: bool = False,
+    ):
+        self.resolution = resolution
+        self.voxel_margin = voxel_margin
+        # cache for a TorchHashMap instance
+        self._torch_hashmap_cache = None
+
+        super().__init__(
+            7,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            up_block_type,
+            block_args,
+            use_fp16,
+        )
+
+    def set_resolution(self, resolution: int) -> None:
+        self.resolution = resolution
+
+    def _build_or_get_hashmap(self, coords: torch.Tensor, grid_size: torch.Tensor):
+        device = coords.device
+        N = coords.shape[0]
+        # compute flat keys for all coords (prepend batch 0 same as original code)
+        b = torch.zeros((N,), dtype=torch.long, device=device)
+        x, y, z = coords[:, 0].to(torch.int32), coords[:, 1].to(torch.int32), coords[:, 2].to(torch.int32)
+        W, H, D = int(grid_size[0].item()), int(grid_size[1].item()), int(grid_size[2].item())
+        flat_keys = b * (W * H * D) + x * (H * D) + y * D + z
+        values = torch.arange(N, dtype=torch.int32, device=device)
+        DEFAULT_VAL = 0xffffffff  # sentinel used in original code
+        return TorchHashMap(flat_keys, values, DEFAULT_VAL)
+
+    def forward(self, x: SparseTensor, gt_intersected: SparseTensor = None, **kwargs):
+        decoded = super().forward(x, **kwargs)
+        out_list = list(decoded) if isinstance(decoded, tuple) else [decoded]
+        h = out_list[0]
+        vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
+        intersected = h.replace(h.feats[..., 3:6] > 0)
+        quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
+        mesh = [Mesh(*flexible_dual_grid_to_mesh(
+            v.coords[:, 1:], v.feats, i.feats, q.feats,
+            aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+            grid_size=self.resolution,
+            train=False,
+            hashmap_builder=self._build_or_get_hashmap,
+        )) for v, i, q in zip(vertices, intersected, quad_lerp)]
+        out_list[0] = mesh
+        return out_list[0] if len(out_list) == 1 else tuple(out_list)
+
+def flexible_dual_grid_to_mesh(
+    coords: torch.Tensor,
+    dual_vertices: torch.Tensor,
+    intersected_flag: torch.Tensor,
+    split_weight: Union[torch.Tensor, None],
+    aabb: Union[list, tuple, np.ndarray, torch.Tensor],
+    voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None,
+    grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None,
+    train: bool = False,
+    hashmap_builder=None,  # optional callable for building/caching a TorchHashMap
+):
+
+    device = coords.device
+    if not hasattr(flexible_dual_grid_to_mesh, "edge_neighbor_voxel_offset") \
+        or flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset.device != device:
+        flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset = torch.tensor([
+            [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]],     # x-axis
+            [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]],     # y-axis
+            [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]],     # z-axis
+        ], dtype=torch.int, device=device).unsqueeze(0)
+    if not hasattr(flexible_dual_grid_to_mesh, "quad_split_1") or flexible_dual_grid_to_mesh.quad_split_1.device != device:
+        flexible_dual_grid_to_mesh.quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device, requires_grad=False)
+    if not hasattr(flexible_dual_grid_to_mesh, "quad_split_2") or flexible_dual_grid_to_mesh.quad_split_2.device != device:
+        flexible_dual_grid_to_mesh.quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device, requires_grad=False)
+    if not hasattr(flexible_dual_grid_to_mesh, "quad_split_train") or flexible_dual_grid_to_mesh.quad_split_train.device != device:
+        flexible_dual_grid_to_mesh.quad_split_train = torch.tensor([0, 1, 4, 1, 2, 4, 2, 3, 4, 3, 0, 4], dtype=torch.long, device=device, requires_grad=False)
+
+    # AABB
+    if isinstance(aabb, (list, tuple)):
+        aabb = np.array(aabb)
+    if isinstance(aabb, np.ndarray):
+        aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
+
+    # Voxel size
+    if voxel_size is not None:
+        if isinstance(voxel_size, float):
+            voxel_size = [voxel_size, voxel_size, voxel_size]
+        if isinstance(voxel_size, (list, tuple)):
+            voxel_size = np.array(voxel_size)
+        if isinstance(voxel_size, np.ndarray):
+            voxel_size = torch.tensor(voxel_size, dtype=torch.float32, device=coords.device)
+        grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
+    else:
+        if isinstance(grid_size, int):
+            grid_size = [grid_size, grid_size, grid_size]
+        if isinstance(grid_size, (list, tuple)):
+            grid_size = np.array(grid_size)
+        if isinstance(grid_size, np.ndarray):
+            grid_size = torch.tensor(grid_size, dtype=torch.int32, device=coords.device)
+        voxel_size = (aabb[1] - aabb[0]) / grid_size
+
+    # Extract mesh
+    N = dual_vertices.shape[0]
+
+    if hashmap_builder is None:
+        # build local TorchHashMap
+        device = coords.device
+        b = torch.zeros((N,), dtype=torch.long, device=device)
+        x, y, z = coords[:, 0].to(torch.int32), coords[:, 1].to(torch.int32), coords[:, 2].to(torch.int32)
+        W, H, D = int(grid_size[0].item()), int(grid_size[1].item()), int(grid_size[2].item())
+        flat_keys = b * (W * H * D) + x * (H * D) + y * D + z
+        values = torch.arange(N, dtype=torch.long, device=device)
+        DEFAULT_VAL = 0xffffffff
+        torch_hashmap = TorchHashMap(flat_keys, values, DEFAULT_VAL)
+    else:
+        torch_hashmap = hashmap_builder(coords, grid_size)
+
+    # Find connected voxels
+    edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset      # (N, 3, 4, 3)
+    connected_voxel = edge_neighbor_voxel[intersected_flag]                           # (M, 4, 3)
+    M = connected_voxel.shape[0]
+    # flatten connected voxel coords and lookup
+    conn_flat_b = torch.zeros((M * 4,), dtype=torch.long, device=coords.device)
+    conn_x = connected_voxel.reshape(-1, 3)[:, 0].to(torch.int32)
+    conn_y = connected_voxel.reshape(-1, 3)[:, 1].to(torch.int32)
+    conn_z = connected_voxel.reshape(-1, 3)[:, 2].to(torch.int32)
+    W, H, D = int(grid_size[0].item()), int(grid_size[1].item()), int(grid_size[2].item())
+    conn_flat = conn_flat_b * (W * H * D) + conn_x * (H * D) + conn_y * D + conn_z
+
+    conn_indices = torch_hashmap.lookup_flat(conn_flat).reshape(M, 4).int()
+    connected_voxel_valid = (conn_indices != 0xffffffff).all(dim=1)
+    quad_indices = conn_indices[connected_voxel_valid].int()                             # (L, 4)
+
+    mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3)
+    if split_weight is None:
+        # if split 1
+        atempt_triangles_0 = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_1]
+        normals0 = torch.cross(mesh_vertices[atempt_triangles_0[:, 1]] - mesh_vertices[atempt_triangles_0[:, 0]], mesh_vertices[atempt_triangles_0[:, 2]] - mesh_vertices[atempt_triangles_0[:, 0]])
+        normals1 = torch.cross(mesh_vertices[atempt_triangles_0[:, 2]] - mesh_vertices[atempt_triangles_0[:, 1]], mesh_vertices[atempt_triangles_0[:, 3]] - mesh_vertices[atempt_triangles_0[:, 1]])
+        align0 = (normals0 * normals1).sum(dim=1, keepdim=True).abs()
+        # if split 2
+        atempt_triangles_1 = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_2]
+        normals0 = torch.cross(mesh_vertices[atempt_triangles_1[:, 1]] - mesh_vertices[atempt_triangles_1[:, 0]], mesh_vertices[atempt_triangles_1[:, 2]] - mesh_vertices[atempt_triangles_1[:, 0]])
+        normals1 = torch.cross(mesh_vertices[atempt_triangles_1[:, 2]] - mesh_vertices[atempt_triangles_1[:, 1]], mesh_vertices[atempt_triangles_1[:, 3]] - mesh_vertices[atempt_triangles_1[:, 1]])
+        align1 = (normals0 * normals1).sum(dim=1, keepdim=True).abs()
+        # select split
+        mesh_triangles = torch.where(align0 > align1, atempt_triangles_0, atempt_triangles_1).reshape(-1, 3)
+    else:
+        split_weight_ws = split_weight[quad_indices]
+        split_weight_ws_02 = split_weight_ws[:, 0] * split_weight_ws[:, 2]
+        split_weight_ws_13 = split_weight_ws[:, 1] * split_weight_ws[:, 3]
+        mesh_triangles = torch.where(
+            split_weight_ws_02 > split_weight_ws_13,
+            quad_indices[:, flexible_dual_grid_to_mesh.quad_split_1],
+            quad_indices[:, flexible_dual_grid_to_mesh.quad_split_2]
+        ).reshape(-1, 3)
+
+    return mesh_vertices, mesh_triangles
+
+class ChannelLayerNorm32(LayerNorm32):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        DIM = x.dim()
+        x = x.permute(0, *range(2, DIM), 1).contiguous()
+        x = super().forward(x)
+        x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+        return x
+
+class UpsampleBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mode = "conv",
+    ):
+        assert mode in ["conv", "nearest"], f"Invalid mode {mode}"
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if mode == "conv":
+            self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1)
+        elif mode == "nearest":
+            assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "conv"):
+            x = self.conv(x)
+            return pixel_shuffle_3d(x, 2)
+        else:
+            return F.interpolate(x, scale_factor=2, mode="nearest")
+
+def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module:
+    return ChannelLayerNorm32(*args, **kwargs)
+
+class ResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        norm_type = "layer",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+
+        self.norm1 = norm_layer(norm_type, channels)
+        self.norm2 = norm_layer(norm_type, self.out_channels)
+        self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
+        self.conv2 = nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1)
+        self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = F.silu(h)
+        dtype = next(self.conv1.parameters()).dtype
+        h = h.to(dtype)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+
+
+class SparseStructureDecoder(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        latent_channels: int,
+        num_res_blocks: int,
+        channels: List[int],
+        num_res_blocks_middle: int = 2,
+        norm_type = "layer",
+        use_fp16: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.latent_channels = latent_channels
+        self.num_res_blocks = num_res_blocks
+        self.channels = channels
+        self.num_res_blocks_middle = num_res_blocks_middle
+        self.norm_type = norm_type
+        self.use_fp16 = use_fp16
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
+
+        self.middle_block = nn.Sequential(*[
+            ResBlock3d(channels[0], channels[0])
+            for _ in range(num_res_blocks_middle)
+        ])
+
+        self.blocks = nn.ModuleList([])
+        for i, ch in enumerate(channels):
+            self.blocks.extend([
+                ResBlock3d(ch, ch)
+                for _ in range(num_res_blocks)
+            ])
+            if i < len(channels) - 1:
+                self.blocks.append(
+                    UpsampleBlock3d(ch, channels[i+1])
+                )
+
+        self.out_layer = nn.Sequential(
+            norm_layer(norm_type, channels[-1]),
+            nn.SiLU(),
+            nn.Conv3d(channels[-1], out_channels, 3, padding=1)
+        )
+
+        if use_fp16:
+            self.convert_to_fp16()
+
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        self.use_fp16 = True
+        self.dtype = torch.float16
+        self.blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = next(self.input_layer.parameters()).dtype
+        x = x.to(dtype)
+        h = self.input_layer(x)
+
+        h = h.type(self.dtype)
+        h = self.middle_block(h)
+        for block in self.blocks:
+            h = block(h)
+
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+        return h
+
+class Vae(nn.Module):
+    def __init__(self, init_txt_model, init_txt_model_only, operations=None):
+        super().__init__()
+        operations = operations or torch.nn
+        if init_txt_model or init_txt_model_only:
+            self.txt_dec = SparseUnetVaeDecoder(
+                out_channels=6,
+                model_channels=[1024, 512, 256, 128, 64],
+                latent_channels=32,
+                num_blocks=[4, 16, 8, 4, 0],
+                block_type=["SparseConvNeXtBlock3d"] * 5,
+                up_block_type=["SparseResBlockC2S3d"] * 4,
+                block_args=[{}, {}, {}, {}, {}],
+                pred_subdiv=False
+            )
+
+        if not init_txt_model_only:
+            self.shape_dec = FlexiDualGridVaeDecoder(
+                resolution=256,
+                model_channels=[1024, 512, 256, 128, 64],
+                latent_channels=32,
+                num_blocks=[4, 16, 8, 4, 0],
+                block_type=["SparseConvNeXtBlock3d"] * 5,
+                up_block_type=["SparseResBlockC2S3d"] * 4,
+                block_args=[{}, {}, {}, {}, {}],
+            )
+
+            self.struct_dec = SparseStructureDecoder(
+                out_channels=1,
+                latent_channels=8,
+                num_res_blocks=2,
+                num_res_blocks_middle=2,
+                channels=[512, 128, 32],
+            )
+        self.register_buffer("resolution", torch.tensor(1024.0), persistent=False)
+
+    @torch.no_grad()
+    def decode_shape_slat(self, slat, resolution: int):
+        self.shape_dec.set_resolution(resolution)
+        return self.shape_dec(slat, return_subs=True)
+
+    @torch.no_grad()
+    def decode_tex_slat(self, slat, subs):
+        if self.txt_dec is None:
+            raise ValueError("Checkpoint doesn't include texture model")
+        return self.txt_dec(slat, guide_subs=subs) * 0.5 + 0.5
+
+    # shouldn't be called (placeholder)
+    @torch.no_grad()
+    def decode(
+        self,
+        shape_slat: SparseTensor,
+        tex_slat: SparseTensor,
+        resolution: int,
+    ):
+        meshes, subs = self.decode_shape_slat(shape_slat, resolution)
+        tex_voxels = self.decode_tex_slat(tex_slat, subs)
+        return tex_voxels
diff --git a/comfy/model_base.py b/comfy/model_base.py
index d81f13c69..2a27e7cca 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -53,6 +53,7 @@ import comfy.ldm.omnigen.omnigen2
 import comfy.ldm.qwen_image.model
 import comfy.ldm.kandinsky5.model
 import comfy.ldm.anima.model
+import comfy.ldm.trellis2.model
 import comfy.ldm.ace.ace_step15
 import comfy.ldm.cogvideo.model
 import comfy.ldm.rt_detr.rtdetr_v4
@@ -1637,6 +1638,16 @@ class WAN22(WAN21):
     def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
         return latent_image
 
+class Trellis2(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None, unet_model=comfy.ldm.trellis2.model.Trellis2):
+        super().__init__(model_config, model_type, device, unet_model)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        embeds = kwargs.get("embeds")
+        out["embeds"] = comfy.conds.CONDRegular(embeds)
+        return out
+
 class WAN21_FlowRVS(WAN21):
     def __init__(self, model_config, model_type=ModelType.IMG_TO_IMG_FLOW, image_to_video=False, device=None):
         model_config.unet_config["model_type"] = "t2v"
@@ -1678,7 +1689,6 @@ class WAN21_SCAIL(WAN21):
         pose_latents = kwargs.get("pose_video_latent", None)
         if pose_latents is not None:
             out['pose_latents'] = [pose_latents.shape[0], 20, *pose_latents.shape[2:]]
-
         return out
 
 class WAN22_WanDancer(WAN21):
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 70b4df8b3..3baf5c501 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -113,6 +113,30 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]
         return unet_config
 
+    if '{}img2shape.blocks.1.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys:
+        unet_config = {}
+        unet_config["image_model"] = "trellis2"
+
+        unet_config["init_txt_model"] = False
+        if '{}shape2txt.blocks.29.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys:
+            unet_config["init_txt_model"] = True
+
+        unet_config["resolution"] = 64
+        if metadata is not None:
+            if "is_512" in metadata:
+                unet_config["resolution"] = 32
+
+        unet_config["num_heads"] = 12
+        return unet_config
+
+    if '{}shape2txt.blocks.29.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys: # trellis2 texture
+        unet_config = {}
+        unet_config["image_model"] = "trellis2"
+        unet_config["resolution"] = 64
+        unet_config["num_heads"] = 12
+        unet_config["txt_only"] = True
+        return unet_config
+
     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"
diff --git a/comfy/sd.py b/comfy/sd.py
index 7bd07ed3a..8bbf77ce9 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -15,6 +15,7 @@ import comfy.ldm.lightricks.vae.causal_video_autoencoder
 import comfy.ldm.lightricks.vae.audio_vae
 import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
+import comfy.ldm.trellis2.vae
 import comfy.ldm.wan.vae2_2
 import comfy.ldm.hunyuan3d.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
@@ -528,6 +529,18 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
+            elif "shape_dec.blocks.1.16.to_subdiv.weight" in sd or "txt_dec.blocks.3.4.conv2.weight" in sd: # trellis2 or trellis2 texture only
+                init_txt_model = False
+                init_txt_model_only = False
+                if "shape_dec.blocks.1.16.to_subdiv.weight" not in sd:
+                    init_txt_model_only = True
+                if "txt_dec.blocks.1.16.norm1.weight" in sd:
+                    init_txt_model = True
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                # TODO
+                self.memory_used_decode = lambda shape, dtype: (2500 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (2500 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.first_stage_model = comfy.ldm.trellis2.vae.Vae(init_txt_model, init_txt_model_only= init_txt_model_only)
             elif "decoder.conv_in.weight" in sd:
                 if sd['decoder.conv_in.weight'].shape[1] == 64:
                     ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 617db4f28..f02a370f0 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -1318,6 +1318,29 @@ class WAN22_T2V(WAN21_T2V):
         out = model_base.WAN22(self, image_to_video=True, device=device)
         return out
 
+class Trellis2(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "trellis2"
+    }
+
+    sampling_settings = {
+        "shift": 3.0,
+    }
+
+    memory_usage_factor = 3.5
+
+    latent_format = latent_formats.Trellis2
+    vae_key_prefix = ["vae."]
+    clip_vision_prefix = "conditioner.main_image_encoder.model."
+    # this is only needed for the texture model
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.Trellis2(self, device=device)
+
+    def clip_target(self, state_dict={}):
+        return None
+
 class WAN21_FlowRVS(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",
@@ -1784,6 +1807,7 @@ class Kandinsky5Image(Kandinsky5):
         return supported_models_base.ClipTarget(comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage, comfy.text_encoders.kandinsky5.te(**hunyuan_detect))
 
 
+
 class ACEStep15(supported_models_base.BASE):
     unet_config = {
         "audio_model": "ace1.5",
@@ -1823,7 +1847,6 @@ class ACEStep15(supported_models_base.BASE):
 
         return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
 
-
 class LongCatImage(supported_models_base.BASE):
     unet_config = {
         "image_model": "flux",
@@ -1901,6 +1924,7 @@ class ErnieImage(supported_models_base.BASE):
         return supported_models_base.ClipTarget(comfy.text_encoders.ernie.ErnieTokenizer, comfy.text_encoders.ernie.te(**hunyuan_detect))
 
 
+
 class SAM3(supported_models_base.BASE):
     unet_config = {"image_model": "SAM3"}
     supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
@@ -2020,7 +2044,6 @@ class CogVideoX_Inpaint(CogVideoX_T2V):
         out = model_base.CogVideoX(self, image_to_video=True, device=device)
         return out
 
-
 models = [
     LotusD,
     Stable_Zero123,
@@ -2107,4 +2130,5 @@ models = [
     CogVideoX_I2V,
     CogVideoX_T2V,
     SVD_img2vid,
+    Trellis2
 ]
diff --git a/comfy_api/latest/_util/geometry_types.py b/comfy_api/latest/_util/geometry_types.py
index cdde60b10..b821fd620 100644
--- a/comfy_api/latest/_util/geometry_types.py
+++ b/comfy_api/latest/_util/geometry_types.py
@@ -7,9 +7,10 @@ import torch
 
 
 class VOXEL:
-    def __init__(self, data: torch.Tensor):
+    def __init__(self, data: torch.Tensor, voxel_colors=None, resolution=None):
         self.data = data
-
+        self.voxel_colors = voxel_colors
+        self.resolution = resolution # each 3d model has its own resolution
 
 class MESH:
     def __init__(self, vertices: torch.Tensor, faces: torch.Tensor,
diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py
new file mode 100644
index 000000000..74eb9e0c1
--- /dev/null
+++ b/comfy_extras/nodes_mesh_postprocess.py
@@ -0,0 +1,845 @@
+import torch
+import numpy as np
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, IO, Types
+import copy
+import comfy.utils
+import logging
+import scipy
+
+def get_mesh_batch_item(mesh, index):
+    if hasattr(mesh, "vertex_counts") and mesh.vertex_counts is not None:
+        vertex_count = int(mesh.vertex_counts[index].item())
+        face_count = int(mesh.face_counts[index].item())
+        vertices = mesh.vertices[index, :vertex_count]
+        faces = mesh.faces[index, :face_count]
+        colors = None
+        if hasattr(mesh, "colors") and mesh.colors is not None:
+            if hasattr(mesh, "color_counts") and mesh.color_counts is not None:
+                color_count = int(mesh.color_counts[index].item())
+                colors = mesh.colors[index, :color_count]
+            else:
+                colors = mesh.colors[index, :vertex_count]
+        return vertices, faces, colors
+
+    colors = None
+    if hasattr(mesh, "colors") and mesh.colors is not None:
+        colors = mesh.colors[index]
+    return mesh.vertices[index], mesh.faces[index], colors
+
+def pack_variable_mesh_batch(vertices, faces, colors=None):
+    batch_size = len(vertices)
+    max_vertices = max(v.shape[0] for v in vertices)
+    max_faces = max(f.shape[0] for f in faces)
+
+    packed_vertices = vertices[0].new_zeros((batch_size, max_vertices, vertices[0].shape[1]))
+    packed_faces = faces[0].new_zeros((batch_size, max_faces, faces[0].shape[1]))
+    vertex_counts = torch.tensor([v.shape[0] for v in vertices], device=vertices[0].device, dtype=torch.int64)
+    face_counts = torch.tensor([f.shape[0] for f in faces], device=faces[0].device, dtype=torch.int64)
+
+    for i, (v, f) in enumerate(zip(vertices, faces)):
+        packed_vertices[i, :v.shape[0]] = v
+        packed_faces[i, :f.shape[0]] = f
+
+    mesh = Types.MESH(packed_vertices, packed_faces)
+    mesh.vertex_counts = vertex_counts
+    mesh.face_counts = face_counts
+
+    if colors is not None:
+        max_colors = max(c.shape[0] for c in colors)
+        packed_colors = colors[0].new_zeros((batch_size, max_colors, colors[0].shape[1]))
+        color_counts = torch.tensor([c.shape[0] for c in colors], device=colors[0].device, dtype=torch.int64)
+        for i, c in enumerate(colors):
+            packed_colors[i, :c.shape[0]] = c
+        mesh.vertex_colors = packed_colors
+        mesh.color_counts = color_counts
+
+    return mesh
+
+
+def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution):
+    """
+    Generic function to paint a mesh using nearest-neighbor colors from a sparse voxel field.
+    """
+    device = comfy.model_management.vae_offload_device()
+
+    origin = torch.tensor([-0.5, -0.5, -0.5], device=device)
+    voxel_size = 1.0 / resolution
+
+    # map voxels
+    voxel_pos = voxel_coords.to(device).float() * voxel_size + origin
+    verts = mesh.vertices.to(device).squeeze(0)
+    voxel_colors = voxel_colors.to(device)
+
+    voxel_pos_np = voxel_pos.numpy()
+    verts_np = verts.numpy()
+
+    tree = scipy.spatial.cKDTree(voxel_pos_np)
+
+    # nearest neighbour k=1
+    _, nearest_idx_np = tree.query(verts_np, k=1, workers=-1)
+
+    nearest_idx = torch.from_numpy(nearest_idx_np).long()
+    v_colors = voxel_colors[nearest_idx]
+
+    # to [0, 1]
+    srgb_colors = v_colors.clamp(0, 1)#(v_colors * 0.5 + 0.5).clamp(0, 1)
+
+    # to Linear RGB (required for GLTF)
+    linear_colors = torch.pow(srgb_colors, 2.2)
+
+    final_colors = linear_colors.unsqueeze(0)
+
+    out_mesh = copy.deepcopy(mesh)
+    out_mesh.vertex_colors = final_colors
+
+    return out_mesh
+
+class PaintMesh(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="PaintMesh",
+            display_name="Paint Mesh",
+            category="latent/3d",
+            description=(
+                "Paints the mesh using colors from the input voxel field by matching each vertex "
+                "to the nearest voxel color."
+            ),
+            inputs=[
+                IO.Mesh.Input("mesh"),
+                IO.Voxel.Input("voxel_colors")
+            ],
+            outputs=[
+                IO.Mesh.Output("mesh"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, mesh, voxel_colors):
+        voxels = voxel_colors
+        coords = voxels.data
+        colors = voxels.voxel_colors
+        resolution = voxels.resolution
+
+        if coords.shape[0] == 0:
+            return IO.NodeOutput(paint_mesh_default_colors(mesh))
+
+        mesh_batch_size = mesh.vertices.shape[0]
+
+        if coords.shape[-1] == 4 and mesh_batch_size > 1:
+            batch_idx = coords[:, 0].long()
+            voxel_coords = coords[:, 1:]
+            mesh_batch_size = mesh.vertices.shape[0]
+
+            out_verts, out_faces, out_colors = [], [], []
+            for i in range(mesh_batch_size):
+                sel = batch_idx == i
+                item_coords = voxel_coords[sel]
+                item_colors = colors[sel]
+                item_vertices, item_faces, _ = get_mesh_batch_item(mesh, i)
+                item_mesh = Types.MESH(vertices=item_vertices.unsqueeze(0), faces=item_faces.unsqueeze(0))
+
+                if item_coords.shape[0] == 0:
+                    painted = paint_mesh_default_colors(item_mesh)
+                else:
+                    painted = paint_mesh_with_voxels(item_mesh, item_coords, item_colors, resolution=resolution)
+
+                out_verts.append(painted.vertices.squeeze(0))
+                out_faces.append(painted.faces.squeeze(0))
+                out_colors.append(painted.vertex_colors.squeeze(0))
+
+            out_mesh = pack_variable_mesh_batch(out_verts, out_faces, out_colors)
+            return IO.NodeOutput(out_mesh)
+
+        if coords.shape[-1] == 4:
+            coords = coords[:, 1:]
+
+        out_mesh = paint_mesh_with_voxels(mesh, coords, colors, resolution=resolution)
+        return IO.NodeOutput(out_mesh)
+
+def paint_mesh_default_colors(mesh):
+    out_mesh = copy.copy(mesh)
+    vertex_count = mesh.vertices.shape[1]
+    out_mesh.vertex_colors = mesh.vertices.new_zeros((1, vertex_count, 3))
+    return out_mesh
+
+
+def fill_holes_fn(vertices, faces, max_perimeter=0.03):
+    is_batched = vertices.ndim == 3
+    if is_batched:
+        v_list, f_list = [], []
+        for i in range(vertices.shape[0]):
+            v_i, f_i = fill_holes_fn(vertices[i], faces[i], max_perimeter)
+            v_list.append(v_i)
+            f_list.append(f_i)
+        max_v = max(v.shape[0] for v in v_list)
+        for i in range(len(v_list)):
+            if v_list[i].shape[0] < max_v:
+                pad = torch.zeros(max_v - v_list[i].shape[0], 3, device=v_list[i].device, dtype=v_list[i].dtype)
+                v_list[i] = torch.cat([v_list[i], pad], dim=0)
+        return torch.stack(v_list), torch.stack(f_list)
+
+    device = vertices.device
+    v = vertices
+    f = faces
+
+    if f.numel() == 0:
+        return v, f
+
+    edges = torch.cat([f[:, [0, 1]], f[:, [1, 2]], f[:, [2, 0]]], dim=0)
+    edges_sorted, _ = torch.sort(edges, dim=1)
+    max_v = v.shape[0]
+    packed_undirected = edges_sorted[:, 0].long() * max_v + edges_sorted[:, 1].long()
+    unique_packed, counts = torch.unique(packed_undirected, return_counts=True)
+    boundary_packed = unique_packed[counts == 1]
+
+    if boundary_packed.numel() == 0:
+        return v, f
+
+    boundary_mask = torch.isin(packed_undirected, boundary_packed)
+    b_edges = edges_sorted[boundary_mask]
+
+    adj = {}
+    for i in range(b_edges.shape[0]):
+        a = b_edges[i, 0].item()
+        b = b_edges[i, 1].item()
+        adj.setdefault(a, []).append(b)
+        adj.setdefault(b, []).append(a)
+
+    # Trace all boundary loops
+    loops = []
+    visited = set()
+    for start_node in adj.keys():
+        if start_node in visited:
+            continue
+        curr = start_node
+        prev = -1
+        loop = []
+        while curr not in visited:
+            visited.add(curr)
+            loop.append(curr)
+            neighbors = adj[curr]
+            candidates = [n for n in neighbors if n != prev]
+            if not candidates:
+                loop = []
+                break
+            next_node = candidates[0]
+            prev, curr = curr, next_node
+            if curr == start_node:
+                loops.append(loop)
+                break
+
+    if not loops:
+        return v, f
+
+    # Mesh normal for winding orientation only
+    face_normals = torch.linalg.cross(
+        v[f[:, 1]] - v[f[:, 0]],
+        v[f[:, 2]] - v[f[:, 0]],
+        dim=-1
+    )
+    mesh_normal = face_normals.mean(dim=0)
+    mesh_normal = mesh_normal / (torch.norm(mesh_normal) + 1e-8)
+
+    # === FIX: Fill ALL boundary loops below perimeter threshold ===
+    new_verts = []
+    new_faces = []
+    v_idx = v.shape[0]
+
+    for loop in loops:
+        loop_t = torch.tensor(loop, device=device, dtype=torch.long)
+        loop_v = v[loop_t]
+
+        # Perimeter check
+        next_v = torch.roll(loop_v, -1, dims=0)
+        diffs = loop_v - next_v
+        perimeter = torch.norm(diffs, dim=1).sum().item()
+
+        if perimeter > max_perimeter:
+            continue
+
+        # Ensure CCW winding consistent with mesh
+        cross = torch.linalg.cross(loop_v, next_v, dim=-1)
+        loop_normal = cross.sum(dim=0)
+        loop_normal = loop_normal / (torch.norm(loop_normal) + 1e-8)
+        if torch.dot(loop_normal, mesh_normal) < 0:
+            loop = loop[::-1]
+            loop_t = torch.tensor(loop, device=device, dtype=torch.long)
+            loop_v = v[loop_t]
+
+        if len(loop) == 3:
+            new_faces.append([loop[0], loop[1], loop[2]])
+        else:
+            centroid = loop_v.mean(dim=0)
+            new_verts.append(centroid)
+            for i in range(len(loop)):
+                new_faces.append([loop[i], loop[(i + 1) % len(loop)], v_idx])
+            v_idx += 1
+
+    if new_verts:
+        v = torch.cat([v, torch.stack(new_verts)], dim=0)
+    if new_faces:
+        f = torch.cat([f, torch.tensor(new_faces, device=device, dtype=torch.long)], dim=0)
+
+    return v, f
+
+def _cleanup_mesh(verts, faces, min_angle_deg=0.5, max_aspect=100.0):
+    if faces.numel() == 0:
+        return verts, faces
+
+    v0 = verts[faces[:, 0]]
+    v1 = verts[faces[:, 1]]
+    v2 = verts[faces[:, 2]]
+    e0 = v1 - v0
+    e1 = v2 - v1
+    e2 = v0 - v2
+    l0 = torch.norm(e0, dim=-1)
+    l1 = torch.norm(e1, dim=-1)
+    l2 = torch.norm(e2, dim=-1)
+    n = torch.cross(e0, e2, dim=-1)
+    area = torch.norm(n, dim=-1)
+
+    max_edge = torch.max(torch.max(l0, l1), l2)
+    aspect = max_edge * max_edge / (2.0 * area + 1e-12)
+
+    cos_a = (l1 * l1 + l2 * l2 - l0 * l0) / (2 * l1 * l2 + 1e-12)
+    cos_b = (l0 * l0 + l2 * l2 - l1 * l1) / (2 * l0 * l2 + 1e-12)
+    cos_c = (l0 * l0 + l1 * l1 - l2 * l2) / (2 * l0 * l1 + 1e-12)
+    cos_all = torch.stack([cos_a, cos_b, cos_c], dim=-1)
+    angles = torch.acos(torch.clamp(cos_all, -1, 1)) * 180 / np.pi
+
+    good = (aspect < max_aspect) & (angles.min(dim=1)[0] > min_angle_deg) & (area > 1e-12)
+    faces = faces[good]
+
+    if faces.numel() == 0:
+        return verts, faces
+
+    used = torch.zeros(verts.shape[0], dtype=torch.bool, device=verts.device)
+    used[faces[:, 0]] = True
+    used[faces[:, 1]] = True
+    used[faces[:, 2]] = True
+
+    remap = torch.full((verts.shape[0],), -1, dtype=torch.int64, device=verts.device)
+    remap[used] = torch.arange(used.sum().item(), device=verts.device)
+    verts = verts[used]
+    faces = remap[faces]
+    return verts, faces
+
+def _pytorch_edge_errors_fast(verts, Q, edges, stabilizer, max_edge_length_sq, mesh_scale_sq):
+    n_edges = edges.shape[0]
+    dtype = verts.dtype
+    if n_edges == 0:
+        return (torch.empty((0, 3), dtype=dtype, device=verts.device),
+                torch.empty((0,), dtype=dtype, device=verts.device),
+                torch.zeros((0,), dtype=torch.bool, device=verts.device))
+
+    device = verts.device
+    mesh_scale = (mesh_scale_sq) ** 0.5
+
+    va = edges[:, 0]
+    vb = edges[:, 1]
+    Q0 = Q[va]
+    Q1 = Q[vb]
+    Qe = Q0 + Q1
+
+    A = Qe[:, :3, :3] + torch.eye(3, device=device, dtype=dtype).unsqueeze(0) * stabilizer
+    b = -Qe[:, :3, 3].unsqueeze(-1)
+
+    dets = torch.det(A)
+    good = dets.abs() > 1e-12
+    opt = torch.zeros((n_edges, 3), dtype=dtype, device=device)
+
+    if good.any():
+        try:
+            sol = torch.linalg.solve(A[good], b[good])
+            opt[good] = sol.squeeze(-1)
+        except Exception:
+            good = torch.zeros_like(good)
+
+    if (~good).any():
+        bad_idx = torch.nonzero(~good, as_tuple=True)[0]
+        opt[bad_idx] = (verts[va[bad_idx]] + verts[vb[bad_idx]]) * 0.5
+
+    pa = verts[va]
+    pb = verts[vb]
+    el = torch.norm(pb - pa, dim=-1)
+    dist_a = torch.norm(opt - pa, dim=-1)
+    dist_b = torch.norm(opt - pb, dim=-1)
+    wander_bad = (dist_a > 4.0 * el) | (dist_b > 4.0 * el)
+
+    if wander_bad.any():
+        bad_idx = torch.nonzero(wander_bad, as_tuple=True)[0]
+        opt[bad_idx] = (verts[va[bad_idx]] + verts[vb[bad_idx]]) * 0.5
+
+    v4 = torch.cat([opt, torch.ones((n_edges, 1), device=device, dtype=dtype)], dim=1)
+    err = torch.abs(torch.einsum("ei,eij,ej->e", v4, Qe, v4))
+
+    length_ok = el > mesh_scale * 1e-5
+    error_ok = err < max_edge_length_sq
+    nan_ok = ~torch.isnan(opt).any(dim=-1) & ~torch.isnan(err)
+    valid = length_ok & error_ok & nan_ok
+
+    return opt, err, valid
+
+
+def _build_quadrics_fast(verts, faces):
+    v0 = verts[faces[:, 0]]
+    v1 = verts[faces[:, 1]]
+    v2 = verts[faces[:, 2]]
+    e1 = v1 - v0
+    e2 = v2 - v0
+    n = torch.cross(e1, e2, dim=-1)
+    area = torch.norm(n, dim=-1)
+    mask = area > 1e-12
+    n_norm = torch.zeros_like(n)
+    n_norm[mask] = n[mask] / area[mask].unsqueeze(-1)
+    d = -(n_norm * v0).sum(dim=-1, keepdim=True)
+    p = torch.cat([n_norm, d], dim=-1)
+    K = torch.einsum("fi,fj->fij", p, p)
+    K = K * area[:, None, None]
+    V = verts.shape[0]
+    Q = torch.zeros((V, 4, 4), dtype=verts.dtype, device=verts.device)
+    K_flat = K.reshape(-1, 16)
+    Q_flat = Q.reshape(V, 16)
+    for corner in range(3):
+        idx = faces[:, corner].unsqueeze(1).expand(-1, 16)
+        Q_flat.scatter_add_(0, idx, K_flat)
+    return Q_flat.reshape(V, 4, 4)
+
+
+def _gpu_greedy_matching_fast(edges, err, v_alive, max_select):
+    """Vectorized greedy matching.
+
+    Selects an independent set of edges (no two share a vertex) preferring
+    lowest error. Replaces _gpu_greedy_sampled's Python per-edge loop with
+    two scatter_reduce calls.
+    """
+    device = edges.device
+    n_edges = edges.shape[0]
+    if n_edges == 0:
+        return torch.empty(0, dtype=torch.int64, device=device)
+
+    va = edges[:, 0]
+    vb = edges[:, 1]
+    num_verts = v_alive.shape[0]
+
+    # Pack (error_bits, edge_idx) into one int64 so amin gives a unique winner.
+    # err is non-negative finite float32 -> IEEE bits are monotonic.
+    err32 = err.to(torch.float32).clamp(min=0).contiguous()
+    err_bits = err32.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    edge_idx = torch.arange(n_edges, device=device, dtype=torch.int64)
+    key = (err_bits << 32) | edge_idx
+
+    INT64_MAX = torch.iinfo(torch.int64).max
+    best_key = torch.full((num_verts,), INT64_MAX, dtype=torch.int64, device=device)
+    best_key.scatter_reduce_(0, va, key, reduce='amin', include_self=True)
+    best_key.scatter_reduce_(0, vb, key, reduce='amin', include_self=True)
+
+    # An edge wins iff it is the min-key edge incident to BOTH its endpoints
+    # AND both endpoints are still alive.
+    is_winner = (key == best_key[va]) & (key == best_key[vb]) & v_alive[va] & v_alive[vb]
+
+    sel = torch.nonzero(is_winner, as_tuple=True)[0]
+
+    if sel.numel() > max_select:
+        sel_err = err[sel]
+        top = torch.topk(sel_err, max_select, largest=False).indices
+        sel = sel[top]
+
+    return sel
+
+
+def _qem_simplify_fast(vertices, faces_in, colors_in, normals_in, target_faces, device, max_edge_length=None):
+    # Use float32 instead of float64. RTX-class consumer GPUs run FP32 ~32-64x
+    # faster than FP64, and QEM only needs the stabilizer for conditioning.
+    # Always copy=True so we can safely mutate verts/colors/normals in-place.
+    verts = vertices.detach().to(device=device, dtype=torch.float32, copy=True)
+    faces = faces_in.detach().to(device=device, dtype=torch.int64)
+    colors = (
+        colors_in.detach().to(device=device, dtype=torch.float32, copy=True)
+        if colors_in is not None
+        else None
+    )
+    # ADDED: Initialize normals
+    normals = (
+        normals_in.detach().to(device=device, dtype=torch.float32, copy=True)
+        if normals_in is not None
+        else None
+    )
+
+    num_verts = verts.shape[0]
+    num_faces = faces.shape[0]
+
+    logging.debug(f"[QEM-fast] Input: {num_verts} verts, {num_faces} faces, target={target_faces}")
+
+    v_alive = torch.ones(num_verts, dtype=torch.bool, device=device)
+    f_alive = torch.ones(num_faces, dtype=torch.bool, device=device)
+
+    Q = _build_quadrics_fast(verts, faces)
+
+    bbox = verts.max(dim=0)[0] - verts.min(dim=0)[0]
+    mesh_scale = torch.norm(bbox).item()
+
+    if max_edge_length is None or max_edge_length <= 0:
+        max_edge_length = mesh_scale * 2.0
+
+    if max_edge_length < 1e-6:
+        max_edge_length = 1.0
+
+    stabilizer = mesh_scale * mesh_scale * 0.001
+    max_edge_length_sq = max_edge_length * max_edge_length
+    mesh_scale_sq = mesh_scale * mesh_scale
+
+    iteration = 0
+    total_collapses = 0
+    last_faces = num_faces
+
+    while True:
+        n_faces = int(f_alive.sum().item())
+
+        if n_faces <= target_faces:
+            break
+
+        alive_v = torch.nonzero(v_alive, as_tuple=True)[0]
+        alive_f = torch.nonzero(f_alive, as_tuple=True)[0]
+
+        if alive_v.numel() <= 4 or alive_f.numel() == 0:
+            break
+
+        # Compact active mesh
+        vmap = torch.full((num_verts,), -1, dtype=torch.int64, device=device)
+        vmap[alive_v] = torch.arange(alive_v.numel(), device=device)
+
+        active_faces = faces[alive_f]
+        remapped = vmap[active_faces]
+
+        # Extract edges
+        e0 = remapped[:, [0, 1]]
+        e1 = remapped[:, [1, 2]]
+        e2 = remapped[:, [2, 0]]
+        edges = torch.cat([e0, e1, e2], dim=0)
+        edges = torch.sort(edges, dim=1)[0]
+        edges = edges[(edges >= 0).all(dim=1)]
+        edges = edges[edges[:, 0] != edges[:, 1]]
+
+        if edges.shape[0] == 0:
+            break
+
+        # Deduplicate edges
+        num_compact = alive_v.numel()
+        packed = edges[:, 0].long() * num_compact + edges[:, 1].long()
+        packed = torch.unique(packed)
+        edges = torch.stack([packed // num_compact, packed % num_compact], dim=1)
+
+        edges_orig = alive_v[edges]
+
+        # Filter by edge length
+        pa = verts[edges_orig[:, 0]]
+        pb = verts[edges_orig[:, 1]]
+        el = torch.norm(pb - pa, dim=-1)
+        short_enough = el < max_edge_length
+
+        if not short_enough.any():
+            max_edge_length = el.max().item() * 2.0
+            max_edge_length_sq = max_edge_length * max_edge_length
+            short_enough = el < max_edge_length
+            if not short_enough.any():
+                break
+
+        edges_orig = edges_orig[short_enough]
+        if edges_orig.shape[0] == 0:
+            break
+
+        # Sample edges for processing
+        n_edges_total = edges_orig.shape[0]
+        max_edges_to_process = 10_000_000
+
+        if n_edges_total > max_edges_to_process:
+            perm = torch.randint(0, n_edges_total, (max_edges_to_process,), device=device)
+            edges_orig = edges_orig[perm]
+            n_edges = max_edges_to_process
+        else:
+            n_edges = n_edges_total
+
+        optimal, err, valid = _pytorch_edge_errors_fast(
+            verts, Q, edges_orig, stabilizer, max_edge_length_sq, mesh_scale_sq
+        )
+
+        if not valid.any():
+            valid = torch.ones(n_edges, dtype=torch.bool, device=device)
+
+        valid_idx = torch.nonzero(valid, as_tuple=True)[0]
+        edges_orig = edges_orig[valid_idx]
+        optimal = optimal[valid_idx]
+        err = err[valid_idx]
+
+        faces_to_remove = n_faces - target_faces
+        max_collapses = min(1_000_000, max(10_000, faces_to_remove // 4))
+
+        sel = _gpu_greedy_matching_fast(edges_orig, err, v_alive, max_collapses)
+
+        if sel.numel() == 0:
+            break
+
+        v_a = edges_orig[sel, 0]
+        v_b = edges_orig[sel, 1]
+
+        # Apply collapses
+        verts[v_a] = optimal[sel]
+        v_alive[v_b] = False
+        Q[v_a] += Q[v_b]
+
+        if colors is not None:
+            colors[v_a] = (colors[v_a] + colors[v_b]) * 0.5
+
+        if normals is not None:
+            normals[v_a] = (normals[v_a] + normals[v_b]) * 0.5
+
+        merge_map = torch.arange(num_verts, device=device)
+        merge_map[v_b] = v_a
+        faces = merge_map[faces]
+
+        bad = (
+            (faces[:, 0] == faces[:, 1])
+            | (faces[:, 1] == faces[:, 2])
+            | (faces[:, 2] == faces[:, 0])
+        )
+        f_alive &= ~bad
+
+        total_collapses += v_a.numel()
+        iteration += 1
+
+        if iteration % 50 == 0 or n_faces < last_faces * 0.9:
+            logging.debug(f"[QEM-fast] Iter {iteration}: {total_collapses} collapses, {int(f_alive.sum().item())} faces, applied {v_a.numel()}")
+            last_faces = n_faces
+
+        if iteration % 5 == 0 and int(f_alive.sum().item()) < num_faces * 0.5:
+            faces = faces[f_alive]
+            f_alive = torch.ones(faces.shape[0], dtype=torch.bool, device=device)
+            num_faces = faces.shape[0]
+
+        if iteration > 5000:
+            break
+
+    # Finalize
+    final_v = verts[v_alive]
+    final_c = colors[v_alive] if colors is not None else None
+    final_n = normals[v_alive] if normals is not None else None
+
+    remap = torch.full((num_verts,), -1, dtype=torch.int64, device=device)
+    remap[v_alive] = torch.arange(int(v_alive.sum().item()), device=device)
+
+    final_f_raw = faces[f_alive]
+    alive_mask = v_alive[final_f_raw].all(dim=1)
+    final_f_raw = final_f_raw[alive_mask]
+    final_f = remap[final_f_raw]
+    valid_faces = (final_f >= 0).all(dim=1)
+    final_f = final_f[valid_faces]
+
+    if final_f.numel() > 0:
+        final_f = torch.unique(torch.sort(final_f, dim=1)[0], dim=0)
+
+    if final_n is not None and final_f.numel() > 0:
+        v0, v1, v2 = final_v[final_f[:, 0]], final_v[final_f[:, 1]], final_v[final_f[:, 2]]
+
+        # calculate the actual normal of the simplified faces
+        face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+
+        # Get the average reference normal for each face
+        n0, n1, n2 = final_n[final_f[:, 0]], final_n[final_f[:, 1]], final_n[final_f[:, 2]]
+        ref_face_normals = (n0 + n1 + n2) / 3.0
+
+        # Dot product to check if they point in the same direction
+        dot_products = (face_normals * ref_face_normals).sum(dim=-1)
+
+        # Flip the indices of ONLY the incorrect faces (swap vertex 1 and 2)
+        wrong_way_mask = dot_products < 0
+        final_f[wrong_way_mask] = final_f[wrong_way_mask][:, [0, 2, 1]]
+
+    final_v, final_f = _cleanup_mesh(final_v, final_f, min_angle_deg=0.5, max_aspect=100.0)
+
+    return final_v, final_f, final_c, final_n
+
+
+def simplify_fn_fast(vertices, faces, colors=None, normals=None, target=100000, max_edge_length=None):
+    if vertices.ndim == 3:
+        v_list, f_list, c_list, n_list = [], [], [], []
+        for i in range(vertices.shape[0]):
+            c_in = colors[i] if colors is not None else None
+            n_in = normals[i] if normals is not None else None
+            v_i, f_i, c_i, n_i = simplify_fn_fast(vertices[i], faces[i], c_in, n_in, target, max_edge_length)
+            v_list.append(v_i)
+            f_list.append(f_i)
+            if c_i is not None:
+                c_list.append(c_i)
+            if n_i is not None:
+                n_list.append(n_i)
+
+        c_out = torch.stack(c_list) if len(c_list) > 0 else None
+        n_out = torch.stack(n_list) if len(n_list) > 0 else None
+        return torch.stack(v_list), torch.stack(f_list), c_out, n_out
+
+    if faces.shape[0] <= target:
+        return vertices, faces, colors, normals
+
+    device = vertices.device
+    dtype = vertices.dtype
+    face_dtype = faces.dtype
+    color_dtype = colors.dtype if colors is not None else None
+    # ADDED: Normal dtype
+    normal_dtype = normals.dtype if normals is not None else None
+
+    # Pass tensors directly; _qem_simplify_fast handles dtype/device + copy.
+    out_v, out_f, out_c, out_n = _qem_simplify_fast(
+        vertices, faces, colors, normals, target, device, max_edge_length
+    )
+
+    final_v = out_v.to(device=device, dtype=dtype)
+    final_f = out_f.to(device=device, dtype=face_dtype)
+    final_c = (
+        out_c.to(device=device, dtype=color_dtype)
+        if out_c is not None
+        else None
+    )
+    final_n = (
+        out_n.to(device=device, dtype=normal_dtype)
+        if out_n is not None
+        else None
+    )
+    return final_v, final_f, final_c, final_n
+
+def compute_vertex_normals(verts, faces):
+    """Computes area-weighted vertex normals."""
+    # QUICK FIX: Ensure indices are int64 for scatter_add_
+    faces_long = faces.to(torch.int64)
+
+    i0, i1, i2 = faces_long[:, 0], faces_long[:, 1], faces_long[:, 2]
+    v0, v1, v2 = verts[i0], verts[i1], verts[i2]
+
+    # calculate unnormalized face normals (magnitude is proportional to area)
+    face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+
+    # accumulate face normals to vertices
+    vertex_normals = torch.zeros_like(verts)
+    vertex_normals.scatter_add_(0, i0.unsqueeze(-1).expand_as(face_normals), face_normals)
+    vertex_normals.scatter_add_(0, i1.unsqueeze(-1).expand_as(face_normals), face_normals)
+    vertex_normals.scatter_add_(0, i2.unsqueeze(-1).expand_as(face_normals), face_normals)
+
+    return torch.nn.functional.normalize(vertex_normals, p=2, dim=-1, eps=1e-6)
+
+def _process_mesh_batch(mesh, per_item_fn):
+    """Handles list/batched/single mesh dispatching, color extraction, and stacking."""
+    mesh = copy.deepcopy(mesh)
+
+    def process_single(v, f, c, bar):
+        v, f, c = per_item_fn(v, f, c)
+        bar.update(1)
+        return v, f, c
+
+    is_list = isinstance(mesh.vertices, list)
+    is_batched_tensor = not is_list and mesh.vertices.ndim == 3
+
+    if is_list or is_batched_tensor:
+        out_v, out_f, out_c = [], [], []
+        bsz = len(mesh.vertices) if is_list else mesh.vertices.shape[0]
+        bar = comfy.utils.ProgressBar(bsz)
+
+        for i in range(bsz):
+            v_i = mesh.vertices[i]
+            f_i = mesh.faces[i]
+            c_i = None
+            if hasattr(mesh, 'vertex_colors') and mesh.vertex_colors is not None:
+                c_i = mesh.vertex_colors[i] if (isinstance(mesh.vertex_colors, list) or mesh.vertex_colors.ndim == 3) else mesh.vertex_colors
+
+            v_i, f_i, c_i = process_single(v_i, f_i, c_i, bar)
+
+            out_v.append(v_i)
+            out_f.append(f_i)
+            if c_i is not None:
+                out_c.append(c_i)
+
+        if all(v.shape == out_v[0].shape for v in out_v) and all(f.shape == out_f[0].shape for f in out_f):
+            mesh.vertices = torch.stack(out_v)
+            mesh.faces = torch.stack(out_f)
+            if out_c:
+                mesh.vertex_colors = torch.stack(out_c)
+        else:
+            mesh.vertices = out_v
+            mesh.faces = out_f
+            if out_c:
+                mesh.vertex_colors = out_c
+    else:
+        c = mesh.vertex_colors if hasattr(mesh, 'vertex_colors') and mesh.vertex_colors is not None else None
+        bar = comfy.utils.ProgressBar(1)
+        v, f, c = process_single(mesh.vertices, mesh.faces, c, bar)
+        mesh.vertices = v
+        mesh.faces = f
+        if c is not None:
+            mesh.vertex_colors = c
+
+    return IO.NodeOutput(mesh)
+
+
+class DecimateMesh(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="DecimateMesh",
+            display_name="Decimate Mesh",
+            category="latent/3d",
+            description="Simplifies a mesh to a target face count using QEM.",
+            inputs=[
+                IO.Mesh.Input("mesh"),
+                IO.Int.Input("target_face_count", default=200_000, min=0, max=50_000_000,
+                             tooltip="Target maximum number of faces. Set to 0 to disable."),
+            ],
+            outputs=[IO.Mesh.Output("mesh")],
+        )
+
+    @classmethod
+    def execute(cls, mesh, target_face_count):
+        def _fn(v, f, c):
+            if target_face_count > 0 and f.shape[0] > target_face_count:
+                n = compute_vertex_normals(v, f)
+                v, f, c, _ = simplify_fn_fast(v, f, colors=c, normals=n, target=target_face_count)
+            return v, f, c
+        return _process_mesh_batch(mesh, _fn)
+
+
+class FillHoles(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="FillHoles",
+            display_name="Fill Holes",
+            category="latent/3d",
+            description="Fills holes in a mesh up to a maximum perimeter threshold.",
+            inputs=[
+                IO.Mesh.Input("mesh"),
+                IO.Float.Input("max_perimeter", default=0.03, min=0.0, step=0.0001,
+                               tooltip="Maximum hole perimeter to fill. Set to 0 to disable."),
+            ],
+            outputs=[IO.Mesh.Output("mesh")],
+        )
+
+    @classmethod
+    def execute(cls, mesh, max_perimeter):
+        def _fn(v, f, c):
+            if max_perimeter > 0:
+                v, f = fill_holes_fn(v, f, max_perimeter=max_perimeter)
+            return v, f, c
+        return _process_mesh_batch(mesh, _fn)
+
+class PostProcessMeshExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            FillHoles,
+            DecimateMesh,
+            PaintMesh
+        ]
+
+
+async def comfy_entrypoint() -> PostProcessMeshExtension:
+    return PostProcessMeshExtension()
diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py
index c03524246..6657c4490 100644
--- a/comfy_extras/nodes_save_3d.py
+++ b/comfy_extras/nodes_save_3d.py
@@ -234,6 +234,12 @@ def save_glb(vertices, faces, filepath, metadata=None,
     textures = []
     samplers = []
     materials = []
+    pbr = {
+        "metallicFactor": 0.0,
+        "roughnessFactor": 0.5,
+        "baseColorFactor": [0.22, 0.22, 0.22, 1.0],
+    }
+
     if texture_png_bytes is not None and "TEXCOORD_0" in primitive_attributes:
         buffer_views.append({
             "buffer": 0,
@@ -243,15 +249,13 @@ def save_glb(vertices, faces, filepath, metadata=None,
         images.append({"bufferView": len(buffer_views) - 1, "mimeType": "image/png"})
         samplers.append({"magFilter": 9729, "minFilter": 9729, "wrapS": 33071, "wrapT": 33071})
         textures.append({"source": 0, "sampler": 0})
-        materials.append({
-            "pbrMetallicRoughness": {
-                "baseColorTexture": {"index": 0, "texCoord": 0},
-                "metallicFactor": 0.0,
-                "roughnessFactor": 1.0,
-            },
-            "doubleSided": True,
-        })
-        primitive["material"] = 0
+        pbr["baseColorTexture"] = {"index": 0, "texCoord": 0}
+
+    materials.append({
+        "pbrMetallicRoughness": pbr,
+        "doubleSided": True,
+    })
+    primitive["material"] = 0
 
     gltf = {
         "asset": {"version": "2.0", "generator": "ComfyUI"},
@@ -373,10 +377,14 @@ class SaveGLB(IO.ComfyNode):
                     continue
                 tex_img = Image.fromarray(texture_np[i], mode="RGB") if texture_np is not None else None
                 f = f"{filename}_{counter:05}_.glb"
-                save_glb(vertices_i, faces_i, os.path.join(full_output_folder, f), metadata,
-                         uvs=uvs_i,
-                         vertex_colors=v_colors,
-                         texture_image=tex_img)
+                save_glb(
+                    vertices_i, faces_i,
+                    os.path.join(full_output_folder, f),
+                    metadata,
+                    uvs=uvs_i,
+                    vertex_colors=v_colors,
+                    texture_image=tex_img,
+                )
                 results.append({
                     "filename": f,
                     "subfolder": subfolder,
diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py
new file mode 100644
index 000000000..c922d58b6
--- /dev/null
+++ b/comfy_extras/nodes_trellis2.py
@@ -0,0 +1,667 @@
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, IO, Types, io
+from comfy.ldm.trellis2.vae import SparseTensor
+from comfy_extras.nodes_mesh_postprocess import pack_variable_mesh_batch
+import comfy.model_management
+from PIL import Image
+import numpy as np
+import torch
+
+ShapeSubdivides = io.Custom("SHAPE_SUBDIVIDES")
+
+def prepare_trellis_vae_for_decode(vae, sample_shape):
+    memory_required = vae.memory_used_decode(sample_shape, vae.vae_dtype)
+    if len(sample_shape) == 5:
+        memory_required *= max(1, int(sample_shape[4]))
+    memory_required = max(1, int(memory_required))
+    device = comfy.model_management.get_torch_device()
+    comfy.model_management.load_models_gpu(
+        [vae.patcher],
+        memory_required=memory_required,
+        force_full_load=getattr(vae, "disable_offload", False),
+    )
+    free_memory = vae.patcher.get_free_memory(device)
+    batch_number = max(1, int(free_memory / memory_required))
+    return batch_number
+
+shape_slat_normalization = {
+    "mean": torch.tensor([
+        0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+        -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+        0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+        -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+    ])[None],
+    "std": torch.tensor([
+        5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+        5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+        4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+        5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+    ])[None]
+}
+
+tex_slat_normalization = {
+    "mean": torch.tensor([
+        3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+        0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+        -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+        1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+    ])[None],
+    "std": torch.tensor([
+        2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+        2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+        2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+        2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+    ])[None]
+}
+
+def shape_norm(shape_latent, coords):
+    std = shape_slat_normalization["std"].to(shape_latent)
+    mean = shape_slat_normalization["mean"].to(shape_latent)
+    samples = SparseTensor(feats = shape_latent, coords=coords)
+    samples = samples * std + mean
+    return samples
+
+
+def infer_batched_coord_layout(coords):
+    if coords.ndim != 2 or coords.shape[1] != 4:
+        raise ValueError(f"Expected Trellis2 coords with shape [N, 4], got {tuple(coords.shape)}")
+
+    if coords.shape[0] == 0:
+        raise ValueError("Trellis2 coords can't be empty")
+
+    batch_ids = coords[:, 0].to(torch.int64)
+    if (batch_ids < 0).any():
+        raise ValueError(f"Trellis2 batch ids must be non-negative, got {batch_ids.unique(sorted=True).tolist()}")
+    batch_size = int(batch_ids.max().item()) + 1
+    counts = torch.bincount(batch_ids, minlength=batch_size)
+
+    if (counts == 0).any():
+        raise ValueError(f"Non-contiguous Trellis2 batch ids in coords: {batch_ids.unique(sorted=True).tolist()}")
+
+    max_tokens = int(counts.max().item())
+    return batch_size, counts, max_tokens
+
+
+def split_batched_coords(coords, coord_counts):
+    if coord_counts.ndim != 1:
+        raise ValueError(f"Trellis2 coord_counts must be 1D, got shape {tuple(coord_counts.shape)}")
+    if (coord_counts < 0).any():
+        raise ValueError(f"Trellis2 coord_counts must be non-negative, got {coord_counts.tolist()}")
+    if int(coord_counts.sum().item()) != coords.shape[0]:
+        raise ValueError(
+            f"Trellis2 coord_counts total {int(coord_counts.sum().item())} does not match coords rows {coords.shape[0]}"
+        )
+
+    batch_ids = coords[:, 0].to(torch.int64)
+    order = torch.argsort(batch_ids, stable=True)
+    sorted_coords = coords.index_select(0, order)
+    sorted_batch_ids = batch_ids.index_select(0, order)
+
+    offsets = coord_counts.cumsum(0) - coord_counts
+    items = []
+    for i in range(coord_counts.shape[0]):
+        count = int(coord_counts[i].item())
+        start = int(offsets[i].item())
+        coords_i = sorted_coords[start:start + count]
+        ids_i = sorted_batch_ids[start:start + count]
+        if coords_i.shape[0] != count or not torch.all(ids_i == i):
+            raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}")
+        items.append(coords_i)
+    return items
+
+def flatten_batched_sparse_latent(samples, coords, coord_counts):
+    samples = samples.squeeze(-1).transpose(1, 2)
+    if coord_counts is None:
+        return samples.reshape(-1, samples.shape[-1]), coords
+
+    coords_items = split_batched_coords(coords, coord_counts)
+    feat_list = []
+    coord_list = []
+    for i, coords_i in enumerate(coords_items):
+        count = int(coord_counts[i].item())
+        feat_list.append(samples[i, :count])
+        coord_list.append(coords_i)
+
+    return torch.cat(feat_list, dim=0), torch.cat(coord_list, dim=0)
+
+
+def split_batched_sparse_latent(samples, coords, coord_counts):
+    samples = samples.squeeze(-1).transpose(1, 2)
+    if coord_counts is None:
+        return [(samples.reshape(-1, samples.shape[-1]), coords)]
+
+    coords_items = split_batched_coords(coords, coord_counts)
+    items = []
+    for i, coords_i in enumerate(coords_items):
+        count = int(coord_counts[i].item())
+        items.append((samples[i, :count], coords_i))
+    return items
+
+class VaeDecodeShapeTrellis(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeShapeTrellis",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+            ],
+            outputs=[
+                IO.Mesh.Output("mesh"),
+                ShapeSubdivides.Output(display_name = "shape_subdivides"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae):
+
+        resolution = int(vae.first_stage_model.resolution.item())
+        sample_tensor = samples["samples"]
+        device = comfy.model_management.get_torch_device()
+        coords = samples["coords"]
+        prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        trellis_vae = vae.first_stage_model
+        coord_counts = samples.get("coord_counts")
+
+        samples = samples["samples"]
+        if coord_counts is None:
+            samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
+            samples = shape_norm(samples.to(device), coords.to(device))
+            mesh, subs = trellis_vae.decode_shape_slat(samples, resolution)
+        else:
+            split_items = split_batched_sparse_latent(samples, coords, coord_counts)
+            mesh = []
+            subs_per_sample = []
+            for feats_i, coords_i in split_items:
+                coords_i = coords_i.to(device).clone()
+                coords_i[:, 0] = 0
+                sample_i = shape_norm(feats_i.to(device), coords_i)
+                mesh_i, subs_i = trellis_vae.decode_shape_slat(sample_i, resolution)
+                mesh.append(mesh_i[0])
+                subs_per_sample.append(subs_i)
+
+            subs = []
+            for stage_index in range(len(subs_per_sample[0])):
+                stage_tensors = [sample_subs[stage_index] for sample_subs in subs_per_sample]
+                feats_list = [stage_tensor.feats for stage_tensor in stage_tensors]
+                coords_list = [stage_tensor.coords for stage_tensor in stage_tensors]
+                subs.append(SparseTensor.from_tensor_list(feats_list, coords_list))
+
+        face_list = [m.faces for m in mesh]
+        vert_list = [m.vertices for m in mesh]
+        if all(v.shape == vert_list[0].shape for v in vert_list) and all(f.shape == face_list[0].shape for f in face_list):
+            mesh = Types.MESH(vertices=torch.stack(vert_list), faces=torch.stack(face_list))
+        else:
+            mesh = pack_variable_mesh_batch(vert_list, face_list)
+        return IO.NodeOutput(mesh, subs)
+
+class VaeDecodeTextureTrellis(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeTextureTrellis",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+                ShapeSubdivides.Input("shape_subdivides",
+                                 tooltip=(
+                                    "Shape information used to guide higher-detail reconstruction during decoding. "
+                                    "Helps preserve structure consistency at higher resolutions."
+                                 )),
+            ],
+            outputs=[
+                IO.Voxel.Output("voxel_colors"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae, shape_subdivides):
+        sample_tensor = samples["samples"]
+        device = comfy.model_management.get_torch_device()
+        coords = samples["coords"]
+        prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        trellis_vae = vae.first_stage_model
+        coord_counts = samples.get("coord_counts")
+
+        samples = samples["samples"]
+        samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
+        samples = samples.to(device)
+        std = tex_slat_normalization["std"].to(samples)
+        mean = tex_slat_normalization["mean"].to(samples)
+        samples = SparseTensor(feats = samples, coords=coords.to(device))
+        samples = samples * std + mean
+
+        voxel = trellis_vae.decode_tex_slat(samples, shape_subdivides)
+        color_feats = voxel.feats[:, :3]
+        voxel_coords = voxel.coords#[:, 1:]
+
+        voxel = Types.VOXEL(voxel_coords, color_feats, 1024)
+        return IO.NodeOutput(voxel)
+
+class VaeDecodeStructureTrellis2(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeStructureTrellis2",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+                IO.Combo.Input("resolution", options=["32", "64"], default="32")
+            ],
+            outputs=[
+                IO.Voxel.Output("voxel"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae, resolution):
+        resolution = int(resolution)
+        sample_tensor = samples["samples"]
+        sample_tensor = sample_tensor[:, :8]
+        batch_number = prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        decoder = vae.first_stage_model.struct_dec
+        load_device = comfy.model_management.get_torch_device()
+        decoded_batches = []
+        for start in range(0, sample_tensor.shape[0], batch_number):
+            sample_chunk = sample_tensor[start:start + batch_number].to(load_device)
+            decoded_batches.append(decoder(sample_chunk) > 0)
+        decoded = torch.cat(decoded_batches, dim=0)
+        current_res = decoded.shape[2]
+
+        if current_res != resolution:
+            ratio = current_res // resolution
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        out = Types.VOXEL(decoded.squeeze(1).float())
+        return IO.NodeOutput(out)
+
+class Trellis2UpsampleCascade(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2UpsampleCascade",
+            category="latent/3d",
+            display_name="Trellis2 Upsample Cascade",
+            description="Upsamples low-resolution Trellis2 shape latents into higher resolution coordinates while respecting the maximum token budget.",
+            inputs=[
+                IO.Latent.Input("shape_latent"),
+                IO.Vae.Input("vae"),
+                IO.Combo.Input("target_resolution", options=["1024", "1536"], default="1024", tooltip="Controls output detail level for upsampling."),
+                IO.Int.Input("max_tokens", default=49152, min=1024, max=100000,
+                             tooltip=(
+                                "Maximum number of output elements (coordinates) allowed after upsampling. "
+                                "Used to limit memory usage and control mesh density."
+                            ))
+            ],
+            outputs=[
+                IO.Voxel.Output(
+                "high_res_voxel",
+                tooltip=(
+                    "High-resolution sparse coordinates produced after cascade upsampling. "
+                    "Represents the refined 3D structure at target resolution."
+                )
+            )
+            ]
+        )
+
+    @classmethod
+    def execute(cls, shape_latent, vae, target_resolution, max_tokens):
+        shape_latent_512 = shape_latent
+        device = comfy.model_management.get_torch_device()
+        prepare_trellis_vae_for_decode(vae, shape_latent_512["samples"].shape)
+
+        coord_counts = shape_latent_512.get("coord_counts")
+        decoder = vae.first_stage_model.shape_dec
+        lr_resolution = 512
+        target_resolution = int(target_resolution)
+
+        if coord_counts is None:
+            feats, coords_512 = flatten_batched_sparse_latent(
+                shape_latent_512["samples"],
+                shape_latent_512["coords"],
+                coord_counts,
+            )
+            feats = feats.to(device)
+            coords_512 = coords_512.to(device)
+            slat = shape_norm(feats, coords_512)
+            slat.feats = slat.feats.to(next(decoder.parameters()).dtype)
+            hr_coords = decoder.upsample(slat, upsample_times=4)
+
+            hr_resolution = target_resolution
+            while True:
+                quant_coords = torch.cat([
+                    hr_coords[:, :1],
+                    ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+                ], dim=1)
+                final_coords = quant_coords.unique(dim=0)
+                num_tokens = final_coords.shape[0]
+
+                if num_tokens < max_tokens or hr_resolution <= 1024:
+                    break
+                hr_resolution -= 128
+
+            return IO.NodeOutput(final_coords,)
+
+        items = split_batched_sparse_latent(
+            shape_latent_512["samples"],
+            shape_latent_512["coords"],
+            coord_counts,
+        )
+        decoder_dtype = next(decoder.parameters()).dtype
+
+        sample_hr_coords = []
+        for feats_i, coords_i in items:
+            feats_i = feats_i.to(device)
+            coords_i = coords_i.to(device).clone()
+            coords_i[:, 0] = 0
+            slat_i = shape_norm(feats_i, coords_i)
+            slat_i.feats = slat_i.feats.to(decoder_dtype)
+            sample_hr_coords.append(decoder.upsample(slat_i, upsample_times=4))
+
+        hr_resolution = target_resolution
+        while True:
+            exceeds_limit = False
+            for hr_coords_i in sample_hr_coords:
+                quant_coords_i = torch.cat([
+                    hr_coords_i[:, :1],
+                    ((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+                ], dim=1)
+                if quant_coords_i.unique(dim=0).shape[0] >= max_tokens:
+                    exceeds_limit = True
+                    break
+            if not exceeds_limit or hr_resolution <= 1024:
+                break
+            hr_resolution -= 128
+
+        final_coords_list = []
+        output_coord_counts = []
+        for sample_offset, hr_coords_i in enumerate(sample_hr_coords):
+            quant_coords_i = torch.cat([
+                hr_coords_i[:, :1],
+                ((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+            ], dim=1)
+            final_coords_i = quant_coords_i.unique(dim=0)
+            final_coords_i = final_coords_i.clone()
+            final_coords_i[:, 0] = sample_offset
+            final_coords_list.append(final_coords_i)
+            output_coord_counts.append(int(final_coords_i.shape[0]))
+
+        coords = torch.cat(final_coords_list, dim=0)
+        output = Types.VOXEL(coords)
+        output.coord_counts = torch.tensor(output_coord_counts, dtype=torch.int64)
+        output.resolutions = torch.full((len(final_coords_list),), int(hr_resolution), dtype=torch.int64)
+        output.upsampled = True
+
+        return IO.NodeOutput(output,)
+
+dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+
+def run_conditioning(model, cropped_img_tensor, include_1024=True):
+    model_internal = model.model
+    device = comfy.model_management.intermediate_device()
+    torch_device = comfy.model_management.get_torch_device()
+
+    def prepare_tensor(pil_img, size):
+        resized_pil = pil_img.resize((size, size), Image.Resampling.LANCZOS)
+        img_np = np.array(resized_pil).astype(np.float32) / 255.0
+        img_t = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(torch_device)
+        return (img_t - dino_mean.to(torch_device)) / dino_std.to(torch_device)
+
+    model_internal.image_size = 512
+    input_512 = prepare_tensor(cropped_img_tensor, 512)
+    cond_512 = model_internal(input_512, skip_norm_elementwise=True)[0]
+
+    cond_1024 = None
+    if include_1024:
+        model_internal.image_size = 1024
+        input_1024 = prepare_tensor(cropped_img_tensor, 1024)
+        cond_1024 = model_internal(input_1024, skip_norm_elementwise=True)[0]
+
+    conditioning = {
+        'cond_512': cond_512.to(device),
+        'neg_cond': torch.zeros_like(cond_512).to(device),
+    }
+    if cond_1024 is not None:
+        conditioning['cond_1024'] = cond_1024.to(device)
+
+    return conditioning
+class Trellis2Conditioning(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2Conditioning",
+            category="conditioning/video_models",
+            inputs=[
+                IO.ClipVision.Input("clip_vision_model"),
+                IO.Image.Input("image"),
+                IO.Mask.Input("mask"),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, clip_vision_model, image, mask) -> IO.NodeOutput:
+        # Normalize to batched form so per-image conditioning loop below is uniform.
+        if image.ndim == 3:
+            image = image.unsqueeze(0)
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        batch_size = image.shape[0]
+        if mask.shape[0] == 1 and batch_size > 1:
+            mask = mask.expand(batch_size, -1, -1)
+        elif mask.shape[0] != batch_size:
+            raise ValueError(f"Trellis2Conditioning mask batch {mask.shape[0]} does not match image batch {batch_size}")
+
+        cond_512_list = []
+        cond_1024_list = []
+
+        for b in range(batch_size):
+            item_image = image[b]
+            item_mask = mask[b] if mask.size(0) > 1 else mask[0]
+
+            img_np = (item_image.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+            mask_np = (item_mask.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+
+            pil_img = Image.fromarray(img_np)
+            pil_mask = Image.fromarray(mask_np)
+
+            max_size = max(pil_img.size)
+            scale = min(1.0, 1024 / max_size)
+            if scale < 1.0:
+                new_w, new_h = int(pil_img.width * scale), int(pil_img.height * scale)
+                pil_img = pil_img.resize((new_w, new_h), Image.Resampling.LANCZOS)
+                pil_mask = pil_mask.resize((new_w, new_h), Image.Resampling.NEAREST)
+
+            rgba_np = np.zeros((pil_img.height, pil_img.width, 4), dtype=np.uint8)
+            rgba_np[:, :, :3] = np.array(pil_img)
+            rgba_np[:, :, 3] = np.array(pil_mask)
+
+            alpha = rgba_np[:, :, 3]
+            bbox_coords = np.argwhere(alpha > 0.8 * 255)
+
+            if len(bbox_coords) > 0:
+                y_min, x_min = np.min(bbox_coords[:, 0]), np.min(bbox_coords[:, 1])
+                y_max, x_max = np.max(bbox_coords[:, 0]), np.max(bbox_coords[:, 1])
+
+                center_y, center_x = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
+                size = max(y_max - y_min, x_max - x_min)
+
+                crop_x1 = int(center_x - size // 2)
+                crop_y1 = int(center_y - size // 2)
+                crop_x2 = int(center_x + size // 2)
+                crop_y2 = int(center_y + size // 2)
+
+                rgba_pil = Image.fromarray(rgba_np)
+                cropped_rgba = rgba_pil.crop((crop_x1, crop_y1, crop_x2, crop_y2))
+                cropped_np = np.array(cropped_rgba).astype(np.float32) / 255.0
+            else:
+                import logging
+                logging.warning("Mask for the image is empty. Trellis2 requires an image with a mask for the best mesh quality.")
+                cropped_np = rgba_np.astype(np.float32) / 255.0
+
+            bg_rgb = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+
+            fg = cropped_np[:, :, :3]
+            alpha_float = cropped_np[:, :, 3:4]
+            composite_np = fg * alpha_float + bg_rgb * (1.0 - alpha_float)
+
+            # to match trellis2 code (quantize -> dequantize)
+            composite_uint8 = (composite_np * 255.0).round().clip(0, 255).astype(np.uint8)
+
+            cropped_pil = Image.fromarray(composite_uint8)
+
+            item_conditioning = run_conditioning(clip_vision_model, cropped_pil, include_1024=True)
+            cond_512_list.append(item_conditioning["cond_512"])
+            cond_1024_list.append(item_conditioning["cond_1024"])
+
+        cond_512_batched = torch.cat(cond_512_list, dim=0)
+        cond_1024_batched = torch.cat(cond_1024_list, dim=0)
+        neg_cond_batched = torch.zeros_like(cond_512_batched)
+        neg_embeds_batched = torch.zeros_like(cond_1024_batched)
+
+        positive = [[cond_512_batched, {"embeds": cond_1024_batched}]]
+        negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]]
+        return IO.NodeOutput(positive, negative)
+
+class EmptyTrellis2ShapeLatent(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2ShapeLatent",
+            category="latent/3d",
+            inputs=[
+                IO.Voxel.Input(
+                    "voxel",
+                    tooltip=(
+                        "Shape structure input. Accepts either a voxel structure "
+                        "or upsampled voxel coordinates from a previous cascade stage."
+                    )
+            )
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, voxel):
+        # to accept the upscaled coords
+        is_512_pass = False
+        upsampled = hasattr(voxel, "upsampled")
+        if upsampled:
+            voxel = voxel.data
+
+        if not upsampled:
+            decoded = voxel.data.unsqueeze(1)
+            coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
+            is_512_pass = True
+
+        else:
+            coords = voxel.int()
+            is_512_pass = False
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+        in_channels = 32
+        # image like format
+        latent = torch.zeros(batch_size, in_channels, max_tokens, 1)
+
+        if is_512_pass:
+            generation_mode = "shape_generation_512"
+        else:
+            generation_mode = "shape_generation"
+        return IO.NodeOutput({"samples": latent, "coords": coords, "coord_counts": counts, "type": "trellis2",
+                              "model_options": {"generation_mode": generation_mode, "coords": coords, "coord_counts": counts}})
+
+class EmptyTrellis2LatentTexture(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2LatentTexture",
+            category="latent/3d",
+            inputs=[
+                IO.Voxel.Input(
+                    "voxel",
+                    tooltip=(
+                        "Shape structure input. Accepts either a voxel structure "
+                        "or upsampled voxel coordinates from a previous cascade stage."
+                    )
+                ),
+                IO.Latent.Input("shape_latent"),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, voxel, shape_latent):
+        channels = 32
+        upsampled = hasattr(voxel, "upsampled")
+        if upsampled:
+            voxel = voxel.data
+
+        if not upsampled:
+            decoded = voxel.data.unsqueeze(1)
+            coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
+        else:
+            coords = voxel.int()
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+
+        shape_latent = shape_latent["samples"]
+        if shape_latent.ndim == 4:
+            shape_latent = shape_latent.squeeze(-1).transpose(1, 2).reshape(-1, channels)
+
+        latent = torch.zeros(batch_size, channels, max_tokens, 1)
+        return IO.NodeOutput({"samples": latent, "type": "trellis2", "coords": coords, "coord_counts": counts,
+                                     "model_options": {"generation_mode": "texture_generation",
+                                                       "coords": coords, "coord_counts": counts, "shape_slat": shape_latent}})
+
+
+class EmptyTrellis2LatentStructure(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2LatentStructure",
+            category="latent/3d",
+            inputs=[
+                IO.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch."),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+    @classmethod
+    def execute(cls, batch_size):
+        in_channels = 8
+        resolution = 16
+        latent = torch.zeros(batch_size, in_channels, resolution, resolution, resolution)
+        output = {
+            "samples": latent,
+            "type": "trellis2",
+        }
+        return IO.NodeOutput(output)
+
+class Trellis2Extension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            Trellis2Conditioning,
+            EmptyTrellis2ShapeLatent,
+            EmptyTrellis2LatentStructure,
+            EmptyTrellis2LatentTexture,
+            VaeDecodeTextureTrellis,
+            VaeDecodeShapeTrellis,
+            VaeDecodeStructureTrellis2,
+            Trellis2UpsampleCascade,
+        ]
+
+
+async def comfy_entrypoint() -> Trellis2Extension:
+    return Trellis2Extension()
diff --git a/nodes.py b/nodes.py
index 13e46ac8a..65413206a 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1537,6 +1537,10 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
+    if "model_options" in latent:
+        inner = model.model.diffusion_model
+        inner.meta = latent["model_options"]
+
     callback = latent_preview.prepare_callback(model, steps)
     disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
     samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
@@ -2430,6 +2434,8 @@ async def init_builtin_extra_nodes():
         "nodes_toolkit.py",
         "nodes_replacements.py",
         "nodes_nag.py",
+        "nodes_trellis2.py",
+        "nodes_mesh_postprocess.py",
         "nodes_sdpose.py",
         "nodes_math.py",
         "nodes_number_convert.py",