Add Trellis 2 Support (CORE-7) (#12183)

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
 

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

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]
+}

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

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)

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')

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

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):

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"

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}

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
 ]

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,

comfy_api_nodes/nodes_bytedance.py

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

comfy_api_nodes/util/__init__.py

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

comfy_api_nodes/util/conversions.py

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

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()

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,

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()

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",