Vertex Clustering, Mask Fix, Normal Fix (#14035)

* Vertex Clustering, Mask Fix, Normal Fix

* detects inverted mask

* update the decimate mesh

comfy/clip_vision.py

@@ -9,6 +9,7 @@ import comfy.model_management
 import comfy.utils
 import comfy.clip_model
 import comfy.image_encoders.dino2
+import comfy.image_encoders.dino3
 
 class Output:
     def __getitem__(self, key):
@@ -23,6 +24,7 @@ IMAGE_ENCODERS = {
     "siglip_vision_model": comfy.clip_model.CLIPVisionModelProjection,
     "siglip2_vision_model": comfy.clip_model.CLIPVisionModelProjection,
     "dinov2": comfy.image_encoders.dino2.Dinov2Model,
+    "dinov3": comfy.image_encoders.dino3.DINOv3ViTModel
 }
 
 class ClipVisionModel():
@@ -134,6 +136,8 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
     elif 'encoder.layer.23.layer_scale2.lambda1' in sd:
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
+    elif 'layer.9.attention.o_proj.bias' in sd: # dinov3
+        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino3_large.json")
     else:
         return None
 

comfy/image_encoders/dino3.py

@@ -0,0 +1,285 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import comfy.model_management
+from comfy.ldm.modules.attention import optimized_attention_for_device
+from comfy.image_encoders.dino2 import LayerScale as DINOv3ViTLayerScale
+
+class DINOv3ViTMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, mlp_bias, device, dtype, operations):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.up_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.down_proj = operations.Linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.act_fn = torch.nn.GELU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+def apply_rotary_pos_emb(q, k, cos, sin, **kwargs):
+    num_tokens = q.shape[-2]
+    num_patches = sin.shape[-2]
+    num_prefix_tokens = num_tokens - num_patches
+
+    q_prefix_tokens, q_patches = q.split((num_prefix_tokens, num_patches), dim=-2)
+    k_prefix_tokens, k_patches = k.split((num_prefix_tokens, num_patches), dim=-2)
+
+    q_patches = (q_patches * cos) + (rotate_half(q_patches) * sin)
+    k_patches = (k_patches * cos) + (rotate_half(k_patches) * sin)
+
+    q = torch.cat((q_prefix_tokens, q_patches), dim=-2)
+    k = torch.cat((k_prefix_tokens, k_patches), dim=-2)
+
+    return q, k
+
+class DINOv3ViTAttention(nn.Module):
+    def __init__(self, hidden_size, num_attention_heads, device, dtype, operations):
+        super().__init__()
+        self.embed_dim = hidden_size
+        self.num_heads = num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+
+        self.k_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=False, device=device, dtype=dtype) # key_bias = False
+        self.v_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+
+        self.q_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+        self.o_proj = operations.Linear(self.embed_dim, self.embed_dim, bias=True, device=device, dtype=dtype)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+
+        batch_size, patches, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is not None:
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        attn = optimized_attention_for_device(query_states.device, mask=False)
+
+        attn_output = attn(
+            query_states, key_states, value_states, self.num_heads, attention_mask, skip_reshape=True, skip_output_reshape=True
+        )
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(batch_size, patches, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output
+
+class DINOv3ViTGatedMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, mlp_bias, device, dtype, operations):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.up_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.down_proj = operations.Linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, device=device, dtype=dtype)
+        self.act_fn = torch.nn.GELU()
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+def get_patches_center_coordinates(
+    num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor:
+
+    coords_h = torch.arange(0.5, num_patches_h, dtype=dtype, device=device)
+    coords_w = torch.arange(0.5, num_patches_w, dtype=dtype, device=device)
+    coords_h = coords_h / num_patches_h
+    coords_w = coords_w / num_patches_w
+    coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)
+    coords = coords.flatten(0, 1)
+    coords = 2.0 * coords - 1.0
+    return coords
+
+class DINOv3ViTRopePositionEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, rope_theta, hidden_size, num_attention_heads, image_size, patch_size, device, dtype):
+        super().__init__()
+        self.base = rope_theta
+        self.head_dim = hidden_size // num_attention_heads
+        self.num_patches_h = image_size // patch_size
+        self.num_patches_w = image_size // patch_size
+        self.patch_size = patch_size
+
+        inv_freq = 1 / self.base ** torch.arange(0, 1, 4 / self.head_dim, dtype=torch.float32, device=device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        _, _, height, width = pixel_values.shape
+        num_patches_h = height // self.patch_size
+        num_patches_w = width // self.patch_size
+
+        device = pixel_values.device
+        device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"
+        with torch.amp.autocast(device_type = device_type, enabled=False):
+            patch_coords = get_patches_center_coordinates(
+                num_patches_h, num_patches_w, dtype=torch.float32, device=device
+            )
+
+            self.inv_freq = self.inv_freq.to(device)
+            angles = 2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :]
+            angles = angles.flatten(1, 2)
+            angles = angles.tile(2)
+
+            cos = torch.cos(angles)
+            sin = torch.sin(angles)
+
+        dtype = pixel_values.dtype
+        return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+class DINOv3ViTEmbeddings(nn.Module):
+    def __init__(self, hidden_size, num_register_tokens, num_channels, patch_size, dtype, device, operations):
+        super().__init__()
+        self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_size, device=device, dtype=dtype))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, hidden_size, device=device, dtype=dtype))
+        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, hidden_size, device=device, dtype=dtype))
+        self.patch_embeddings = operations.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
+        )
+
+    def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None):
+        batch_size = pixel_values.shape[0]
+        target_dtype = self.patch_embeddings.weight.dtype
+
+        patch_embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
+        patch_embeddings = patch_embeddings.flatten(2).transpose(1, 2)
+
+        if bool_masked_pos is not None:
+            mask_token = self.mask_token.to(patch_embeddings.dtype)
+            patch_embeddings = torch.where(bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings)
+
+        cls_token = self.cls_token.expand(batch_size, -1, -1)
+        register_tokens = self.register_tokens.expand(batch_size, -1, -1)
+        device = patch_embeddings.device
+        cls_token = cls_token.to(device)
+        register_tokens = register_tokens.to(device)
+        embeddings = torch.cat([cls_token, register_tokens, patch_embeddings], dim=1)
+
+        return embeddings
+
+class DINOv3ViTLayer(nn.Module):
+
+    def __init__(self, hidden_size, layer_norm_eps, use_gated_mlp, mlp_bias, intermediate_size, num_attention_heads,
+                 device, dtype, operations):
+        super().__init__()
+
+        self.norm1 = operations.LayerNorm(hidden_size, eps=layer_norm_eps, device=device, dtype=dtype)
+        self.attention = DINOv3ViTAttention(hidden_size, num_attention_heads, device=device, dtype=dtype, operations=operations)
+        self.layer_scale1 = DINOv3ViTLayerScale(hidden_size, device=device, dtype=dtype, operations=None)
+
+        self.norm2 = operations.LayerNorm(hidden_size, eps=layer_norm_eps, device=device, dtype=dtype)
+
+        if use_gated_mlp:
+            self.mlp = DINOv3ViTGatedMLP(hidden_size, intermediate_size, mlp_bias, device=device, dtype=dtype, operations=operations)
+        else:
+            self.mlp = DINOv3ViTMLP(hidden_size, intermediate_size=intermediate_size, mlp_bias=mlp_bias, device=device, dtype=dtype, operations=operations)
+        self.layer_scale2 = DINOv3ViTLayerScale(hidden_size, device=device, dtype=dtype, operations=None)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+        )
+        hidden_states = self.layer_scale1(hidden_states)
+        hidden_states = hidden_states + residual
+
+        residual = hidden_states
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.layer_scale2(hidden_states)
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class DINOv3ViTModel(nn.Module):
+    def __init__(self, config, dtype, device, operations):
+        super().__init__()
+        use_bf16 = comfy.model_management.should_use_bf16(device, prioritize_performance=True)
+        if dtype == torch.float16 and use_bf16:
+            dtype = torch.bfloat16
+        elif dtype == torch.float16 and not use_bf16:
+            dtype = torch.float32
+        num_hidden_layers = config["num_hidden_layers"]
+        hidden_size = config["hidden_size"]
+        num_attention_heads = config["num_attention_heads"]
+        num_register_tokens = config["num_register_tokens"]
+        intermediate_size = config["intermediate_size"]
+        layer_norm_eps = config["layer_norm_eps"]
+        num_channels = config["num_channels"]
+        patch_size = config["patch_size"]
+        rope_theta = config["rope_theta"]
+
+        self.embeddings = DINOv3ViTEmbeddings(
+            hidden_size, num_register_tokens, num_channels=num_channels, patch_size=patch_size, dtype=dtype, device=device, operations=operations
+        )
+        self.rope_embeddings = DINOv3ViTRopePositionEmbedding(
+            rope_theta, hidden_size, num_attention_heads, image_size=512, patch_size=patch_size, dtype=dtype, device=device
+        )
+        self.layer = nn.ModuleList(
+            [DINOv3ViTLayer(hidden_size, layer_norm_eps, use_gated_mlp=False, mlp_bias=True,
+                            intermediate_size=intermediate_size,num_attention_heads = num_attention_heads,
+                            dtype=dtype, device=device, operations=operations)
+            for _ in range(num_hidden_layers)])
+        self.norm = operations.LayerNorm(hidden_size, eps=layer_norm_eps, dtype=dtype, device=device)
+
+    def get_input_embeddings(self):
+        return self.embeddings.patch_embeddings
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        bool_masked_pos: torch.Tensor | None = None,
+        **kwargs,
+    ):
+
+        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
+        position_embeddings = self.rope_embeddings(pixel_values)
+
+        for i, layer_module in enumerate(self.layer):
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        if kwargs.get("skip_norm_elementwise", False):
+            sequence_output= F.layer_norm(hidden_states, hidden_states.shape[-1:])
+        else:
+            norm = self.norm.to(hidden_states.device)
+            sequence_output = norm(hidden_states)
+        pooled_output = sequence_output[:, 0, :]
+
+        return sequence_output, None, pooled_output, None

comfy/image_encoders/dino3_large.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "dinov3",
+  "hidden_size": 1024,
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "key_bias": false,
+  "layer_norm_eps": 1e-05,
+  "mlp_bias": true,
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "num_register_tokens": 4,
+  "patch_size": 16,
+  "pos_embed_rescale": 2.0,
+  "proj_bias": true,
+  "query_bias": true,
+  "rope_theta": 100.0,
+  "use_gated_mlp": false,
+  "value_bias": true,
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_std": [0.229, 0.224, 0.225]
+}

comfy/latent_formats.py

@@ -760,6 +760,8 @@ class Hunyuan3Dv2_1(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1
 
+class Trellis2(LatentFormat): # TODO
+    latent_channels = 32
 class Hunyuan3Dv2mini(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1

comfy/ldm/trellis2/attention.py

@@ -0,0 +1,282 @@
+import torch
+import math
+from comfy.ldm.modules.attention import optimized_attention
+from typing import Tuple, Union, List
+from comfy.ldm.trellis2.vae import VarLenTensor
+import comfy.ops
+
+
+# replica of the seedvr2 code
+def var_attn_arg(kwargs):
+    cu_seqlens_q = kwargs.get("cu_seqlens_q", None)
+    max_seqlen_q = kwargs.get("max_seqlen_q", None)
+    cu_seqlens_k = kwargs.get("cu_seqlens_kv", cu_seqlens_q)
+    max_seqlen_k = kwargs.get("max_kv_seqlen", max_seqlen_q)
+    assert cu_seqlens_q is not None, "cu_seqlens_q shouldn't be None when var_length is True"
+    return cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k
+
+def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    var_length = True
+    if var_length:
+        cu_seqlens_q, cu_seqlens_k, _, _ = var_attn_arg(kwargs)
+        if not skip_reshape:
+            # assumes 2D q, k,v [total_tokens, embed_dim]
+            total_tokens, embed_dim = q.shape
+            head_dim = embed_dim // heads
+            q = q.view(total_tokens, heads, head_dim)
+            k = k.view(k.shape[0], heads, head_dim)
+            v = v.view(v.shape[0], heads, head_dim)
+
+        b = q.size(0)
+        dim_head = q.shape[-1]
+        q = torch.nested.nested_tensor_from_jagged(q, offsets=cu_seqlens_q.long())
+        k = torch.nested.nested_tensor_from_jagged(k, offsets=cu_seqlens_k.long())
+        v = torch.nested.nested_tensor_from_jagged(v, offsets=cu_seqlens_k.long())
+
+        mask = None
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+    if mask is not None:
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+
+    out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+    if var_length:
+        return out.transpose(1, 2).values()
+    if not skip_output_reshape:
+        out = (
+            out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+        )
+    return out
+
+def scaled_dot_product_attention(*args, **kwargs):
+    num_all_args = len(args) + len(kwargs)
+
+    q = None
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs.get('qkv')
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs.get('q')
+        kv = args[1] if len(args) > 1 else kwargs.get('kv')
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs.get('q')
+        k = args[1] if len(args) > 1 else kwargs.get('k')
+        v = args[2] if len(args) > 2 else kwargs.get('v')
+
+    if q is not None:
+        heads = q.shape[2]
+    else:
+        heads = qkv.shape[3]
+
+    if num_all_args == 1:
+        q, k, v = qkv.unbind(dim=2)
+    elif num_all_args == 2:
+        k, v = kv.unbind(dim=2)
+
+    q = q.permute(0, 2, 1, 3)
+    k = k.permute(0, 2, 1, 3)
+    v = v.permute(0, 2, 1, 3)
+
+    out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True, **kwargs)
+
+    out = out.permute(0, 2, 1, 3)
+
+    return out
+
+def sparse_windowed_scaled_dot_product_self_attention(
+    qkv,
+    window_size: int,
+    shift_window: Tuple[int, int, int] = (0, 0, 0)
+):
+
+    serialization_spatial_cache_name = f'windowed_attention_{window_size}_{shift_window}'
+    serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+    if serialization_spatial_cache is None:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = calc_window_partition(qkv, window_size, shift_window)
+        qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, attn_func_args))
+    else:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = serialization_spatial_cache
+
+    qkv_feats = qkv.feats[fwd_indices]      # [M, 3, H, C]
+    heads = qkv_feats.shape[2]
+
+    if optimized_attention.__name__ == 'attention_xformers':
+        q, k, v = qkv_feats.unbind(dim=1)
+        q = q.unsqueeze(0)                                                              # [1, M, H, C]
+        k = k.unsqueeze(0)                                                              # [1, M, H, C]
+        v = v.unsqueeze(0)                                                              # [1, M, H, C]
+        #out = xops.memory_efficient_attention(q, k, v, **attn_func_args)[0]             # [M, H, C]
+        out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True)
+    elif optimized_attention.__name__ == 'attention_flash':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args)  # [M, H, C]
+    else:
+        out = optimized_attention(q, k, v, heads, skip_output_reshape=True, skip_reshape=True)
+
+    out = out[bwd_indices]      # [T, H, C]
+
+    return qkv.replace(out)
+
+def calc_window_partition(
+    tensor,
+    window_size: Union[int, Tuple[int, ...]],
+    shift_window: Union[int, Tuple[int, ...]] = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
+
+    DIM = tensor.coords.shape[1] - 1
+    shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
+    window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
+    shifted_coords = tensor.coords.clone().detach()
+    shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+
+    MAX_COORDS = [i + j for i, j in zip(tensor.spatial_shape, shift_window)]
+    NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
+    OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
+
+    shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+    shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
+    fwd_indices = torch.argsort(shifted_indices)
+    bwd_indices = torch.empty_like(fwd_indices)
+    bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
+    seq_lens = torch.bincount(shifted_indices)
+    mask = seq_lens != 0
+    seq_lens = seq_lens[mask]
+
+    if optimized_attention.__name__ == 'attention_xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        attn_func_args = {
+            'attn_bias': xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+        }
+    elif optimized_attention.__name__ == 'attention_flash':
+        attn_func_args = {
+            'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'max_seqlen': torch.max(seq_lens)
+        }
+
+    return fwd_indices, bwd_indices, seq_lens, attn_func_args
+
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+    q=None
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        device = qkv.device
+
+        s = qkv
+        q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+        kv_seqlen = q_seqlen
+        qkv = qkv.feats     # [T, 3, H, C]
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, C]
+        else:
+            s = None
+            N, L, H, C = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, C)   # [T_Q, H, C]
+
+        if isinstance(kv, VarLenTensor):
+            kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+            kv = kv.feats     # [T_KV, 2, H, C]
+        else:
+            N, L, _, H, C = kv.shape
+            kv_seqlen = [L] * N
+            kv = kv.reshape(N * L, 2, H, C)   # [T_KV, 2, H, C]
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, Ci]
+        else:
+            s = None
+            N, L, H, CI = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, CI)  # [T_Q, H, Ci]
+
+        if isinstance(k, VarLenTensor):
+            kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+            k = k.feats     # [T_KV, H, Ci]
+            v = v.feats     # [T_KV, H, Co]
+        else:
+            N, L, H, CI, CO = *k.shape, v.shape[-1]
+            kv_seqlen = [L] * N
+            k = k.reshape(N * L, H, CI)     # [T_KV, H, Ci]
+            v = v.reshape(N * L, H, CO)     # [T_KV, H, Co]
+
+    # TODO: change
+    if q is not None:
+        heads = q
+    else:
+        heads = qkv
+    heads = heads.shape[2]
+    if optimized_attention.__name__ == 'attention_xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        q = q.unsqueeze(0)
+        k = k.unsqueeze(0)
+        v = v.unsqueeze(0)
+        mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+        out = xops.memory_efficient_attention(q, k, v, mask)[0]
+    elif optimized_attention.__name__ == 'attention_flash':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+
+    elif optimized_attention.__name__ == "attention_pytorch":
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        else:
+            cu_seqlens_kv = cu_seqlens_q
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        out = attention_pytorch(q, k, v, heads=heads,cu_seqlens_q=cu_seqlens_q,
+                                cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=max(q_seqlen), max_kv_seqlen=max(kv_seqlen),
+                                skip_reshape=True, skip_output_reshape=True)
+
+    if s is not None:
+        return s.replace(out)
+    else:
+        return out.reshape(N, L, H, -1)

comfy/ldm/trellis2/flexgemm.py

@@ -0,0 +1,298 @@
+# will contain every cuda -> pytorch operation
+
+from typing import Optional, Tuple
+import torch
+
+UINT32_SENTINEL = 0xFFFFFFFF
+
+
+def compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device):
+    """Kernel spatial offsets in the same order as the CUDA/Triton kernels."""
+    offsets = []
+    for vx in range(Kw):
+        for vy in range(Kh):
+            for vz in range(Kd):
+                offsets.append((vx * Dw, vy * Dh, vz * Dd))
+    return torch.tensor(offsets, device=device, dtype=torch.int32)
+
+
+class TorchHashMap:
+    """Sorted-array hashmap backed by torch.searchsorted."""
+
+    def __init__(self, keys: torch.Tensor, values: torch.Tensor, default_value: int):
+        device = keys.device
+        self.sorted_keys, order = torch.sort(keys.to(torch.long))
+        self.sorted_vals = values.to(torch.long)[order]
+        self.default_value = torch.tensor(default_value, dtype=torch.long, device=device)
+        self._n = self.sorted_keys.numel()
+
+    def lookup_flat(self, flat_keys: torch.Tensor) -> torch.Tensor:
+        flat = flat_keys.to(torch.long)
+        if self._n == 0:
+            return torch.full((flat.shape[0],), -1, device=flat.device, dtype=torch.int32)
+        idx = torch.searchsorted(self.sorted_keys, flat)
+        idx_safe = torch.clamp(idx, max=self._n - 1)
+        found = (idx < self._n) & (self.sorted_keys[idx_safe] == flat)
+        out = torch.full((flat.shape[0],), -1, device=flat.device, dtype=torch.int32)
+        if found.any():
+            out[found] = self.sorted_vals[idx_safe[found]].to(torch.int32)
+        return out
+
+
+def build_submanifold_neighbor_map(
+    hashmap,
+    coords: torch.Tensor,
+    W, H, D,
+    Kw, Kh, Kd,
+    Dw, Dh, Dd,
+):
+    device = coords.device
+    M = coords.shape[0]
+    V = Kw * Kh * Kd
+    half_V = V // 2 + 1
+    INVALID = -1
+
+    # int32 neighbour map: 4 bytes/elem vs 8 bytes for int64
+    neighbor = torch.full((M, V), INVALID, device=device, dtype=torch.int32)
+
+    b = coords[:, 0].long()
+    x = coords[:, 1].long()
+    y = coords[:, 2].long()
+    z = coords[:, 3].long()
+
+    offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device)
+
+    ox = x - (Kw // 2) * Dw
+    oy = y - (Kh // 2) * Dh
+    oz = z - (Kd // 2) * Dd
+
+    for v in range(half_V):
+        if v == half_V - 1:
+            # Center voxel always maps to itself
+            neighbor[:, v] = torch.arange(M, device=device, dtype=torch.int32)
+            continue
+
+        dx, dy, dz = offsets[v]
+
+        kx = ox + dx
+        ky = oy + dy
+        kz = oz + dz
+
+        valid = (
+            (kx >= 0) & (kx < W) &
+            (ky >= 0) & (ky < H) &
+            (kz >= 0) & (kz < D)
+        )
+
+        flat = (
+            b[valid] * (W * H * D) +
+            kx[valid] * (H * D) +
+            ky[valid] * D +
+            kz[valid]
+        )
+
+        if flat.numel() > 0:
+            found = hashmap.lookup_flat(flat)
+            idx_in_M = torch.where(valid)[0]
+            neighbor[idx_in_M, v] = found.to(torch.int32)
+
+            # BUG FIX: old code used  found != hashmap.default_value  which
+            # compared int32 -1 against int64 4294967295 → always True.
+            # We now explicitly check for valid indices.
+            valid_found_mask = found >= 0
+            if valid_found_mask.any():
+                src_points = idx_in_M[valid_found_mask]
+                dst_points = found[valid_found_mask].long()
+                neighbor[dst_points, V - 1 - v] = src_points.to(torch.int32)
+
+    return neighbor
+
+def get_recommended_chunk_mem(
+    device=None,
+    safety_fraction: float = 0.4,
+    min_gb: float = 0.25,
+    max_gb: float = 8.0,
+):
+
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    else:
+        device = torch.device(device)
+
+    if device.type == 'cuda':
+        try:
+            idx = device.index if device.index is not None else 0
+            free_bytes, total_bytes = torch.cuda.mem_get_info(idx)
+            free_gb = free_bytes / (1024 ** 3)
+            total_gb = total_bytes / (1024 ** 3)
+
+            recommended = free_gb * safety_fraction
+            result = max(min_gb, min(recommended, max_gb))
+            return result
+
+        except Exception:
+            try:
+                idx = device.index if device.index is not None else 0
+                total_gb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 3)
+            except Exception:
+                total_gb = 16.0
+
+            if total_gb < 12:
+                result = 0.5
+            elif total_gb < 16:
+                result = 0.75
+            elif total_gb < 24:
+                result = 1.0
+            elif total_gb < 32:
+                result = 2.0
+            elif total_gb < 48:
+                result = 4.0
+            else:
+                result = 6.0
+            return result
+
+    else:
+        try:
+            import psutil
+            avail_gb = psutil.virtual_memory().available / (1024 ** 3)
+            recommended = avail_gb * safety_fraction
+            result = max(min_gb, min(recommended, max_gb))
+            return result
+        except ImportError:
+            return min_gb
+
+def sparse_submanifold_conv3d(
+    feats: torch.Tensor,
+    coords: torch.Tensor,
+    shape: tuple,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    neighbor_cache: Optional[torch.Tensor],
+    dilation: tuple,
+    max_chunk_mem_gb: float = 6.0,
+    accumulate_f32: bool = True,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+    if feats.shape[0] == 0:
+        Co = weight.shape[0]
+        return torch.empty((0, Co), device=feats.device, dtype=feats.dtype), None
+
+    if len(shape) == 5:
+        _, _, W, H, D = shape
+    else:
+        W, H, D = shape
+
+    Co, Kw, Kh, Kd, Ci = weight.shape
+    V = Kw * Kh * Kd
+    device = feats.device
+    sentinel = -1
+    max_chunk_mem_gb = get_recommended_chunk_mem(device)
+
+    if neighbor_cache is None:
+        b_stride = W * H * D
+        x_stride = H * D
+        y_stride = D
+        z_stride = 1
+
+        flat_keys = (coords[:, 0].long() * b_stride +
+                     coords[:, 1].long() * x_stride +
+                     coords[:, 2].long() * y_stride +
+                     coords[:, 3].long() * z_stride)
+        vals = torch.arange(coords.shape[0], dtype=torch.int32, device=device)
+        hashmap = TorchHashMap(flat_keys, vals, UINT32_SENTINEL)
+
+        neighbor = build_submanifold_neighbor_map(
+            hashmap, coords, W, H, D, Kw, Kh, Kd,
+            dilation[0], dilation[1], dilation[2]
+        )
+    else:
+        neighbor = neighbor_cache
+
+    N_pts = feats.shape[0]
+
+    if accumulate_f32:
+        weight_T = weight.view(Co, V * Ci).to(torch.float32).T.contiguous()
+        output = torch.zeros(N_pts, Co, device=device, dtype=torch.float32)
+    else:
+        weight_T = weight.view(Co, V * Ci).to(feats.dtype).T.contiguous()
+        output = torch.zeros(N_pts, Co, device=device, dtype=feats.dtype)
+
+    # ------------------------------------------------------------------
+    # Chunk size from memory budget
+    # ------------------------------------------------------------------
+    bytes_per_elem = 4 if accumulate_f32 else feats.element_size()
+    mem_per_row = V * Ci * bytes_per_elem
+    max_chunk_mem = max_chunk_mem_gb * (1024 ** 3)
+    chunk_size = max(1, int(max_chunk_mem / mem_per_row))
+    chunk_size = min(chunk_size, N_pts)
+
+    # ------------------------------------------------------------------
+    # Chunked forward pass
+    #   Each iteration:
+    #     1. gather   (chunk, V, Ci)     – memory bound
+    #     2. mask     zero invalids       – in-place, no extra alloc
+    #     3. reshape  (chunk, V*Ci)
+    #     4. GEMM     (chunk, V*Ci) @ (V*Ci, Co) → (chunk, Co)  – cuBLAS
+    #        written directly into output slice via out= argument
+    # ------------------------------------------------------------------
+    for start in range(0, N_pts, chunk_size):
+        end = min(start + chunk_size, N_pts)
+        actual_chunk = end - start
+
+        # (chunk, V) int32
+        chunk_neighbor = neighbor[start:end]
+        chunk_valid = chunk_neighbor != sentinel
+
+        # Clamp sentinel -1 → 0 for safe indexing.  No clone of the full map.
+        chunk_idx = chunk_neighbor.clamp(min=0).long()
+
+        # Gather: (chunk, V, Ci).  Memory-bound, single index_select.
+        gathered = feats[chunk_idx]
+
+        # Zero invalid neighbours in-place.  gathered is a fresh tensor from
+        # advanced indexing, so in-place mutation is safe.
+        gathered.mul_(chunk_valid.unsqueeze(-1))
+
+        # Reshape to (chunk, V*Ci)
+        gathered_flat = gathered.view(actual_chunk, V * Ci)
+        if accumulate_f32:
+            gathered_flat = gathered_flat.to(torch.float32)
+
+        # Single GEMM call per chunk, written directly into output.
+        # This avoids allocating a temporary (chunk, Co) tensor.
+        torch.matmul(gathered_flat, weight_T, out=output[start:end])
+
+    if accumulate_f32:
+        output = output.to(feats.dtype)
+
+    if bias is not None:
+        output = output + bias.unsqueeze(0).to(output.dtype)
+
+    return output, neighbor
+
+class Mesh:
+    def __init__(self,
+        vertices,
+        faces,
+        vertex_attrs=None
+    ):
+        self.vertices = vertices.float()
+        self.faces = faces.int()
+        self.vertex_attrs = vertex_attrs
+
+    @property
+    def device(self):
+        return self.vertices.device
+
+    def to(self, device, non_blocking=False):
+        return Mesh(
+            self.vertices.to(device, non_blocking=non_blocking),
+            self.faces.to(device, non_blocking=non_blocking),
+            self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None,
+        )
+
+    def cuda(self, non_blocking=False):
+        return self.to('cuda', non_blocking=non_blocking)
+
+    def cpu(self):
+        return self.to('cpu')

comfy/ldm/trellis2/model.py

@@ -0,0 +1,935 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from comfy.ldm.trellis2.vae import SparseTensor, SparseLinear, sparse_cat, VarLenTensor
+from typing import Optional, Tuple, Literal, Union, List
+from comfy.ldm.trellis2.attention import (
+    sparse_windowed_scaled_dot_product_self_attention, sparse_scaled_dot_product_attention, scaled_dot_product_attention
+)
+from comfy.ldm.genmo.joint_model.layers import TimestepEmbedder
+from comfy.ldm.flux.math import apply_rope, apply_rope1
+
+class SparseGELU(nn.GELU):
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
+
+class SparseFeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            SparseLinear(channels, int(channels * mlp_ratio), device=device, dtype=dtype, operations=operations),
+            SparseGELU(approximate="tanh"),
+            SparseLinear(int(channels * mlp_ratio), channels, device=device, dtype=dtype, operations=operations),
+        )
+
+    def forward(self, x: VarLenTensor) -> VarLenTensor:
+        return self.mlp(x)
+
+def manual_cast(obj, dtype):
+    return obj.to(dtype=dtype)
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int, device, dtype):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim, device=device, dtype=dtype))
+
+    def forward(self, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        x_type = x.dtype
+        x = x.float()
+        if isinstance(x, VarLenTensor):
+            x = x.replace(F.normalize(x.feats, dim=-1) * self.gamma * self.scale)
+        else:
+            x = F.normalize(x, dim=-1) * self.gamma * self.scale
+        return x.to(x_type)
+
+class SparseRotaryPositionEmbedder(nn.Module):
+    def __init__(
+        self,
+        head_dim: int,
+        dim: int = 3,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        device=None
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.dim = dim
+        self.rope_freq = rope_freq
+        self.freq_dim = head_dim // 2 // dim
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32, device=device) / self.freq_dim
+        self.freqs = rope_freq[0] / (rope_freq[1] ** (self.freqs))
+
+    def _get_freqs_cis(self, coords: torch.Tensor) -> torch.Tensor:
+        phases_list = []
+        for i in range(self.dim):
+            phases_list.append(torch.outer(coords[..., i], self.freqs.to(coords.device)))
+
+        phases = torch.cat(phases_list, dim=-1)
+
+        if phases.shape[-1] < self.head_dim // 2:
+            padn = self.head_dim // 2 - phases.shape[-1]
+            phases = torch.cat([phases, torch.zeros(*phases.shape[:-1], padn, device=phases.device)], dim=-1)
+
+        cos = torch.cos(phases)
+        sin = torch.sin(phases)
+
+        f_cis_0 = torch.stack([cos, sin], dim=-1)
+        f_cis_1 = torch.stack([-sin, cos], dim=-1)
+        freqs_cis = torch.stack([f_cis_0, f_cis_1], dim=-1)
+
+        return freqs_cis
+
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+
+    def forward(self, q, k=None):
+        cache_name = f'rope_cis_{self.dim}d_f{self.rope_freq[1]}_hd{self.head_dim}'
+        freqs_cis = q.get_spatial_cache(cache_name)
+
+        if freqs_cis is None:
+            coords = q.coords[..., 1:].to(torch.float32)
+            freqs_cis = self._get_freqs_cis(coords)
+            q.register_spatial_cache(cache_name, freqs_cis)
+
+        if q.feats.ndim == 3:
+            f_cis = freqs_cis.unsqueeze(1)
+        else:
+            f_cis = freqs_cis
+
+        if k is None:
+            return q.replace(apply_rope1(q.feats, f_cis))
+
+        q_feats, k_feats = apply_rope(q.feats, k.feats, f_cis)
+        return q.replace(q_feats), k.replace(k_feats)
+
+    @staticmethod
+    def apply_rotary_embedding(x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases.unsqueeze(-2)
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+
+class RotaryPositionEmbedder(SparseRotaryPositionEmbedder):
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+        if torch.is_complex(phases):
+            phases = phases.to(torch.complex64)
+        else:
+            phases = phases.to(torch.float32)
+        if phases.shape[-1] < self.head_dim // 2:
+                padn = self.head_dim // 2 - phases.shape[-1]
+                phases = torch.cat([phases, torch.polar(
+                    torch.ones(*phases.shape[:-1], padn, device=phases.device, dtype=torch.float32),
+                    torch.zeros(*phases.shape[:-1], padn, device=phases.device, dtype=torch.float32)
+                )], dim=-1)
+        return phases
+
+class SparseMultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int] = None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed", "double_windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = operations.Linear(channels, channels * 3, bias=qkv_bias, device=device, dtype=dtype)
+        else:
+            self.to_q = operations.Linear(channels, channels, bias=qkv_bias, device=device, dtype=dtype)
+            self.to_kv = operations.Linear(self.ctx_channels, channels * 2, bias=qkv_bias, device=device, dtype=dtype)
+
+        if self.qk_rms_norm:
+            self.q_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+            self.k_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+
+        self.to_out = operations.Linear(channels, channels, device=device, dtype=dtype)
+
+        if use_rope:
+            self.rope = SparseRotaryPositionEmbedder(self.head_dim, rope_freq=rope_freq, device=device)
+
+    @staticmethod
+    def _linear(module: nn.Linear, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.replace(module(x.feats))
+        else:
+            return module(x)
+
+    @staticmethod
+    def _reshape_chs(x: Union[VarLenTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.reshape(*shape)
+        else:
+            return x.reshape(*x.shape[:2], *shape)
+
+    def _fused_pre(self, x: Union[VarLenTensor, torch.Tensor], num_fused: int) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            x_feats = x.feats.unsqueeze(0)
+        else:
+            x_feats = x
+        x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+        return x.replace(x_feats.squeeze(0)) if isinstance(x, VarLenTensor) else x_feats
+
+    def forward(self, x: SparseTensor, context: Optional[Union[VarLenTensor, torch.Tensor]] = None) -> SparseTensor:
+        if self._type == "self":
+            dtype = next(self.to_qkv.parameters()).dtype
+            x = x.to(dtype)
+            qkv = self._linear(self.to_qkv, x)
+            qkv = self._fused_pre(qkv, num_fused=3)
+            if self.qk_rms_norm or self.use_rope:
+                q, k, v = qkv.unbind(dim=-3)
+                if self.qk_rms_norm:
+                    q = self.q_rms_norm(q)
+                    k = self.k_rms_norm(k)
+                if self.use_rope:
+                    q, k = self.rope(q, k)
+                qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+            if self.attn_mode == "full":
+                h = sparse_scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                h = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv, self.window_size, shift_window=self.shift_window
+                )
+            elif self.attn_mode == "double_windowed":
+                qkv0 = qkv.replace(qkv.feats[:, :, self.num_heads//2:])
+                qkv1 = qkv.replace(qkv.feats[:, :, :self.num_heads//2])
+                h0 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv0, self.window_size, shift_window=(0, 0, 0)
+                )
+                h1 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv1, self.window_size, shift_window=tuple([self.window_size//2] * 3)
+                )
+                h = qkv.replace(torch.cat([h0.feats, h1.feats], dim=1))
+        else:
+            q = self._linear(self.to_q, x)
+            q = self._reshape_chs(q, (self.num_heads, -1))
+            dtype = next(self.to_kv.parameters()).dtype
+            context = context.to(dtype)
+            kv = self._linear(self.to_kv, context)
+            kv = self._fused_pre(kv, num_fused=2)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=-3)
+                k = self.k_rms_norm(k)
+                h = sparse_scaled_dot_product_attention(q, k, v)
+            else:
+                h = sparse_scaled_dot_product_attention(q, kv)
+        h = self._reshape_chs(h, (-1,))
+        h = self._linear(self.to_out, h)
+        return h
+
+class ModulatedSparseTransformerCrossBlock(nn.Module):
+    """
+    Sparse Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, device=device)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.self_attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.cross_attn = SparseMultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+            device=device, dtype=dtype, operations=operations
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(channels, 6 * channels, bias=True, device=device, dtype=dtype)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels, device=device, dtype=dtype) / channels ** 0.5)
+
+    def _forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = x.replace(self.norm1(x.feats))
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.self_attn(h)
+        h = h * gate_msa
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = x.replace(self.norm3(x.feats))
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        return self._forward(x, mod, context)
+
+
+class SLatFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "rope",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        dtype = None,
+        device = None,
+        operations = None,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = dtype
+
+        self.t_embedder = TimestepEmbedder(model_channels, device=device, dtype=dtype, operations=operations)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(model_channels, 6 * model_channels, bias=True, device=device, dtype=dtype)
+            )
+
+        self.input_layer = SparseLinear(in_channels, model_channels, device=device, dtype=dtype, operations=operations)
+
+        self.blocks = nn.ModuleList([
+            ModulatedSparseTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=self.share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                device=device, dtype=dtype, operations=operations
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = SparseLinear(model_channels, out_channels, device=device, dtype=dtype, operations=operations)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        x: SparseTensor,
+        t: torch.Tensor,
+        cond: Union[torch.Tensor, List[torch.Tensor]],
+        concat_cond: Optional[SparseTensor] = None,
+        **kwargs
+    ) -> SparseTensor:
+        if concat_cond is not None:
+            x = sparse_cat([x, concat_cond], dim=-1)
+        if isinstance(cond, list):
+            cond = VarLenTensor.from_tensor_list(cond)
+
+        dtype = next(self.input_layer.parameters()).dtype
+        x = x.to(dtype)
+        h = self.input_layer(x)
+        h = manual_cast(h, self.dtype)
+        t = t.to(dtype)
+        t_embedder = self.t_embedder.to(dtype)
+        t_emb = t_embedder(t, out_dtype = t.dtype)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+
+        h = manual_cast(h, x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return h
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            operations.Linear(channels, int(channels * mlp_ratio), device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(int(channels * mlp_ratio), channels, device=device, dtype=dtype),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int, device=None, dtype=None):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim, device=device, dtype=dtype))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int]=None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = operations.Linear(channels, channels * 3, bias=qkv_bias, dtype=dtype, device=device)
+        else:
+            self.to_q = operations.Linear(channels, channels, bias=qkv_bias, device=device, dtype=dtype)
+            self.to_kv = operations.Linear(self.ctx_channels, channels * 2, bias=qkv_bias, device=device, dtype=dtype)
+
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads, device=device, dtype=dtype)
+
+        self.to_out = operations.Linear(channels, channels, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self._type == "self":
+            x = x.to(next(self.to_qkv.parameters()).dtype)
+            qkv = self.to_qkv(x)
+            qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+
+            if self.attn_mode == "full":
+                if self.qk_rms_norm or self.use_rope:
+                    q, k, v = qkv.unbind(dim=2)
+                    if self.qk_rms_norm:
+                        q = self.q_rms_norm(q)
+                        k = self.k_rms_norm(k)
+                    if self.use_rope:
+                        assert phases is not None, "Phases must be provided for RoPE"
+                        q = RotaryPositionEmbedder.apply_rotary_embedding(q, phases)
+                        k = RotaryPositionEmbedder.apply_rotary_embedding(k, phases)
+                    h = scaled_dot_product_attention(q, k, v)
+                else:
+                    h = scaled_dot_product_attention(qkv)
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x)
+            context = context.to(next(self.to_kv.parameters()).dtype)
+            kv = self.to_kv(context)
+            q = q.reshape(B, L, self.num_heads, -1)
+            kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q, k, v)
+            else:
+                h = scaled_dot_product_attention(q, kv)
+        h = h.reshape(B, L, -1)
+        h = self.to_out(h)
+        return h
+
+class ModulatedTransformerCrossBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, device=device)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6, device=device)
+        self.self_attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.cross_attn = MultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+            device=device, dtype=dtype, operations=operations
+        )
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+            device=device, dtype=dtype, operations=operations
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(channels, 6 * channels, bias=True, dtype=dtype, device=device)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels, device=device, dtype=dtype) / channels ** 0.5)
+
+    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = self.norm1(x)
+        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        h = self.self_attn(h, phases=phases)
+        h = h * gate_msa.unsqueeze(1)
+        x = x + h
+        h = self.norm2(x)
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = self.norm3(x)
+        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        h = self.mlp(h)
+        h = h * gate_mlp.unsqueeze(1)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return self._forward(x, mod, context, phases)
+
+
+class SparseStructureFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "rope",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        operations=None,
+        device = None,
+        dtype = torch.float32,
+        **kwargs
+    ):
+        super().__init__()
+        self.device = device
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = dtype
+        self.device = device
+
+        self.t_embedder = TimestepEmbedder(model_channels, dtype=dtype, device=device, operations=operations)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(model_channels, 6 * model_channels, bias=True, device=device, dtype=dtype)
+            )
+
+        pos_embedder = RotaryPositionEmbedder(self.model_channels // self.num_heads, 3, device=device)
+        coords = torch.meshgrid(*[torch.arange(res, device=self.device, dtype=dtype) for res in [resolution] * 3], indexing='ij')
+        coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+        rope_phases = pos_embedder(coords)
+        self.register_buffer("rope_phases", rope_phases, persistent=False)
+
+        if pe_mode != "rope":
+            self.rope_phases = None
+
+        self.input_layer = operations.Linear(in_channels, model_channels, device=device, dtype=dtype)
+
+        self.blocks = nn.ModuleList([
+            ModulatedTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                device=device, dtype=dtype, operations=operations
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = operations.Linear(model_channels, out_channels, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        x = x.view(x.shape[0], self.in_channels, *[self.resolution] * 3)
+
+        h = x.view(*x.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+        h = h.to(next(self.input_layer.parameters()).dtype)
+        h = self.input_layer(h)
+        t_emb = self.t_embedder(t, out_dtype = t.dtype)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        h = manual_cast(h, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond, self.rope_phases)
+        h = manual_cast(h, x.dtype)
+        h = F.layer_norm(h, h.shape[-1:])
+        h = h.to(next(self.out_layer.parameters()).dtype)
+        h = self.out_layer(h)
+
+        h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution] * 3).contiguous()
+
+        return h
+
+def timestep_reshift(t_shifted, old_shift=3.0, new_shift=5.0):
+    t_shifted = t_shifted / 1000.0
+    t_linear = t_shifted / (old_shift - t_shifted * (old_shift - 1))
+    t_new = (new_shift * t_linear) / (1 + (new_shift - 1) * t_linear)
+    t_new *= 1000.0
+    return t_new
+
+class Trellis2(nn.Module):
+    def __init__(self, resolution,
+                 in_channels = 32,
+                 out_channels = 32,
+                 model_channels = 1536,
+                 cond_channels = 1024,
+                 num_blocks = 30,
+                 num_heads = 12,
+                 mlp_ratio = 5.3334,
+                 share_mod = True,
+                 qk_rms_norm = True,
+                 qk_rms_norm_cross = True,
+                 init_txt_model=False, # for now
+                 dtype=None, device=None, operations=None, **kwargs):
+
+        super().__init__()
+        self.dtype = dtype
+        operations = operations or nn
+        # for some reason it passes num_heads = -1
+        if num_heads == -1:
+            num_heads = 12
+        args = {
+            "out_channels":out_channels, "num_blocks":num_blocks, "cond_channels" :cond_channels,
+            "model_channels":model_channels, "num_heads":num_heads, "mlp_ratio": mlp_ratio, "share_mod": share_mod,
+            "qk_rms_norm": qk_rms_norm, "qk_rms_norm_cross": qk_rms_norm_cross, "device": device, "dtype": dtype, "operations": operations
+        }
+        txt_only = kwargs.get("txt_only", False)
+        if not txt_only:
+            self.img2shape = SLatFlowModel(resolution=resolution, in_channels=in_channels, **args)
+            self.shape2txt = None
+            if init_txt_model:
+                self.shape2txt = SLatFlowModel(resolution=resolution, in_channels=in_channels*2, **args)
+            self.img2shape_512 = SLatFlowModel(resolution=32, in_channels=in_channels, **args)
+            args.pop("out_channels")
+            self.structure_model = SparseStructureFlowModel(resolution=16, in_channels=8, out_channels=8, **args)
+        else:
+            self.shape2txt = SLatFlowModel(resolution=resolution, in_channels=in_channels*2, **args)
+        self.guidance_interval = [0.6, 1.0]
+        self.guidance_interval_txt = [0.6, 0.9]
+
+    def forward(self, x, timestep, context, **kwargs):
+        transformer_options = kwargs.get("transformer_options", {})
+        model_options = {}
+        if hasattr(self, "meta"):
+            model_options = self.meta
+        timestep = timestep.to(x.dtype)
+        embeds = kwargs.get("embeds")
+        if embeds is None:
+            raise ValueError("Trellis2.forward requires 'embeds' in kwargs")
+
+        is_1024 = True#self.img2shape.resolution == 1024
+        coords = model_options.get("coords", None)
+        coord_counts = model_options.get("coord_counts", None)
+        mode = model_options.get("generation_mode", "structure_generation")
+
+        is_512_run = False
+        if mode == "shape_generation_512":
+            is_512_run = True
+            mode = "shape_generation"
+
+        if coords is not None:
+            if x.ndim == 4:
+                x = x.squeeze(-1).transpose(1, 2)
+            not_struct_mode = True
+        else:
+            mode = "structure_generation"
+            not_struct_mode = False
+
+        if x.size(-1) == 16 and x.size(-2) == 16:
+            mode = "structure_generation"
+            not_struct_mode = False
+
+        if not not_struct_mode:
+            bsz = x.size(0)
+            x = x[:, :8]
+            x = x.view(bsz, 8, 16, 16, 16)
+
+        if is_1024 and not_struct_mode and not is_512_run:
+            context = embeds
+
+        sigmas = transformer_options.get("sigmas")[0].item()
+        if sigmas < 1.00001:
+            timestep *= 1000.0
+
+        if context.size(0) > 1:
+            cond = context.chunk(2)[1]
+        else:
+            cond = context
+
+        shape_rule = sigmas < self.guidance_interval[0] or sigmas > self.guidance_interval[1]
+        txt_rule = sigmas < self.guidance_interval_txt[0] or sigmas > self.guidance_interval_txt[1]
+
+        if not_struct_mode:
+            orig_bsz = x.shape[0]
+            rule = txt_rule if mode == "texture_generation" else shape_rule
+
+            # CFG Bypass Slicing
+            if rule and orig_bsz > 1:
+                half = orig_bsz // 2
+                x_eval = x[half:]
+                t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep
+                c_eval = cond
+            else:
+                x_eval = x
+                t_eval = timestep
+                c_eval = context
+
+            B, N, C = x_eval.shape
+
+            # Vectorized SparseTensor Construction
+            if mode in ["shape_generation", "texture_generation"]:
+                if coord_counts is not None:
+                    logical_batch = coord_counts.shape[0]
+                    # Duplicate coords if CFG is active
+                    if B > logical_batch:
+                        c_pos = coords.clone()
+                        c_pos[:, 0] += logical_batch
+                        batched_coords = torch.cat([coords, c_pos], dim=0)
+                        counts_eval = torch.cat([coord_counts, coord_counts], dim=0)
+                    else:
+                        batched_coords = coords
+                        counts_eval = coord_counts
+
+                    # Create boolean mask [B, N] to drop the padded zeros instantly
+                    mask = torch.arange(N, device=x.device).unsqueeze(0) < counts_eval.unsqueeze(1)
+                    feats_flat = x_eval[mask]
+                else:
+                    feats_flat = x_eval.reshape(-1, C)
+                    coords_list =[]
+                    for i in range(B):
+                        c = coords.clone()
+                        c[:, 0] = i
+                        coords_list.append(c)
+                    batched_coords = torch.cat(coords_list, dim=0)
+                    mask = None
+            else:
+                batched_coords = coords
+                feats_flat = x_eval
+                mask = None
+
+            x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32))
+
+        if mode == "shape_generation":
+            if is_512_run:
+                out = self.img2shape_512(x_st, t_eval, c_eval)
+            else:
+                out = self.img2shape(x_st, t_eval, c_eval)
+
+        elif mode == "texture_generation":
+            if self.shape2txt is None:
+                raise ValueError("Checkpoint for Trellis2 doesn't include texture generation!")
+            slat = model_options.get("shape_slat")
+            if slat is None:
+                raise ValueError("shape_slat can't be None")
+
+            slat_feats = slat
+            # Duplicate shape context if CFG is active
+            if coord_counts is not None and B > coord_counts.shape[0]:
+                slat_feats = torch.cat([slat_feats, slat_feats], dim=0)
+            elif coord_counts is None:
+                slat_feats = slat_feats[:N].repeat(B, 1)
+
+            x_st = x_st.replace(feats=torch.cat([x_st.feats, slat_feats.to(x_st.feats.device)], dim=-1))
+            out = self.shape2txt(x_st, t_eval, c_eval)
+
+        else: # structure
+            orig_bsz = x.shape[0]
+            if shape_rule and orig_bsz > 1:
+                half = orig_bsz // 2
+                x_eval = x[half:]
+                t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep
+                out = self.structure_model(x_eval, t_eval, cond)
+                out = out.repeat(2, 1, 1, 1, 1)
+            else:
+                out = self.structure_model(x, timestep, context)
+
+        if not_struct_mode:
+            if mask is not None:
+                # Instantly scatter the valid tokens back into a padded rectangular tensor
+                padded_out = torch.zeros((B, N, out.feats.shape[-1]), device=x.device, dtype=out.feats.dtype)
+                padded_out[mask] = out.feats
+                out_tensor = padded_out.transpose(1, 2).unsqueeze(-1)
+            else:
+                out_tensor = out.feats.view(B, N, -1).transpose(1, 2).unsqueeze(-1)
+
+            if rule and orig_bsz > 1:
+                out_tensor = out_tensor.repeat(2, 1, 1, 1)
+            return out_tensor
+        else:
+            out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, 0, 0, 24))
+
+        return out

comfy/model_base.py

@@ -53,6 +53,7 @@ import comfy.ldm.omnigen.omnigen2
 import comfy.ldm.qwen_image.model
 import comfy.ldm.kandinsky5.model
 import comfy.ldm.anima.model
+import comfy.ldm.trellis2.model
 import comfy.ldm.ace.ace_step15
 import comfy.ldm.cogvideo.model
 import comfy.ldm.rt_detr.rtdetr_v4
@@ -1637,6 +1638,16 @@ class WAN22(WAN21):
     def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
         return latent_image
 
+class Trellis2(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None, unet_model=comfy.ldm.trellis2.model.Trellis2):
+        super().__init__(model_config, model_type, device, unet_model)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        embeds = kwargs.get("embeds")
+        out["embeds"] = comfy.conds.CONDRegular(embeds)
+        return out
+
 class WAN21_FlowRVS(WAN21):
     def __init__(self, model_config, model_type=ModelType.IMG_TO_IMG_FLOW, image_to_video=False, device=None):
         model_config.unet_config["model_type"] = "t2v"
@@ -1678,7 +1689,6 @@ class WAN21_SCAIL(WAN21):
         pose_latents = kwargs.get("pose_video_latent", None)
         if pose_latents is not None:
             out['pose_latents'] = [pose_latents.shape[0], 20, *pose_latents.shape[2:]]
-
         return out
 
 class WAN22_WanDancer(WAN21):

comfy/model_detection.py

@@ -113,6 +113,30 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]
         return unet_config
 
+    if '{}img2shape.blocks.1.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys:
+        unet_config = {}
+        unet_config["image_model"] = "trellis2"
+
+        unet_config["init_txt_model"] = False
+        if '{}shape2txt.blocks.29.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys:
+            unet_config["init_txt_model"] = True
+
+        unet_config["resolution"] = 64
+        if metadata is not None:
+            if "is_512" in metadata:
+                unet_config["resolution"] = 32
+
+        unet_config["num_heads"] = 12
+        return unet_config
+
+    if '{}shape2txt.blocks.29.cross_attn.k_rms_norm.gamma'.format(key_prefix) in state_dict_keys: # trellis2 texture
+        unet_config = {}
+        unet_config["image_model"] = "trellis2"
+        unet_config["resolution"] = 64
+        unet_config["num_heads"] = 12
+        unet_config["txt_only"] = True
+        return unet_config
+
     if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit
         unet_config = {}
         unet_config["audio_model"] = "dit1.0"

comfy/sd.py

@@ -15,6 +15,7 @@ import comfy.ldm.lightricks.vae.causal_video_autoencoder
 import comfy.ldm.lightricks.vae.audio_vae
 import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
+import comfy.ldm.trellis2.vae
 import comfy.ldm.wan.vae2_2
 import comfy.ldm.hunyuan3d.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
@@ -528,6 +529,18 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
+            elif "shape_dec.blocks.1.16.to_subdiv.weight" in sd or "txt_dec.blocks.3.4.conv2.weight" in sd: # trellis2 or trellis2 texture only
+                init_txt_model = False
+                init_txt_model_only = False
+                if "shape_dec.blocks.1.16.to_subdiv.weight" not in sd:
+                    init_txt_model_only = True
+                if "txt_dec.blocks.1.16.norm1.weight" in sd:
+                    init_txt_model = True
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                # TODO
+                self.memory_used_decode = lambda shape, dtype: (2500 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (2500 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.first_stage_model = comfy.ldm.trellis2.vae.Vae(init_txt_model, init_txt_model_only= init_txt_model_only)
             elif "decoder.conv_in.weight" in sd:
                 if sd['decoder.conv_in.weight'].shape[1] == 64:
                     ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}

comfy/supported_models.py

@@ -1318,6 +1318,29 @@ class WAN22_T2V(WAN21_T2V):
         out = model_base.WAN22(self, image_to_video=True, device=device)
         return out
 
+class Trellis2(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "trellis2"
+    }
+
+    sampling_settings = {
+        "shift": 3.0,
+    }
+
+    memory_usage_factor = 3.5
+
+    latent_format = latent_formats.Trellis2
+    vae_key_prefix = ["vae."]
+    clip_vision_prefix = "conditioner.main_image_encoder.model."
+    # this is only needed for the texture model
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.Trellis2(self, device=device)
+
+    def clip_target(self, state_dict={}):
+        return None
+
 class WAN21_FlowRVS(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",
@@ -1784,6 +1807,7 @@ class Kandinsky5Image(Kandinsky5):
         return supported_models_base.ClipTarget(comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage, comfy.text_encoders.kandinsky5.te(**hunyuan_detect))
 
 
+
 class ACEStep15(supported_models_base.BASE):
     unet_config = {
         "audio_model": "ace1.5",
@@ -1823,7 +1847,6 @@ class ACEStep15(supported_models_base.BASE):
 
         return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
 
-
 class LongCatImage(supported_models_base.BASE):
     unet_config = {
         "image_model": "flux",
@@ -1901,6 +1924,7 @@ class ErnieImage(supported_models_base.BASE):
         return supported_models_base.ClipTarget(comfy.text_encoders.ernie.ErnieTokenizer, comfy.text_encoders.ernie.te(**hunyuan_detect))
 
 
+
 class SAM3(supported_models_base.BASE):
     unet_config = {"image_model": "SAM3"}
     supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
@@ -2020,7 +2044,6 @@ class CogVideoX_Inpaint(CogVideoX_T2V):
         out = model_base.CogVideoX(self, image_to_video=True, device=device)
         return out
 
-
 models = [
     LotusD,
     Stable_Zero123,
@@ -2107,4 +2130,5 @@ models = [
     CogVideoX_I2V,
     CogVideoX_T2V,
     SVD_img2vid,
+    Trellis2
 ]

comfy_api/latest/_caching.py

@@ -21,10 +21,6 @@ class CacheProvider(ABC):
     Exceptions from provider methods are caught by the caller and never break execution.
     """
 
-    async def on_set_prompt(self) -> None:
-        """Called after prompt cache keys are prepared. Dispatched via asyncio.create_task."""
-        pass
-
     @abstractmethod
     async def on_lookup(self, context: CacheContext) -> Optional[CacheValue]:
         """Called on local cache miss. Return CacheValue if found, None otherwise."""

comfy_api/latest/_util/geometry_types.py

@@ -7,9 +7,10 @@ import torch
 
 
 class VOXEL:
-    def __init__(self, data: torch.Tensor):
+    def __init__(self, data: torch.Tensor, voxel_colors=None, resolution=None):
         self.data = data
-
+        self.voxel_colors = voxel_colors
+        self.resolution = resolution # each 3d model has its own resolution
 
 class MESH:
     def __init__(self, vertices: torch.Tensor, faces: torch.Tensor,

comfy_execution/caching.py

@@ -164,7 +164,6 @@ class BasicCache:
         await self.cache_key_set.add_keys(node_ids)
         self.is_changed_cache = is_changed_cache
         self.initialized = True
-        await self._notify_providers_set_prompt()
 
     def all_node_ids(self):
         assert self.initialized
@@ -264,24 +263,6 @@ class BasicCache:
             except Exception as e:
                 _logger.warning(f"Cache provider {provider.__class__.__name__} error on store: {e}")
 
-    async def _notify_providers_set_prompt(self):
-        from comfy_execution.cache_provider import (
-            _has_cache_providers, _get_cache_providers, _logger
-        )
-
-        if not self.enable_providers:
-            return
-        if not _has_cache_providers():
-            return
-
-        for provider in _get_cache_providers():
-            try:
-                task = asyncio.create_task(self._safe_provider_set_prompt(provider))
-                self._pending_store_tasks.add(task)
-                task.add_done_callback(self._pending_store_tasks.discard)
-            except Exception as e:
-                _logger.warning(f"Cache provider {provider.__class__.__name__} error on set_prompt: {e}")
-
     @staticmethod
     async def _safe_provider_store(provider, context, cache_value):
         from comfy_execution.cache_provider import _logger
@@ -290,14 +271,6 @@ class BasicCache:
         except Exception as e:
             _logger.warning(f"Cache provider {provider.__class__.__name__} async store error: {e}")
 
-    @staticmethod
-    async def _safe_provider_set_prompt(provider):
-        from comfy_execution.cache_provider import _logger
-        try:
-            await provider.on_set_prompt()
-        except Exception as e:
-            _logger.warning(f"Cache provider {provider.__class__.__name__} async set_prompt error: {e}")
-
     async def _check_providers_lookup(self, node_id, cache_key):
         from comfy_execution.cache_provider import (
             _has_cache_providers, _get_cache_providers,

comfy_extras/nodes_audio.py

@@ -543,7 +543,7 @@ class AudioConcat(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioConcat",
             search_aliases=["join audio", "combine audio", "append audio"],
-            display_name="Audio Concat",
+            display_name="Concatenate Audio",
             description="Concatenates the audio1 to audio2 in the specified direction.",
             category="audio",
             inputs=[
@@ -597,7 +597,7 @@ class AudioMerge(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioMerge",
             search_aliases=["mix audio", "overlay audio", "layer audio"],
-            display_name="Audio Merge",
+            display_name="Merge Audio",
             description="Combine two audio tracks by overlaying their waveforms.",
             category="audio",
             inputs=[
@@ -667,8 +667,9 @@ class AudioAdjustVolume(IO.ComfyNode):
         return IO.Schema(
             node_id="AudioAdjustVolume",
             search_aliases=["audio gain", "loudness", "audio level"],
-            display_name="Audio Adjust Volume",
+            display_name="Adjust Audio Volume",
             category="audio",
+            description="Adjust the volume of the audio by a specified amount in decibels (dB).",
             inputs=[
                 IO.Audio.Input("audio"),
                 IO.Int.Input(

comfy_extras/nodes_dataset.py

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

comfy_extras/nodes_hunyuan3d.py

@@ -419,15 +419,17 @@ class VoxelToMeshBasic(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="VoxelToMeshBasic",
-            display_name="Voxel to Mesh (Basic)",
+            display_name="Voxel to Mesh (Basic) (DEPRECATED)",
             category="3d",
+            description="Converts a voxel grid to a mesh.",
+            is_deprecated=True, # This node is superseded by the Voxel To Mesh node
             inputs=[
                 IO.Voxel.Input("voxel"),
                 IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
             ],
             outputs=[
                 IO.Mesh.Output(),
-            ]
+            ],
         )
 
     @classmethod
@@ -453,9 +455,10 @@ class VoxelToMesh(IO.ComfyNode):
             node_id="VoxelToMesh",
             display_name="Voxel to Mesh",
             category="3d",
+            description="Converts a voxel grid to a mesh.",
             inputs=[
                 IO.Voxel.Input("voxel"),
-                IO.Combo.Input("algorithm", options=["surface net", "basic"], advanced=True),
+                IO.Combo.Input("algorithm", options=["surface net", "basic"]),
                 IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
             ],
             outputs=[

comfy_extras/nodes_images.py

@@ -55,9 +55,10 @@ class ImageCropV2(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="ImageCropV2",
-            search_aliases=["trim"],
+            search_aliases=["crop", "cut", "trim"],
             display_name="Crop Image",
             category="image/transform",
+            description = "Crop an image to the specified dimensions.",
             essentials_category="Image Tools",
             has_intermediate_output=True,
             inputs=[

comfy_extras/nodes_lt_audio.py

@@ -11,8 +11,8 @@ class LTXVAudioVAELoader(io.ComfyNode):
     def define_schema(cls) -> io.Schema:
         return io.Schema(
             node_id="LTXVAudioVAELoader",
-            display_name="LTXV Audio VAE Loader",
-            category="audio",
+            display_name="Load LTXV Audio VAE",
+            category="loaders",
             inputs=[
                 io.Combo.Input(
                     "ckpt_name",
@@ -40,7 +40,7 @@ class LTXVAudioVAEEncode(VAEEncodeAudio):
         return io.Schema(
             node_id="LTXVAudioVAEEncode",
             display_name="LTXV Audio VAE Encode",
-            category="audio",
+            category="latent/audio",
             inputs=[
                 io.Audio.Input("audio", tooltip="The audio to be encoded."),
                 io.Vae.Input(
@@ -63,7 +63,7 @@ class LTXVAudioVAEDecode(io.ComfyNode):
         return io.Schema(
             node_id="LTXVAudioVAEDecode",
             display_name="LTXV Audio VAE Decode",
-            category="audio",
+            category="latent/audio",
             inputs=[
                 io.Latent.Input("samples", tooltip="The latent to be decoded."),
                 io.Vae.Input(

comfy_extras/nodes_mediapipe.py

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

comfy_extras/nodes_moge.py

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

comfy_extras/nodes_save_3d.py

@@ -234,6 +234,12 @@ def save_glb(vertices, faces, filepath, metadata=None,
     textures = []
     samplers = []
     materials = []
+    pbr = {
+        "metallicFactor": 0.0,
+        "roughnessFactor": 0.5,
+        "baseColorFactor": [0.22, 0.22, 0.22, 1.0],
+    }
+
     if texture_png_bytes is not None and "TEXCOORD_0" in primitive_attributes:
         buffer_views.append({
             "buffer": 0,
@@ -243,15 +249,13 @@ def save_glb(vertices, faces, filepath, metadata=None,
         images.append({"bufferView": len(buffer_views) - 1, "mimeType": "image/png"})
         samplers.append({"magFilter": 9729, "minFilter": 9729, "wrapS": 33071, "wrapT": 33071})
         textures.append({"source": 0, "sampler": 0})
-        materials.append({
-            "pbrMetallicRoughness": {
-                "baseColorTexture": {"index": 0, "texCoord": 0},
-                "metallicFactor": 0.0,
-                "roughnessFactor": 1.0,
-            },
-            "doubleSided": True,
-        })
-        primitive["material"] = 0
+        pbr["baseColorTexture"] = {"index": 0, "texCoord": 0}
+
+    materials.append({
+        "pbrMetallicRoughness": pbr,
+        "doubleSided": True,
+    })
+    primitive["material"] = 0
 
     gltf = {
         "asset": {"version": "2.0", "generator": "ComfyUI"},
@@ -373,10 +377,14 @@ class SaveGLB(IO.ComfyNode):
                     continue
                 tex_img = Image.fromarray(texture_np[i], mode="RGB") if texture_np is not None else None
                 f = f"{filename}_{counter:05}_.glb"
-                save_glb(vertices_i, faces_i, os.path.join(full_output_folder, f), metadata,
-                         uvs=uvs_i,
-                         vertex_colors=v_colors,
-                         texture_image=tex_img)
+                save_glb(
+                    vertices_i, faces_i,
+                    os.path.join(full_output_folder, f),
+                    metadata,
+                    uvs=uvs_i,
+                    vertex_colors=v_colors,
+                    texture_image=tex_img,
+                )
                 results.append({
                     "filename": f,
                     "subfolder": subfolder,

comfy_extras/nodes_trellis2.py

@@ -0,0 +1,712 @@
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, IO, Types, io
+from comfy.ldm.trellis2.vae import SparseTensor
+from comfy_extras.nodes_mesh_postprocess import pack_variable_mesh_batch
+import comfy.model_management
+from PIL import Image
+import numpy as np
+import torch
+
+ShapeSubdivides = io.Custom("SHAPE_SUBDIVIDES")
+
+def prepare_trellis_vae_for_decode(vae, sample_shape):
+    memory_required = vae.memory_used_decode(sample_shape, vae.vae_dtype)
+    if len(sample_shape) == 5:
+        memory_required *= max(1, int(sample_shape[4]))
+    memory_required = max(1, int(memory_required))
+    device = comfy.model_management.get_torch_device()
+    comfy.model_management.load_models_gpu(
+        [vae.patcher],
+        memory_required=memory_required,
+        force_full_load=getattr(vae, "disable_offload", False),
+    )
+    free_memory = vae.patcher.get_free_memory(device)
+    batch_number = max(1, int(free_memory / memory_required))
+    return batch_number
+
+shape_slat_normalization = {
+    "mean": torch.tensor([
+        0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+        -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+        0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+        -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+    ])[None],
+    "std": torch.tensor([
+        5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+        5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+        4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+        5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+    ])[None]
+}
+
+tex_slat_normalization = {
+    "mean": torch.tensor([
+        3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+        0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+        -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+        1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+    ])[None],
+    "std": torch.tensor([
+        2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+        2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+        2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+        2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+    ])[None]
+}
+
+def shape_norm(shape_latent, coords):
+    std = shape_slat_normalization["std"].to(shape_latent)
+    mean = shape_slat_normalization["mean"].to(shape_latent)
+    samples = SparseTensor(feats = shape_latent, coords=coords)
+    samples = samples * std + mean
+    return samples
+
+
+def infer_batched_coord_layout(coords):
+    if coords.ndim != 2 or coords.shape[1] != 4:
+        raise ValueError(f"Expected Trellis2 coords with shape [N, 4], got {tuple(coords.shape)}")
+
+    if coords.shape[0] == 0:
+        raise ValueError("Trellis2 coords can't be empty")
+
+    batch_ids = coords[:, 0].to(torch.int64)
+    if (batch_ids < 0).any():
+        raise ValueError(f"Trellis2 batch ids must be non-negative, got {batch_ids.unique(sorted=True).tolist()}")
+    batch_size = int(batch_ids.max().item()) + 1
+    counts = torch.bincount(batch_ids, minlength=batch_size)
+
+    if (counts == 0).any():
+        raise ValueError(f"Non-contiguous Trellis2 batch ids in coords: {batch_ids.unique(sorted=True).tolist()}")
+
+    max_tokens = int(counts.max().item())
+    return batch_size, counts, max_tokens
+
+
+def split_batched_coords(coords, coord_counts):
+    if coord_counts.ndim != 1:
+        raise ValueError(f"Trellis2 coord_counts must be 1D, got shape {tuple(coord_counts.shape)}")
+    if (coord_counts < 0).any():
+        raise ValueError(f"Trellis2 coord_counts must be non-negative, got {coord_counts.tolist()}")
+    if int(coord_counts.sum().item()) != coords.shape[0]:
+        raise ValueError(
+            f"Trellis2 coord_counts total {int(coord_counts.sum().item())} does not match coords rows {coords.shape[0]}"
+        )
+
+    batch_ids = coords[:, 0].to(torch.int64)
+    order = torch.argsort(batch_ids, stable=True)
+    sorted_coords = coords.index_select(0, order)
+    sorted_batch_ids = batch_ids.index_select(0, order)
+
+    offsets = coord_counts.cumsum(0) - coord_counts
+    items = []
+    for i in range(coord_counts.shape[0]):
+        count = int(coord_counts[i].item())
+        start = int(offsets[i].item())
+        coords_i = sorted_coords[start:start + count]
+        ids_i = sorted_batch_ids[start:start + count]
+        if coords_i.shape[0] != count or not torch.all(ids_i == i):
+            raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}")
+        items.append(coords_i)
+    return items
+
+def flatten_batched_sparse_latent(samples, coords, coord_counts):
+    samples = samples.squeeze(-1).transpose(1, 2)
+    if coord_counts is None:
+        return samples.reshape(-1, samples.shape[-1]), coords
+
+    coords_items = split_batched_coords(coords, coord_counts)
+    feat_list = []
+    coord_list = []
+    for i, coords_i in enumerate(coords_items):
+        count = int(coord_counts[i].item())
+        feat_list.append(samples[i, :count])
+        coord_list.append(coords_i)
+
+    return torch.cat(feat_list, dim=0), torch.cat(coord_list, dim=0)
+
+
+def split_batched_sparse_latent(samples, coords, coord_counts):
+    samples = samples.squeeze(-1).transpose(1, 2)
+    if coord_counts is None:
+        return [(samples.reshape(-1, samples.shape[-1]), coords)]
+
+    coords_items = split_batched_coords(coords, coord_counts)
+    items = []
+    for i, coords_i in enumerate(coords_items):
+        count = int(coord_counts[i].item())
+        items.append((samples[i, :count], coords_i))
+    return items
+
+class VaeDecodeShapeTrellis(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeShapeTrellis",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+            ],
+            outputs=[
+                IO.Mesh.Output("mesh"),
+                ShapeSubdivides.Output(display_name = "shape_subdivides"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae):
+
+        resolution = int(vae.first_stage_model.resolution.item())
+        sample_tensor = samples["samples"]
+        device = comfy.model_management.get_torch_device()
+        coords = samples["coords"]
+        prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        trellis_vae = vae.first_stage_model
+        coord_counts = samples.get("coord_counts")
+
+        samples = samples["samples"]
+        if coord_counts is None:
+            samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
+            samples = shape_norm(samples.to(device), coords.to(device))
+            mesh, subs = trellis_vae.decode_shape_slat(samples, resolution)
+        else:
+            split_items = split_batched_sparse_latent(samples, coords, coord_counts)
+            mesh = []
+            subs_per_sample = []
+            for feats_i, coords_i in split_items:
+                coords_i = coords_i.to(device).clone()
+                coords_i[:, 0] = 0
+                sample_i = shape_norm(feats_i.to(device), coords_i)
+                mesh_i, subs_i = trellis_vae.decode_shape_slat(sample_i, resolution)
+                mesh.append(mesh_i[0])
+                subs_per_sample.append(subs_i)
+
+            subs = []
+            for stage_index in range(len(subs_per_sample[0])):
+                stage_tensors = [sample_subs[stage_index] for sample_subs in subs_per_sample]
+                feats_list = [stage_tensor.feats for stage_tensor in stage_tensors]
+                coords_list = [stage_tensor.coords for stage_tensor in stage_tensors]
+                subs.append(SparseTensor.from_tensor_list(feats_list, coords_list))
+
+        face_list = [m.faces for m in mesh]
+        vert_list = [m.vertices for m in mesh]
+        if all(v.shape == vert_list[0].shape for v in vert_list) and all(f.shape == face_list[0].shape for f in face_list):
+            mesh = Types.MESH(vertices=torch.stack(vert_list), faces=torch.stack(face_list))
+        else:
+            mesh = pack_variable_mesh_batch(vert_list, face_list)
+        return IO.NodeOutput(mesh, subs)
+
+class VaeDecodeTextureTrellis(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeTextureTrellis",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+                ShapeSubdivides.Input("shape_subdivides",
+                                 tooltip=(
+                                    "Shape information used to guide higher-detail reconstruction during decoding. "
+                                    "Helps preserve structure consistency at higher resolutions."
+                                 )),
+            ],
+            outputs=[
+                IO.Voxel.Output("voxel_colors"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae, shape_subdivides):
+        sample_tensor = samples["samples"]
+        device = comfy.model_management.get_torch_device()
+        coords = samples["coords"]
+        prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        trellis_vae = vae.first_stage_model
+        coord_counts = samples.get("coord_counts")
+
+        samples = samples["samples"]
+        samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
+        samples = samples.to(device)
+        std = tex_slat_normalization["std"].to(samples)
+        mean = tex_slat_normalization["mean"].to(samples)
+        samples = SparseTensor(feats = samples, coords=coords.to(device))
+        samples = samples * std + mean
+
+        voxel = trellis_vae.decode_tex_slat(samples, shape_subdivides)
+        color_feats = voxel.feats[:, :3]
+        voxel_coords = voxel.coords#[:, 1:]
+
+        voxel = Types.VOXEL(voxel_coords, color_feats, 1024)
+        return IO.NodeOutput(voxel)
+
+class VaeDecodeStructureTrellis2(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="VaeDecodeStructureTrellis2",
+            category="latent/3d",
+            inputs=[
+                IO.Latent.Input("samples"),
+                IO.Vae.Input("vae"),
+                IO.Combo.Input("resolution", options=["32", "64"], default="32")
+            ],
+            outputs=[
+                IO.Voxel.Output("voxel"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, samples, vae, resolution):
+        resolution = int(resolution)
+        sample_tensor = samples["samples"]
+        sample_tensor = sample_tensor[:, :8]
+        batch_number = prepare_trellis_vae_for_decode(vae, sample_tensor.shape)
+        decoder = vae.first_stage_model.struct_dec
+        load_device = comfy.model_management.get_torch_device()
+        decoded_batches = []
+        for start in range(0, sample_tensor.shape[0], batch_number):
+            sample_chunk = sample_tensor[start:start + batch_number].to(load_device)
+            decoded_batches.append(decoder(sample_chunk) > 0)
+        decoded = torch.cat(decoded_batches, dim=0)
+        current_res = decoded.shape[2]
+
+        if current_res != resolution:
+            ratio = current_res // resolution
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        out = Types.VOXEL(decoded.squeeze(1).float())
+        return IO.NodeOutput(out)
+
+class Trellis2UpsampleCascade(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2UpsampleCascade",
+            category="latent/3d",
+            display_name="Trellis2 Upsample Cascade",
+            description="Upsamples low-resolution Trellis2 shape latents into higher resolution coordinates while respecting the maximum token budget.",
+            inputs=[
+                IO.Latent.Input("shape_latent"),
+                IO.Vae.Input("vae"),
+                IO.Combo.Input("target_resolution", options=["1024", "1536"], default="1024", tooltip="Controls output detail level for upsampling."),
+                IO.Int.Input("max_tokens", default=49152, min=1024, max=100000,
+                             tooltip=(
+                                "Maximum number of output elements (coordinates) allowed after upsampling. "
+                                "Used to limit memory usage and control mesh density."
+                            ))
+            ],
+            outputs=[
+                IO.Voxel.Output(
+                "high_res_voxel",
+                tooltip=(
+                    "High-resolution sparse coordinates produced after cascade upsampling. "
+                    "Represents the refined 3D structure at target resolution."
+                )
+            )
+            ]
+        )
+
+    @classmethod
+    def execute(cls, shape_latent, vae, target_resolution, max_tokens):
+        shape_latent_512 = shape_latent
+        device = comfy.model_management.get_torch_device()
+        prepare_trellis_vae_for_decode(vae, shape_latent_512["samples"].shape)
+
+        coord_counts = shape_latent_512.get("coord_counts")
+        decoder = vae.first_stage_model.shape_dec
+        lr_resolution = 512
+        target_resolution = int(target_resolution)
+
+        if coord_counts is None:
+            feats, coords_512 = flatten_batched_sparse_latent(
+                shape_latent_512["samples"],
+                shape_latent_512["coords"],
+                coord_counts,
+            )
+            feats = feats.to(device)
+            coords_512 = coords_512.to(device)
+            slat = shape_norm(feats, coords_512)
+            slat.feats = slat.feats.to(next(decoder.parameters()).dtype)
+            hr_coords = decoder.upsample(slat, upsample_times=4)
+
+            hr_resolution = target_resolution
+            while True:
+                quant_coords = torch.cat([
+                    hr_coords[:, :1],
+                    ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+                ], dim=1)
+                final_coords = quant_coords.unique(dim=0)
+                num_tokens = final_coords.shape[0]
+
+                if num_tokens < max_tokens or hr_resolution <= 1024:
+                    break
+                hr_resolution -= 128
+
+            return IO.NodeOutput(final_coords,)
+
+        items = split_batched_sparse_latent(
+            shape_latent_512["samples"],
+            shape_latent_512["coords"],
+            coord_counts,
+        )
+        decoder_dtype = next(decoder.parameters()).dtype
+
+        sample_hr_coords = []
+        for feats_i, coords_i in items:
+            feats_i = feats_i.to(device)
+            coords_i = coords_i.to(device).clone()
+            coords_i[:, 0] = 0
+            slat_i = shape_norm(feats_i, coords_i)
+            slat_i.feats = slat_i.feats.to(decoder_dtype)
+            sample_hr_coords.append(decoder.upsample(slat_i, upsample_times=4))
+
+        hr_resolution = target_resolution
+        while True:
+            exceeds_limit = False
+            for hr_coords_i in sample_hr_coords:
+                quant_coords_i = torch.cat([
+                    hr_coords_i[:, :1],
+                    ((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+                ], dim=1)
+                if quant_coords_i.unique(dim=0).shape[0] >= max_tokens:
+                    exceeds_limit = True
+                    break
+            if not exceeds_limit or hr_resolution <= 1024:
+                break
+            hr_resolution -= 128
+
+        final_coords_list = []
+        output_coord_counts = []
+        for sample_offset, hr_coords_i in enumerate(sample_hr_coords):
+            quant_coords_i = torch.cat([
+                hr_coords_i[:, :1],
+                ((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+            ], dim=1)
+            final_coords_i = quant_coords_i.unique(dim=0)
+            final_coords_i = final_coords_i.clone()
+            final_coords_i[:, 0] = sample_offset
+            final_coords_list.append(final_coords_i)
+            output_coord_counts.append(int(final_coords_i.shape[0]))
+
+        coords = torch.cat(final_coords_list, dim=0)
+        output = Types.VOXEL(coords)
+        output.coord_counts = torch.tensor(output_coord_counts, dtype=torch.int64)
+        output.resolutions = torch.full((len(final_coords_list),), int(hr_resolution), dtype=torch.int64)
+        output.upsampled = True
+
+        return IO.NodeOutput(output,)
+
+dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+
+def run_conditioning(model, cropped_img_tensor, include_1024=True):
+    model_internal = model.model
+    device = comfy.model_management.intermediate_device()
+    torch_device = comfy.model_management.get_torch_device()
+
+    def prepare_tensor(pil_img, size):
+        resized_pil = pil_img.resize((size, size), Image.Resampling.LANCZOS)
+        img_np = np.array(resized_pil).astype(np.float32) / 255.0
+        img_t = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(torch_device)
+        return (img_t - dino_mean.to(torch_device)) / dino_std.to(torch_device)
+
+    model_internal.image_size = 512
+    input_512 = prepare_tensor(cropped_img_tensor, 512)
+    cond_512 = model_internal(input_512, skip_norm_elementwise=True)[0]
+
+    cond_1024 = None
+    if include_1024:
+        model_internal.image_size = 1024
+        input_1024 = prepare_tensor(cropped_img_tensor, 1024)
+        cond_1024 = model_internal(input_1024, skip_norm_elementwise=True)[0]
+
+    conditioning = {
+        'cond_512': cond_512.to(device),
+        'neg_cond': torch.zeros_like(cond_512).to(device),
+    }
+    if cond_1024 is not None:
+        conditioning['cond_1024'] = cond_1024.to(device)
+
+    return conditioning
+class Trellis2Conditioning(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Trellis2Conditioning",
+            category="conditioning/video_models",
+            inputs=[
+                IO.ClipVision.Input("clip_vision_model"),
+                IO.Image.Input("image"),
+                IO.Mask.Input("mask"),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ]
+        )
+
+    @classmethod
+    @classmethod
+    def execute(cls, clip_vision_model, image, mask) -> IO.NodeOutput:
+        # Normalize to batched form so per-image conditioning loop below is uniform.
+        if image.ndim == 3:
+            image = image.unsqueeze(0)
+        elif image.ndim == 4:
+            if image.shape[1] in [1, 3, 4] and image.shape[-1] not in [1, 3, 4]:
+                image = image.permute(0, 2, 3, 1)
+
+        # normalize mask to standard [B, H, W] (handling 2D, 3D, and 4D variants)
+        if mask.ndim == 4:
+            if mask.shape[1] == 1:
+                mask = mask.squeeze(1)
+            elif mask.shape[-1] == 1:
+                mask = mask.squeeze(-1)
+            else:
+                mask = mask[:, :, :, 0] # take first channel as fallback
+
+        if mask.ndim == 3:
+            if mask.shape[-1] == 1:
+                mask = mask.squeeze(-1).unsqueeze(0)
+        elif mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+
+        batch_size = image.shape[0]
+        if mask.shape[0] == 1 and batch_size > 1:
+            mask = mask.expand(batch_size, -1, -1)
+        elif mask.shape[0] != batch_size:
+            raise ValueError(f"Trellis2Conditioning mask batch {mask.shape[0]} does not match image batch {batch_size}")
+
+        cond_512_list = []
+        cond_1024_list = []
+
+        for b in range(batch_size):
+            item_image = image[b]
+            item_mask = mask[b] if mask.size(0) > 1 else mask[0]
+
+            img_np = (item_image.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+            mask_np = (item_mask.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+
+            # Ensure img_np is either 2D (grayscale) or 3D (RGB/RGBA)
+            if img_np.ndim == 3 and img_np.shape[-1] == 1:
+                img_np = img_np.squeeze(-1)
+
+            mask_np = mask_np.squeeze()
+
+            # detect inverted mask
+            border_pixels = np.concatenate([
+                mask_np[0, :], mask_np[-1, :], mask_np[:, 0], mask_np[:, -1]
+            ])
+            if np.mean(border_pixels) > 127:
+                mask_np = 255 - mask_np
+
+            mask_np[mask_np < 35] = 0
+
+            border_shave = 4
+            mask_np[:border_shave, :] = 0
+            mask_np[-border_shave:, :] = 0
+            mask_np[:, :border_shave] = 0
+            mask_np[:, -border_shave:] = 0
+
+            pil_img = Image.fromarray(img_np)
+            pil_mask = Image.fromarray(mask_np)
+
+            max_size = max(pil_img.size)
+            scale = min(1.0, 1024 / max_size)
+            if scale < 1.0:
+                new_w, new_h = int(pil_img.width * scale), int(pil_img.height * scale)
+                pil_img = pil_img.resize((new_w, new_h), Image.Resampling.LANCZOS)
+                pil_mask = pil_mask.resize((new_w, new_h), Image.Resampling.NEAREST)
+
+            rgba_np = np.zeros((pil_img.height, pil_img.width, 4), dtype=np.uint8)
+            rgba_np[:, :, :3] = np.array(pil_img.convert("RGB"))
+            rgba_np[:, :, 3] = np.array(pil_mask)
+
+            alpha = rgba_np[:, :, 3]
+            bbox_coords = np.argwhere(alpha > 0.8 * 255)
+
+            if len(bbox_coords) > 0:
+                y_min, x_min = np.min(bbox_coords[:, 0]), np.min(bbox_coords[:, 1])
+                y_max, x_max = np.max(bbox_coords[:, 0]), np.max(bbox_coords[:, 1])
+
+                center_y, center_x = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
+                size = max(y_max - y_min, x_max - x_min)
+
+                crop_x1 = int(center_x - size // 2)
+                crop_y1 = int(center_y - size // 2)
+                crop_x2 = int(center_x + size // 2)
+                crop_y2 = int(center_y + size // 2)
+
+                rgba_pil = Image.fromarray(rgba_np)
+                cropped_rgba = rgba_pil.crop((crop_x1, crop_y1, crop_x2, crop_y2))
+                cropped_np = np.array(cropped_rgba).astype(np.float32) / 255.0
+            else:
+                import logging
+                logging.warning("Mask for the image is empty. Trellis2 requires an image with a mask for the best mesh quality.")
+                cropped_np = rgba_np.astype(np.float32) / 255.0
+
+            bg_rgb = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+
+            fg = cropped_np[:, :, :3]
+            alpha_float = cropped_np[:, :, 3:4]
+            composite_np = fg * alpha_float + bg_rgb * (1.0 - alpha_float)
+
+            # Keep the image as 4-channel RGBA to force TRELLIS to bypass its internal background remover
+            rgb_uint8 = (composite_np * 255.0).round().clip(0, 255).astype(np.uint8)
+            alpha_uint8 = (alpha_float.squeeze(-1) * 255.0).round().clip(0, 255).astype(np.uint8)
+
+            rgba_composite = np.zeros((cropped_np.shape[0], cropped_np.shape[1], 4), dtype=np.uint8)
+            rgba_composite[:, :, :3] = rgb_uint8
+            rgba_composite[:, :, 3] = alpha_uint8
+
+            cropped_pil = Image.fromarray(rgba_composite, mode="RGBA")
+
+            # Convert to RGB to ensure the CLIP/DINO model receives a 3-channel image
+            item_conditioning = run_conditioning(clip_vision_model, cropped_pil.convert("RGB"), include_1024=True)
+            cond_512_list.append(item_conditioning["cond_512"])
+            cond_1024_list.append(item_conditioning["cond_1024"])
+
+        cond_512_batched = torch.cat(cond_512_list, dim=0)
+        cond_1024_batched = torch.cat(cond_1024_list, dim=0)
+        neg_cond_batched = torch.zeros_like(cond_512_batched)
+        neg_embeds_batched = torch.zeros_like(cond_1024_batched)
+
+        positive = [[cond_512_batched, {"embeds": cond_1024_batched}]]
+        negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]]
+        return IO.NodeOutput(positive, negative)
+
+class EmptyTrellis2ShapeLatent(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2ShapeLatent",
+            category="latent/3d",
+            inputs=[
+                IO.Voxel.Input(
+                    "voxel",
+                    tooltip=(
+                        "Shape structure input. Accepts either a voxel structure "
+                        "or upsampled voxel coordinates from a previous cascade stage."
+                    )
+            )
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, voxel):
+        # to accept the upscaled coords
+        is_512_pass = False
+        upsampled = hasattr(voxel, "upsampled")
+        if upsampled:
+            voxel = voxel.data
+
+        if not upsampled:
+            decoded = voxel.data.unsqueeze(1)
+            coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
+            is_512_pass = True
+
+        else:
+            coords = voxel.int()
+            is_512_pass = False
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+        in_channels = 32
+        # image like format
+        latent = torch.zeros(batch_size, in_channels, max_tokens, 1)
+
+        if is_512_pass:
+            generation_mode = "shape_generation_512"
+        else:
+            generation_mode = "shape_generation"
+        return IO.NodeOutput({"samples": latent, "coords": coords, "coord_counts": counts, "type": "trellis2",
+                              "model_options": {"generation_mode": generation_mode, "coords": coords, "coord_counts": counts}})
+
+class EmptyTrellis2LatentTexture(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2LatentTexture",
+            category="latent/3d",
+            inputs=[
+                IO.Voxel.Input(
+                    "voxel",
+                    tooltip=(
+                        "Shape structure input. Accepts either a voxel structure "
+                        "or upsampled voxel coordinates from a previous cascade stage."
+                    )
+                ),
+                IO.Latent.Input("shape_latent"),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, voxel, shape_latent):
+        channels = 32
+        upsampled = hasattr(voxel, "upsampled")
+        if upsampled:
+            voxel = voxel.data
+
+        if not upsampled:
+            decoded = voxel.data.unsqueeze(1)
+            coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
+        else:
+            coords = voxel.int()
+
+        batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
+
+        shape_latent = shape_latent["samples"]
+        if shape_latent.ndim == 4:
+            shape_latent = shape_latent.squeeze(-1).transpose(1, 2).reshape(-1, channels)
+
+        latent = torch.zeros(batch_size, channels, max_tokens, 1)
+        return IO.NodeOutput({"samples": latent, "type": "trellis2", "coords": coords, "coord_counts": counts,
+                                     "model_options": {"generation_mode": "texture_generation",
+                                                       "coords": coords, "coord_counts": counts, "shape_slat": shape_latent}})
+
+
+class EmptyTrellis2LatentStructure(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="EmptyTrellis2LatentStructure",
+            category="latent/3d",
+            inputs=[
+                IO.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch."),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
+    @classmethod
+    def execute(cls, batch_size):
+        in_channels = 8
+        resolution = 16
+        latent = torch.zeros(batch_size, in_channels, resolution, resolution, resolution)
+        output = {
+            "samples": latent,
+            "type": "trellis2",
+        }
+        return IO.NodeOutput(output)
+
+class Trellis2Extension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            Trellis2Conditioning,
+            EmptyTrellis2ShapeLatent,
+            EmptyTrellis2LatentStructure,
+            EmptyTrellis2LatentTexture,
+            VaeDecodeTextureTrellis,
+            VaeDecodeShapeTrellis,
+            VaeDecodeStructureTrellis2,
+            Trellis2UpsampleCascade,
+        ]
+
+
+async def comfy_entrypoint() -> Trellis2Extension:
+    return Trellis2Extension()

folder_paths.py

@@ -60,7 +60,7 @@ folder_names_and_paths["geometry_estimation"] = ([os.path.join(models_dir, "geom
 
 folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flow")], supported_pt_extensions)
 
-folder_names_and_paths["mediapipe"] = ([os.path.join(models_dir, "mediapipe")], supported_pt_extensions)
+folder_names_and_paths["detection"] = ([os.path.join(models_dir, "detection")], supported_pt_extensions)
 
 output_directory = os.path.join(base_path, "output")
 temp_directory = os.path.join(base_path, "temp")

nodes.py

@@ -1537,6 +1537,10 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
+    if "model_options" in latent:
+        inner = model.model.diffusion_model
+        inner.meta = latent["model_options"]
+
     callback = latent_preview.prepare_callback(model, steps)
     disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
     samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
@@ -2430,6 +2434,8 @@ async def init_builtin_extra_nodes():
         "nodes_toolkit.py",
         "nodes_replacements.py",
         "nodes_nag.py",
+        "nodes_trellis2.py",
+        "nodes_mesh_postprocess.py",
         "nodes_sdpose.py",
         "nodes_math.py",
         "nodes_number_convert.py",

openapi.yaml

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