Merge 4cf6a777f6 into 7f287b705e

comfy/clip_vision.py

@@ -10,6 +10,7 @@ import comfy.utils
 import comfy.clip_model
 import comfy.image_encoders.dino2
 import comfy.image_encoders.dino3
+from comfy.image_encoders.naf import NAF
 
 class Output:
     def __getitem__(self, key):
@@ -53,6 +54,7 @@ class ClipVisionModel():
         self.model.eval()
 
         self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
+        self.naf = None
 
     def load_sd(self, sd):
         return self.model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
@@ -141,6 +143,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_large.json")
     elif 'layer.0.mlp.gate_proj.weight' in sd and 'layer.31.norm1.weight' in sd: # Dinov3 ViT-H/16+ (SwiGLU gated MLP, 32 layers)
         json_config = comfy.image_encoders.dino3.DINOV3_VITH_CONFIG
+    elif 'layer.9.attention.o_proj.bias' in sd: # dinov3 large (24 layers); generic o_proj.bias key, so must come after the ViT-H check
+        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino3_large.json")
     else:
         return None
 
@@ -153,6 +157,14 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
     for k in keys:
         if k not in u:
             sd.pop(k)
+    # NAF feature upsampler bundled into the DINOv3 file under the `naf.` prefix.
+    naf_keys = [k for k in sd if k.startswith("naf.")]
+    if naf_keys:
+        naf_sd = {k[len("naf."):]: sd.pop(k) for k in naf_keys}
+        naf = NAF().eval()
+        naf.load_state_dict(naf_sd, strict=False)
+        naf.to(comfy.model_management.text_encoder_dtype(clip.load_device))
+        clip.naf = comfy.model_patcher.CoreModelPatcher(naf, load_device=clip.load_device, offload_device=comfy.model_management.text_encoder_offload_device())
     return clip
 
 def load(ckpt_path):

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/image_encoders/naf.py

@@ -0,0 +1,283 @@
+"""NAF (Neighborhood Attention Filtering) feature upsampler.
+
+Vendored from valeoai/NAF (Apache-2.0):
+  https://github.com/valeoai/NAF — src/model/naf.py + src/layers/{convolutions,attentions,rope}.py
+Used by Pixal3D's shape/texture conditioning to produce
+the 2x-upsampled half of the 2048-channel proj feature map.
+"""
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Pure-torch neighborhood attention (replaces natten.na2d / na2d_qk + na2d_av).
+
+def upsample_lr_slice(src_lr: torch.Tensor, lr_dh: int, lr_dw: int,
+                       hr_h_range: Tuple[int, int], hr_w_range: Tuple[int, int]) -> torch.Tensor:
+    """Slice a LR-layout tensor [B, h_lr, w_lr, n, C], permute to BCHW, and
+    nearest-exact upsample only the region covering [hr_h_range, hr_w_range].
+    Returns BCHW at hr_h_end-hr_h_start x hr_w_end-hr_w_start (no padding for
+    out-of-bounds regions)."""
+    B = src_lr.shape[0]
+    n = src_lr.shape[-2]
+    C = src_lr.shape[-1]
+    h_hr_start, h_hr_end = hr_h_range
+    w_hr_start, w_hr_end = hr_w_range
+    # LR positions covering [h_hr_start, h_hr_end). Nearest-exact maps HR p → p // D.
+    lr_h_start = h_hr_start // lr_dh
+    lr_h_end   = (h_hr_end - 1) // lr_dh + 1
+    lr_w_start = w_hr_start // lr_dw
+    lr_w_end   = (w_hr_end - 1) // lr_dw + 1
+    lr_slice = src_lr[:, lr_h_start:lr_h_end, lr_w_start:lr_w_end]
+    lh, lw = lr_slice.shape[1], lr_slice.shape[2]
+    lr_bcd = lr_slice.permute(0, 3, 4, 1, 2).reshape(B * n, C, lh, lw).contiguous()
+    up = F.interpolate(lr_bcd, scale_factor=(lr_dh, lr_dw), mode="nearest-exact")
+    offset_h = h_hr_start - lr_h_start * lr_dh
+    offset_w = w_hr_start - lr_w_start * lr_dw
+    return up[:, :, offset_h:offset_h + (h_hr_end - h_hr_start),
+                    offset_w:offset_w + (w_hr_end - w_hr_start)]
+
+
+def na2d_pure(
+        q: torch.Tensor,                # [B, H, W, n_heads, d_qk] at HR.
+        k_lr: torch.Tensor,             # [B, h_lr, w_lr, n_heads, d_qk] at LR
+        v_lr: torch.Tensor,             # [B, h_lr, w_lr, n_heads, d_v] at LR
+        kernel_size: Tuple[int, int],   # (Kh, Kw) attention window.
+        dilation: Tuple[int, int],      # (Dh, Dw) stride within the unrolled K/V grid; also the LR→HR upsample factor.
+        scale: float,                   # 1 / sqrt(d_qk) scaling for the Q·K scores.
+        tile: int = 128,                # Spatial tile size (output positions per tile)
+        v_chunk: int = 64,              # Sub-divide d_v into chunks of this size when computing attn·V. None disables chunking.
+        output: torch.Tensor = None,    # Pre-allocated [B, n_heads, d_v, H, W] buffer (may be on CPU).
+    ) -> torch.Tensor:                  # [B, n_heads, d_v, H, W] (caller views as BCHW).
+    """Neighborhood attention in pure torch via F.unfold + per-tile slicing.
+
+    K and V are passed at LR resolution and upsampled (nearest-exact) per-tile only
+    for the slice the unfold needs. Avoids the [B, n*d, H, W] HR allocations for K
+    (512 MB) and V (2 GB) at tex_1024 fp16. Spatial tiling bounds the per-tile
+    F.unfold blob; `v_chunk` further slices d_v so attn·V is computed in C-sized
+    chunks (attn is reused, computed once from Q/K).
+
+    """
+    B, H, W, n, d_qk = q.shape
+    d_v = v_lr.shape[-1]
+    Kh, Kw = kernel_size
+    Dh, Dw = dilation
+    pad_h, pad_w = (Kh // 2) * Dh, (Kw // 2) * Dw
+
+    out = output if output is not None else torch.empty((B, n, d_v, H, W), device=q.device, dtype=q.dtype)
+
+    th = min(tile, H) if tile else H
+    tw = min(tile, W) if tile else W
+    chunk = v_chunk if (v_chunk and v_chunk < d_v) else d_v
+
+    for h0 in range(0, H, th):
+        for w0 in range(0, W, tw):
+            h1, w1 = min(h0 + th, H), min(w0 + tw, W)
+            t_h, t_w = h1 - h0, w1 - w0
+
+            # Padded HR region the unfold needs (kernel span = (K-1)*D + 1).
+            h_src_start = max(0, h0 - pad_h)
+            h_src_end   = min(H, h1 + pad_h)
+            w_src_start = max(0, w0 - pad_w)
+            w_src_end   = min(W, w1 + pad_w)
+            pad_top = max(0, pad_h - h0)
+            pad_bot = max(0, (h1 + pad_h) - H)
+            pad_lft = max(0, pad_w - w0)
+            pad_rgt = max(0, (w1 + pad_w) - W)
+
+            # Upsample only the tile region from k_lr / v_lr.
+            k_tile = upsample_lr_slice(k_lr, Dh, Dw,
+                                        (h_src_start, h_src_end),
+                                        (w_src_start, w_src_end))
+            v_tile = upsample_lr_slice(v_lr, Dh, Dw,
+                                        (h_src_start, h_src_end),
+                                        (w_src_start, w_src_end))
+            if pad_top or pad_bot or pad_lft or pad_rgt:
+                k_tile = F.pad(k_tile, [pad_lft, pad_rgt, pad_top, pad_bot])
+                v_tile = F.pad(v_tile, [pad_lft, pad_rgt, pad_top, pad_bot])
+
+            # Q·K → attention weights (small: KK=81 per output position).
+            KK = Kh * Kw
+            k_w = F.unfold(k_tile, kernel_size=(Kh, Kw), dilation=(Dh, Dw), padding=0)
+            k_w = k_w.view(B, n, d_qk, KK, t_h * t_w).permute(0, 1, 4, 3, 2)  # [B, n, t, KK, d_qk]
+            # q is [B, H, W, n, d_qk]; per-tile slice + permute -> [B, n, t_h*t_w, 1, d_qk].
+            q_tile = q[:, h0:h1, w0:w1].permute(0, 3, 1, 2, 4).reshape(B, n, t_h * t_w, 1, d_qk)
+            scores = torch.matmul(q_tile, k_w.transpose(-1, -2)) * scale
+            attn = scores.softmax(dim=-1)
+            del k_w, scores, q_tile, k_tile
+
+            # attn · V, chunked over d_v.
+            for c0 in range(0, d_v, chunk):
+                c1 = min(c0 + chunk, d_v)
+                v_w = F.unfold(v_tile[:, c0:c1], kernel_size=(Kh, Kw),dilation=(Dh, Dw), padding=0) # [B*n, (c1-c0)*KK, t]
+                v_w = v_w.view(B, n, c1 - c0, KK, t_h * t_w).permute(0, 1, 4, 3, 2)
+                out_chunk = torch.matmul(attn, v_w).squeeze(-2) # [B, n, t, c1-c0]
+                out_chunk = out_chunk.view(B, n, t_h, t_w, c1 - c0).permute(0, 1, 4, 2, 3)
+                out[:, :, c0:c1, h0:h1, w0:w1] = out_chunk
+                del v_w, out_chunk
+            del attn, v_tile
+
+    return out  # [B, n, d_v, H, W] — sole caller (CrossAttention) views it as BCHW directly.
+
+
+class CrossAttention(nn.Module):
+    """Window-restricted cross-attention. No learnable parameters; the model's
+    capacity lives entirely in the ImageEncoder convs."""
+
+    def __init__(self, dim: int, num_heads: int, kernel_size: Tuple[int, int] = (9, 9)):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads = num_heads
+        self.kernel_size = kernel_size
+        self.scale = (dim // num_heads) ** -0.5
+
+    @staticmethod
+    def _split_heads_lr(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+        """[B, n*d, h, w] -> [B, h, w, n, d] at the input resolution (no upsample)."""
+        B, C, H, W = x.shape
+        return x.view(B, num_heads, C // num_heads, H, W).permute(0, 3, 4, 1, 2).contiguous()
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                output_device=None) -> torch.Tensor:
+        hq, wq = q.shape[-2:]
+        hk, wk = k.shape[-2:]
+        dilation = (hq // hk, wq // wk)
+        B, C, _, _ = q.shape
+        q = q.view(B, self.num_heads, C // self.num_heads, hq, wq).permute(0, 3, 4, 1, 2).contiguous()
+        k_lr = self._split_heads_lr(k, self.num_heads).to(q.dtype)
+        v_lr = self._split_heads_lr(v, self.num_heads).to(q.dtype)
+        out_buf = None
+        if output_device is not None:
+            n = self.num_heads
+            d_v = v.shape[1] // n
+            out_buf = torch.empty(B, n, d_v, hq, wq, device=output_device, dtype=q.dtype)
+        out = na2d_pure(q, k_lr, v_lr, self.kernel_size, dilation, self.scale, output=out_buf)
+        return out.view(B, -1, hq, wq)
+
+
+# RoPE positional embedding
+
+def rope_rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+class RoPE(nn.Module):
+    def __init__(self, embed_dim: int, num_heads: int, base: float = 100.0):
+        super().__init__()
+        assert embed_dim % (4 * num_heads) == 0
+        self.num_heads = num_heads
+        self.D_head = embed_dim // num_heads
+        self.base = base
+        self.register_buffer("periods", torch.empty(self.D_head // 4), persistent=True) # loaded from the checkpoint
+
+    def _cos_sin(self, H: int, W: int, dtype: torch.dtype):
+        """cos/sin depend only on (H, W, dtype) and the checkpoint-fixed periods; recomputed per forward."""
+        device = self.periods.device
+        coords_h = torch.arange(0.5, H, device=device, dtype=torch.float32) / H
+        coords_w = torch.arange(0.5, W, device=device, dtype=torch.float32) / W
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)  # [H, W, 2]
+        coords = coords.flatten(0, 1) * 2.0 - 1.0  # [HW, 2]
+        angles = 2 * math.pi * coords[:, :, None] / self.periods.to(coords.dtype)[None, None, :]  # [HW, 2, D//4]
+        angles = angles.flatten(1, 2).tile(2)  # [HW, D]
+        cos = torch.cos(angles).to(dtype)
+        sin = torch.sin(angles).to(dtype)
+        return cos, sin
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, n*D_head, H, W]
+        B, C, H, W = x.shape
+        n = self.num_heads
+        D = C // n
+        x = x.view(B, n, D, H, W).permute(0, 1, 3, 4, 2).reshape(B, n, H * W, D)
+        cos, sin = self._cos_sin(H, W, x.dtype)
+        x = (x * cos) + (rope_rotate_half(x) * sin)
+        x = x.view(B, n, H, W, D).permute(0, 1, 4, 2, 3).reshape(B, n * D, H, W)
+        return x
+
+
+# Image encoder
+
+class EncBlock(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, num_groups: int = 8):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=kernel_size,
+                               padding=kernel_size // 2, padding_mode="reflect", bias=True)
+        self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=kernel_size,
+                               padding=kernel_size // 2, padding_mode="reflect", bias=True)
+        self.activation_fn = nn.SiLU()
+
+    def forward(self, x):
+        x = self.norm1(x)
+        x = self.activation_fn(x)
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.activation_fn(x)
+        x = self.conv2(x)
+        return x  # no skip connection
+
+
+def _encoder(in_dim: int, hidden_dim: int, kernel_size: int = 1, ks_res: int = 1, num_layers: int = 2) -> nn.Sequential:
+    return nn.Sequential(
+        nn.Conv2d(in_dim, hidden_dim, kernel_size=kernel_size, padding=kernel_size // 2, padding_mode="reflect", bias=True),
+        *[EncBlock(hidden_dim, kernel_size=ks_res) for _ in range(num_layers)],
+    )
+
+
+class ImageEncoder(nn.Module):
+    """Two parallel conv stacks (1x1 + 3x3) producing dim/2 channels each, then concat,
+    spatial average-pool to target size, RoPE-embed positions."""
+
+    def __init__(self, in_channels: int = 3, out_channels: int = 256,
+                 heads_rope: int = 4, rope_base: float = 100.0, img_layers: int = 2):
+        super().__init__()
+        half = out_channels // 2
+        self.encoder = _encoder(in_channels, half, kernel_size=1, ks_res=1, num_layers=img_layers)
+        self.sem_encoder = _encoder(in_channels, half, kernel_size=3, ks_res=3, num_layers=img_layers)
+        self.rope = RoPE(embed_dim=out_channels, num_heads=heads_rope, base=rope_base)
+
+    def forward(self, x: torch.Tensor, output_size: Tuple[int, int]) -> torch.Tensor:
+        # Avoid running the conv stacks on >4× the target resolution.
+        out_h, out_w = output_size
+        if x.shape[-2] > 4 * out_h or x.shape[-1] > 4 * out_w:
+            x = F.interpolate(x, size=(min(x.shape[-2], 4 * out_h),
+                                       min(x.shape[-1], 4 * out_w)),
+                              mode="bilinear", align_corners=False)
+        x = torch.cat([self.encoder(x), self.sem_encoder(x)], dim=1)
+        x = F.adaptive_avg_pool2d(x, output_size=output_size)
+        x = self.rope(x)
+        return x
+
+
+class NAF(nn.Module):
+    """NAF feature upsampler."""
+
+    def __init__(
+            self, dim: int = 256,       # internal channel dimension of the ImageEncoder
+            heads_attn: int = 4,        # attention heads in the windowed cross-attn
+            heads_rope: int = 4,        # heads for RoPE position encoding (must divide dim)
+            kernel_size: int = 9,       # square kernel for the neighborhood attention window
+            rope_base: float = 100.0,   # base for RoPE frequency periods
+            img_layers: int = 2         # number of EncBlocks in each conv stack
+        ):
+        super().__init__()
+        self.image_encoder = ImageEncoder(in_channels=3, out_channels=dim, heads_rope=heads_rope, rope_base=rope_base, img_layers=img_layers)
+        self.upsampler = CrossAttention(dim=dim, num_heads=heads_attn, kernel_size=(kernel_size, kernel_size))
+
+    def forward(
+            self,
+            image: torch.Tensor,            # [B, 3, H_img, W_img] in [0, 1].
+            features: torch.Tensor,         # [B, C, H_feat, W_feat] low-resolution features (any C).
+            output_size: Tuple[int, int],   # (H_out, W_out) target spatial resolution for the upsampled features.
+            output_device=None,
+        ) -> torch.Tensor:                  # [B, C, H_out, W_out] upsampled features.
+        """Upsample low-res feature map to output_size, guided by the image."""
+        q = self.image_encoder(image, output_size=output_size)
+        k = F.adaptive_avg_pool2d(q, output_size=features.shape[-2:])
+        return self.upsampler(q, k, features, output_device=output_device)

comfy/latent_formats.py

@@ -249,6 +249,53 @@ class TripoSplat(LatentFormat):
     def process_out(self, latent):
         return latent
 
+class Trellis2(LatentFormat):
+    latent_channels = 32
+
+class Trellis2SLAT(Trellis2):
+    # Sparse structured latent: per-token feats [N, 32]. process_out denormalizes
+    # the decoded feats (latent * std + mean); subclasses carry each space's stats.
+    latents_mean = None
+    latents_std = None
+
+    def process_in(self, latent):
+        mean = self.latents_mean.to(latent.device, latent.dtype)
+        std = self.latents_std.to(latent.device, latent.dtype)
+        return (latent - mean) / std
+
+    def process_out(self, latent):
+        mean = self.latents_mean.to(latent.device, latent.dtype)
+        std = self.latents_std.to(latent.device, latent.dtype)
+        return latent * std + mean
+
+class Trellis2ShapeSLAT(Trellis2SLAT):
+    latents_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]
+    latents_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]
+
+class Trellis2TexSLAT(Trellis2SLAT):
+    latents_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]
+    latents_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]
+
 class Mochi(LatentFormat):
     latent_channels = 12
     latent_dimensions = 3

comfy/ldm/trellis2/flexgemm.py

@@ -0,0 +1,154 @@
+from typing import Optional, Tuple
+import torch
+
+import comfy.model_management
+
+
+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):
+        self.sorted_keys, order = torch.sort(keys.to(torch.long))
+        self.sorted_vals = values.to(torch.long)[order]
+        self._n = self.sorted_keys.numel()
+
+    # Chunk size for lookup_flat, caps each transient to ~CHUNK rows.
+    _LOOKUP_CHUNK = 1 << 23   # 8M rows ≈ 64 MB per int64 temp
+
+    def lookup_flat(self, flat_keys: torch.Tensor) -> torch.Tensor:
+        N = flat_keys.shape[0]
+        out = torch.full((N,), -1, device=flat_keys.device, dtype=torch.int32)
+        if self._n == 0 or N == 0:
+            return out
+        for s in range(0, N, self._LOOKUP_CHUNK):
+            e = min(s + self._LOOKUP_CHUNK, N)
+            flat_chunk = flat_keys[s:e].to(torch.long)
+            idx = torch.searchsorted(self.sorted_keys, flat_chunk)
+            in_range = idx < self._n
+            idx.clamp_(max=self._n - 1)  # reuse idx as the "safe" index
+            found = in_range & (self.sorted_keys[idx] == flat_chunk)
+            if found.any():
+                found_idx = found.nonzero(as_tuple=True)[0]
+                out[s + found_idx] = self.sorted_vals[idx[found_idx]].to(torch.int32)
+        return out
+
+
+def build_submanifold_neighbor_map(
+    hashmap,
+    coords: torch.Tensor,
+    W, H, D,
+    Kw, Kh, Kd,
+    Dw, Dh, Dd,
+):
+    # neighbor[i, v] = index of the voxel at voxel i's coord + kernel-offset v, or -1.
+    # Chunked over voxels so the [chunk, V, 3] candidate transient stays bounded.
+    device = coords.device
+    M = coords.shape[0]
+    offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device).long()      # [V, 3]
+    V = offsets.shape[0]
+    center = torch.tensor([(Kw // 2) * Dw, (Kh // 2) * Dh, (Kd // 2) * Dd], device=device)
+    WHD, HD = W * H * D, H * D
+
+    neighbor = torch.empty((M, V), dtype=torch.int32, device=device)
+    # ~V*40 bytes/voxel of transient (int64 cand + flat + masks); cap at ~0.5 GB.
+    chunk = max(1, min(M, int(0.5 * (1024 ** 3) / (V * 40))))
+
+    for s in range(0, M, chunk):
+        e = min(s + chunk, M)
+        b = coords[s:e, 0].long()
+        cand = coords[s:e, 1:4].long()[:, None, :] + offsets[None, :, :] - center  # [c, V, 3]
+        x, y, z = cand[..., 0], cand[..., 1], cand[..., 2]
+        in_bounds = (x >= 0) & (x < W) & (y >= 0) & (y < H) & (z >= 0) & (z < D)   # [c, V]
+        flat = b[:, None] * WHD + x * HD + y * D + z                               # [c, V]
+        flat = torch.where(in_bounds, flat, torch.full_like(flat, -1))            # OOB -> guaranteed miss
+        neighbor[s:e] = hashmap.lookup_flat(flat.reshape(-1)).view(e - s, V)
+    return neighbor
+
+def get_recommended_chunk_mem(
+    device=None,
+    safety_fraction: float = 0.2,
+    min_gb: float = 0.25,
+    max_gb: float = 2.0,
+):
+    """Pick a chunk-memory budget (in GB) for sparse conv batching."""
+    free_gb = comfy.model_management.get_free_memory(device) / (1024 ** 3)
+    return max(min_gb, min(free_gb * safety_fraction, max_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,
+) -> 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
+
+    W, H, D = shape
+
+    Co, Kw, Kh, Kd, Ci = weight.shape
+    V = Kw * Kh * Kd
+    device = feats.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)
+
+        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]
+
+    weight_T = weight.view(Co, V * Ci).T
+
+    output = torch.empty(N_pts, Co, device=device, dtype=feats.dtype)
+
+    # Zero row at index N_pts; missing neighbors (-1) gather it -> no separate masking.
+    feats_padded = torch.cat([feats, feats.new_zeros(1, Ci)], dim=0)
+
+    # Chunk over voxels to bound the (chunk, V, Ci) gather.
+    max_chunk_mem_gb = get_recommended_chunk_mem(device)
+    mem_per_row = V * Ci * feats.element_size()
+    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)
+
+    for start in range(0, N_pts, chunk_size):
+        end = min(start + chunk_size, N_pts)
+        actual_chunk = end - start
+
+        chunk_idx = torch.where(neighbor[start:end] < 0, N_pts, neighbor[start:end])  # -1 -> zero row
+        gathered = feats_padded[chunk_idx]                  # (chunk, V, Ci)
+        gathered_flat = gathered.view(actual_chunk, V * Ci)
+        torch.matmul(gathered_flat, weight_T, out=output[start:end])  # (chunk, V*Ci) @ (V*Ci, Co)
+
+    if bias is not None:
+        output += bias.unsqueeze(0).to(output.dtype)
+
+    return output, neighbor

comfy/model_base.py

@@ -61,6 +61,7 @@ import comfy.ldm.ideogram4.model
 import comfy.ldm.krea2.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
@@ -1853,6 +1854,27 @@ 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)
+        # CONDConstant: shared across pos/neg
+        for k in ("trellis2_coords", "trellis2_coord_counts",
+                  "trellis2_generation_mode", "trellis2_shape_slat",
+                  "trellis2_proj_feats", "trellis2_model_frame"):
+            v = kwargs.get(k)
+            if v is not None:
+                out[k] = comfy.conds.CONDConstant(v)
+        # Pixal3D's per-stage feature maps + camera params travel as a dict
+        proj_feat_pack = kwargs.get("proj_feat_pack")
+        if proj_feat_pack is not None:
+            out["proj_feat_pack"] = comfy.conds.CONDConstant(proj_feat_pack)
+        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"

comfy/model_detection.py

@@ -113,6 +113,25 @@ 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
 
+    shape_key = '{}img2shape.t_embedder.mlp.0.weight'.format(key_prefix)
+    tex_key = '{}shape2txt.t_embedder.mlp.0.weight'.format(key_prefix)
+    if shape_key in state_dict_keys or tex_key in state_dict_keys:  # trellis2 / pixal3d
+        has_shape = shape_key in state_dict_keys
+        has_tex = tex_key in state_dict_keys
+        unet_config = {
+            "image_model": "trellis2",
+            "resolution": 32 if (metadata is not None and "is_512" in metadata) else 64,
+            "init_txt_model": has_tex,
+            "txt_only": has_tex and not has_shape,
+        }
+        # Per-submodel projection head (Pixal3D adds `proj_linear`; Trellis2 doesn't).
+        for sub, name in (("img2shape", "shape"), ("shape2txt", "texture"), ("structure_model", "structure")):
+            key = '{}{}.blocks.0.cross_attn.proj_linear.weight'.format(key_prefix, sub)
+            if key in state_dict_keys:
+                unet_config["image_attn_mode_{}".format(name)] = "proj"
+                unet_config["proj_in_channels_{}".format(name)] = int(state_dict[key].shape[1])
+        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

@@ -14,6 +14,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.triposplat.vae
@@ -543,6 +544,16 @@ 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: # trellis2 shape vae (struct_dec + shape_dec)
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                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.ShapeVae()
+            elif "txt_dec.blocks.3.4.conv2.weight" in sd: # trellis2 texture vae
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                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.TextureVae()
             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}
@@ -1101,6 +1112,15 @@ class VAE:
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
 
+    def prepare_decode(self, sample_shape, memory_required=None):
+        """For VAEs whose real decode entry point bypasses decode()"""
+        if memory_required is None:
+            memory_required = self.memory_used_decode(sample_shape, self.vae_dtype)
+        memory_required = max(1, int(memory_required))
+        model_management.load_models_gpu([self.patcher], memory_required=memory_required, force_full_load=self.disable_offload)
+        free_memory = self.patcher.get_free_memory(self.device)
+        return max(1, int(free_memory / memory_required))
+
     def decode_tiled(self, samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
         self.throw_exception_if_invalid()
         memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile

comfy/supported_models.py

@@ -1432,6 +1432,31 @@ 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"
+    }
+
+    unet_extra_config = {"num_heads": 12}
+
+    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",
@@ -2369,5 +2394,6 @@ models = [
     CogVideoX_I2V,
     CogVideoX_T2V,
     SVD_img2vid,
+    Trellis2,
     DepthAnything3,
 ]

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 SPLAT:
     """A batch of 3D Gaussian splats in render-ready (activated, world-space) form.
@@ -34,9 +35,16 @@ class MESH:
                  uvs: torch.Tensor | None = None,
                  vertex_colors: torch.Tensor | None = None,
                  texture: torch.Tensor | None = None,
+                 metallic_roughness: torch.Tensor | None = None,
                  vertex_counts: torch.Tensor | None = None,
                  face_counts: torch.Tensor | None = None,
-                 unlit: bool = False):
+                 unlit: bool = False,
+                 normals: torch.Tensor | None = None,
+                 tangents: torch.Tensor | None = None,
+                 normal_map: torch.Tensor | None = None,
+                 occlusion_in_mr: bool = False,
+                 material: dict | None = None,
+                 emissive: torch.Tensor | None = None):
 
         assert (vertex_counts is None) == (face_counts is None), \
             "vertex_counts and face_counts must be provided together (both or neither)"
@@ -44,13 +52,25 @@ class MESH:
         self.faces = faces                  # faces: (B, M, 3)
         self.uvs = uvs                      # uvs: (B, N, 2)
         self.vertex_colors = vertex_colors  # vertex_colors: (B, N, 3 or 4)
-        self.texture = texture              # texture: (B, H, W, 3)
+        # Optional per-vertex normals: (B, N, 3). When None, SaveGLB computes smooth
+        # area-weighted normals so viewers don't fall back to flat (per-face) shading.
+        self.normals = normals
+        self.texture = texture              # texture (baseColor): (B, H, W, 3)
+        # glTF metallicRoughness texture: (B, H, W, 3), R unused, G=roughness, B=metallic
+        self.metallic_roughness = metallic_roughness
         # When vertices/faces are zero-padded to a common N/M across the batch (variable-size mesh batch),
         # these hold the real per-item lengths (B,). None means rows are uniform and no slicing is needed.
         self.vertex_counts = vertex_counts
         self.face_counts = face_counts
         # Render flat / emissive (no scene lighting) when saved, e.g. for gaussian-splat-derived meshes.
         self.unlit = unlit
+        # Extra maps / material overrides attached by bake, normal/AO, and SetMeshMaterial nodes;
+        # consumed by SaveGLB. Declared here (with defaults) so consumers read them directly.
+        self.tangents = tangents            # (B, N, 4) per-vertex tangents for normal mapping
+        self.normal_map = normal_map        # tangent-space normal map: (B, H, W, 3)
+        self.occlusion_in_mr = occlusion_in_mr  # True = R channel of metallic_roughness holds AO (ORM)
+        self.material = material             # SetMeshMaterial scalar/factor overrides
+        self.emissive = emissive             # emissive map: (B, H, W, 3)
 
 
 class File3D:

comfy_extras/mesh3d/uv_unwrap/mesh.py

@@ -0,0 +1,162 @@
+"""Mesh container, edge/face adjacency, manifold cleanup."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List
+
+import numpy as np
+import torch
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import connected_components
+from torch import Tensor
+
+
+# ---- Per-face / per-vertex geometry ----
+
+def face_normals(vertices: Tensor, faces: Tensor) -> Tensor:
+    """[F,3] unit face normals (degenerate faces -> zero)."""
+    v0 = vertices[faces[:, 0]]
+    v1 = vertices[faces[:, 1]]
+    v2 = vertices[faces[:, 2]]
+    n = torch.linalg.cross(v1 - v0, v2 - v0)
+    return n / n.norm(dim=1, keepdim=True).clamp_min(1e-20)
+
+
+def face_areas(vertices: Tensor, faces: Tensor) -> Tensor:
+    """[F] triangle areas."""
+    v0 = vertices[faces[:, 0]]
+    v1 = vertices[faces[:, 1]]
+    v2 = vertices[faces[:, 2]]
+    return 0.5 * torch.linalg.cross(v1 - v0, v2 - v0).norm(dim=1)
+
+
+def face_centroids(vertices: Tensor, faces: Tensor) -> Tensor:
+    """[F,3] triangle centroids."""
+    return vertices[faces].mean(dim=1)
+
+
+def face_edge_lengths(vertices: Tensor, faces: Tensor) -> Tensor:
+    """[F,3] edge lengths; column e = |v[faces[:,e]] - v[faces[:,(e+1)%3]]|."""
+    va = vertices[faces]
+    vb = vertices[faces.roll(shifts=-1, dims=1)]
+    return (vb - va).norm(dim=-1).to(torch.float32)
+
+
+def chart_3d_areas(face_area: Tensor, face_chart: Tensor, n_charts: int) -> Tensor:
+    """[n_charts] sum of face areas per chart."""
+    out = torch.zeros(n_charts, dtype=face_area.dtype, device=face_area.device)
+    out.scatter_add_(0, face_chart, face_area)
+    return out
+
+
+@dataclass
+class MeshData:
+    """Cleaned mesh with adjacency; face_face[f, i] = face sharing edge (faces[f,i], faces[f,(i+1)%3]) or -1 if boundary."""
+
+    vertices: Tensor                 # [V, 3] float
+    faces: Tensor                    # [F, 3] long
+    face_face: Tensor                # [F, 3] long, neighbor face id or -1
+    face_normal: Tensor              # [F, 3] float
+    face_area: Tensor                # [F] float
+    face_centroid: Tensor            # [F, 3] float
+    component: Tensor                # [F] long, connected-component id
+    n_components: int
+
+
+def build_mesh(vertices: Tensor, faces: Tensor) -> MeshData:
+    """Build adjacency; non-manifold edges (>2 incident faces) get no neighbor and act as boundary."""
+    if vertices.dtype != torch.float32:
+        vertices = vertices.to(torch.float32)
+    if faces.dtype != torch.long:
+        faces = faces.to(torch.long)
+
+    device = faces.device
+    V = vertices.shape[0]
+    F = faces.shape[0]
+
+    # Per directed face-edge; flat layout p = f*3+i.
+    a = faces.flatten()
+    b = faces.roll(shifts=-1, dims=1).flatten()
+    lo = torch.minimum(a, b)
+    hi = torch.maximum(a, b)
+    edge_key = lo * (V + 1) + hi
+
+    # Pair manifold (count==2) face-edges; others get no neighbor.
+    _, inverse, counts = torch.unique(edge_key, return_inverse=True, return_counts=True)
+    edge_count = counts[inverse]
+    manifold_mask = edge_count == 2
+
+    sort_idx = torch.argsort(edge_key, stable=True)
+    sorted_manifold = manifold_mask[sort_idx]
+    pair_positions = sort_idx[sorted_manifold]
+    pair_a = pair_positions[0::2]
+    pair_b = pair_positions[1::2]
+
+    face_id_flat = torch.arange(F, device=device).repeat_interleave(3)
+    face_face_flat = torch.full((3 * F,), -1, dtype=torch.long, device=device)
+    face_face_flat[pair_a] = face_id_flat[pair_b]
+    face_face_flat[pair_b] = face_id_flat[pair_a]
+    face_face = face_face_flat.view(F, 3)
+
+    face_face_np = face_face.cpu().numpy()
+    rows_mask = face_face_np >= 0
+    if rows_mask.any():
+        rows = np.broadcast_to(np.arange(F)[:, None], (F, 3))[rows_mask]
+        cols = face_face_np[rows_mask]
+        adj = csr_matrix(
+            (np.ones(rows.size, dtype=np.int8), (rows, cols)),
+            shape=(F, F),
+        )
+    else:
+        adj = csr_matrix((F, F), dtype=np.int8)
+    n_components, labels = connected_components(adj, directed=False)
+
+    face_normal = face_normals(vertices, faces)
+    face_area = face_areas(vertices, faces)
+    face_centroid = face_centroids(vertices, faces)
+
+    return MeshData(
+        vertices=vertices,
+        faces=faces,
+        face_face=face_face,
+        face_normal=face_normal,
+        face_area=face_area,
+        face_centroid=face_centroid,
+        component=torch.from_numpy(labels.astype(np.int64)).to(device),
+        n_components=int(n_components),
+    )
+
+
+def chart_boundary_loops(
+    faces_subset: Tensor, face_face_subset: Tensor
+) -> List[List[int]]:
+    """Return ordered boundary vertex loops for a chart submesh (face_face_subset[f,i]==-1 marks a boundary edge)."""
+    F = faces_subset.shape[0]
+    faces_np = faces_subset.cpu().numpy()
+    ff = face_face_subset.cpu().numpy()
+
+    next_v: Dict[int, int] = {}
+    for f in range(F):
+        for i in range(3):
+            if ff[f, i] == -1:
+                a = int(faces_np[f, i])
+                b = int(faces_np[f, (i + 1) % 3])
+                next_v[a] = b
+
+    loops: List[List[int]] = []
+    visited = set()
+    for start in list(next_v.keys()):
+        if start in visited:
+            continue
+        loop = [start]
+        visited.add(start)
+        cur = next_v.get(start)
+        while cur is not None and cur != start:
+            if cur in visited:
+                break
+            loop.append(cur)
+            visited.add(cur)
+            cur = next_v.get(cur)
+        if len(loop) >= 3:
+            loops.append(loop)
+    return loops

comfy_extras/mesh3d/uv_unwrap/pack.py

@@ -0,0 +1,861 @@
+"""Atlas packing via bitmap rasterize-and-place."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.nn.functional import max_pool1d
+
+import comfy.model_management
+
+# Numba is optional, but ~5x faster than torch on these operations, potential TODO: comfy-kitchen cuda/triton kernels as even faster alternative
+try:
+    from numba import njit as _njit, prange as _prange, get_num_threads as _nb_threads
+    _HAVE_NUMBA_PACK = True
+except ImportError:
+    _HAVE_NUMBA_PACK = False
+    _prange = range
+    def _nb_threads(): return 1
+    def _njit(*args, **kwargs):
+        def deco(fn): return fn
+        return deco if not args else args[0]
+
+
+# Cap on deterministic sweep density: tiny charts on a large atlas would otherwise enumerate every texel column.
+_SWEEP_CAP = 1024
+
+
+@dataclass
+class ChartPlacement:
+    chart_id: int
+    offset: Tuple[float, float]       # in texels
+    scale: float                      # texels per UV unit
+    rotation: float = 0.0             # radians
+    swap_xy: bool = False             # extra 90° bitmap rotation chosen at place time
+    chart_h: float = 0.0              # unswapped bitmap height in texels (rotation pivot)
+
+
+@_njit(cache=True, boundscheck=False, parallel=True)
+def _prepare_dims_jit(uvs, uv_off, a3, auv, tpu, padding, theta, scale, bw, bh, rot_uv):
+    """Pass 1: per-chart best rotation, texel scale, rotated/scaled UVs, padded bitmap dims."""
+    n = uv_off.shape[0] - 1
+    half_pi = math.pi * 0.5
+    for c in _prange(n):
+        v0, v1 = uv_off[c], uv_off[c + 1]
+        best_area = 1e30
+        best_t = 0.0
+        for k in range(36):
+            th = half_pi * k / 36.0
+            co = math.cos(th)
+            si = math.sin(th)
+            xmin = 1e30
+            xmax = -1e30
+            ymin = 1e30
+            ymax = -1e30
+            for i in range(v0, v1):
+                xr = uvs[i, 0] * co - uvs[i, 1] * si
+                yr = uvs[i, 0] * si + uvs[i, 1] * co
+                if xr < xmin:
+                    xmin = xr
+                if xr > xmax:
+                    xmax = xr
+                if yr < ymin:
+                    ymin = yr
+                if yr > ymax:
+                    ymax = yr
+            area = (xmax - xmin) * (ymax - ymin)
+            if area < best_area:
+                best_area = area
+                best_t = th
+        theta[c] = best_t
+        co = math.cos(best_t)
+        si = math.sin(best_t)
+        xmin = 1e30
+        xmax = -1e30
+        ymin = 1e30
+        ymax = -1e30
+        for i in range(v0, v1):
+            xr = uvs[i, 0] * co - uvs[i, 1] * si
+            yr = uvs[i, 0] * si + uvs[i, 1] * co
+            rot_uv[i, 0] = xr
+            rot_uv[i, 1] = yr
+            if xr < xmin:
+                xmin = xr
+            if xr > xmax:
+                xmax = xr
+            if yr < ymin:
+                ymin = yr
+            if yr > ymax:
+                ymax = yr
+        if v1 == v0:
+            xmin = 0.0
+            xmax = 0.0
+            ymin = 0.0
+            ymax = 0.0
+        s = math.sqrt(max(a3[c], 1e-12) / max(auv[c], 1e-12)) * tpu
+        nominal = math.sqrt(max(a3[c], 1e-12)) * tpu
+        max_bbox = max(8.0, 4.0 * nominal)
+        bbox_max = max(max(xmax - xmin, ymax - ymin), 1e-12)
+        if s * bbox_max > max_bbox:
+            s = max_bbox / bbox_max
+        scale[c] = s
+        wmax = 0.0
+        hmax = 0.0
+        for i in range(v0, v1):
+            rot_uv[i, 0] = (rot_uv[i, 0] - xmin) * s
+            rot_uv[i, 1] = (rot_uv[i, 1] - ymin) * s
+            if rot_uv[i, 0] > wmax:
+                wmax = rot_uv[i, 0]
+            if rot_uv[i, 1] > hmax:
+                hmax = rot_uv[i, 1]
+        bw[c] = int(math.ceil(wmax)) + padding + 1
+        bh[c] = int(math.ceil(hmax)) + padding + 1
+
+
+@_njit(cache=True, boundscheck=False, parallel=True)
+def _raster_all_jit(rot_uv, uv_off, faces, f_off, bw, bh, boff, buf, padding,
+                    tw, th_out, perim):
+    """Pass 2: rasterize + dilate each chart into the flat buffer; records trimmed dims
+    (origin kept) and the perimeter used for placement ordering."""
+    n = uv_off.shape[0] - 1
+    eps = 1e-7
+    for c in _prange(n):
+        f0, f1 = f_off[c], f_off[c + 1]
+        v0 = uv_off[c]
+        V = uv_off[c + 1] - v0
+        w = bw[c]
+        h = bh[c]
+        o = boff[c]
+        for fi in range(f0, f1):
+            i0 = faces[fi, 0] + v0
+            i1 = faces[fi, 1] + v0
+            i2 = faces[fi, 2] + v0
+            x0 = rot_uv[i0, 0]
+            y0 = rot_uv[i0, 1]
+            x1 = rot_uv[i1, 0]
+            y1 = rot_uv[i1, 1]
+            x2 = rot_uv[i2, 0]
+            y2 = rot_uv[i2, 1]
+            xmin_f = min(x0, min(x1, x2))
+            xmax_f = max(x0, max(x1, x2))
+            ymin_f = min(y0, min(y1, y2))
+            ymax_f = max(y0, max(y1, y2))
+            xmin = max(int(math.floor(xmin_f)), 0)
+            xmax = min(int(math.ceil(xmax_f)), w - 1)
+            ymin = max(int(math.floor(ymin_f)), 0)
+            ymax = min(int(math.ceil(ymax_f)), h - 1)
+            if xmax < xmin or ymax < ymin:
+                continue
+            denom = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2)
+            if abs(denom) < 1e-20:
+                continue
+            inv_denom = 1.0 / denom
+            for py in range(ymin, ymax + 1):
+                yc = py + 0.5
+                for px in range(xmin, xmax + 1):
+                    xc = px + 0.5
+                    aa = ((y1 - y2) * (xc - x2) + (x2 - x1) * (yc - y2)) * inv_denom
+                    bb = ((y2 - y0) * (xc - x2) + (x0 - x2) * (yc - y2)) * inv_denom
+                    cc = 1.0 - aa - bb
+                    if aa >= -eps and bb >= -eps and cc >= -eps:
+                        buf[o + py * w + px] = True
+        # Manhattan dilation by `padding` steps (ping-pong on a scratch copy)
+        if padding > 0 and f1 > f0:
+            tmp = np.empty(h * w, dtype=np.bool_)
+            for _ in range(padding):
+                for j in range(h * w):
+                    tmp[j] = buf[o + j]
+                for py in range(h):
+                    for px in range(w):
+                        if tmp[py * w + px]:
+                            continue
+                        hit = False
+                        if py > 0 and tmp[(py - 1) * w + px]:
+                            hit = True
+                        elif py < h - 1 and tmp[(py + 1) * w + px]:
+                            hit = True
+                        elif px > 0 and tmp[py * w + px - 1]:
+                            hit = True
+                        elif px < w - 1 and tmp[py * w + px + 1]:
+                            hit = True
+                        if hit:
+                            buf[o + py * w + px] = True
+        # trimmed dims (keep origin; 1x1 empty bitmap when nothing was rasterized)
+        rmax = -1
+        cmax = -1
+        for py in range(h):
+            for px in range(w):
+                if buf[o + py * w + px]:
+                    if py > rmax:
+                        rmax = py
+                    if px > cmax:
+                        cmax = px
+        if rmax < 0:
+            for j in range(h * w):
+                buf[o + j] = False
+            tw[c] = 1
+            th_out[c] = 1
+        else:
+            tw[c] = cmax + 1
+            th_out[c] = rmax + 1
+        # unique-edge perimeter via sorted int64 keys
+        Fc = f1 - f0
+        if Fc > 0 and V > 0:
+            keys = np.empty(Fc * 3, dtype=np.int64)
+            for fi in range(f0, f1):
+                for j in range(3):
+                    a = faces[fi, j]
+                    b = faces[fi, (j + 1) % 3]
+                    if a < b:
+                        keys[(fi - f0) * 3 + j] = a * V + b
+                    else:
+                        keys[(fi - f0) * 3 + j] = b * V + a
+            keys = np.sort(keys)
+            p = 0.0
+            for i in range(keys.shape[0]):
+                if i > 0 and keys[i] == keys[i - 1]:
+                    continue
+                a = keys[i] // V + v0
+                b = keys[i] % V + v0
+                dx = rot_uv[a, 0] - rot_uv[b, 0]
+                dy = rot_uv[a, 1] - rot_uv[b, 1]
+                p += math.sqrt(dx * dx + dy * dy)
+            perim[c] = p
+
+
+@_njit(cache=True, boundscheck=False, parallel=True)
+def _place_all_jit(buf, boff, stride_w, tw, th, order, start, stop,
+                   atlas, skyline, pool, attempts, sweep_cap, margin,
+                   n_threads, cur_wh, out_x, out_y, out_sw):
+    """Place charts order[start:stop]; returns the first index NOT processed (== stop when
+    done, earlier when the atlas must grow — the caller resizes and resumes). The candidate
+    scan is striped with a (score, index) min-reduction: deterministic for any thread count,
+    and no thread intrinsics (dynamic globals would defeat cache=True)."""
+    aw = atlas.shape[1]
+    ah = atlas.shape[0]
+    cur_w = cur_wh[0]
+    cur_h = cur_wh[1]
+    n_pool = pool.shape[0]
+    big = np.int64(1) << 62
+    nt = n_threads
+    t_score = np.empty(nt, dtype=np.int64)
+    t_k = np.empty(nt, dtype=np.int64)
+    t_x = np.empty(nt, dtype=np.int64)
+    t_y = np.empty(nt, dtype=np.int64)
+    t_sw = np.empty(nt, dtype=np.int64)
+    for oi in range(start, stop):
+        ci = order[oi]
+        if cur_h + margin > ah or cur_w + margin > aw:
+            cur_wh[0] = cur_w
+            cur_wh[1] = cur_h
+            return oi
+        w0 = tw[ci]                                   # unswapped trimmed dims
+        h0 = th[ci]
+        W = stride_w[ci]                              # row stride of the untrimmed block
+        o = boff[ci]
+        step = min(w0, h0) // 8
+        if step < 1:
+            step = 1
+        cap_step = max(cur_w, cur_h) // sweep_cap
+        if cap_step > step:
+            step = cap_step
+
+        poff = (oi * attempts) % (n_pool - attempts + 1)
+        x_range = cur_w + 1 if cur_w > 0 else 1
+        y_range = cur_h + 1 if cur_h > 0 else 1
+        # candidate groups per orientation: skyline-flush sweep, y=0 / y=cur_h sweeps,
+        # x=0 / x=cur_w sweeps; then the shared random pool
+        nx = max(cur_w, 1) // step + 2
+        ny = max(cur_h, 1) // step + 2
+        n_det = nx * 3 + ny * 2
+        total = n_det * 2 + attempts
+        for t in range(nt):
+            t_score[t] = big
+            t_k[t] = big
+        for t2 in _prange(nt):
+            for k in range(t2, total, nt):
+                x = 0                                      # int inits and no body-level continue:
+                y = 0                                      # parfor lowering types undef-path
+                swap = 0                                   # variables as f64
+                valid = True
+                if k < 2 * n_det:
+                    if k >= n_det:
+                        swap = 1
+                    kk = k - n_det if swap == 1 else k
+                    cw = w0 if swap == 0 else h0
+                    if kk < nx:                            # skyline-flush sweep
+                        x = kk * step
+                        if x > cur_w:
+                            valid = False
+                        else:
+                            x_end = x + cw
+                            if x_end > skyline.shape[0]:
+                                x_end = skyline.shape[0]
+                            for xs in range(x, x_end):
+                                if skyline[xs] > y:
+                                    y = int(skyline[xs])
+                    elif kk < 3 * nx:                      # y=0 and y=cur_h sweeps
+                        kk2 = kk - nx
+                        x = (kk2 % nx) * step
+                        if x > cur_w:
+                            valid = False
+                        elif kk2 >= nx:
+                            y = cur_h
+                    else:                                  # x=0 and x=cur_w sweeps
+                        kk2 = kk - 3 * nx
+                        if kk2 >= 2 * ny:
+                            valid = False
+                        else:
+                            y = (kk2 % ny) * step
+                            if y > cur_h:
+                                valid = False
+                            elif kk2 >= ny:
+                                x = cur_w
+                else:
+                    r = k - 2 * n_det
+                    x = int(pool[poff + r, 0] % x_range)
+                    y = int(pool[poff + r, 1] % y_range)
+                    swap = int(r & 1)
+
+                if valid:
+                    ch = h0 if swap == 0 else w0
+                    cw = w0 if swap == 0 else h0
+                    nw = cur_w if cur_w > x + cw else x + cw
+                    nh = cur_h if cur_h > y + ch else y + ch
+                    ext = nw if nw > nh else nh
+                    score = ext * ext + nw * nh
+                    if score < t_score[t2] or (score == t_score[t2] and k < t_k[t2]):
+                        ok = True
+                        for j in range(ch):
+                            yy = int(y + j)
+                            if yy >= ah:
+                                continue
+                            for i in range(cw):
+                                if swap == 0:
+                                    bit = buf[o + j * W + i]
+                                else:
+                                    # 90deg rotation: bm_rot[j, i] = bm[h0-1-i, j]
+                                    bit = buf[o + (h0 - 1 - i) * W + j]
+                                if not bit:
+                                    continue
+                                xx = int(x + i)
+                                if xx >= aw:
+                                    continue
+                                if atlas[yy, xx]:
+                                    ok = False
+                                    break
+                            if not ok:
+                                break
+                        if ok:
+                            t_score[t2] = score
+                            t_k[t2] = k
+                            t_x[t2] = x
+                            t_y[t2] = y
+                            t_sw[t2] = swap
+
+        best_x = -1
+        best_y = -1
+        best_swap = 0
+        bs = big
+        bk = big
+        for t in range(nt):
+            if t_score[t] < bs or (t_score[t] == bs and t_k[t] < bk):
+                bs = t_score[t]
+                bk = t_k[t]
+                best_x = t_x[t]
+                best_y = t_y[t]
+                best_swap = t_sw[t]
+
+        if best_x < 0:                                 # fallback: extension corner
+            best_x = cur_w
+            best_y = 0
+            best_swap = 0
+        bh_ = h0 if best_swap == 0 else w0
+        bw_ = w0 if best_swap == 0 else h0
+        # blit + extents + skyline lift
+        for j in range(bh_):
+            for i in range(bw_):
+                if best_swap == 0:
+                    bit = buf[o + j * W + i]
+                else:
+                    bit = buf[o + (h0 - 1 - i) * W + j]
+                if bit:
+                    atlas[best_y + j, best_x + i] = True
+        if best_x + bw_ > cur_w:
+            cur_w = best_x + bw_
+        if best_y + bh_ > cur_h:
+            cur_h = best_y + bh_
+        for i in range(bw_):
+            col_x = best_x + i
+            if col_x >= skyline.shape[0]:
+                continue
+            col_top = -1
+            for j in range(bh_ - 1, -1, -1):
+                if best_swap == 0:
+                    bit = buf[o + j * W + i]
+                else:
+                    bit = buf[o + (h0 - 1 - i) * W + j]
+                if bit:
+                    col_top = j
+                    break
+            if col_top >= 0:
+                nh2 = best_y + col_top + 1
+                if nh2 > skyline[col_x]:
+                    skyline[col_x] = nh2
+        out_x[ci] = best_x
+        out_y[ci] = best_y
+        out_sw[ci] = best_swap
+    cur_wh[0] = cur_w
+    cur_wh[1] = cur_h
+    return stop
+
+
+# Torch fallback (used when numba is unavailable; runs on GPU if present)
+
+def _dilate_local(x: Tensor, p: int) -> Tensor:
+    """4-connectivity dilation by p over a batch of (cnt,g,g) bitmaps. Dilation distributes
+    over union, so dilating per-triangle then OR-scattering equals dilating the chart."""
+    for _ in range(p):
+        y = x.clone()
+        y[:, 1:, :] |= x[:, :-1, :]
+        y[:, :-1, :] |= x[:, 1:, :]
+        y[:, :, 1:] |= x[:, :, :-1]
+        y[:, :, :-1] |= x[:, :, 1:]
+        x = y
+    return x
+
+
+def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device):
+    """Rasterize every chart into one flat bool buffer; buf[cbase[i]:cbase[i+1]].view(bh,bw)
+    is chart i's bitmap. Triangles are bucketed by next-pow2 bbox size to bound memory."""
+    n = uvs_tex_pad.shape[0]
+    fmax = faces_pad.shape[1]
+    bwL, bhL = bw_t.long(), bh_t.long()
+    cbase = torch.zeros(n + 1, dtype=torch.long, device=device)
+    torch.cumsum(bwL * bhL, 0, out=cbase[1:])
+    buf = torch.zeros(int(cbase[-1].item()), dtype=torch.bool, device=device)
+
+    # gather all triangle coords, keep only valid faces -> (Ttot,3,2) + chart id per triangle
+    fp = faces_pad.reshape(n, fmax * 3)
+    tri = torch.gather(uvs_tex_pad, 1, fp[..., None].expand(-1, -1, 2)).reshape(n * fmax, 3, 2)
+    fm = fmask.reshape(-1)
+    tri_f = tri[fm]
+    if tri_f.shape[0] == 0:
+        return buf, cbase
+    cid = torch.arange(n, device=device).repeat_interleave(fmax)[fm]
+
+    # per-triangle pixel bbox, inflated by padding (origin >= 0); bucket by next-pow2 max-dim
+    tmin = tri_f.amin(1)
+    tmax = tri_f.amax(1)
+    x0 = (tmin[:, 0].floor().long() - padding).clamp_min(0)
+    y0 = (tmin[:, 1].floor().long() - padding).clamp_min(0)
+    bbw = (tmax[:, 0].ceil().long() + padding) - x0 + 1
+    bbh = (tmax[:, 1].ceil().long() + padding) - y0 + 1
+    mxd = torch.maximum(bbw, bbh).clamp_min(1)
+    bsz = (2 ** torch.ceil(torch.log2(mxd.float())).long()).long()
+
+    a = tri_f[:, 0]
+    b = tri_f[:, 1]
+    c = tri_f[:, 2]
+    v0 = b - a
+    v1 = c - a
+    d00 = (v0 * v0).sum(-1)
+    d01 = (v0 * v1).sum(-1)
+    d11 = (v1 * v1).sum(-1)
+    den = (d00 * d11 - d01 * d01).clamp(min=1e-20)
+
+    for g in sorted(set(bsz.tolist())):                                  # one batch per pow2 grid
+        sel = (bsz == g).nonzero(as_tuple=True)[0]
+        m = sel.shape[0]
+        xs0 = x0[sel].view(m, 1, 1)
+        ys0 = y0[sel].view(m, 1, 1)
+        cc = cid[sel]
+        bwp = bwL[cc].view(m, 1, 1)
+        bhp = bhL[cc].view(m, 1, 1)
+        gi = torch.arange(g, device=device)
+        px = xs0 + gi.view(1, 1, g)
+        py = ys0 + gi.view(1, g, 1)                                        # (m,g,g) int
+        pxf = px.float() + 0.5
+        pyf = py.float() + 0.5
+        v2x = pxf - a[sel, 0].view(m, 1, 1)
+        v2y = pyf - a[sel, 1].view(m, 1, 1)
+        d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1)
+        d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1)
+        idn = den[sel].view(m, 1, 1).reciprocal()
+        vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn
+        ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn
+        uu = 1.0 - vv - ww
+        inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6)
+        if padding > 0:
+            inside = _dilate_local(inside, padding)
+        valid = inside & (px < bwp) & (py < bhp)
+        flat = (cbase[cc].view(m, 1, 1) + py * bwp + px)[valid]
+        buf[flat] = True
+    return buf, cbase
+
+
+def _build_candidates_gpu(sky_t, ar, cur_w, cur_h, bw0, bw1, step, rand01, device):
+    """Candidate (x, y) positions as a (2, M, 2) tensor (dim 0 = orientation). The first
+    n_sky rows per orientation are skyline-flush and collision-free by construction.
+    rand01 is (2, rand_n, 2) pre-drawn uniforms; ar a preallocated arange."""
+    hi_x = max(cur_w, 1) + 1
+    hi_y = max(cur_h, 1) + 1
+    xs = ar[0:hi_x:step]
+    ys = ar[0:hi_y:step]
+    n_sky = (hi_x + step - 1) // step
+    zx = torch.zeros_like(xs)
+    zy = torch.zeros_like(ys)
+    common = torch.cat([
+        torch.stack([xs, zx], 1), torch.stack([xs, zx + cur_h], 1),
+        torch.stack([zy, ys], 1), torch.stack([zy + cur_w, ys], 1)])
+    wm = []
+    for cw in (bw0, bw1):
+        span = (n_sky - 1) * step + cw
+        wm.append(max_pool1d(sky_t[:span].view(1, 1, -1).float(), kernel_size=cw,
+                             stride=step).view(-1))
+    sky = torch.stack([torch.stack([xs, wm[0].long()], 1),
+                       torch.stack([xs, wm[1].long()], 1)])
+    lim = torch.tensor([hi_x, hi_y], dtype=rand01.dtype, device=device)
+    rnd = (rand01 * lim).long()
+    return torch.cat([sky, common.expand(2, -1, -1), rnd], 1), n_sky
+
+
+def _best_placement_torch(atlas, pix0, dim0, dim1, cands, n_sky, cur_w, cur_h, device):
+    """Lowest-score non-colliding placement as a (3,) int tensor [x, y, swap]. The best
+    skyline candidate bounds the score; only strictly better candidates are pixel-tested."""
+    m = cands.shape[1]
+    chw = torch.tensor([[dim0[0], dim0[1]], [dim1[0], dim1[1]]], device=device)
+    nw = torch.clamp(cands[..., 0] + chw[:, 1:], min=cur_w)             # (2,M)
+    nh = torch.clamp(cands[..., 1] + chw[:, :1], min=cur_h)
+    ext = torch.maximum(nw, nh)
+    sc = ext * ext + nw * nh
+    js = sc[:, :n_sky].reshape(-1).argmin()                             # best skyline candidate
+    sky_o = js // n_sky
+    s_star = sc[:, :n_sky].reshape(-1)[js]
+    sky = torch.cat([cands[sky_o, js % n_sky], sky_o.reshape(1)])
+    cflat = cands.reshape(-1, 2)
+    surv = (sc.reshape(-1) < s_star).nonzero(as_tuple=True)[0]          # compact once
+    total = surv.shape[0]
+    if total == 0:
+        return sky
+
+    k = pix0.shape[0]
+    if k == 0:                                                          # empty chart: anywhere free
+        j = surv[sc.reshape(-1)[surv].argmin()]
+        return torch.cat([cflat[j], (j // m).reshape(1)])
+    ordr = surv[torch.argsort(sc.reshape(-1)[surv], stable=True)]
+
+    # flattened-index collision test: one int32 gather index instead of two int64 rows/cols
+    aw = atlas.shape[1]
+    idt = torch.int32 if atlas.numel() < (1 << 31) else torch.long
+    lin0 = (pix0[:, 0] * aw + pix0[:, 1]).to(idt)                       # (y, x)
+    lin1 = (pix0[:, 1] * aw + (dim0[0] - 1 - pix0[:, 0])).to(idt)       # rotated: (x, h-1-y)
+    linp = torch.stack([lin0, lin1])
+    aflat = atlas.view(-1)
+    og = (ordr >= m).long()
+    base = (cflat[ordr, 1] * aw + cflat[ordr, 0]).to(idt)
+
+    # prescreen survivors on ~128 strided pixels: a sampled hit proves collision, so only
+    # subsample-clean candidates need the exact test
+    stride = (k + 127) // 128
+    linp_sub = linp[:, ::stride].contiguous()
+    maybe = ~aflat[base[:, None] + linp_sub[og]].any(1)
+    passers = maybe.nonzero(as_tuple=True)[0]                           # ascending = score-sorted
+    npass = passers.shape[0]
+    if npass == 0:
+        return sky
+    if stride == 1:                                                     # prescreen was already exact
+        j = ordr[passers[0]]
+        return torch.cat([cflat[j], (j // m).reshape(1)])
+
+    budget = 1 << 22                                                    # pixel-tests per chunk
+    start = 0
+    while start < npass:
+        take = max(1, budget // k)
+        pi = passers[start:start + take]
+        free = ~aflat[base[pi][:, None] + linp[og[pi]]].any(1)          # (t,k) True-pixel gather
+        # single host read per chunk: whether a free hit exists and where
+        has, first = torch.stack([free.any().long(), free.long().argmax()]).tolist()
+        if has:
+            j = ordr[pi[first]]                                         # lowest score: sorted order
+            return torch.cat([cflat[j], (j // m).reshape(1)])
+        start += take
+        budget = min(budget * 4, 1 << 25)
+    return sky
+
+
+def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces,
+                       texels_per_unit, padding_texels, attempts=4096, rng_seed=0,
+                       progress_callback=None):
+    """Torch rasterize-and-place packer (numba-free fallback). Returns (placements, atlas_w, atlas_h)."""
+    n = len(chart_uvs)
+    if n == 0:
+        return [], 1, 1
+    device = comfy.model_management.get_torch_device()
+    ang = torch.linspace(0.0, math.pi / 2.0, 37, device=device)[:-1]
+    cos_a, sin_a = ang.cos(), ang.sin()
+
+    # ---- Prepare pass 1: best-rotation + scale + bbox for ALL charts at once (batched) ----
+    vcount = [int(u.shape[0]) for u in chart_uvs]
+    fcount = [int(f.shape[0]) for f in chart_faces]
+    vmax = max(vcount)
+    fmax = max(fcount)
+    uvs_pad = torch.zeros(n, vmax, 2, device=device)
+    vmask = torch.zeros(n, vmax, dtype=torch.bool, device=device)
+    faces_pad = torch.zeros(n, fmax, 3, dtype=torch.long, device=device)
+    fmask = torch.zeros(n, fmax, dtype=torch.bool, device=device)
+    for i in range(n):
+        uvs_pad[i, :vcount[i]] = chart_uvs[i].to(device=device, dtype=torch.float32)
+        vmask[i, :vcount[i]] = True
+        if fcount[i]:
+            faces_pad[i, :fcount[i]] = chart_faces[i].to(device=device, dtype=torch.long)
+            fmask[i, :fcount[i]] = True
+    u0, u1 = uvs_pad[..., 0], uvs_pad[..., 1]                            # (N,Vmax)
+    BIG = 1e30
+    mlo = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), BIG))
+    mhi = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), -BIG))
+    xr = torch.addcmul(u0[:, :, None] * cos_a, u1[:, :, None], sin_a, value=-1)   # (N,Vmax,A)
+    yr = torch.addcmul(u0[:, :, None] * sin_a, u1[:, :, None], cos_a)
+    xsp = (xr + mhi[:, :, None]).amax(1) - (xr + mlo[:, :, None]).amin(1)         # (N,A) masked span
+    ysp = (yr + mhi[:, :, None]).amax(1) - (yr + mlo[:, :, None]).amin(1)
+    ti = (xsp * ysp).argmin(1)                                          # (N,) best angle per chart
+    cc, ss = cos_a[ti][:, None], sin_a[ti][:, None]                     # (N,1)
+    rx = torch.addcmul(u0 * cc, u1, ss, value=-1)                      # (N,Vmax)
+    ry = torch.addcmul(u0 * ss, u1, cc)
+    rxmin = (rx + mlo).amin(1)                                          # (N,)
+    rxmax = (rx + mhi).amax(1)
+    rymin = (ry + mlo).amin(1)
+    rymax = (ry + mhi).amax(1)
+    a3 = torch.tensor([max(a, 1e-12) for a in chart_3d_areas], device=device)
+    au = torch.tensor([max(a, 1e-12) for a in chart_uv_areas], device=device)
+    base = (a3 / au).sqrt() * texels_per_unit
+    maxb = (4.0 * a3.sqrt() * texels_per_unit).clamp_min(8.0)
+    bbm = torch.maximum(rxmax - rxmin, rymax - rymin).clamp_min(1e-12)
+    scale = torch.minimum(base, maxb / bbm)                            # (N,)
+    uvs_tex_pad = torch.stack([(rx - rxmin[:, None]) * scale[:, None],
+                               (ry - rymin[:, None]) * scale[:, None]], dim=-1)   # (N,Vmax,2)
+    bw_t = ((rxmax - rxmin) * scale).ceil().int() + padding_texels + 1
+    bh_t = ((rymax - rymin) * scale).ceil().int() + padding_texels + 1
+
+    # one sync: pull all per-chart scalars
+    thetas = ang[ti].cpu().tolist()
+    scales = scale.cpu().tolist()
+
+    # ---- Prepare pass 2: rasterize ALL charts at once, then derive per-chart sparse data ----
+    buf, cbase = _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding_texels, device)
+
+    # nonzero over the flat buffer is ascending, so pixels come out grouped by chart
+    nz = buf.nonzero(as_tuple=True)[0]
+    del buf
+    cid = torch.searchsorted(cbase, nz, right=True) - 1
+    bwl = bw_t.long()
+    local = nz - cbase[cid]
+    py = local // bwl[cid]
+    px = local - py * bwl[cid]
+    del nz, local
+    counts = torch.bincount(cid, minlength=n)
+    rmax = torch.full((n,), -1, dtype=torch.long, device=device)
+    cmax = torch.full((n,), -1, dtype=torch.long, device=device)
+    rmax.scatter_reduce_(0, cid, py, reduce="amax")
+    cmax.scatter_reduce_(0, cid, px, reduce="amax")
+    ht = (rmax + 1).clamp_min(1)                    # trimmed bitmap dims (1x1 when empty)
+    wt = (cmax + 1).clamp_min(1)
+    pix_all = torch.stack([py, px], 1)              # True-pixel (row, col) offsets, sparse
+    pixr_all = torch.stack([px, rmax[cid] - py], 1)  # 90deg rotation: (y, x) -> (x, h-1-y)
+    meta = torch.stack([ht, wt, counts.cumsum(0)], 1).cpu().tolist()     # one sync for all charts
+    dim_l = [(m[0], m[1]) for m in meta]
+    dimr_l = [(w, h) for (h, w) in dim_l]
+    offs = [0] + [m[2] for m in meta]
+    pix_l = [pix_all[offs[i]:offs[i + 1]] for i in range(n)]
+    pixr_l = [pixr_all[offs[i]:offs[i + 1]] for i in range(n)]
+
+    # column tops (skyline lift), batched via flat scatter-amax over (chart, column) keys
+    wmax = max(max(h, w) for (h, w) in dim_l)
+    ct_pad = torch.full((n * wmax,), -1, dtype=torch.long, device=device)
+    ctr_pad = torch.full((n * wmax,), -1, dtype=torch.long, device=device)
+    ct_pad.scatter_reduce_(0, cid * wmax + px, py, reduce="amax")
+    ctr_pad.scatter_reduce_(0, cid * wmax + (rmax[cid] - py), px, reduce="amax")
+    ct_pad = ct_pad.view(n, wmax)
+    ctr_pad = ctr_pad.view(n, wmax)
+    del cid, py, px, rmax, cmax
+
+    # ---- Placement: skyline bin-pack on GPU ----
+    order = sorted(range(n), key=lambda i: -(dim_l[i][0] * dim_l[i][1]))   # biggest bitmap first
+    max_b = max(max(d) for d in dim_l)
+    margin = max_b + 8
+    side_guess = int(math.sqrt(sum(d[0] * d[1] for d in dim_l)) * 2) + 16
+    cap = side_guess + margin
+    atlas = torch.zeros((cap, cap), dtype=torch.bool, device=device)
+    sky_t = torch.zeros(cap, dtype=torch.long, device=device)
+    ar = torch.arange(cap + 1, device=device)
+    cur_w = cur_h = 0
+    placements = [None] * n
+    gen = torch.Generator(device=device).manual_seed(rng_seed)
+    rand_n = min(512, attempts)                     # random samples per orientation
+    # no _SWEEP_CAP here: the skyline-bound pruning depends on the dense sweep
+    rand01 = torch.rand(n, 2, rand_n, 2, generator=gen, device=device)   # all draws upfront
+
+    for t_i, ci in enumerate(order):
+        if progress_callback is not None and (t_i & 255) == 0:
+            progress_callback(n + t_i, 2 * n)
+        if cur_h + margin > atlas.shape[0] or cur_w + margin > atlas.shape[1]:
+            ns = max(atlas.shape[0], cur_h + margin, cur_w + margin)
+            na = torch.zeros((ns, ns), dtype=torch.bool, device=device)
+            na[:atlas.shape[0], :atlas.shape[1]] = atlas
+            atlas = na
+            nsk = torch.zeros(ns, dtype=torch.long, device=device)
+            nsk[:sky_t.shape[0]] = sky_t
+            sky_t = nsk
+            ar = torch.arange(ns + 1, device=device)
+        dim, dimr = dim_l[ci], dimr_l[ci]
+        step = max(1, min(dim[0], dim[1]) // 8)
+        cands, n_sky = _build_candidates_gpu(
+            sky_t, ar, cur_w, cur_h, dim[1], dimr[1], step, rand01[t_i], device)
+        res = _best_placement_torch(atlas, pix_l[ci], dim, dimr,
+                                    cands, n_sky, cur_w, cur_h, device)
+        bx, by, swap = (int(v) for v in res.tolist())
+        if bx < 0:
+            bx, by, swap = cur_w, 0, 0
+        pix = pixr_l[ci] if swap else pix_l[ci]
+        bh_, bw_ = (dimr if swap else dim)
+        atlas[by + pix[:, 0], bx + pix[:, 1]] = True                    # sparse blit
+        cur_w = max(cur_w, bx + bw_)
+        cur_h = max(cur_h, by + bh_)
+        ct = (ctr_pad if swap else ct_pad)[ci, :bw_]                    # GPU skyline lift
+        ix = ar[bx:bx + bw_]
+        sky_t[ix] = torch.where(ct >= 0, torch.maximum(sky_t[ix], by + ct + 1), sky_t[ix])
+        placements[ci] = ChartPlacement(chart_id=ci, offset=(float(bx), float(by)),
+                                        scale=scales[ci], rotation=thetas[ci], swap_xy=bool(swap),
+                                        chart_h=float(dim_l[ci][0]))
+    return placements, cur_w, cur_h
+
+
+def pack_bitmap_concat(
+    uvs_cat: np.ndarray,          # (sumV, 2) per-chart concatenated UVs
+    uv_offsets: np.ndarray,       # (n+1,)
+    faces_cat: np.ndarray,        # (sumF, 3) local vert ids per chart
+    face_offsets: np.ndarray,     # (n+1,)
+    chart_3d_areas: np.ndarray,
+    chart_uv_areas: np.ndarray,
+    texels_per_unit: float = 256.0,
+    padding_texels: int = 2,
+    attempts: int = 4096,
+    rng_seed: int = 0,
+    progress_callback=None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int]:
+    """Rasterize-and-place packer over concatenated chart arrays (no per-chart python).
+    Returns (x, y, swap, rotation, scale, chart_h, atlas_w, atlas_h) with one entry per chart.
+    progress_callback(done, total) is invoked periodically; total is 2*n_charts."""
+    n = int(uv_offsets.shape[0]) - 1
+    empty = np.zeros(n, dtype=np.int64)
+    if n == 0:
+        return empty, empty, empty, empty.astype(np.float64), empty.astype(np.float64), empty, 1, 1
+    if not _HAVE_NUMBA_PACK:
+        chart_uvs = [torch.from_numpy(np.ascontiguousarray(uvs_cat[uv_offsets[c]:uv_offsets[c + 1]]))
+                     for c in range(n)]
+        chart_faces = [torch.from_numpy(np.ascontiguousarray(faces_cat[face_offsets[c]:face_offsets[c + 1]]))
+                       for c in range(n)]
+        placements, w, h = _pack_bitmap_torch(
+            chart_uvs, [float(a) for a in chart_3d_areas], [float(a) for a in chart_uv_areas],
+            chart_faces, texels_per_unit, padding_texels, attempts=attempts,
+            rng_seed=rng_seed, progress_callback=progress_callback)
+        px = np.array([p.offset[0] for p in placements], dtype=np.int64)
+        py = np.array([p.offset[1] for p in placements], dtype=np.int64)
+        sw = np.array([1 if p.swap_xy else 0 for p in placements], dtype=np.int64)
+        th = np.array([p.rotation for p in placements], dtype=np.float64)
+        sc = np.array([p.scale for p in placements], dtype=np.float64)
+        chh = np.array([p.chart_h for p in placements], dtype=np.int64)
+        return px, py, sw, th, sc, chh, w, h
+
+    uvs64 = np.ascontiguousarray(uvs_cat, dtype=np.float64)
+    faces64 = np.ascontiguousarray(faces_cat, dtype=np.int64)
+    uv_off = np.ascontiguousarray(uv_offsets, dtype=np.int64)
+    f_off = np.ascontiguousarray(face_offsets, dtype=np.int64)
+    a3 = np.ascontiguousarray(chart_3d_areas, dtype=np.float64)
+    auv = np.ascontiguousarray(chart_uv_areas, dtype=np.float64)
+
+    theta = np.zeros(n, dtype=np.float64)
+    scale = np.zeros(n, dtype=np.float64)
+    bw = np.zeros(n, dtype=np.int64)
+    bh = np.zeros(n, dtype=np.int64)
+    rot_uv = np.empty_like(uvs64)
+    _prepare_dims_jit(uvs64, uv_off, a3, auv, float(texels_per_unit), int(padding_texels),
+                      theta, scale, bw, bh, rot_uv)
+    boff = np.zeros(n + 1, dtype=np.int64)
+    np.cumsum(bw * bh, out=boff[1:])
+    buf = np.zeros(int(boff[-1]), dtype=np.bool_)
+    tw = np.zeros(n, dtype=np.int64)
+    th_arr = np.zeros(n, dtype=np.int64)
+    perim = np.zeros(n, dtype=np.float64)
+    _raster_all_jit(rot_uv, uv_off, faces64, f_off, bw, bh, boff, buf,
+                    int(padding_texels), tw, th_arr, perim)
+    if progress_callback is not None:
+        progress_callback(n, 2 * n)
+
+    order = np.argsort(-perim, kind="stable")
+    max_b = int(max(int(tw.max()), int(th_arr.max())))
+    margin = max_b + 8
+    side_guess = int(math.sqrt(float((tw * th_arr).sum()))) * 2 + 16
+    cap = side_guess + margin
+    atlas = np.zeros((cap, cap), dtype=np.bool_)
+    skyline = np.zeros(cap, dtype=np.int64)
+    rng = np.random.default_rng(rng_seed)
+    # shared random pool, sliced at a rotating offset per chart
+    pool = rng.integers(0, 1 << 31, size=(attempts * 8, 2)).astype(np.int64)
+    out_x = np.full(n, -1, dtype=np.int64)
+    out_y = np.full(n, -1, dtype=np.int64)
+    out_sw = np.zeros(n, dtype=np.int64)
+    cur_wh = np.zeros(2, dtype=np.int64)
+    start = 0
+    while start < n:
+        stop = min(n, start + 1024)
+        nxt = _place_all_jit(buf, boff, bw, tw, th_arr, order, start, stop,
+                             atlas, skyline, pool, int(attempts), int(_SWEEP_CAP),
+                             int(margin), int(_nb_threads()), cur_wh, out_x, out_y, out_sw)
+        if nxt < stop:                               # atlas must grow before this chart fits
+            ns = max(atlas.shape[0], int(cur_wh[1]) + margin, int(cur_wh[0]) + margin)
+            na = np.zeros((ns, ns), dtype=np.bool_)
+            na[:atlas.shape[0], :atlas.shape[1]] = atlas
+            atlas = na
+            nsk = np.zeros(ns, dtype=np.int64)
+            nsk[:skyline.shape[0]] = skyline
+            skyline = nsk
+        start = nxt
+        if progress_callback is not None:
+            progress_callback(n + start, 2 * n)
+    return out_x, out_y, out_sw, theta, scale, th_arr, int(cur_wh[0]), int(cur_wh[1])
+
+
+def apply_placements_concat(
+    uvs_cat: np.ndarray, uv_offsets: np.ndarray,
+    px: np.ndarray, py: np.ndarray, sw: np.ndarray,
+    theta: np.ndarray, scale: np.ndarray, chart_h: np.ndarray,
+    atlas_w: int, atlas_h: int,
+) -> np.ndarray:
+    """apply_placements over concatenated charts, fully vectorized. Returns (sumV, 2) float32."""
+    n = int(uv_offsets.shape[0]) - 1
+    side = float(max(atlas_w, atlas_h, 1))
+    cov = np.repeat(np.arange(n), np.diff(uv_offsets))
+    u_in = uvs_cat[:, 0].astype(np.float64)
+    v_in = uvs_cat[:, 1].astype(np.float64)
+    c = np.cos(theta)[cov]
+    s = np.sin(theta)[cov]
+    u = u_in * c - v_in * s
+    v = u_in * s + v_in * c
+    umin = np.full(n, np.inf)
+    vmin = np.full(n, np.inf)
+    np.minimum.at(umin, cov, u)
+    np.minimum.at(vmin, cov, v)
+    u = (u - umin[cov]) * scale[cov]
+    v = (v - vmin[cov]) * scale[cov]
+    swv = sw[cov].astype(bool)
+    # 90 deg rotation matching the rotated-bitmap access: (u, v) -> (chart_h - v, u)
+    u2 = np.where(swv, chart_h[cov] - v, u) + px[cov]
+    v2 = np.where(swv, u, v) + py[cov]
+    out = np.stack([u2, v2], axis=1) / side
+    np.clip(out, 0.0, 1.0, out=out)                  # slivers can stick sub-texel past extents
+    return out.astype(np.float32)

comfy_extras/mesh3d/uv_unwrap/parameterize.py

@@ -0,0 +1,565 @@
+"""Chart parameterization: ortho PCA projection, falling back to ABF/LSCM."""
+from __future__ import annotations
+
+import warnings
+from typing import List, Tuple
+
+import numpy as np
+import scipy.sparse as sp
+import scipy.sparse.linalg as spla
+import torch
+from torch import Tensor
+
+from . import mesh as _mesh
+
+LSCM_BATCH_MAX_VERTS = 256      # charts above this solve per-chart sparse (lscm_chart)
+
+
+def solve_least_squares(A: sp.csr_matrix, b: np.ndarray) -> np.ndarray:
+    """Solve ||Ax - b||^2 by factorizing AtA."""
+    At = A.T.tocsr()
+    AtA = (At @ A).tocsc()
+    Atb = At @ b
+    return spla.spsolve(AtA, Atb)
+
+
+def _triangle_local_2d(verts_3d: np.ndarray, faces: np.ndarray) -> np.ndarray:
+    """Per-triangle 2D coords [F, 3, 2] with v0 at origin, v1 along +x."""
+    v0 = verts_3d[faces[:, 0]]
+    v1 = verts_3d[faces[:, 1]]
+    v2 = verts_3d[faces[:, 2]]
+    e01 = v1 - v0
+    e02 = v2 - v0
+    L01 = np.linalg.norm(e01, axis=1).clip(min=1e-20)
+    x_axis = e01 / L01[:, None]
+    n = np.cross(e01, e02)
+    n /= np.linalg.norm(n, axis=1, keepdims=True).clip(min=1e-20)
+    y_axis = np.cross(n, x_axis)
+
+    out = np.zeros((faces.shape[0], 3, 2), dtype=np.float64)
+    out[:, 1, 0] = L01
+    out[:, 2, 0] = (e02 * x_axis).sum(axis=1)
+    out[:, 2, 1] = (e02 * y_axis).sum(axis=1)
+    return out
+
+
+def _pick_pins(loops: List[List[int]], verts_3d: np.ndarray) -> Tuple[int, int]:
+    """Pick the longest-diameter axis-extremal boundary vertex pair across all boundary verts."""
+    if not loops:
+        # Closed surface: two far verts via two-pass farthest.
+        d2 = np.sum((verts_3d - verts_3d[0]) ** 2, axis=1)
+        a = int(np.argmax(d2))
+        d2 = np.sum((verts_3d - verts_3d[a]) ** 2, axis=1)
+        b = int(np.argmax(d2))
+        return a, b
+    boundary_verts: List[int] = []
+    for loop in loops:
+        boundary_verts.extend(loop)
+    seen = set()
+    uniq = []
+    for v in boundary_verts:
+        if v not in seen:
+            seen.add(v)
+            uniq.append(v)
+    bv = np.asarray(uniq, dtype=np.int64)
+    pts = verts_3d[bv]
+    pin_pairs = []
+    for axis in range(3):
+        i_min = int(bv[int(np.argmin(pts[:, axis]))])
+        i_max = int(bv[int(np.argmax(pts[:, axis]))])
+        d = float(np.linalg.norm(verts_3d[i_min] - verts_3d[i_max]))
+        pin_pairs.append((d, i_min, i_max))
+    d0, _, _ = pin_pairs[0]
+    d1, _, _ = pin_pairs[1]
+    d2, _, _ = pin_pairs[2]
+    if d0 > d1 and d0 > d2:
+        _, a, b = pin_pairs[0]
+    elif d1 > d2:
+        _, a, b = pin_pairs[1]
+    else:
+        _, a, b = pin_pairs[2]
+    return a, b
+
+
+def _ortho_project(verts_3d: np.ndarray) -> np.ndarray:
+    """PCA-fit plane normal, axis-aligned tangent, project verts to 2D."""
+    centroid = verts_3d.mean(axis=0)
+    pts = verts_3d - centroid
+    cov = pts.T @ pts
+    _w, ev = np.linalg.eigh(cov)
+    normal = ev[:, 0]
+    a = np.abs(normal)
+    if a[0] < a[1] and a[0] < a[2]:
+        t = np.array([1.0, 0.0, 0.0])
+    elif a[1] < a[2]:
+        t = np.array([0.0, 1.0, 0.0])
+    else:
+        t = np.array([0.0, 0.0, 1.0])
+    t = t - normal * float(np.dot(normal, t))
+    t /= max(float(np.linalg.norm(t)), 1e-20)
+    b = np.cross(normal, t)
+    return np.stack([verts_3d @ t, verts_3d @ b], axis=1)
+
+
+def ortho_project_concat(verts: np.ndarray, chart_of_vert: np.ndarray, n_charts: int) -> np.ndarray:
+    """_ortho_project for every chart at once over concatenated per-chart vertices."""
+    cnt = np.bincount(chart_of_vert, minlength=n_charts).clip(min=1).astype(np.float64)
+    cen = np.stack([np.bincount(chart_of_vert, weights=verts[:, i], minlength=n_charts)
+                    for i in range(3)], axis=1) / cnt[:, None]
+    d = verts - cen[chart_of_vert]
+    cov = np.zeros((n_charts, 3, 3), dtype=np.float64)
+    for i in range(3):
+        for j in range(i, 3):
+            s = np.bincount(chart_of_vert, weights=d[:, i] * d[:, j], minlength=n_charts)
+            cov[:, i, j] = s
+            cov[:, j, i] = s
+    _w, ev = np.linalg.eigh(cov)
+    normal = ev[:, :, 0]
+    t = np.eye(3, dtype=np.float64)[np.argmin(np.abs(normal), axis=1)]
+    t = t - normal * (normal * t).sum(axis=1, keepdims=True)
+    t /= np.linalg.norm(t, axis=1, keepdims=True).clip(min=1e-20)
+    b = np.cross(normal, t)
+    tt, bb = t[chart_of_vert], b[chart_of_vert]
+    return np.stack([(verts * tt).sum(1), (verts * bb).sum(1)], axis=1)
+
+
+def stretch_metrics_concat(
+    verts: np.ndarray, uvs: np.ndarray, faces: np.ndarray,
+    chart_of_face: np.ndarray, n_charts: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Per-chart Sander stretch metrics (rms, max, n_flipped, n_zero_area); rms/max inf where undefined."""
+    p = verts[faces]
+    t = uvs[faces]
+    pa_signed = 0.5 * (
+        (t[:, 1, 1] - t[:, 0, 1]) * (t[:, 2, 0] - t[:, 0, 0])
+        - (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0]))
+    n_flip = np.bincount(chart_of_face[pa_signed < -1e-12], minlength=n_charts)
+    n_zero = np.bincount(chart_of_face[np.abs(pa_signed) < 1e-12], minlength=n_charts)
+    pa = np.abs(pa_signed).clip(min=1e-20)
+    ga = 0.5 * np.linalg.norm(np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1)
+    keep = (ga > 1e-12) & (np.abs(pa_signed) > 1e-12)
+    t1, s1 = t[:, 0, 0], t[:, 0, 1]
+    t2, s2 = t[:, 1, 0], t[:, 1, 1]
+    t3, s3 = t[:, 2, 0], t[:, 2, 1]
+    inv_2pa = 1.0 / (2.0 * pa)
+    Ss = (p[:, 0] * (t2 - t3)[:, None] + p[:, 1] * (t3 - t1)[:, None]
+          + p[:, 2] * (t1 - t2)[:, None]) * inv_2pa[:, None]
+    St = (p[:, 0] * (s3 - s2)[:, None] + p[:, 1] * (s1 - s3)[:, None]
+          + p[:, 2] * (s2 - s1)[:, None]) * inv_2pa[:, None]
+    a = (Ss * Ss).sum(axis=1)
+    bb = (Ss * St).sum(axis=1)
+    c = (St * St).sum(axis=1)
+    sigma2_sq = 0.5 * (a + c + np.sqrt(np.maximum(0.0, (a - c) ** 2 + 4 * bb ** 2)))
+    rms_sq = (a + c) * 0.5
+    cf = chart_of_face[keep]
+    tg = np.bincount(cf, weights=ga[keep], minlength=n_charts)
+    tp = np.bincount(cf, weights=pa[keep], minlength=n_charts)
+    rs = np.bincount(cf, weights=(rms_sq * ga)[keep], minlength=n_charts)
+    smax = np.zeros(n_charts, dtype=np.float64)
+    np.maximum.at(smax, cf, sigma2_sq[keep])
+    ok = tg > 0.0
+    tg_safe = np.where(ok, tg, 1.0)
+    norm = np.sqrt(tp / tg_safe)
+    rms = np.where(ok, np.sqrt(rs / tg_safe) * norm, np.inf)
+    mx = np.where(ok, np.sqrt(smax) * norm, np.inf)
+    return rms, mx, n_flip, n_zero
+
+
+def _segment_argmax(vals: np.ndarray, seg: np.ndarray, n: int) -> np.ndarray:
+    """Index of the (first) max element per segment; -1 for empty segments."""
+    amax = np.full(n, -np.inf)
+    np.maximum.at(amax, seg, vals)
+    hit = vals == amax[seg]
+    out = np.full(n, np.iinfo(np.int64).max, dtype=np.int64)
+    np.minimum.at(out, seg[hit], np.nonzero(hit)[0])
+    return np.where(out == np.iinfo(np.int64).max, -1, out)
+
+
+def lscm_charts_batch(
+    verts: np.ndarray,            # (sumV, 3) float64, per-chart concatenated
+    uv_pins: np.ndarray,          # (sumV, 2) float64, ortho UVs (pin values + fallback)
+    faces_gl: np.ndarray,         # (sumF, 3) global-local ids into verts
+    face_pos: np.ndarray,         # (sumF,) row index of each face within its chart
+    chart_of_face: np.ndarray,    # (sumF,)
+    chart_of_vert: np.ndarray,    # (sumV,)
+    vert_offsets: np.ndarray,     # (n_charts+1,)
+    chart_ids: np.ndarray,        # charts to solve (each with >=3 verts, >=1 face)
+    n_charts: int,
+    max_bucket_verts: int = LSCM_BATCH_MAX_VERTS,
+    device: "torch.device | None" = None,
+) -> dict:
+    """Batched dense ABF/LSCM; returns {chart_id: (Vc, 2) float32}. Charts larger than
+    max_bucket_verts are left out (the caller solves those sparse)."""
+    out: dict = {}
+    if chart_ids.size == 0:
+        return out
+    sel = np.zeros(n_charts, dtype=bool)
+    sel[chart_ids] = True
+    vcounts = np.diff(vert_offsets)
+
+    # ABF coefficients for all selected faces in one shot
+    fmask = sel[chart_of_face]
+    f_ids = np.nonzero(fmask)[0]
+    abf_ids, abf_cos, abf_sin, abf_valid = _abf_face_coefficients(verts, faces_gl[f_ids])
+
+    # farthest-point pin pair per chart (two passes)
+    vmask = sel[chart_of_vert]
+    v_ids = np.nonzero(vmask)[0]
+    cv = chart_of_vert[v_ids]
+    first = vert_offsets[:-1]
+    d0 = ((verts[v_ids] - verts[first[cv]]) ** 2).sum(1)
+    pin_a = _segment_argmax(d0, cv, n_charts)          # global vert index (into v_ids space)
+    pin_a = np.where(pin_a >= 0, v_ids[pin_a.clip(min=0)], -1)
+    d1 = ((verts[v_ids] - verts[pin_a.clip(min=0)[cv]]) ** 2).sum(1)
+    pin_b = _segment_argmax(d1, cv, n_charts)
+    pin_b = np.where(pin_b >= 0, v_ids[pin_b.clip(min=0)], -1)
+    # degenerate (all verts coincide): any distinct vert within the chart (Vc >= 3 guaranteed)
+    alt = np.where(pin_a == first, first + 1, first)
+    pin_b = np.where(pin_a == pin_b, alt, pin_b)
+
+    fcounts = np.bincount(chart_of_face[f_ids], minlength=n_charts)
+    # size-sorted chunks padded to their own max, bounded by an element budget so one
+    # face-heavy chart can't inflate a whole chunk
+    small = chart_ids[vcounts[chart_ids] <= max_bucket_verts]
+    sorted_ids = small[np.argsort(vcounts[small], kind="stable")]
+    budget = (96 << 20) // 8                            # float64 elements in a chunk's A
+    chunks = []
+    cs = 0
+    fmax_r = vmax_r = 0
+    for idx in range(sorted_ids.size):
+        c2 = sorted_ids[idx]
+        fm2 = max(fmax_r, int(fcounts[c2]))
+        vm2 = max(vmax_r, int(vcounts[c2]))
+        nb = idx - cs + 1
+        if nb > 1 and (nb > 128 or nb * 4 * fm2 * vm2 > budget):
+            chunks.append((cs, idx))
+            cs = idx
+            fmax_r, vmax_r = int(fcounts[c2]), int(vcounts[c2])
+        else:
+            fmax_r, vmax_r = fm2, vm2
+    if sorted_ids.size:
+        chunks.append((cs, sorted_ids.size))
+    for s, e in chunks:
+        cids = sorted_ids[s:e]
+        B = cids.size
+        Vmax = int(vcounts[cids].max())
+        Fmax = int(fcounts[cids].max())
+        N = 2 * Vmax
+        R = 2 * Fmax
+        compact = np.full(n_charts, -1, dtype=np.int64)
+        compact[cids] = np.arange(B)
+        fm = compact[chart_of_face[f_ids]] >= 0
+        fi = f_ids[fm]                               # face rows for this chunk
+        bi = compact[chart_of_face[fi]]              # chart slot per face
+        frow = face_pos[fi]
+        v0 = vert_offsets[chart_of_face[fi]]         # local id = global-local - v0
+        am = fm.nonzero()[0]                         # index into abf_* arrays
+
+        pieces_i: list = []
+        pieces_v: list = []
+
+        def scatter(rows, cols, vals, bsel):
+            pieces_i.append((bsel * R + rows) * N + cols)
+            pieces_v.append(vals)
+
+        val = abf_valid[am]
+        ii = am[val]
+        ids = abf_ids[ii] - v0[val, None]            # local vert ids, reordered
+        cosf, sinf = abf_cos[ii], abf_sin[ii]
+        rr, bsel = frow[val] * 2, bi[val]
+        ones = np.ones(ii.size)
+        for cc2, vv in ((ids[:, 0], cosf - 1.0), (ids[:, 0] + Vmax, -sinf),
+                        (ids[:, 1], -cosf), (ids[:, 1] + Vmax, sinf), (ids[:, 2], ones)):
+            scatter(rr, cc2, vv, bsel)
+        for cc2, vv in ((ids[:, 0], sinf), (ids[:, 0] + Vmax, cosf - 1.0),
+                        (ids[:, 1], -sinf), (ids[:, 1] + Vmax, -cosf), (ids[:, 2] + Vmax, ones)):
+            scatter(rr + 1, cc2, vv, bsel)
+
+        inv = ~val
+        if inv.any():
+            jj = fi[inv]
+            tri2d = _triangle_local_2d(verts, faces_gl[jj])
+            twice = tri2d[:, 1, 0] * tri2d[:, 2, 1] - tri2d[:, 1, 1] * tri2d[:, 2, 0]
+            w = 1.0 / np.sqrt(2.0 * np.abs(twice).clip(min=1e-20))
+            rr2, bs2 = frow[inv] * 2, bi[inv]
+            lids = faces_gl[jj] - v0[inv, None]
+            for j in range(3):
+                jp1, jp2 = (j + 1) % 3, (j + 2) % 3
+                aj = (tri2d[:, jp1, 0] - tri2d[:, jp2, 0]) * w
+                bj = (tri2d[:, jp1, 1] - tri2d[:, jp2, 1]) * w
+                vc2 = lids[:, j]
+                scatter(rr2, vc2, aj, bs2)
+                scatter(rr2, vc2 + Vmax, -bj, bs2)
+                scatter(rr2 + 1, vc2, bj, bs2)
+                scatter(rr2 + 1, vc2 + Vmax, aj, bs2)
+
+        flat = np.concatenate(pieces_i)
+        A = np.bincount(flat, weights=np.concatenate(pieces_v),
+                        minlength=B * R * N).reshape(B, R, N)
+
+        # pins: move their columns to the RHS, then constrain via identity rows
+        voff = vert_offsets[cids]
+        pa_l = pin_a[cids] - voff
+        pb_l = pin_b[cids] - voff
+        pin_cols = np.stack([pa_l, pb_l, pa_l + Vmax, pb_l + Vmax], 1)       # (B,4)
+        pin_vals = np.stack([uv_pins[pin_a[cids], 0], uv_pins[pin_b[cids], 0],
+                             uv_pins[pin_a[cids], 1], uv_pins[pin_b[cids], 1]], 1)
+        rhs = np.zeros((B, R), dtype=np.float64)
+        barange = np.arange(B)
+        for k in range(4):
+            rhs -= A[barange, :, pin_cols[:, k]] * pin_vals[:, k, None]
+            A[barange, :, pin_cols[:, k]] = 0.0
+
+        # constrained columns: the 4 pins + padding beyond each chart's vert count
+        vcs = vcounts[cids]
+        padm = np.arange(Vmax)[None, :] >= vcs[:, None]
+        con = np.concatenate([padm, padm], axis=1)                            # (B,N)
+        np.put_along_axis(con, pin_cols, True, axis=1)
+        cval = np.zeros((B, N), dtype=np.float64)
+        np.put_along_axis(cval, pin_cols, pin_vals, axis=1)
+
+        # normal equations + batched solve; the fp64 dense algebra goes to the GPU when available
+        use_gpu = device is not None and device.type == "cuda"
+        if use_gpu:
+            A_t = torch.from_numpy(A).to(device)
+            At = A_t.transpose(1, 2)
+            AtA = At @ A_t
+            Atb = (At @ torch.from_numpy(rhs).to(device).unsqueeze(2)).squeeze(2)
+            con_t = torch.from_numpy(con).to(device)
+            free2 = (~con_t[:, :, None]) & (~con_t[:, None, :])
+            AtA = AtA * free2
+            diag = torch.diagonal(AtA, dim1=1, dim2=2)
+            # median (not max) positive diagonal: a degenerate face's ~1e19 squared row
+            # weight would blow a max-scaled eps past the unit ABF rows
+            dpos = torch.where(diag > 0, diag, torch.full_like(diag, float("nan")))
+            dsc = 1e-12 * torch.nan_to_num(dpos.nanmedian(dim=1).values, nan=1e-8).clamp_min(1e-20)
+            diag += torch.where(con_t, torch.ones_like(diag), dsc[:, None].expand_as(diag))
+            Atb = torch.where(con_t, torch.from_numpy(cval).to(device), Atb)
+            x = torch.linalg.solve(AtA, Atb).cpu().numpy()
+        else:
+            At = A.transpose(0, 2, 1)
+            AtA = At @ A                             # batched BLAS dgemm
+            Atb = (At @ rhs[:, :, None])[:, :, 0]
+            AtA *= (~con[:, :, None]) & (~con[:, None, :])
+            dg = AtA.reshape(B, -1)[:, ::N + 1]
+            # median positive diagonal (see GPU branch): robust to degenerate-face weights
+            dpos = np.where(dg > 0, dg, np.nan)
+            with np.errstate(all="ignore"):
+                dsc = 1e-12 * np.nan_to_num(np.nanmedian(dpos, axis=1), nan=1e-8).clip(min=1e-20)
+            dg += np.where(con, 1.0, dsc[:, None])
+            Atb2 = np.where(con, cval, Atb)
+            x = np.linalg.solve(AtA, Atb2)
+        for i2, c2 in enumerate(cids):
+            vc3 = int(vcs[i2])
+            out[int(c2)] = np.stack([x[i2, :vc3], x[i2, Vmax:Vmax + vc3]], 1).astype(np.float32)
+    return out
+
+
+def _uv_boundary_self_intersects(
+    uvs: np.ndarray, faces: np.ndarray, face_face: np.ndarray, eps: float = 1e-9
+) -> bool:
+    """True if any chart-boundary edge pair crosses in 2D (ortho folded the chart)."""
+    fi, ei = np.nonzero(face_face < 0)
+    n = fi.size
+    if n < 2:
+        return False
+    a = uvs[faces[fi, ei]].astype(np.float64)
+    b = uvs[faces[fi, (ei + 1) % 3]].astype(np.float64)
+    d = b - a
+    # Pairwise segment crossings, row-chunked to bound memory at chunk*n.
+    chunk = max(1, min(n, 1_000_000 // max(n, 1)))
+    for s in range(0, n, chunk):
+        e = min(s + chunk, n)
+        d1 = d[s:e, None, :]
+        denom = d1[:, :, 0] * d[None, :, 1] - d1[:, :, 1] * d[None, :, 0]
+        rx = a[None, :, 0] - a[s:e, None, 0]
+        ry = a[None, :, 1] - a[s:e, None, 1]
+        with np.errstate(divide="ignore", invalid="ignore"):
+            t = (rx * d[None, :, 1] - ry * d[None, :, 0]) / denom
+            u = (rx * d1[:, :, 1] - ry * d1[:, :, 0]) / denom
+        cross = (
+            (np.abs(denom) >= eps)
+            & (t > eps) & (t < 1.0 - eps)
+            & (u > eps) & (u < 1.0 - eps)
+        )
+        if bool(cross.any()):
+            return True
+    return False
+
+
+def _abf_face_coefficients(
+    verts_3d: np.ndarray, faces: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Per-face ABF constraint (largest-sine vertex at local index 2); returns (faces_reordered, cosine, sine, valid_mask) with valid_mask False for degenerate tris."""
+    p0 = verts_3d[faces[:, 0]]
+    p1 = verts_3d[faces[:, 1]]
+    p2 = verts_3d[faces[:, 2]]
+    e01 = p1 - p0
+    e12 = p2 - p1
+    e20 = p0 - p2
+    L01 = np.linalg.norm(e01, axis=1).clip(min=1e-20)
+    L12 = np.linalg.norm(e12, axis=1).clip(min=1e-20)
+    L20 = np.linalg.norm(e20, axis=1).clip(min=1e-20)
+    cos_a0 = ((-e20) * e01).sum(axis=1) / (L20 * L01)
+    cos_a1 = ((-e01) * e12).sum(axis=1) / (L01 * L12)
+    cos_a2 = ((-e12) * e20).sum(axis=1) / (L12 * L20)
+    cos_a0 = cos_a0.clip(-1.0, 1.0)
+    cos_a1 = cos_a1.clip(-1.0, 1.0)
+    cos_a2 = cos_a2.clip(-1.0, 1.0)
+    a = np.arccos(cos_a0)
+    b_ang = np.arccos(cos_a1)
+    c_ang = np.arccos(cos_a2)
+    angles = np.stack([a, b_ang, c_ang], axis=1)
+    sines = np.stack([np.sin(a), np.sin(b_ang), np.sin(c_ang)], axis=1)
+    valid = (angles > 1e-12).all(axis=1)
+    ids = faces.astype(np.int64).copy()
+
+    s0, s1, s2 = sines[:, 0], sines[:, 1], sines[:, 2]
+    pattA = (s1 > s0) & (s1 > s2)
+    pattB = (~pattA) & (s0 > s1) & (s0 > s2)
+
+    if pattA.any():
+        old_a = angles[pattA].copy()
+        old_s = sines[pattA].copy()
+        old_id = ids[pattA].copy()
+        angles[pattA] = old_a[:, [2, 0, 1]]
+        sines[pattA] = old_s[:, [2, 0, 1]]
+        ids[pattA] = old_id[:, [2, 0, 1]]
+    if pattB.any():
+        old_a = angles[pattB].copy()
+        old_s = sines[pattB].copy()
+        old_id = ids[pattB].copy()
+        angles[pattB] = old_a[:, [1, 2, 0]]
+        sines[pattB] = old_s[:, [1, 2, 0]]
+        ids[pattB] = old_id[:, [1, 2, 0]]
+
+    a0 = angles[:, 0]
+    s0 = sines[:, 0]
+    s1 = sines[:, 1]
+    s2 = sines[:, 2]
+    c0 = np.cos(a0)
+    ratio = np.where(s2 > 0.0, s1 / s2.clip(min=1e-20), 1.0)
+    cosine = c0 * ratio
+    sine = s0 * ratio
+    return ids, cosine, sine, valid
+
+
+def lscm_chart(
+    local_verts: Tensor,
+    local_faces: Tensor,
+    local_face_face: Tensor,
+    pin_positions: "np.ndarray | None" = None,
+) -> Tensor:
+    """ABF parameterization on one chart (degenerate faces use plain LSCM rows; two pins fix gauge at pin_positions)."""
+    verts_np = local_verts.detach().cpu().numpy().astype(np.float64)
+    faces_np = local_faces.detach().cpu().numpy().astype(np.int64)
+    Vc = verts_np.shape[0]
+    Fc = faces_np.shape[0]
+
+    if Vc < 3 or Fc == 0:
+        return torch.zeros((Vc, 2), dtype=torch.float32, device=local_verts.device)
+
+    loops = _mesh.chart_boundary_loops(local_faces, local_face_face)
+    pin_a, pin_b = _pick_pins(loops, verts_np)
+
+    if pin_positions is not None and pin_positions.shape == (Vc, 2):
+        pa = pin_positions[pin_a]
+        pb = pin_positions[pin_b]
+        u_a, v_a = float(pa[0]), float(pa[1])
+        u_b, v_b = float(pb[0]), float(pb[1])
+    else:
+        u_a, v_a = 0.0, 0.0
+        u_b, v_b = 1.0, 0.0
+
+    abf_ids, abf_cos, abf_sin, abf_valid = _abf_face_coefficients(verts_np, faces_np)
+
+    rows_list: List[np.ndarray] = []
+    cols_list: List[np.ndarray] = []
+    vals_list: List[np.ndarray] = []
+
+    # ABF rows for valid faces.
+    valid_idx = np.nonzero(abf_valid)[0]
+    if valid_idx.size:
+        Nv = valid_idx.size
+        id0 = abf_ids[valid_idx, 0]
+        id1 = abf_ids[valid_idx, 1]
+        id2 = abf_ids[valid_idx, 2]
+        cosf = abf_cos[valid_idx]
+        sinf = abf_sin[valid_idx]
+        r_real = valid_idx * 2
+        r_imag = valid_idx * 2 + 1
+        ones = np.ones(Nv, dtype=np.float64)
+        rows_list.extend([r_real] * 5)
+        cols_list.extend([id0, id0 + Vc, id1, id1 + Vc, id2])
+        vals_list.extend([cosf - 1.0, -sinf, -cosf, sinf, ones])
+        rows_list.extend([r_imag] * 5)
+        cols_list.extend([id0, id0 + Vc, id1, id1 + Vc, id2 + Vc])
+        vals_list.extend([sinf, cosf - 1.0, -sinf, -cosf, ones])
+
+    # Plain-LSCM rows for invalid (degenerate) faces.
+    invalid_idx = np.nonzero(~abf_valid)[0]
+    if invalid_idx.size:
+        tri2d_inv = _triangle_local_2d(verts_np, faces_np[invalid_idx])
+        twice_area_inv = (
+            tri2d_inv[:, 1, 0] * tri2d_inv[:, 2, 1]
+            - tri2d_inv[:, 1, 1] * tri2d_inv[:, 2, 0]
+        )
+        weight_inv = 1.0 / np.sqrt(2.0 * np.abs(twice_area_inv).clip(min=1e-20))
+        r_real_inv = invalid_idx * 2
+        r_imag_inv = invalid_idx * 2 + 1
+        for j in range(3):
+            jp1 = (j + 1) % 3
+            jp2 = (j + 2) % 3
+            a_j = (tri2d_inv[:, jp1, 0] - tri2d_inv[:, jp2, 0]) * weight_inv
+            b_j = (tri2d_inv[:, jp1, 1] - tri2d_inv[:, jp2, 1]) * weight_inv
+            v_idx = faces_np[invalid_idx, j]
+            rows_list.extend([r_real_inv, r_real_inv, r_imag_inv, r_imag_inv])
+            cols_list.extend([v_idx, v_idx + Vc, v_idx, v_idx + Vc])
+            vals_list.extend([a_j, -b_j, b_j, a_j])
+
+    rows = np.concatenate(rows_list) if rows_list else np.empty(0, dtype=np.int64)
+    cols = np.concatenate(cols_list) if cols_list else np.empty(0, dtype=np.int64)
+    vals = np.concatenate(vals_list) if vals_list else np.empty(0, dtype=np.float64)
+
+    A_full = sp.csr_matrix((vals, (rows, cols)), shape=(2 * Fc, 2 * Vc))
+
+    pin_cols = np.array([pin_a, pin_b, pin_a + Vc, pin_b + Vc], dtype=np.int64)
+    pin_vals = np.array([u_a, u_b, v_a, v_b], dtype=np.float64)
+
+    free_mask = np.ones(2 * Vc, dtype=bool)
+    free_mask[pin_cols] = False
+    free_cols = np.nonzero(free_mask)[0]
+
+    A_pinned = A_full[:, pin_cols]
+    A_free = A_full[:, free_cols]
+    b = -(A_pinned @ pin_vals)
+
+    # Singular system (under-constrained chart) falls back to ortho.
+    fallback_to_ortho = False
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", category=sp.linalg.MatrixRankWarning)
+            x_free = solve_least_squares(A_free, b)
+        if not np.all(np.isfinite(x_free)):
+            fallback_to_ortho = True
+    except (sp.linalg.MatrixRankWarning, RuntimeError):
+        fallback_to_ortho = True  # singular / under-constrained system
+
+    if fallback_to_ortho:
+        if pin_positions is not None and pin_positions.shape == (Vc, 2):
+            uvs = pin_positions.astype(np.float32)
+        else:
+            uvs = _ortho_project(verts_np).astype(np.float32)
+        return torch.from_numpy(uvs).to(local_verts.device)
+
+    full = np.zeros(2 * Vc, dtype=np.float64)
+    full[free_cols] = x_free
+    full[pin_cols] = pin_vals
+    uvs = np.stack([full[:Vc], full[Vc:]], axis=1).astype(np.float32)
+    if not np.all(np.isfinite(uvs)):
+        if pin_positions is not None and pin_positions.shape == (Vc, 2):
+            uvs = pin_positions.astype(np.float32)
+        else:
+            uvs = _ortho_project(verts_np).astype(np.float32)
+
+    return torch.from_numpy(uvs).to(local_verts.device)

comfy_extras/mesh3d/uv_unwrap/segment.py

@@ -0,0 +1,414 @@
+"""Adaptive cost-grow chart segmentation (vectorized torch, CPU or GPU)."""
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+from torch import Tensor
+from tqdm import tqdm
+
+from .mesh import MeshData, face_edge_lengths
+
+
+DEFAULT_W_NORMAL_DEVIATION = 2.0
+DEFAULT_W_ROUNDNESS = 0.01
+DEFAULT_W_STRAIGHTNESS = 6.0
+DEFAULT_MAX_COST = 2.0
+NORMAL_DEVIATION_HARD_CUTOFF = 0.707  # ~75°
+
+
+def _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum, area, perim, K,
+               nd_cutoff, tau, w_nd, w_round, w_straight):
+    """One grow pass: each frontier face joins its lowest-cost adjacent chart if cost <= tau;
+    returns the number of faces assigned."""
+    u = frontier.nonzero(as_tuple=True)[0]
+    if u.numel() == 0:
+        return 0
+    nb = ff[u]                                       # (U,3) neighbor face ids
+    nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
+    valid = nbc >= 0
+    d = (fn[u][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
+    nd = (1.0 - d).clamp(0.0, 1.0)
+    valid &= nd < nd_cutoff
+    el = fel[u]                                      # (U,3)
+    # l_in per candidate chart j: edge k counts if its (assigned) neighbor is in chart j
+    inm = (nbc[:, :, None] == nbc[:, None, :]) & valid[:, None, :]
+    l_in = (el[:, None, :] * inm).sum(-1)            # (U,3)
+    tot = el.sum(-1, keepdim=True)
+    l_out = tot - l_in
+    ca = area[nbc.clamp_min(0)]
+    cp = perim[nbc.clamp_min(0)]
+    new_perim = cp - l_in + l_out
+    new_r = new_perim * new_perim / (ca + fa[u][:, None]).clamp_min(1e-20)
+    round_cost = torch.where((cp <= 1e-20) | (ca <= 1e-20) | (new_r <= 1e-20),
+                             torch.zeros_like(new_r),
+                             1.0 - (cp * cp / ca.clamp_min(1e-20)) / new_r.clamp_min(1e-20))
+    straight_cost = ((l_out - l_in) / tot.clamp_min(1e-20)).clamp(max=0.0)
+    cost = w_nd * nd + w_round * round_cost + w_straight * straight_cost
+    cost = torch.where(valid, cost, cost.new_full((), float("inf")))
+    best_cost, best_j = cost.min(1)
+    acc = best_cost <= tau
+    n_acc = int(acc.sum())
+    if n_acc == 0:
+        return 0
+
+    f_acc = u[acc]
+    c_acc = nbc.gather(1, best_j[:, None]).squeeze(1)[acc]
+    nbc_old = nbc[acc]                               # neighbor charts before this commit
+    face_chart[f_acc] = c_acc
+    nb_acc = nb[acc]
+    nbs_acc = nb_acc.clamp_min(0)
+    nbc_post = torch.where(nb_acc >= 0, face_chart[nbs_acc], nb_acc.new_full((), -1))
+    # frontier update: committed faces leave; their still-unassigned neighbors enter
+    frontier[f_acc] = False
+    grow_nb = nbs_acc[(nb_acc >= 0) & (nbc_post < 0)]
+    frontier[grow_nb] = True
+    el_acc = el[acc]
+    cx = c_acc[:, None]
+    dper = torch.where(nbc_old == cx, -el_acc,       # was member: edge turns interior
+                       torch.where(nbc_post == cx, torch.zeros_like(el_acc),  # co-committer
+                                   el_acc)).sum(1)   # boundary / other chart
+    perim.scatter_add_(0, c_acc, dper)
+    area.scatter_add_(0, c_acc, fa[f_acc])
+    nsum.index_add_(0, c_acc, fn[f_acc] * fa[f_acc, None])
+    nl = nsum[:K].norm(dim=1, keepdim=True)
+    basis[:K] = torch.where(nl > 1e-20, nsum[:K] / nl.clamp_min(1e-20), basis[:K])
+    return n_acc
+
+
+def segment_charts(
+    mesh: MeshData,
+    max_cost: float = DEFAULT_MAX_COST,
+    w_normal_deviation: float = DEFAULT_W_NORMAL_DEVIATION,
+    w_roundness: float = DEFAULT_W_ROUNDNESS,
+    w_straightness: float = DEFAULT_W_STRAIGHTNESS,
+    progress_callback=None,
+) -> Tensor:
+    """Segment mesh into charts (parallel batch cost-grow). Returns face -> chart_id."""
+    F = mesh.faces.shape[0]
+    device = mesh.faces.device
+    if F == 0:
+        return torch.zeros(0, dtype=torch.long, device=device)
+
+    fn = mesh.face_normal.detach().to(torch.float32)
+    fa = mesh.face_area.detach().to(torch.float32)
+    fc = mesh.face_centroid.detach().to(torch.float32)
+    ff = mesh.face_face.detach().long()
+    fel = face_edge_lengths(mesh.vertices, mesh.faces).detach().to(torch.float32)
+    nd_cutoff = NORMAL_DEVIATION_HARD_CUTOFF
+
+    # one seed per connected component (first face of each)
+    comp = mesh.component.detach().long().to(device)
+    ncomp = int(comp.max()) + 1 if comp.numel() else 0
+    if ncomp:
+        seeds = torch.full((ncomp,), F, dtype=torch.long, device=device)
+        seeds.scatter_reduce_(0, comp, torch.arange(F, device=device), reduce="amin")
+    else:
+        seeds = torch.zeros(1, dtype=torch.long, device=device)
+    K = seeds.shape[0]
+
+    max_total_charts = max(F, 8000)
+    cap = K + F + 1                                  # every re-seed assigns a face, so K < K0 + F
+    face_chart = torch.full((F,), -1, dtype=torch.long, device=device)
+    basis = torch.zeros(cap, 3, dtype=torch.float32, device=device)
+    nsum = torch.zeros(cap, 3, dtype=torch.float32, device=device)
+    area = torch.zeros(cap, dtype=torch.float32, device=device)
+    perim = torch.zeros(cap, dtype=torch.float32, device=device)
+    face_chart[seeds] = torch.arange(K, device=device)
+    basis[:K] = fn[seeds]
+    nsum[:K] = fn[seeds] * fa[seeds, None]
+    area[:K] = fa[seeds]
+    perim[:K] = fel[seeds].sum(1)
+    frontier = torch.zeros(F, dtype=torch.bool, device=device)
+    seed_nb = ff[seeds]
+    seed_nb = seed_nb[seed_nb >= 0]
+    frontier[seed_nb] = True
+    frontier &= face_chart < 0
+
+    min_d2 = torch.full((F,), float("inf"), dtype=torch.float32, device=device)
+    for i in range(0, K, 32):                        # chunked: (F, <=32, 3) stays small
+        d2 = ((fc[:, None, :] - fc[seeds[i:i + 32]][None, :, :]) ** 2).sum(-1)
+        min_d2 = torch.minimum(min_d2, d2.amin(1))
+
+    # Multi-pass threshold schedule (low-cost first); tau cap 0.5 keeps cones ~30deg.
+    tau_final = min(max_cost * 0.25, 0.5)
+    thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final]
+    max_inner = max(64, int(F ** 0.5) * 2)
+    outer_iter = 0
+    assigned = 0
+    tq = tqdm(total=F, desc="unwrap: segment (adaptive)", unit="face", leave=False)
+    while True:
+        outer_iter += 1
+        if outer_iter > F + 16:
+            break
+        for tau in thresholds:
+            for _ in range(max_inner):
+                n_added = _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum,
+                                     area, perim, K, nd_cutoff, tau, w_normal_deviation,
+                                     w_roundness, w_straightness)
+                if n_added == 0:
+                    break
+                tq.update(n_added)
+                assigned += n_added
+                if progress_callback is not None:
+                    progress_callback(assigned, F)
+        unassigned = face_chart < 0
+        if int(unassigned.sum()) == 0:
+            break
+        if K >= max_total_charts:
+            break
+        # re-seed at the unassigned face farthest from every existing seed
+        new_seed = int(torch.where(unassigned, min_d2,
+                                   min_d2.new_full((), float("-inf"))).argmax())
+        face_chart[new_seed] = K
+        basis[K] = fn[new_seed]
+        nsum[K] = fn[new_seed] * fa[new_seed]
+        area[K] = fa[new_seed]
+        perim[K] = fel[new_seed].sum()
+        K += 1
+        min_d2 = torch.minimum(min_d2, ((fc - fc[new_seed]) ** 2).sum(-1))
+        tq.update(1)
+        frontier[new_seed] = False
+        ns_nb = ff[new_seed]
+        ns_nb = ns_nb[ns_nb >= 0]
+        frontier[ns_nb[face_chart[ns_nb] < 0]] = True
+
+    tq.close()
+
+    # Orphan cleanup: leftover faces join their best-matching neighbor's chart.
+    while True:
+        orphans = (face_chart < 0).nonzero(as_tuple=True)[0]
+        if orphans.numel() == 0:
+            break
+        nb = ff[orphans]
+        nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
+        valid = nbc >= 0
+        assignable = valid.any(1)
+        if not bool(assignable.any()):
+            break
+        d = (fn[orphans][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
+        ndv = torch.where(valid, 1.0 - d, d.new_full((), float("inf")))
+        best_c = nbc.gather(1, ndv.argmin(1, keepdim=True)).squeeze(1)
+        face_chart[orphans[assignable]] = best_c[assignable]
+    leftover = (face_chart < 0).nonzero(as_tuple=True)[0]
+    if leftover.numel():                             # isolated faces become singleton charts
+        face_chart[leftover] = K + torch.arange(leftover.numel(), device=device)
+
+    _, inverse = torch.unique(face_chart, sorted=True, return_inverse=True)
+    return inverse
+
+
+# Parallel edge-collapse (PEC) chart clustering (GPU)
+def _combine_normal_cones(
+    axis_a: Tensor, half_a: Tensor,
+    axis_b: Tensor, half_b: Tensor,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """Merge two normal cones along the great circle from axis_a; returns (combined_axis, combined_half_angle, axis_angle)."""
+    cos_angle = (axis_a * axis_b).sum(dim=-1).clamp(-1.0, 1.0)
+    axis_angle = torch.acos(cos_angle)
+    new_low = torch.minimum(-half_a, axis_angle - half_b)
+    new_high = torch.maximum(half_a, axis_angle + half_b)
+    new_half = (new_high - new_low) * 0.5
+    rot_angle = (new_high + new_low) * 0.5
+    b_perp = axis_b - axis_a * cos_angle.unsqueeze(-1)
+    b_perp_norm = b_perp.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+    b_perp_unit = b_perp / b_perp_norm
+    new_axis = (
+        axis_a * torch.cos(rot_angle).unsqueeze(-1)
+        + b_perp_unit * torch.sin(rot_angle).unsqueeze(-1)
+    )
+    new_axis_norm = new_axis.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+    new_axis = new_axis / new_axis_norm
+    return new_axis, new_half, axis_angle
+
+
+def _build_chart_edges(
+    face_face: Tensor,
+    chart_id: Tensor,
+    face_edge_len: Tensor,
+) -> Tuple[Tensor, Tensor]:
+    """Build chart-edge list (chart_pairs[E,2] with a<b, edge_length[E]); same-chart edges dropped, duplicates summed."""
+    F = face_face.shape[0]
+    device = face_face.device
+    f_idx = torch.arange(F, device=device).repeat_interleave(3)
+    nb = face_face.flatten()
+    valid = nb >= 0
+    f_idx = f_idx[valid]
+    nb = nb[valid]
+    el = face_edge_len.flatten()[valid]
+
+    ca = chart_id[f_idx]
+    cb = chart_id[nb]
+    diff = ca != cb
+    ca = ca[diff]
+    cb = cb[diff]
+    el = el[diff]
+    if ca.numel() == 0:
+        return (
+            torch.empty((0, 2), dtype=torch.long, device=device),
+            torch.empty(0, device=device),
+        )
+
+    lo = torch.minimum(ca, cb)
+    hi = torch.maximum(ca, cb)
+    V = int(chart_id.max().item()) + 1
+    key = lo * V + hi
+    sort_idx = torch.argsort(key)
+    sorted_key = key[sort_idx]
+    sorted_lo = lo[sort_idx]
+    sorted_hi = hi[sort_idx]
+    sorted_el = el[sort_idx]
+    unique_key, inverse, counts = torch.unique(
+        sorted_key, return_inverse=True, return_counts=True
+    )
+    n_unique = unique_key.shape[0]
+    reduced_el = torch.zeros(n_unique, device=device, dtype=el.dtype)
+    reduced_el.scatter_add_(0, inverse, sorted_el)
+    first_idx = torch.cat([
+        torch.zeros(1, dtype=torch.long, device=device),
+        counts.cumsum(0)[:-1],
+    ])
+    pair_lo = sorted_lo[first_idx]
+    pair_hi = sorted_hi[first_idx]
+    chart_pairs = torch.stack([pair_lo, pair_hi], dim=1)
+    return chart_pairs, reduced_el
+
+
+def _merge_small_charts(
+    chart_id: Tensor, face_normal: Tensor, face_area: Tensor,
+    face_face: Tensor, face_edge_len: Tensor,
+    min_faces: int, cost_cap: float,
+) -> Tensor:
+    """Absorb charts under min_faces faces into their lowest-cone-cost neighbor (capped at cost_cap)."""
+    if chart_id.numel() == 0:
+        return chart_id
+    device = chart_id.device
+    for _ in range(16):
+        N = int(chart_id.max().item()) + 1
+        sizes = torch.bincount(chart_id, minlength=N)
+        # recompute cones from scratch: area-weighted mean axis, max deviation as half-angle
+        axis = torch.zeros(N, 3, dtype=torch.float32, device=device)
+        axis.index_add_(0, chart_id, face_normal * face_area[:, None])
+        axis = axis / axis.norm(dim=1, keepdim=True).clamp_min(1e-12)
+        dev = torch.acos((face_normal * axis[chart_id]).sum(1).clamp(-1.0, 1.0))
+        half = torch.zeros(N, dtype=torch.float32, device=device)
+        half.scatter_reduce_(0, chart_id, dev, reduce="amax")
+
+        edges, _ = _build_chart_edges(face_face, chart_id, face_edge_len)
+        if edges.shape[0] == 0:
+            break
+        a, b = edges[:, 0], edges[:, 1]
+        _, new_half, _ = _combine_normal_cones(axis[a], half[a], axis[b], half[b])
+        ok = new_half <= cost_cap
+        E = edges.shape[0]
+        key = (torch.clamp(new_half * 1e6, max=2e9).to(torch.int64) << 32) \
+            | torch.arange(E, dtype=torch.long, device=device)
+        best = torch.full((N,), 1 << 62, dtype=torch.long, device=device)
+        va = (sizes[a] < min_faces) & ok
+        vb = (sizes[b] < min_faces) & ok
+        best.scatter_reduce_(0, a[va], key[va], reduce="amin")
+        best.scatter_reduce_(0, b[vb], key[vb], reduce="amin")
+        src = (best < (1 << 62)).nonzero(as_tuple=True)[0]
+        if src.numel() == 0:
+            break
+        eid = best[src] & 0xFFFFFFFF
+        ea, eb = a[eid], b[eid]
+        tgt = torch.where(ea == src, eb, ea)
+        # cycle break: keep src->tgt only if tgt merges nowhere itself or src > tgt;
+        # the kept graph is then a DAG, so the pointer-doubling below terminates
+        prop = torch.arange(N, dtype=torch.long, device=device)
+        prop[src] = tgt
+        keepm = (prop[tgt] == tgt) | (src > tgt)
+        remap = torch.arange(N, dtype=torch.long, device=device)
+        remap[src[keepm]] = tgt[keepm]
+        for _ in range(32):
+            nr = remap[remap]
+            if torch.equal(nr, remap):
+                break
+            remap = nr
+        chart_id = remap[chart_id]
+        _, chart_id = torch.unique(chart_id, return_inverse=True)
+    return chart_id
+
+
+def cluster_charts_pec(
+    mesh: MeshData,
+    max_cost: float = 0.7,
+    max_iters: int = 1024,
+    min_faces: int = 8,
+    progress_callback=None,
+) -> Tensor:
+    """Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge
+    cutoff (~0.7 rad ~ 40deg); charts under min_faces are then absorbed at a relaxed 2x cutoff."""
+    device = mesh.faces.device
+    F = mesh.faces.shape[0]
+    faces = mesh.faces.to(torch.long)
+    vertices = mesh.vertices.to(torch.float32)
+    face_normal = mesh.face_normal.to(torch.float32)
+    face_face = mesh.face_face.to(torch.long)
+
+    face_edge_len = face_edge_lengths(vertices, faces)
+
+    chart_id = torch.arange(F, dtype=torch.long, device=device)
+    chart_axis = face_normal.clone()
+    chart_half = torch.zeros(F, dtype=torch.float32, device=device)
+
+    for it in range(max_iters):
+        edges, _ = _build_chart_edges(face_face, chart_id, face_edge_len)
+        if edges.shape[0] == 0:
+            break
+
+        a = edges[:, 0]
+        b = edges[:, 1]
+        axis_a = chart_axis[a]
+        axis_b = chart_axis[b]
+        half_a = chart_half[a]
+        half_b = chart_half[b]
+        _, new_half, _ = _combine_normal_cones(axis_a, half_a, axis_b, half_b)
+        cost = new_half.clone()
+
+        # Pack (cost, edge_id) so scatter_reduce amin picks the right edge.
+        E = edges.shape[0]
+        N = int(chart_id.max().item()) + 1
+        edge_ids = torch.arange(E, dtype=torch.long, device=device)
+        cost_i32 = torch.clamp(cost * 1e6, max=2e9).to(torch.int64)
+        key = (cost_i32 << 32) | edge_ids
+        chart_min = torch.full((N,), (2**62), dtype=torch.long, device=device)
+        chart_min.scatter_reduce_(0, a, key, reduce="amin", include_self=True)
+        chart_min.scatter_reduce_(0, b, key, reduce="amin", include_self=True)
+
+        # Mutual-min collapse: each chart in at most one merge per iter (winners are disjoint pairs).
+        is_a_min = chart_min[a] == key
+        is_b_min = chart_min[b] == key
+        mutual = is_a_min & is_b_min
+        within = cost <= max_cost
+        winners = mutual & within
+
+        n_merge = int(winners.sum().item())
+        if n_merge == 0:
+            break
+        if progress_callback is not None:
+            progress_callback(F - N + n_merge, F)      # saturating: charts remaining vs faces
+
+        win_a = a[winners]
+        win_b = b[winners]
+
+        axis_a_w = chart_axis[win_a]
+        half_a_w = chart_half[win_a]
+        axis_b_w = chart_axis[win_b]
+        half_b_w = chart_half[win_b]
+        new_axis, new_half_w, _ = _combine_normal_cones(
+            axis_a_w, half_a_w, axis_b_w, half_b_w,
+        )
+        chart_axis[win_a] = new_axis
+        chart_half[win_a] = new_half_w
+
+        remap = torch.arange(N, dtype=torch.long, device=device)
+        remap[win_b] = win_a
+        chart_id = remap[chart_id]
+
+    if min_faces > 1:
+        chart_id = _merge_small_charts(chart_id, face_normal, mesh.face_area.to(torch.float32),
+                                       face_face, face_edge_len, min_faces, 2.0 * max_cost)
+    _, inverse = torch.unique(chart_id, sorted=True, return_inverse=True)
+    return inverse

comfy_extras/nodes_model_advanced.py

@@ -277,13 +277,30 @@ class RescaleCFG:
     CATEGORY = "model/patch"
 
     def patch(self, model, multiplier):
+        model_sampling = model.get_model_object("model_sampling")
+        is_x0_space = not isinstance(model_sampling, comfy.model_sampling.EPS)
+
         def rescale_cfg(args):
+            x_orig = args["input"]
+            cond_scale = args["cond_scale"]
+
+            if is_x0_space:
+                # Flow-matching / X0 models: cond_denoised/uncond_denoised are x_0 estimates,
+                # so the eps↔v conversion below would be wrong. Rescale directly in x_0 space.
+                x_0_cond = args["cond_denoised"]
+                x_0_uncond = args["uncond_denoised"]
+                x_0_cfg = x_0_uncond + cond_scale * (x_0_cond - x_0_uncond)
+                dims = tuple(range(1, x_0_cond.ndim))
+                ro_pos = x_0_cond.std(dim=dims, keepdim=True)
+                ro_cfg = x_0_cfg.std(dim=dims, keepdim=True).clamp(min=1e-8)
+                x_0_rescaled = x_0_cfg * (ro_pos / ro_cfg)
+                x_0_final = multiplier * x_0_rescaled + (1.0 - multiplier) * x_0_cfg
+                return x_orig - x_0_final
+
             cond = args["cond"]
             uncond = args["uncond"]
-            cond_scale = args["cond_scale"]
             sigma = args["sigma"]
             sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
-            x_orig = args["input"]
 
             #rescale cfg has to be done on v-pred model output
             x = x_orig / (sigma * sigma + 1.0)

comfy_extras/nodes_moge.py

@@ -2,6 +2,7 @@
 
 
 import torch
+import math
 
 import comfy.utils
 import folder_paths
@@ -391,10 +392,57 @@ class MoGePointMapToMesh(io.ComfyNode):
         return io.NodeOutput(mesh)
 
 
+class MoGeGeometryToFOV(io.ComfyNode):
+    """Extract horizontal/vertical FOV from MoGe intrinsics, e.g. fov_y to feed SAM3DBody_Predict."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MoGeGeometryToFOV",
+            search_aliases=["moge", "fov", "geometry", "intrinsics", "field of view"],
+            display_name="Get FoV from MoGe Geometry",
+            description="Derive the field of view and focal length from MoGe intrinsics.",
+            category="image/geometry estimation",
+            inputs=[
+                MoGeGeometry.Input("moge_geometry"),
+                io.Combo.Input("axis", options=["vertical", "horizontal", "diagonal"], default="vertical",
+                               tooltip="'vertical' (fov_y), 'horizontal' (fov_x), or 'diagonal'."),
+                io.Combo.Input("unit", options=["degrees", "radians"], default="degrees",
+                               tooltip="Output unit for the FOV."),
+            ],
+            outputs=[
+                io.Float.Output(display_name="fov"),
+                io.Float.Output(display_name="focal_pixels"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, moge_geometry, axis, unit) -> io.NodeOutput:
+        K = moge_geometry.get("intrinsics") if isinstance(moge_geometry, dict) else None
+        if K is None:
+            raise ValueError("moge_geometry has no intrinsics (panorama geometry has none).")
+        if K.ndim == 3:
+            K = K[0]
+        # MoGe normalizes fx by width and fy by height; with cx=cy=0.5 the half-extent
+        # in normalized units is 0.5, so fov = 2*atan(0.5 / f) per axis (hypot for diagonal).
+        hx = 0.5 / float(K[0, 0].item())
+        hy = 0.5 / float(K[1, 1].item())
+        half_tan = {"horizontal": hx, "vertical": hy, "diagonal": math.hypot(hx, hy)}[axis]
+        fov_radians = 2.0 * math.atan(half_tan)
+        fov = fov_radians if unit == "radians" else math.degrees(fov_radians)
+        # Pixels are square here, so fy*H == fx*W is the single lens focal in pixels.
+        src = next((moge_geometry[k] for k in ("image", "points", "depth") if k in moge_geometry), None)
+        if src is None:
+            raise ValueError("moge_geometry has no image/points/depth to read the pixel height from.")
+        H = int(src.shape[1])
+        focal_pixels = float(K[1, 1].item()) * H
+        return io.NodeOutput(fov, focal_pixels)
+
+
 class MoGeExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [LoadMoGeModel, MoGeInference, MoGePanoramaInference, MoGeRender, MoGePointMapToMesh]
+        return [LoadMoGeModel, MoGeInference, MoGePanoramaInference, MoGeRender, MoGePointMapToMesh, MoGeGeometryToFOV]
 
 
 async def comfy_entrypoint() -> MoGeExtension:

comfy_extras/nodes_save_3d.py

@@ -1,10 +1,13 @@
 """Save-side 3D nodes: mesh packing/slicing helpers + GLB writer + SaveGLB node."""
 
+import copy
 import json
 import logging
+import math
 import os
 import struct
 from io import BytesIO
+from typing import TypedDict
 
 import numpy as np
 from PIL import Image
@@ -13,14 +16,16 @@ from typing_extensions import override
 
 import folder_paths
 from comfy.cli_args import args
-from comfy_api.latest import ComfyExtension, IO, Types
+from comfy_api.latest import ComfyExtension, IO, Types, UI
+from server import PromptServer
 
 
-def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
+def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False,
-    # Pack lists of (Nᵢ, *) vertex/face/color/uv tensors into padded batched tensors,
+                             normals=None, metallic_roughness=None, tangents=None, normal_map=None,
-    # stashing per-item lengths as runtime attrs so consumers can recover the real slice.
+                             occlusion_in_mr=False, material=None, emissive=None):
-    # colors and uvs are 1:1 with vertices, so they're padded to max_vertices and read with vertex_counts.
+    # Pack per-item tensors into padded batches, stashing per-item lengths as runtime attrs.
-    # texture is (B, H, W, 3) — passed through unchanged
+    # colors/uvs/normals/tangents are 1:1 with vertices (padded to max_vertices); texture/
+    # metallic_roughness/normal_map are (B,H,W,*) image stacks passed through unchanged.
     batch_size = len(vertices)
     max_vertices = max(v.shape[0] for v in vertices)
     max_faces = max(f.shape[0] for f in faces)
@@ -52,43 +57,86 @@ def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=Non
             )
             packed_uvs[i, :u.shape[0]] = u
 
+    packed_normals = None
+    if normals is not None:
+        packed_normals = normals[0].new_zeros((batch_size, max_vertices, normals[0].shape[1]))
+        for i, nrm in enumerate(normals):
+            assert nrm.shape[0] == vertices[i].shape[0], (
+                f"normals[{i}] has {nrm.shape[0]} entries, expected {vertices[i].shape[0]} (1:1 with vertices)"
+            )
+            packed_normals[i, :nrm.shape[0]] = nrm
+
+    packed_tangents = None
+    if tangents is not None:
+        packed_tangents = tangents[0].new_zeros((batch_size, max_vertices, tangents[0].shape[1]))
+        for i, tn in enumerate(tangents):
+            assert tn.shape[0] == vertices[i].shape[0], (
+                f"tangents[{i}] has {tn.shape[0]} entries, expected {vertices[i].shape[0]} (1:1 with vertices)"
+            )
+            packed_tangents[i, :tn.shape[0]] = tn
+
     return Types.MESH(packed_vertices, packed_faces,
                       uvs=packed_uvs, vertex_colors=packed_colors, texture=texture,
-                      vertex_counts=vertex_counts, face_counts=face_counts, unlit=unlit)
+                      metallic_roughness=metallic_roughness,
+                      vertex_counts=vertex_counts, face_counts=face_counts, unlit=unlit,
+                      normals=packed_normals, tangents=packed_tangents,
+                      normal_map=normal_map, occlusion_in_mr=occlusion_in_mr,
+                      material=material, emissive=emissive)
 
 
 def get_mesh_batch_item(mesh, index):
     # Returns (vertices, faces, colors, uvs) for batch index, slicing to real lengths
     # if the mesh carries per-item counts (variable-size batch).
-    v_colors = getattr(mesh, "vertex_colors", None)
+    v_colors = mesh.vertex_colors
-    v_uvs = getattr(mesh, "uvs", None)
+    v_uvs = mesh.uvs
-    if getattr(mesh, "vertex_counts", None) is not None:
+    v_normals = mesh.normals
+    if 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 = v_colors[index, :vertex_count] if v_colors is not None else None
         uvs = v_uvs[index, :vertex_count] if v_uvs is not None else None
-        return vertices, faces, colors, uvs
+        normals = v_normals[index, :vertex_count] if v_normals is not None else None
+        return vertices, faces, colors, uvs, normals
 
     colors = v_colors[index] if v_colors is not None else None
     uvs = v_uvs[index] if v_uvs is not None else None
-    return mesh.vertices[index], mesh.faces[index], colors, uvs
+    normals = v_normals[index] if v_normals is not None else None
+    return mesh.vertices[index], mesh.faces[index], colors, uvs, normals
 
 
-def save_glb(vertices, faces, filepath, metadata=None,
+def save_glb(vertices, faces, filepath=None, metadata=None,
-             uvs=None, vertex_colors=None, texture_image=None, unlit=False):
+             uvs=None, vertex_colors=None, texture_image=None,
+             metallic_roughness_image=None, unlit=False,
+             normals=None, normal_map_image=None, tangents=None, occlusion_in_mr=False,
+             material=None, emissive_image=None):
     """
     Save PyTorch tensor vertices and faces as a GLB file without external dependencies.
 
     Parameters:
     vertices: torch.Tensor of shape (N, 3) - The vertex coordinates
     faces: torch.Tensor of shape (M, 3) - The face indices (triangle faces)
-    filepath: str - Output filepath (should end with .glb)
+    filepath: str - Output filepath (should end with .glb). None returns the GLB bytes instead of writing.
     metadata: dict - Optional asset.extras metadata
     uvs: torch.Tensor of shape (N, 2) - Optional per-vertex texture coordinates
     vertex_colors: torch.Tensor of shape (N, 3) or (N, 4) - Optional per-vertex colors in [0, 1]
     texture_image: PIL.Image - Optional baseColor texture, embedded as PNG
+    metallic_roughness_image: PIL.Image - Optional glTF metallicRoughness texture
+        (R unused, G=roughness, B=metallic), embedded as PNG
+    normals: torch.Tensor of shape (N, 3) - Optional per-vertex normals, written as the
+        glTF NORMAL attribute. When omitted, NO normals are written and viewers fall back
+        to flat (per-face) shading — use the MeshSmoothNormals node to generate them.
+    normal_map_image: PIL.Image - Optional tangent-space normal map (glTF/OpenGL +Y),
+        written as the material normalTexture. Needs TEXCOORD_0.
+    tangents: torch.Tensor of shape (N, 4) - Optional per-vertex tangents (xyz + handedness w),
+        written as the glTF TANGENT attribute. Without it viewers derive tangents in-shader.
+    occlusion_in_mr: bool - When True, R of metallic_roughness_image holds AO (ORM packing) and
+        occlusionTexture is pointed at that same image.
+    material: dict - Optional scalar overrides from SetMeshMaterial (base_color_factor,
+        metallic/roughness_factor with <0 = auto, emissive_factor/strength, normal_scale,
+        occlusion_strength, double_sided).
+    emissive_image: PIL.Image - Optional emissive (glow) texture, written as emissiveTexture.
     """
 
     # Convert tensors to numpy arrays
@@ -117,44 +165,82 @@ def save_glb(vertices, faces, filepath, metadata=None,
         raise ValueError(
             f"save_glb: vertex_colors has {colors_np.shape[0]} entries but vertex count is {n_verts}"
         )
+
+    normals_np = normals.cpu().numpy().astype(np.float32) if normals is not None else None
+    if normals_np is not None and normals_np.shape[0] != n_verts:
+        raise ValueError(
+            f"save_glb: normals has {normals_np.shape[0]} entries but vertex count is {n_verts}"
+        )
+    tangents_np = tangents.cpu().numpy().astype(np.float32) if tangents is not None else None
+    if tangents_np is not None and tangents_np.shape != (n_verts, 4):
+        raise ValueError(
+            f"save_glb: tangents must be (N, 4) with N={n_verts}, got {tuple(tangents_np.shape)}"
+        )
     faces_np = faces_signed.astype(np.uint32)
     texture_png_bytes = None
     if texture_image is not None:
         buf = BytesIO()
         texture_image.save(buf, format="PNG")
         texture_png_bytes = buf.getvalue()
+    mr_png_bytes = None
+    if metallic_roughness_image is not None:
+        buf = BytesIO()
+        metallic_roughness_image.save(buf, format="PNG")
+        mr_png_bytes = buf.getvalue()
+    nm_png_bytes = None
+    if normal_map_image is not None:
+        buf = BytesIO()
+        normal_map_image.save(buf, format="PNG")
+        nm_png_bytes = buf.getvalue()
+    em_png_bytes = None
+    if emissive_image is not None:
+        buf = BytesIO()
+        emissive_image.save(buf, format="PNG")
+        em_png_bytes = buf.getvalue()
 
     vertices_buffer = vertices_np.tobytes()
     indices_buffer = faces_np.tobytes()
     uvs_buffer = uvs_np.tobytes() if uvs_np is not None else b""
     colors_buffer = colors_np.tobytes() if colors_np is not None else b""
+    normals_buffer = normals_np.tobytes() if normals_np is not None else b""
+    tangents_buffer = tangents_np.tobytes() if tangents_np is not None else b""
     texture_buffer = texture_png_bytes if texture_png_bytes is not None else b""
+    mr_buffer = mr_png_bytes if mr_png_bytes is not None else b""
+    nm_buffer = nm_png_bytes if nm_png_bytes is not None else b""
+    em_buffer = em_png_bytes if em_png_bytes is not None else b""
 
     def pad_to_4_bytes(buffer):
         padding_length = (4 - (len(buffer) % 4)) % 4
         return buffer + b'\x00' * padding_length
 
-    vertices_buffer_padded = pad_to_4_bytes(vertices_buffer)
+    # Blob order in one place; offsets accumulated in a pass so adding a buffer is one entry.
-    indices_buffer_padded = pad_to_4_bytes(indices_buffer)
+    _blobs = [
-    uvs_buffer_padded = pad_to_4_bytes(uvs_buffer)
+        ("vertices", vertices_buffer), ("indices", indices_buffer), ("uvs", uvs_buffer),
-    colors_buffer_padded = pad_to_4_bytes(colors_buffer)
+        ("colors", colors_buffer), ("normals", normals_buffer), ("tangents", tangents_buffer),
-    texture_buffer_padded = pad_to_4_bytes(texture_buffer)
+        ("texture", texture_buffer), ("mr", mr_buffer), ("nm", nm_buffer), ("em", em_buffer),
-
+    ]
-    buffer_data = b"".join([
+    byte_offset = {}
-        vertices_buffer_padded,
+    acc = 0
-        indices_buffer_padded,
+    parts = []
-        uvs_buffer_padded,
+    for name, b in _blobs:
-        colors_buffer_padded,
+        padded = pad_to_4_bytes(b)
-        texture_buffer_padded,
+        byte_offset[name] = acc
-    ])
+        acc += len(padded)
+        parts.append(padded)
+    buffer_data = b"".join(parts)
 
     vertices_byte_length = len(vertices_buffer)
-    vertices_byte_offset = 0
     indices_byte_length = len(indices_buffer)
-    indices_byte_offset = len(vertices_buffer_padded)
+    vertices_byte_offset = byte_offset["vertices"]
-    uvs_byte_offset = indices_byte_offset + len(indices_buffer_padded)
+    indices_byte_offset = byte_offset["indices"]
-    colors_byte_offset = uvs_byte_offset + len(uvs_buffer_padded)
+    uvs_byte_offset = byte_offset["uvs"]
-    texture_byte_offset = colors_byte_offset + len(colors_buffer_padded)
+    colors_byte_offset = byte_offset["colors"]
+    normals_byte_offset = byte_offset["normals"]
+    tangents_byte_offset = byte_offset["tangents"]
+    texture_byte_offset = byte_offset["texture"]
+    mr_byte_offset = byte_offset["mr"]
+    nm_byte_offset = byte_offset["nm"]
+    em_byte_offset = byte_offset["em"]
 
     buffer_views = [
         {
@@ -224,6 +310,40 @@ def save_glb(vertices, faces, filepath, metadata=None,
         })
         primitive_attributes["COLOR_0"] = accessor_idx
 
+    if normals_np is not None and len(normals_np) > 0:
+        buffer_views.append({
+            "buffer": 0,
+            "byteOffset": normals_byte_offset,
+            "byteLength": len(normals_buffer),
+            "target": 34962
+        })
+        accessor_idx = len(accessors)
+        accessors.append({
+            "bufferView": len(buffer_views) - 1,
+            "byteOffset": 0,
+            "componentType": 5126,  # FLOAT
+            "count": len(normals_np),
+            "type": "VEC3",
+        })
+        primitive_attributes["NORMAL"] = accessor_idx
+
+    if tangents_np is not None and len(tangents_np) > 0:
+        buffer_views.append({
+            "buffer": 0,
+            "byteOffset": tangents_byte_offset,
+            "byteLength": len(tangents_buffer),
+            "target": 34962
+        })
+        accessor_idx = len(accessors)
+        accessors.append({
+            "bufferView": len(buffer_views) - 1,
+            "byteOffset": 0,
+            "componentType": 5126,  # FLOAT
+            "count": len(tangents_np),
+            "type": "VEC4",  # xyz tangent + w handedness (glTF TANGENT)
+        })
+        primitive_attributes["TANGENT"] = accessor_idx
+
     primitive = {
         "attributes": primitive_attributes,
         "indices": 1,
@@ -235,9 +355,24 @@ def save_glb(vertices, faces, filepath, metadata=None,
     samplers = []
     materials = []
     extensions_used = []
+
+    def add_image_texture(png_byte_offset, png_byte_length):
+        """Append an embedded PNG image + a texture referencing it; return the texture index."""
+        buffer_views.append({"buffer": 0, "byteOffset": png_byte_offset, "byteLength": png_byte_length})
+        images.append({"bufferView": len(buffer_views) - 1, "mimeType": "image/png"})
+        if not samplers:
+            samplers.append({"magFilter": 9729, "minFilter": 9729, "wrapS": 33071, "wrapT": 33071})
+        textures.append({"source": len(images) - 1, "sampler": 0})
+        return len(textures) - 1
+
+    has_uv = "TEXCOORD_0" in primitive_attributes
     if unlit and texture_png_bytes is None:
         # Flat, light-independent shading (KHR_materials_unlit): COLOR_0 is shown as-is, matching how a
         # gaussian splat renders (emissive). Without this the viewer lights the mesh and washes the colours.
+        if nm_png_bytes is not None or em_png_bytes is not None or occlusion_in_mr or material is not None:
+            logging.warning(
+                "save_glb: unlit material ignores normal/occlusion/emissive maps and SetMeshMaterial "
+                "overrides — those are PBR-lit features. Disable unlit to export them.")
         materials.append({
             "pbrMetallicRoughness": {"baseColorFactor": [1.0, 1.0, 1.0, 1.0], "metallicFactor": 0.0, "roughnessFactor": 1.0},
             "extensions": {"KHR_materials_unlit": {}},
@@ -245,23 +380,67 @@ def save_glb(vertices, faces, filepath, metadata=None,
         })
         extensions_used.append("KHR_materials_unlit")
         primitive["material"] = 0
-    if texture_png_bytes is not None and "TEXCOORD_0" in primitive_attributes:
+    else:
-        buffer_views.append({
+        pbr = {
-            "buffer": 0,
+            "metallicFactor": 0.0,
-            "byteOffset": texture_byte_offset,
+            "roughnessFactor": 0.5,
-            "byteLength": len(texture_buffer),
+            "baseColorFactor": [0.22, 0.22, 0.22, 1.0],   # neutral-gray fallback for bare geometry only
-        })
+        }
-        images.append({"bufferView": len(buffer_views) - 1, "mimeType": "image/png"})
+        if texture_png_bytes is not None and has_uv:
-        samplers.append({"magFilter": 9729, "minFilter": 9729, "wrapS": 33071, "wrapT": 33071})
+            pbr["baseColorTexture"] = {"index": add_image_texture(texture_byte_offset, len(texture_buffer)), "texCoord": 0}
-        textures.append({"source": 0, "sampler": 0})
+
-        materials.append({
+        if (texture_png_bytes is not None and has_uv) or "COLOR_0" in primitive_attributes:
-            "pbrMetallicRoughness": {
+            pbr["baseColorFactor"] = [1.0, 1.0, 1.0, 1.0]
-                "baseColorTexture": {"index": 0, "texCoord": 0},
+            pbr["roughnessFactor"] = 1.0
-                "metallicFactor": 0.0,
+
-                "roughnessFactor": 1.0,
+        if mr_png_bytes is not None and has_uv:
-            },
+            mr_texture_index = add_image_texture(mr_byte_offset, len(mr_buffer))
-            "doubleSided": True,
+            pbr["metallicRoughnessTexture"] = {"index": mr_texture_index, "texCoord": 0}
-        })
+            # When a metallicRoughness texture is present, the factors scale it; use 1.0
+            # so the texture values pass through unchanged (glTF convention).
+            pbr["metallicFactor"] = 1.0
+            pbr["roughnessFactor"] = 1.0
+
+        mat = material if isinstance(material, dict) else {}
+        # Scalar overrides from SetMeshMaterial (factor < 0 means "leave auto").
+        if mat.get("base_color_factor") is not None:
+            pbr["baseColorFactor"] = [float(x) for x in mat["base_color_factor"]]
+        if mat.get("metallic_factor", -1.0) >= 0.0:
+            pbr["metallicFactor"] = float(mat["metallic_factor"])
+        if mat.get("roughness_factor", -1.0) >= 0.0:
+            pbr["roughnessFactor"] = float(mat["roughness_factor"])
+
+        material = {
+            "pbrMetallicRoughness": pbr,
+            "doubleSided": bool(mat.get("double_sided", True)),
+        }
+        if occlusion_in_mr and mr_png_bytes is not None and has_uv:
+            # ORM packing: occlusionTexture reuses the MR image (glTF reads its R channel).
+            material["occlusionTexture"] = {"index": mr_texture_index, "texCoord": 0,
+                                            "strength": float(mat.get("occlusion_strength", 1.0))}
+        if nm_png_bytes is not None and has_uv:
+            material["normalTexture"] = {"index": add_image_texture(nm_byte_offset, len(nm_buffer)),
+                                         "texCoord": 0, "scale": float(mat.get("normal_scale", 1.0))}
+
+        emissive_factor = [float(x) for x in mat.get("emissive_factor", [0.0, 0.0, 0.0])]
+        emissive_strength = float(mat.get("emissive_strength", 1.0))
+        has_em_tex = em_png_bytes is not None and has_uv
+        if any(c > 0.0 for c in emissive_factor) or has_em_tex:
+            # glTF multiplies emissiveFactor × texture, so a texture with no color would go black;
+            # default the factor to white in that case.
+            if has_em_tex and not any(c > 0.0 for c in emissive_factor):
+                emissive_factor = [1.0, 1.0, 1.0]
+            material["emissiveFactor"] = [min(1.0, c) for c in emissive_factor]
+            if has_em_tex:
+                material["emissiveTexture"] = {"index": add_image_texture(em_byte_offset, len(em_buffer)),
+                                               "texCoord": 0}
+            if emissive_strength != 1.0:
+                material.setdefault("extensions", {})["KHR_materials_emissive_strength"] = {
+                    "emissiveStrength": emissive_strength}
+                if "KHR_materials_emissive_strength" not in extensions_used:
+                    extensions_used.append("KHR_materials_emissive_strength")
+
+        materials.append(material)
         primitive["material"] = 0
 
     gltf = {
@@ -306,17 +485,48 @@ def save_glb(vertices, faces, filepath, metadata=None,
     # Create BIN chunk header (chunk type 1)
     bin_chunk_header = struct.pack('<II', len(buffer_data), 0x004E4942)  # "BIN\0" in little endian
 
-    # Write the GLB file
+    glb = b"".join([glb_header, json_chunk_header, gltf_json_padded, bin_chunk_header, buffer_data])
+    if filepath is None:
+        return glb                       # in-memory GLB bytes (e.g. for a File3D object)
     with open(filepath, 'wb') as f:
-        f.write(glb_header)
+        f.write(glb)
-        f.write(json_chunk_header)
-        f.write(gltf_json_padded)
-        f.write(bin_chunk_header)
-        f.write(buffer_data)
-
     return filepath
 
 
+def mesh_item_to_glb_bytes(mesh, index, metadata=None):
+    """Serialize one batch item of a MESH to in-memory GLB bytes, carrying every PBR attribute
+    (uvs, colors, normals, texture, ORM/occlusion, normal map + tangents, emissive, material).
+    Returns None for an empty item. Shared by SaveGLB (per item) and MeshToFile3D."""
+    vertices_i, faces_i, v_colors, uvs_i, normals_i = get_mesh_batch_item(mesh, index)
+    if vertices_i.shape[0] == 0 or faces_i.shape[0] == 0:
+        return None
+
+    def _img(attr):
+        t = getattr(mesh, attr, None)
+        if t is None:
+            return None
+        a = (t[index].clamp(0.0, 1.0).cpu().numpy() * 255).astype(np.uint8)
+        assert a.ndim == 3 and a.shape[-1] == 3, f"{attr} must be (B, H, W, 3), got {tuple(t.shape)}"
+        return Image.fromarray(a, mode="RGB")
+
+    tangents_b = mesh.tangents
+    tangents_i = tangents_b[index, :vertices_i.shape[0]] if tangents_b is not None else None
+    return save_glb(
+        vertices_i, faces_i, None, metadata,
+        uvs=uvs_i,
+        vertex_colors=v_colors,
+        texture_image=_img("texture"),
+        metallic_roughness_image=_img("metallic_roughness"),
+        unlit=mesh.unlit,
+        normals=normals_i,
+        normal_map_image=_img("normal_map"),
+        tangents=tangents_i,
+        occlusion_in_mr=mesh.occlusion_in_mr,
+        material=mesh.material,
+        emissive_image=_img("emissive"),
+    )
+
+
 class SaveGLB(IO.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -378,25 +588,14 @@ class SaveGLB(IO.ComfyNode):
             counter += 1
         else:
             # Handle Mesh input - save vertices and faces as GLB; carry optional UVs / colors / texture.
-            texture_b = getattr(mesh, "texture", None)
-            texture_np = None
-            if texture_b is not None:
-                texture_np = (texture_b.clamp(0.0, 1.0).cpu().numpy() * 255).astype(np.uint8)
-                assert texture_np.ndim == 4 and texture_np.shape[-1] == 3, (
-                    f"texture must be (B, H, W, 3) RGB, got shape {tuple(texture_np.shape)}"
-                )
             for i in range(mesh.vertices.shape[0]):
-                vertices_i, faces_i, v_colors, uvs_i = get_mesh_batch_item(mesh, i)
+                glb = mesh_item_to_glb_bytes(mesh, i, metadata)
-                if vertices_i.shape[0] == 0 or faces_i.shape[0] == 0:
+                if glb is None:
                     logging.warning(f"SaveGLB: skipping empty mesh at batch index {i}")
                     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,
+                with open(os.path.join(full_output_folder, f), "wb") as fh:
-                         uvs=uvs_i,
+                    fh.write(glb)
-                         vertex_colors=v_colors,
-                         texture_image=tex_img,
-                         unlit=getattr(mesh, "unlit", False))
                 results.append({
                     "filename": f,
                     "subfolder": subfolder,
@@ -406,10 +605,273 @@ class SaveGLB(IO.ComfyNode):
         return IO.NodeOutput(ui={"3d": results})
 
 
+class MeshToFile3D(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="MeshToFile3D",
+            display_name="Create 3D File (from Mesh)",
+            search_aliases=["mesh to glb", "mesh to file", "export mesh"],
+            category="3d",
+            description="Serialize a mesh to a GLB File3D object for Save / Preview 3D nodes, "
+                        "carrying its UVs, colors, normals, texture, normal/occlusion/emissive "
+                        "maps and material. Supports one item per batch only.",
+            inputs=[IO.Mesh.Input("mesh")],
+            outputs=[IO.File3DGLB.Output(display_name="model_3d")],
+        )
+
+    @classmethod
+    def execute(cls, mesh) -> IO.NodeOutput:
+        if mesh.vertices.shape[0] > 1:
+            logging.warning("MeshToFile3D supports one item per batch only. Got %d; using first.",
+                            mesh.vertices.shape[0])
+        glb = mesh_item_to_glb_bytes(mesh, 0)
+        if glb is None:
+            raise ValueError("MeshToFile3D: mesh is empty (no vertices/faces).")
+        return IO.NodeOutput(Types.File3D(BytesIO(glb), file_format="glb"))
+
+
+class RotateMesh(IO.ComfyNode):
+    class ModeValues(TypedDict, total=False):
+        mode: str
+        angle_x: float
+        angle_y: float
+        angle_z: float
+        qw: float
+        qx: float
+        qy: float
+        qz: float
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="RotateMesh",
+            display_name="Rotate Mesh",
+            category="3d/mesh",
+            description=(
+                "Rotate a mesh. Euler XYZ applies X then Y then Z about the world axes (degrees). "
+                "Quaternion is (w, x, y, z), auto-normalized."
+            ),
+            inputs=[
+                IO.Mesh.Input("mesh"),
+                IO.DynamicCombo.Input(
+                    "mode",
+                    options=[
+                        IO.DynamicCombo.Option("euler_xyz", [
+                            IO.Float.Input("angle_x", default=0.0, min=-360.0, max=360.0, step=0.1,
+                                           tooltip="Rotation around the X axis in degrees."),
+                            IO.Float.Input("angle_y", default=0.0, min=-360.0, max=360.0, step=0.1,
+                                           tooltip="Rotation around the Y axis in degrees."),
+                            IO.Float.Input("angle_z", default=0.0, min=-360.0, max=360.0, step=0.1,
+                                           tooltip="Rotation around the Z axis in degrees."),
+                        ]),
+                        IO.DynamicCombo.Option("quaternion", [
+                            IO.Float.Input("qw", default=1.0, min=-1.0, max=1.0, step=0.001),
+                            IO.Float.Input("qx", default=0.0, min=-1.0, max=1.0, step=0.001),
+                            IO.Float.Input("qy", default=0.0, min=-1.0, max=1.0, step=0.001),
+                            IO.Float.Input("qz", default=0.0, min=-1.0, max=1.0, step=0.001),
+                        ]),
+                    ],
+                ),
+            ],
+            outputs=[IO.Mesh.Output("mesh")],
+        )
+
+    @classmethod
+    def execute(cls, mesh: Types.MESH, mode: ModeValues) -> IO.NodeOutput:
+        mode_name = mode["mode"]
+        if mode_name == "euler_xyz":
+            ax = math.radians(mode["angle_x"])
+            ay = math.radians(mode["angle_y"])
+            az = math.radians(mode["angle_z"])
+            if ax == 0.0 and ay == 0.0 and az == 0.0:
+                return IO.NodeOutput(mesh)
+            cx, sx = math.cos(ax), math.sin(ax)
+            cy, sy = math.cos(ay), math.sin(ay)
+            cz, sz = math.cos(az), math.sin(az)
+            R_rows = [
+                [cy * cz, sx * sy * cz - cx * sz, cx * sy * cz + sx * sz],
+                [cy * sz, sx * sy * sz + cx * cz, cx * sy * sz - sx * cz],
+                [-sy,     sx * cy,                cx * cy],
+            ]
+        elif mode_name == "quaternion":
+            qw, qx, qy, qz = mode["qw"], mode["qx"], mode["qy"], mode["qz"]
+            n = math.sqrt(qw * qw + qx * qx + qy * qy + qz * qz)
+            if n < 1e-8:
+                raise ValueError("RotateMesh: quaternion has zero magnitude")
+            qw, qx, qy, qz = qw / n, qx / n, qy / n, qz / n
+            if qw == 1.0 and qx == 0.0 and qy == 0.0 and qz == 0.0:
+                return IO.NodeOutput(mesh)
+            R_rows = [
+                [1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw),     2 * (qx * qz + qy * qw)],
+                [2 * (qx * qy + qz * qw),     1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
+                [2 * (qx * qz - qy * qw),     2 * (qy * qz + qx * qw),     1 - 2 * (qx * qx + qy * qy)],
+            ]
+        else:
+            raise ValueError(f"RotateMesh: unknown mode {mode_name!r}")
+
+        def rotate(v: torch.Tensor) -> torch.Tensor:
+            R = torch.tensor(R_rows, device=v.device, dtype=v.dtype)
+            return v @ R.T
+
+        out = copy.copy(mesh)
+        if isinstance(mesh.vertices, list):
+            out.vertices = [rotate(v) for v in mesh.vertices]
+        else:
+            out.vertices = rotate(mesh.vertices)
+        # Normals are directions; rotate them too (R is orthogonal) so they stay valid.
+        nrm = mesh.normals
+        if nrm is not None:
+            out.normals = [rotate(n) for n in nrm] if isinstance(nrm, list) else rotate(nrm)
+        return IO.NodeOutput(out)
+
+
+class MergeMeshes(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        autogrow_template = IO.Autogrow.TemplatePrefix(
+            IO.Mesh.Input("mesh"), prefix="mesh", min=2, max=50,
+        )
+        return IO.Schema(
+            node_id="MergeMeshes",
+            display_name="Merge Meshes",
+            category="3d/mesh",
+            description=(
+                "Concatenate N meshes into one by offsetting face indices and stacking verts, "
+                "faces, uvs, and colors."
+            ),
+            inputs=[
+                IO.Autogrow.Input("meshes", template=autogrow_template),
+            ],
+            outputs=[IO.Mesh.Output("mesh")],
+        )
+
+    @classmethod
+    def execute(cls, meshes: IO.Autogrow.Type) -> IO.NodeOutput:
+        # Concatenate the input meshes into one (B=1) mesh: cumulative face-index offset,
+        # missing uvs/colors padded (zeros/white), texture from the first input that has one
+        # (later dropped — a single-primitive glb can't carry multiple atlases).
+        meshes = list(meshes.values())
+        if not meshes:
+            raise ValueError("MergeMeshes: need at least one mesh")
+
+        def _b0(t):
+            return t[0] if t.ndim == 3 else t
+
+        any_uvs = any(m.uvs is not None for m in meshes)
+        any_colors = any(m.vertex_colors is not None for m in meshes)
+
+        verts_list, faces_list, uvs_list, colors_list = [], [], [], []
+        texture = None
+        offset = 0
+        for m in meshes:
+            # Coerce to CPU so CUDA-side (MoGe) meshes merge cleanly with our outputs.
+            v = _b0(m.vertices).cpu()
+            f = _b0(m.faces).cpu()
+            verts_list.append(v)
+            faces_list.append(f + offset)
+            offset += v.shape[0]
+            if any_uvs:
+                mu = m.uvs
+                uvs_list.append(_b0(mu).cpu() if mu is not None else v.new_zeros((v.shape[0], 2)))
+            if any_colors:
+                mc = m.vertex_colors
+                c = _b0(mc).cpu() if mc is not None else v.new_ones((v.shape[0], 3))
+                colors_list.append(c)
+            mt = m.texture
+            if mt is not None:
+                if texture is None:
+                    texture = mt.cpu()
+                else:
+                    logging.warning("MergeMeshes: dropping extra texture from input; only one texture is kept.")
+
+        merged_verts = torch.cat(verts_list, dim=0).unsqueeze(0)
+        merged_faces = torch.cat(faces_list, dim=0).unsqueeze(0)
+        merged_uvs = torch.cat(uvs_list, dim=0).unsqueeze(0) if any_uvs else None
+        merged_colors = torch.cat(colors_list, dim=0).unsqueeze(0) if any_colors else None
+
+        return IO.NodeOutput(Types.MESH(
+            vertices=merged_verts,
+            faces=merged_faces,
+            uvs=merged_uvs,
+            vertex_colors=merged_colors,
+            texture=texture,
+        ))
+
+
+class GetMeshInfo(IO.ComfyNode):
+    """Report vertex / face counts and attributes for a MESH, displayed on the
+    node (and as a string output). Counts are comma-formatted since meshes can
+    run into the millions of faces. Passes the mesh through unchanged."""
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="GetMeshInfo",
+            display_name="Get Mesh Info",
+            category="3d/mesh",
+            inputs=[IO.Mesh.Input("mesh")],
+            outputs=[
+                IO.Mesh.Output(display_name="mesh"),
+                IO.String.Output(display_name="info"),
+            ],
+            hidden=[IO.Hidden.unique_id],
+        )
+
+    @staticmethod
+    def _fmt(n: int) -> str:
+        # e.g. 1234567 -> "1,234,567 (1.23M)"; small numbers stay plain.
+        s = f"{n:,}"
+        if n >= 1_000_000:
+            s += f" ({n / 1_000_000:.2f}M)"
+        elif n >= 10_000:
+            s += f" ({n / 1_000:.1f}K)"
+        return s
+
+    @classmethod
+    def execute(cls, mesh):
+        B = mesh.vertices.shape[0]
+        # Honour per-item counts when the batch is zero-padded; else use the row sizes.
+        if mesh.vertex_counts is not None:
+            v_counts = [int(x) for x in mesh.vertex_counts.tolist()]
+            f_counts = [int(x) for x in mesh.face_counts.tolist()]
+        else:
+            v_counts = [int(mesh.vertices.shape[1])] * B
+            f_counts = [int(mesh.faces.shape[1])] * B
+
+        attrs = []
+        for name in ("uvs", "vertex_colors", "normals", "tangents", "texture", "metallic_roughness", "normal_map"):
+            t = getattr(mesh, name, None)
+            if t is not None:
+                if name in ("texture", "metallic_roughness", "normal_map"):
+                    attrs.append(f"{name} {int(t.shape[-3])}×{int(t.shape[-2])}")  # H×W
+                else:
+                    attrs.append(name)
+
+        lines = []
+        if B > 1:
+            lines.append(f"Batch:      {B} meshes")
+            lines.append(f"Vertices:   {cls._fmt(sum(v_counts))} total")
+            lines.append(f"Faces:      {cls._fmt(sum(f_counts))} total")
+            for i in range(B):
+                lines.append(f"  [{i}]  {v_counts[i]:>10,} verts  ·  {f_counts[i]:>10,} faces")
+        else:
+            lines.append(f"Vertices:   {cls._fmt(v_counts[0])}")
+            lines.append(f"Faces:      {cls._fmt(f_counts[0])}")
+        lines.append(f"Attributes: {', '.join(attrs) if attrs else 'none'}")
+
+        info = "\n".join(lines)
+        logging.info("[GetMeshInfo]\n%s", info)
+
+        if cls.hidden.unique_id:
+            PromptServer.instance.send_progress_text(info, cls.hidden.unique_id)
+        return IO.NodeOutput(mesh, info, ui=UI.PreviewText(info))
+
+
 class Save3DExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
-        return [SaveGLB]
+        return [SaveGLB, MeshToFile3D, RotateMesh, MergeMeshes, GetMeshInfo]
 
 
 async def comfy_entrypoint() -> Save3DExtension:

comfy_extras/nodes_trellis2.py

@@ -0,0 +1,947 @@
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, IO, Types, io
+from comfy.ldm.trellis2.vae import SparseTensor
+from comfy.ldm.trellis2.model import build_proj_transform_matrix, compute_stage_proj_feats
+
+from comfy_extras.nodes_mesh_postprocess import pack_variable_mesh_batch
+import comfy.latent_formats
+import comfy.model_management
+import comfy.utils
+import logging
+import math
+import torch
+
+ShapeSubdivides = io.Custom("SHAPE_SUBDIVIDES")
+
+
+shape_slat_format = comfy.latent_formats.Trellis2ShapeSLAT()
+tex_slat_format = comfy.latent_formats.Trellis2TexSLAT()
+
+def shape_norm(shape_latent, coords):
+    feats = shape_slat_format.process_out(shape_latent)
+    return SparseTensor(feats=feats, coords=coords)
+
+
+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):
+        # Mesh grid_size must match the actual coord resolution the upstream
+        # stage was run at (1024 cascade -> 64, 1536 cascade -> 96). The VAE's
+        # built-in `.resolution` buffer defaults to 1024 and is otherwise stale;
+        # take coord_resolution from the latent dict if the stage node set it.
+        coord_resolution = samples.get("coord_resolution")
+        if coord_resolution is not None:
+            resolution = int(coord_resolution) * 16
+        else:
+            resolution = int(vae.first_stage_model.resolution.item())
+        model_frame = samples.get("model_frame", "y_up")
+        sample_tensor = samples["samples"]
+        device = comfy.model_management.get_torch_device()
+        coords = samples["coords"]
+        vae.prepare_decode(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.to(vae.vae_dtype), 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.to(vae.vae_dtype), 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))
+
+        # Rotate Z-up (Trellis2 training frame) vertices to glTF Y-up. Pixal3D outputs are already Y-up.
+        if model_frame == "z_up":
+            vert_list = [torch.stack([v[..., 0], v[..., 2], -v[..., 1]], dim=-1).float().cpu()
+                         for v, _ in mesh]
+        else:
+            vert_list = [v.float().cpu() for v, _ in mesh]
+        face_list = [f.int().cpu() for _, f 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"]
+        vae.prepare_decode(sample_tensor.shape)
+        trellis_vae = vae.first_stage_model
+        coord_counts = samples.get("coord_counts")
+        model_frame = samples.get("model_frame", "y_up")
+        coord_resolution = samples.get("coord_resolution")
+
+        samples = samples["samples"]
+        samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
+        samples = samples.to(device)
+        feats = tex_slat_format.process_out(samples)
+        samples = SparseTensor(feats=feats, coords=coords.to(device))
+
+        voxel = trellis_vae.decode_tex_slat(samples.to(vae.vae_dtype), shape_subdivides)
+        # Keep all decoded channels. The texture VAE emits 6: base_color (0:3),
+        # metallic (3), roughness (4), alpha (5) — all in [0, 1]. Vertex-color
+        # consumers (PaintMesh) slice [:3]
+        color_feats = voxel.feats
+        voxel_coords = voxel.coords
+
+        if coord_resolution is not None:
+            tex_resolution = int(coord_resolution) * 16
+        elif voxel_coords.numel() > 0 and voxel_coords.shape[-1] >= 3:
+            spatial = voxel_coords[:, -3:] if voxel_coords.shape[-1] == 4 else voxel_coords
+            max_idx = int(spatial.max().item()) + 1
+            tex_resolution = next((r for r in (256, 512, 1024, 1536, 2048) if r >= max_idx), max_idx)
+        else:
+            tex_resolution = 1024
+
+        # Remap Z-up voxel coords to Y-up: (x, y, z) -> (x, z, R-1-y), matching the
+        # R_x(-90°) applied to mesh vertices in VaeDecodeShapeTrellis. Keeps PaintMesh's
+        # NN lookup correctly aligned without it needing to know the source frame.
+        if model_frame == "z_up" and voxel_coords.numel() > 0 and voxel_coords.shape[-1] >= 3:
+            R = tex_resolution
+            if voxel_coords.shape[-1] == 4:
+                batch_col = voxel_coords[:, :1]
+                spatial = voxel_coords[:, 1:]
+                spatial_yup = torch.stack(
+                    [spatial[:, 0], spatial[:, 2], (R - 1) - spatial[:, 1]], dim=-1
+                )
+                voxel_coords = torch.cat([batch_col, spatial_yup], dim=-1)
+            else:
+                voxel_coords = torch.stack(
+                    [voxel_coords[:, 0], voxel_coords[:, 2], (R - 1) - voxel_coords[:, 1]],
+                    dim=-1,
+                )
+
+        voxel = Types.VOXEL(voxel_coords, color_feats, tex_resolution)
+        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 = vae.prepare_decode(sample_tensor.shape)
+        shape_vae = vae.first_stage_model
+        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(shape_vae.decode_structure(sample_chunk.to(vae.vae_dtype)) > 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
+        voxel_data = decoded.squeeze(1).float()
+        return IO.NodeOutput(Types.VOXEL(voxel_data))
+
+class Trellis2UpsampleStage(IO.ComfyNode):
+    """Cascade-upsamples a 512-resolution shape latent into high-resolution
+    sparse coords and sets up the second shape-stage sampling pass at the
+    target resolution, attaching per-stage metadata to the conditioning for
+    the model to consume via extra_conds. target_resolution is reduced in
+    128-step decrements until the unique upsampled coord count fits under
+    max_tokens (floor 1024)."""
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2UpsampleStage",
+            category="model/conditioning/trellis2",
+            display_name="Trellis2 Upsample Stage",
+            inputs=[
+                IO.Conditioning.Input("positive"),
+                IO.Conditioning.Input("negative"),
+                IO.Latent.Input("shape_latent", tooltip="The 512-resolution shape latent output from the first shape-stage KSampler."),
+                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.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+                IO.Latent.Output(),
+            ]
+        )
+
+    @staticmethod
+    def _quantize_unique(hr_coords: torch.Tensor, lr_resolution: int, hr_resolution: int, pixal3d_mode: bool = False) -> torch.Tensor:
+        # Trellis2 uses `floor((c+0.5) * grid_res / lr_res)
+        # Pixal3D uses `round((c+0.5) * (grid_res-1) / lr_res)`
+        # this is a half-cell spatial shift. Branch so each upstream is matched bit-for-bit.
+        grid_res = hr_resolution // 16
+        spatial = hr_coords[:, 1:].float()
+        if pixal3d_mode:
+            spatial.add_(0.5).mul_((grid_res - 1) / lr_resolution).round_()
+        else:
+            spatial.add_(0.5).mul_(grid_res / lr_resolution)
+        quant = torch.cat([hr_coords[:, :1], spatial.int()], dim=1)
+        return quant.unique(dim=0)
+
+    @classmethod
+    def execute(cls, positive, negative, shape_latent, vae, target_resolution, max_tokens):
+        device = comfy.model_management.get_torch_device()
+        vae.prepare_decode(shape_latent["samples"].shape)
+
+        coord_counts = shape_latent.get("coord_counts")
+        shape_vae = vae.first_stage_model
+        lr_resolution = 512
+        target_resolution = int(target_resolution)
+        proj_pack = _proj_pack_from_conditioning(positive)
+        pixal3d_mode = proj_pack is not None
+
+        # Decode each sample's HR coords, then search for the largest hr_resolution
+        # that fits under max_tokens across all samples.
+        if coord_counts is None:
+            feats, coords_512 = flatten_batched_sparse_latent(
+                shape_latent["samples"], shape_latent["coords"], coord_counts,
+            )
+            slat = shape_norm(feats.to(device), coords_512.to(device))
+            sample_hr_coords = [shape_vae.upsample_shape(slat.to(vae.vae_dtype), upsample_times=4)]
+        else:
+            items = split_batched_sparse_latent(
+                shape_latent["samples"], shape_latent["coords"], coord_counts,
+            )
+            sample_hr_coords = []
+            for feats_i, coords_i in items:
+                coords_i = coords_i.to(device).clone()
+                coords_i[:, 0] = 0
+                slat_i = shape_norm(feats_i.to(device), coords_i)
+                sample_hr_coords.append(shape_vae.upsample_shape(slat_i.to(vae.vae_dtype), upsample_times=4))
+
+        # Resolution search — cache the final iteration's quantized unique tensors
+        # so we don't recompute .unique() per sample after picking hr_resolution.
+        hr_resolution = target_resolution
+        quant_unique_list = []
+        while True:
+            quant_unique_list = []
+            exceeds_limit = False
+            for hr_coords_i in sample_hr_coords:
+                qu = cls._quantize_unique(hr_coords_i, lr_resolution, hr_resolution, pixal3d_mode)
+                quant_unique_list.append(qu)
+                if qu.shape[0] >= max_tokens:
+                    exceeds_limit = True
+                    break
+            if not exceeds_limit:
+                break
+            if hr_resolution <= 1024:
+                for k in range(len(quant_unique_list), len(sample_hr_coords)):
+                    quant_unique_list.append(
+                        cls._quantize_unique(sample_hr_coords[k], lr_resolution, hr_resolution, pixal3d_mode)
+                    )
+                break
+            hr_resolution -= 128
+
+        # Rewrite batch column to match per-sample offset and concat.
+        per_sample_counts = []
+        for sample_offset, qu in enumerate(quant_unique_list):
+            qu[:, 0] = sample_offset
+            per_sample_counts.append(int(qu.shape[0]))
+        coords = torch.cat(quant_unique_list, dim=0)
+        counts = torch.tensor(per_sample_counts, dtype=torch.int64)
+        coord_resolution = hr_resolution // 16
+
+        batch_size, _, max_tokens_out = infer_batched_coord_layout(coords)
+        latent = torch.zeros(batch_size, 32, max_tokens_out, 1)
+
+        extras = {
+            "trellis2_generation_mode": "shape_generation",
+            "trellis2_coords": coords,
+            "trellis2_coord_counts": counts,
+        }
+        if proj_pack is not None:
+            extras["trellis2_proj_feats"] = compute_stage_proj_feats(
+                proj_pack, "shape_1024", coords=coords, coord_resolution=coord_resolution,
+            )
+        positive_out = _conditioning_set_extras(positive, extras)
+        negative_out = _conditioning_set_extras(negative, extras)
+        out_latent = {"samples": latent, "coords": coords, "coord_counts": counts,
+                      "coord_resolution": coord_resolution, "type": "trellis2",
+                      "model_frame": shape_latent.get("model_frame",
+                                                       "y_up" if proj_pack is not None else "z_up")}
+        return IO.NodeOutput(positive_out, negative_out, out_latent)
+
+def _dinov3_encode(model, image_bchw, image_size, want_patches=False):
+    """Run DINOv3 once at the requested resolution.
+
+    image_bchw: [B, 3, H, W] float in [0, 1] (any source resolution; resized here).
+    Returns the full sequence tensor (Trellis2 path) or a dict with the global
+    tokens split out + a 2D patch grid (Pixal3D path) when `want_patches=True`.
+    """
+    model_internal = model.model
+    device = comfy.model_management.get_torch_device()
+    img_t = comfy.utils.common_upscale(image_bchw, image_size, image_size, "lanczos", "disabled").to(device)
+    mean = torch.tensor(model.image_mean or [0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)
+    std = torch.tensor(model.image_std or [0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)
+    img_t = (img_t - mean) / std
+    tokens = model_internal(img_t, skip_norm_elementwise=True)[0]
+    if not want_patches:
+        return tokens
+    h_p = w_p = image_size // 16
+    n_reg = tokens.shape[1] - 1 - h_p * w_p
+    return {"tokens": tokens[:, :1 + n_reg], "patches_2d": _dinov3_patches_to_2d(tokens, image_size)}
+
+
+class Trellis2Conditioning(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2Conditioning",
+            category="model/conditioning/trellis2",
+            inputs=[
+                IO.ClipVision.Input("clip_vision_model"),
+                IO.Image.Input("image", tooltip="Preprocessed image from ImageCropToMask (pad_factor=1.0 for TRELLIS.2)."),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, clip_vision_model, image) -> IO.NodeOutput:
+        out_device = comfy.model_management.intermediate_device()
+        cond = _dino_encode_batch(clip_vision_model, image, out_device)
+        cond_512_batched, cond_1024_batched = cond["global_512"], cond["global_1024"]
+        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)
+
+def _proj_pack_from_conditioning(conditioning):
+    """Return the proj_feat_pack dict embedded in a Pixal3D conditioning (or None
+    for vanilla Trellis2 / no conditioning connected). Pixal3DConditioning ships
+    the pack in cond[0][1]["proj_feat_pack"]; Trellis2Conditioning doesn't set it."""
+    if not conditioning:
+        return None
+    entry = conditioning[0]
+    if not isinstance(entry, (list, tuple)) or len(entry) < 2 or not isinstance(entry[1], dict):
+        return None
+    return entry[1].get("proj_feat_pack")
+
+
+def _conditioning_set_extras(conditioning, extras: dict):
+    """Return a copy of `conditioning` with `extras` merged into each entry's
+    dict — same shallow-copy pattern ControlNetApplyAdvanced uses. The dicts
+    are copied so we don't mutate upstream conditioning."""
+    out = []
+    for entry in conditioning:
+        if isinstance(entry, (list, tuple)) and len(entry) >= 2 and isinstance(entry[1], dict):
+            new_dict = entry[1].copy()
+            new_dict.update(extras)
+            out.append([entry[0], new_dict])
+        else:
+            out.append(entry)
+    return out
+
+
+class Trellis2ShapeStage(IO.ComfyNode):
+    """Sets up the first shape-stage sampling pass: extracts sparse coords from
+    the dense structure voxel produced by VaeDecodeStructureTrellis2, builds an
+    empty sparse latent, and attaches per-stage metadata to the conditioning so
+    the model reads it via extra_conds at sample time. For the second shape pass
+    (post-upsample), use Trellis2UpsampleStage instead — it combines the cascade
+    and the second-pass stage setup."""
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2ShapeStage",
+            category="model/conditioning/trellis2",
+            inputs=[
+                IO.Conditioning.Input("positive"),
+                IO.Conditioning.Input("negative"),
+                IO.Voxel.Input(
+                    "voxel",
+                    tooltip="Dense structure voxel from VaeDecodeStructureTrellis2.",
+                ),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, voxel):
+        decoded = voxel.data.unsqueeze(1)
+        coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
+        coord_resolution = int(decoded.shape[-1])
+
+        # Dispatch based on the upstream voxel resolution, mirroring upstream's
+        # pipeline_type → ss_res table:
+        #   coord_res == 32 → first cascade shape pass OR pure-512 pipeline
+        #                     (img2shape_512 + shape_512 proj stage, 512 DINO).
+        #   coord_res >  32 → pure-1024 non-cascade pipeline
+        #                     (img2shape + shape_1024 proj stage, 1024 DINO).
+        if coord_resolution <= 32:
+            mode = "shape_generation_512"
+            stage = "shape_512"
+        else:
+            mode = "shape_generation"
+            stage = "shape_1024"
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+        latent = torch.zeros(batch_size, 32, max_tokens, 1)
+
+        extras = {
+            "trellis2_generation_mode": mode,
+            "trellis2_coords": coords,
+            "trellis2_coord_counts": counts,
+        }
+        proj_pack = _proj_pack_from_conditioning(positive)
+        if proj_pack is not None:
+            extras["trellis2_proj_feats"] = compute_stage_proj_feats(
+                proj_pack, stage, coords=coords, coord_resolution=coord_resolution,
+            )
+        positive_out = _conditioning_set_extras(positive, extras)
+        negative_out = _conditioning_set_extras(negative, extras)
+        out_latent = {"samples": latent, "coords": coords, "coord_counts": counts,
+                      "coord_resolution": coord_resolution, "type": "trellis2",
+                      "model_frame": "y_up" if proj_pack is not None else "z_up"}
+        return IO.NodeOutput(positive_out, negative_out, out_latent)
+
+
+class Trellis2TextureStage(IO.ComfyNode):
+    """Sets up the texture-stage sampling pass. Reads coords / coord_counts /
+    coord_resolution and the shape_slat (the per-voxel shape latent) from the
+    incoming shape_latent dict — set there by Trellis2ShapeStage or
+    Trellis2UpsampleStage. Builds an empty sparse latent at the same coord
+    layout and attaches per-stage metadata to the conditioning."""
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2TextureStage",
+            category="model/conditioning/trellis2",
+            inputs=[
+                IO.Conditioning.Input("positive"),
+                IO.Conditioning.Input("negative"),
+                IO.Latent.Input("shape_latent"),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, shape_latent):
+        channels = 32
+        coords = shape_latent["coords"]
+        coord_resolution = shape_latent.get("coord_resolution")
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+
+        shape_slat = shape_latent["samples"]
+        if shape_slat.ndim == 4:
+            shape_slat = shape_slat.squeeze(-1).transpose(1, 2).reshape(-1, channels)
+
+        latent = torch.zeros(batch_size, channels, max_tokens, 1)
+        proj_pack = _proj_pack_from_conditioning(positive)
+        model_frame = shape_latent.get("model_frame",
+                                       "y_up" if proj_pack is not None else "z_up")
+        extras = {
+            "trellis2_generation_mode": "texture_generation",
+            "trellis2_coords": coords,
+            "trellis2_coord_counts": counts,
+            "trellis2_shape_slat": shape_slat,
+            "trellis2_model_frame": model_frame,
+        }
+        if proj_pack is not None and coord_resolution is not None:
+            extras["trellis2_proj_feats"] = compute_stage_proj_feats(
+                proj_pack, "tex_1024", coords=coords, coord_resolution=coord_resolution,
+            )
+
+        positive_out = _conditioning_set_extras(positive, extras)
+        negative_out = _conditioning_set_extras(negative, extras)
+        out_latent = {"samples": latent, "type": "trellis2", "coords": coords, "coord_counts": counts,
+                      "model_frame": shape_latent.get("model_frame",
+                                                       "y_up" if proj_pack is not None else "z_up")}
+        if coord_resolution is not None:
+            out_latent["coord_resolution"] = coord_resolution
+        return IO.NodeOutput(positive_out, negative_out, out_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 = 32
+        resolution = 16
+        latent = torch.zeros(batch_size, in_channels, resolution, resolution, resolution)
+        return IO.NodeOutput({"samples": latent, "type": "trellis2"})
+
+def _dinov3_patches_to_2d(tokens, image_size, patch_size=16):
+    h_p = w_p = image_size // patch_size
+    n_patches = h_p * w_p
+    n_reg = tokens.shape[1] - 1 - n_patches
+    if n_reg < 0 or tokens.shape[1] != 1 + n_reg + n_patches:
+        raise ValueError(
+            f"_dinov3_patches_to_2d: got {tokens.shape[1]} tokens, expected "
+            f"1 (CLS) + N_reg + {h_p}*{w_p}={n_patches} patches at image_size={image_size}, "
+            f"patch_size={patch_size}. Inferred N_reg={n_reg} which is invalid."
+        )
+    start = 1 + n_reg
+    patches = tokens[:, start:start + n_patches]
+    return patches.transpose(1, 2).reshape(tokens.shape[0], -1, h_p, w_p).contiguous()
+
+
+def _crop_image_with_mask(item_image, item_mask, max_image_size=1024, pad_factor=1.1,
+                          mask_offset=0, mask_threshold=0.05, bg_rgb=(0.0, 0.0, 0.0),
+                          aspect_ratio=1.0):
+    img = item_image.permute(2, 0, 1).unsqueeze(0).cpu().float().clamp(0, 1)
+    mask = item_mask.unsqueeze(0).unsqueeze(0).cpu().float().clamp(0, 1)
+
+    # Detect and correct an inverted mask, only when border and center have opposite polarity.
+    m2d = mask[0, 0]
+    h, w = m2d.shape
+    border = torch.cat([m2d[0, :], m2d[-1, :], m2d[:, 0], m2d[:, -1]])
+    center = m2d[h // 4:h - h // 4, w // 4:w - w // 4]
+    if float(border.mean()) > 0.5 and float(center.mean()) < 0.5:
+        mask = 1.0 - mask
+
+    if mask_offset > 0:
+        r = mask_offset
+        mask = torch.nn.functional.max_pool2d(mask, kernel_size=2 * r + 1, stride=1, padding=r)
+    elif mask_offset < 0:
+        r = -mask_offset
+        mask = 1.0 - torch.nn.functional.max_pool2d(1.0 - mask, kernel_size=2 * r + 1, stride=1, padding=r)
+
+    if mask_threshold > 0.0:
+        mask = torch.where(mask < mask_threshold, torch.zeros_like(mask), mask)
+
+    H, W = img.shape[-2:]
+    if max(H, W) > max_image_size:
+        scale = max_image_size / max(H, W)
+        new_w, new_h = int(W * scale), int(H * scale)
+        img  = comfy.utils.common_upscale(img,  new_w, new_h, "lanczos", "disabled")
+        mask = comfy.utils.common_upscale(mask, new_w, new_h, "lanczos", "disabled")
+        # common_upscale's lanczos path drops the singleton channel dim for masks (utils.py:1062).
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+        H, W = new_h, new_w
+    scene_size = (W, H)
+
+    alpha_u8 = (mask[0, 0].clamp(0, 1) * 255.0).to(torch.uint8)
+    fg_pixels = (alpha_u8 > 204).nonzero()
+    if fg_pixels.numel() == 0:
+        # Try the inverted mask — auto-invert above may have been too conservative.
+        inv_fg = ((255 - alpha_u8) > 204).nonzero()
+        if inv_fg.numel() > 0:
+            logging.info("Trellis2 preprocess: mask bbox empty, using inverted mask.")
+            mask = 1.0 - mask
+            fg_pixels = inv_fg
+    if fg_pixels.numel() > 0:
+        y_min, x_min = fg_pixels.min(dim=0).values.tolist()
+        y_max, x_max = fg_pixels.max(dim=0).values.tolist()
+        center_y, center_x = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
+        bw = x_max - x_min
+        bh = y_max - y_min
+        # Grow the bbox so its aspect matches `aspect_ratio` (width/height),
+        # anchored on the max side. Then apply pad_factor.
+        if bw / max(bh, 1) >= aspect_ratio:
+            crop_w = int(bw * pad_factor)
+            crop_h = int(bw / aspect_ratio * pad_factor)
+        else:
+            crop_h = int(bh * pad_factor)
+            crop_w = int(bh * aspect_ratio * pad_factor)
+        half_w, half_h = crop_w // 2, crop_h // 2
+        crop_x1 = int(center_x - half_w)
+        crop_y1 = int(center_y - half_h)
+        crop_x2 = crop_x1 + 2 * half_w
+        crop_y2 = crop_y1 + 2 * half_h
+    else:
+        logging.warning("Mask for the image is empty; a clean foreground mask is required for best quality.")
+        crop_x1, crop_y1, crop_x2, crop_y2 = 0, 0, W, H
+    crop_bbox = (crop_x1, crop_y1, crop_x2, crop_y2)
+
+    # Zero-pad out-of-bounds slice (PIL.crop semantics).
+    pad_l = max(0, -crop_x1)
+    pad_t = max(0, -crop_y1)
+    pad_r = max(0, crop_x2 - W)
+    pad_b = max(0, crop_y2 - H)
+    if pad_l or pad_t or pad_r or pad_b:
+        img  = torch.nn.functional.pad(img,  (pad_l, pad_r, pad_t, pad_b), value=0.0)
+        mask = torch.nn.functional.pad(mask, (pad_l, pad_r, pad_t, pad_b), value=0.0)
+        crop_x1 += pad_l
+        crop_x2 += pad_l
+        crop_y1 += pad_t
+        crop_y2 += pad_t
+    cropped_img  = img [..., crop_y1:crop_y2, crop_x1:crop_x2]
+    cropped_mask = mask[..., crop_y1:crop_y2, crop_x1:crop_x2]
+
+    bg = torch.tensor(bg_rgb, dtype=cropped_img.dtype, device=cropped_img.device).view(1, 3, 1, 1)
+    composite = (cropped_img * cropped_mask + bg * (1.0 - cropped_mask)).clamp(0, 1)
+    return composite, crop_bbox, scene_size
+
+
+def _dino_encode_batch(clip_vision_model, image, out_device, *, want_patches=False):
+    """Encode an already-preprocessed image through DINOv3 at 512 and 1024.
+
+    Expects `image` to be a comfy IMAGE tensor [B, H, W, 3] of squared composites
+    (from ImageCropToMask). Returns batched global tokens; with want_patches also
+    the 2D patch grids and the per-item BCHW composites that the Pixal3D NAF path needs."""
+    image = image[..., :3]
+    batch_size = image.shape[0]
+
+    cond_512_list, cond_1024_list = [], []
+    patches_512_list, patches_1024_list = [], []
+    composite_list = []
+    for b in range(batch_size):
+        item = image[b].movedim(-1, -3).unsqueeze(0).contiguous().float().clamp(0, 1)
+        c512 = _dinov3_encode(clip_vision_model, item, 512, want_patches=want_patches)
+        c1024 = _dinov3_encode(clip_vision_model, item, 1024, want_patches=want_patches)
+        if want_patches:
+            cond_512_list.append(c512["tokens"].to(out_device))
+            cond_1024_list.append(c1024["tokens"].to(out_device))
+            patches_512_list.append(c512["patches_2d"].to(out_device))
+            patches_1024_list.append(c1024["patches_2d"].to(out_device))
+            composite_list.append(item)
+        else:
+            cond_512_list.append(c512.to(out_device))
+            cond_1024_list.append(c1024.to(out_device))
+
+    out = {
+        "batch_size": batch_size,
+        "global_512": torch.cat(cond_512_list, dim=0),
+        "global_1024": torch.cat(cond_1024_list, dim=0),
+    }
+    if want_patches:
+        out["patches_512"] = torch.cat(patches_512_list, dim=0)
+        out["patches_1024"] = torch.cat(patches_1024_list, dim=0)
+        out["composites"] = composite_list
+    return out
+
+
+class ImageCropToMask(IO.ComfyNode):
+    """Crop an image to its mask's bounding box (centered square, with pad_factor
+    margin), then composite `img * mask` and resize to a square. Handles OOB crops
+    with zero-padding. Useful for 3D pipelines that expect a centered, background-free
+    subject at a fixed input resolution (Trellis2, Pixal3D, Hunyuan3D, TripoSR, etc.)."""
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="ImageCropToMask",
+            display_name="Image Crop to Mask",
+            category="image/transform",
+            search_aliases=["crop to mask", "mask crop", "crop mask", "mask crop resize", "crop mask resize", "trellis2", "pixal3d"],
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Mask.Input("mask"),
+                IO.Int.Input("width", default=1024, min=64, max=4096, step=8, tooltip="Output width in pixels."),
+                IO.Int.Input("height", default=1024, min=64, max=4096, step=8, tooltip="Output height in pixels."),
+                IO.Float.Input("pad_factor", default=1.0, min=1.0, max=2.0, step=0.01, tooltip="Extra margin around the mask bbox as a multiplier."),
+                IO.Int.Input("mask_offset", default=0, min=-32, max=32, step=1, tooltip="Grow or shrink the mask by this many pixels before cropping."),
+                IO.Color.Input("background", default="#000000", tooltip="Fill color behind the masked subject."),
+            ],
+            outputs=[IO.Image.Output(display_name="image")],
+        )
+
+    @classmethod
+    def execute(cls, image, mask, width, height, pad_factor, mask_offset, background) -> IO.NodeOutput:
+        h = background.lstrip("#")
+        bg_rgb = (int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0) if len(h) == 6 else (0.0, 0.0, 0.0)
+        image = image[..., :3]
+        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"Mask batch {mask.shape[0]} does not match image batch {batch_size}")
+
+        out_images = []
+        for b in range(batch_size):
+            composite, _, _ = _crop_image_with_mask(
+                image[b], mask[b], max_image_size=max(width, height), pad_factor=pad_factor,
+                mask_offset=mask_offset, bg_rgb=bg_rgb, aspect_ratio=width / height,
+            )
+            composite = comfy.utils.common_upscale(composite, width, height, "lanczos", "disabled")
+            out_images.append(composite.movedim(-3, -1))
+
+        result = torch.cat(out_images, dim=0).to(
+            device=comfy.model_management.intermediate_device(),
+            dtype=comfy.model_management.intermediate_dtype(),
+        )
+        return IO.NodeOutput(result)
+
+
+class Pixal3DConditioning(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Pixal3DConditioning",
+            category="model/conditioning/trellis2",
+            inputs=[
+                IO.ClipVision.Input("clip_vision_model", tooltip="DINOv3 ViT-L/16 ClipVision."),
+                IO.Image.Input("image", tooltip="Preprocessed image from ImageCropToMask (pad_factor=1.1 for Pixal3D)."),
+                IO.Float.Input(
+                    "camera_angle_x", display_name="fov",
+                    default=49.13, min=1.0, max=170.0, step=0.01, advanced=True,
+                    tooltip="Horizontal FOV in degrees. Wire a MoGeGeometryToFOV "
+                            "(axis='horizontal', unit='degrees') for a per-image FoV (matches upstream default).",
+                ),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip_vision_model, image, camera_angle_x) -> IO.NodeOutput:
+        naf_model = clip_vision_model.naf
+        out_device = comfy.model_management.intermediate_device()
+        compute_device = comfy.model_management.get_torch_device()
+
+        cond = _dino_encode_batch(clip_vision_model, image, out_device, want_patches=True)
+        batch_size = cond["batch_size"]
+        global_512, global_1024 = cond["global_512"], cond["global_1024"]
+        fm_512_dino, fm_1024_dino = cond["patches_512"], cond["patches_1024"]
+        composite_list = cond["composites"]
+
+        # The LR DINO grid AND the NAF HR grid are sampled separately
+        # NAF targets per stage: shape_512=512, shape_1024=512, tex_1024=1024.
+        def _naf_hr(lr_feat, composites, image_size, naf_target):
+            if naf_model is None or naf_target is None:
+                return None
+            comfy.model_management.load_model_gpu(naf_model)
+            inner = naf_model.model
+            model_dtype = next(inner.parameters()).dtype  # set at load time (see clip_vision NAF)
+            hrs = []
+            for i, c in enumerate(composites):
+                img_i = comfy.utils.common_upscale(c, image_size, image_size, "lanczos", "disabled")\
+                    .to(compute_device).to(model_dtype)
+                lr_i = lr_feat[i:i + 1].to(compute_device).to(model_dtype)
+                hr_i = inner(img_i, lr_i, naf_target, output_device=out_device)
+                hrs.append(hr_i)
+            return torch.cat(hrs, dim=0)
+
+        hr_shape_512  = _naf_hr(fm_512_dino,  composite_list, 512,  (512, 512))
+        hr_shape_1024 = _naf_hr(fm_1024_dino, composite_list, 1024, (512, 512))
+        hr_tex_1024   = _naf_hr(fm_1024_dino, composite_list, 1024, (1024, 1024))
+
+        # distance_from_fov: grid_point (-1, 0, 0) projects to pixel (0, image_resolution-1).
+        # FOV widget is in degrees for UX; trig + downstream projection expect radians.
+        camera_angle_x = math.radians(float(camera_angle_x))
+        distance = 0.5 / math.tan(camera_angle_x / 2.0)
+        cam_angle_t = torch.tensor([camera_angle_x] * batch_size, device=out_device, dtype=torch.float32)
+        dist_t = torch.tensor([distance] * batch_size, device=out_device, dtype=torch.float32)
+        scale_t = torch.ones(batch_size, device=out_device, dtype=torch.float32)
+        T = build_proj_transform_matrix(dist_t, batch_size, device=out_device, dtype=torch.float32)
+
+        proj_pack = {
+            "stages": {
+                "ss":         {"feature_map": fm_512_dino,    "feature_map_hr": None,         "image_resolution": 512},
+                "shape_512":  {"feature_map": fm_512_dino,    "feature_map_hr": hr_shape_512, "image_resolution": 512},
+                "shape_1024": {"feature_map": fm_1024_dino,   "feature_map_hr": hr_shape_1024,"image_resolution": 1024},
+                "tex_1024":   {"feature_map": fm_1024_dino,   "feature_map_hr": hr_tex_1024,  "image_resolution": 1024},
+            },
+            "transform_matrix": T,
+            "camera_angle_x": cam_angle_t,
+            "mesh_scale": scale_t,
+            "distance": dist_t,
+            "patch_size": 16,
+        }
+
+        # global_512 → SS/shape_512 cross-attn; global_1024 → shape_1024/tex_1024.
+        ss_proj_feats = compute_stage_proj_feats(
+            proj_pack, "ss", dense_grid_resolution=16, batch_size=batch_size,
+            device=compute_device,
+        )
+        neg_global = torch.zeros_like(global_512)
+        neg_embeds = torch.zeros_like(global_1024)
+        base_extras = {
+            "embeds": global_1024, "proj_feat_pack": proj_pack,
+            "trellis2_proj_feats": ss_proj_feats,
+        }
+        neg_extras = {
+            "embeds": neg_embeds, "proj_feat_pack": proj_pack,
+            "trellis2_proj_feats": ss_proj_feats,
+        }
+        positive = [[global_512, base_extras]]
+        negative = [[neg_global, neg_extras]]
+        return IO.NodeOutput(positive, negative)
+
+
+class Trellis2Extension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            ImageCropToMask,
+            Trellis2Conditioning,
+            Pixal3DConditioning,
+            Trellis2ShapeStage,
+            EmptyTrellis2LatentStructure,
+            Trellis2TextureStage,
+            VaeDecodeTextureTrellis,
+            VaeDecodeShapeTrellis,
+            VaeDecodeStructureTrellis2,
+            Trellis2UpsampleStage,
+        ]
+
+
+async def comfy_entrypoint() -> Trellis2Extension:
+    return Trellis2Extension()

nodes.py

@@ -2482,6 +2482,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",