Trellis2: inline batched mesh helpers

comfy_extras/mesh_batch_utils.py

@@ -1,53 +0,0 @@
-import torch
-from comfy_api.latest import Types
-
-
-def pack_variable_mesh_batch(vertices, faces, colors=None):
-    batch_size = len(vertices)
-    max_vertices = max(v.shape[0] for v in vertices)
-    max_faces = max(f.shape[0] for f in faces)
-
-    packed_vertices = vertices[0].new_zeros((batch_size, max_vertices, vertices[0].shape[1]))
-    packed_faces = faces[0].new_zeros((batch_size, max_faces, faces[0].shape[1]))
-    vertex_counts = torch.tensor([v.shape[0] for v in vertices], device=vertices[0].device, dtype=torch.int64)
-    face_counts = torch.tensor([f.shape[0] for f in faces], device=faces[0].device, dtype=torch.int64)
-
-    for i, (v, f) in enumerate(zip(vertices, faces)):
-        packed_vertices[i, :v.shape[0]] = v
-        packed_faces[i, :f.shape[0]] = f
-
-    mesh = Types.MESH(packed_vertices, packed_faces)
-    mesh.vertex_counts = vertex_counts
-    mesh.face_counts = face_counts
-
-    if colors is not None:
-        max_colors = max(c.shape[0] for c in colors)
-        packed_colors = colors[0].new_zeros((batch_size, max_colors, colors[0].shape[1]))
-        color_counts = torch.tensor([c.shape[0] for c in colors], device=colors[0].device, dtype=torch.int64)
-        for i, c in enumerate(colors):
-            packed_colors[i, :c.shape[0]] = c
-        mesh.colors = packed_colors
-        mesh.color_counts = color_counts
-
-    return mesh
-
-
-def get_mesh_batch_item(mesh, index):
-    if hasattr(mesh, "vertex_counts"):
-        vertex_count = int(mesh.vertex_counts[index].item())
-        face_count = int(mesh.face_counts[index].item())
-        vertices = mesh.vertices[index, :vertex_count]
-        faces = mesh.faces[index, :face_count]
-        colors = None
-        if hasattr(mesh, "colors") and mesh.colors is not None:
-            if hasattr(mesh, "color_counts"):
-                color_count = int(mesh.color_counts[index].item())
-                colors = mesh.colors[index, :color_count]
-            else:
-                colors = mesh.colors[index, :vertex_count]
-        return vertices, faces, colors
-
-    colors = None
-    if hasattr(mesh, "colors") and mesh.colors is not None:
-        colors = mesh.colors[index]
-    return mesh.vertices[index], mesh.faces[index], colors

comfy_extras/nodes_hunyuan3d.py

@@ -10,7 +10,6 @@ from comfy.cli_args import args
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, IO, Types
 from comfy_api.latest._util import MESH, VOXEL  # only for backward compatibility if someone import it from this file (will be removed later) # noqa
-from comfy_extras.mesh_batch_utils import pack_variable_mesh_batch, get_mesh_batch_item
 
 
 class EmptyLatentHunyuan3Dv2(IO.ComfyNode):
@@ -632,6 +631,57 @@ def save_glb(vertices, faces, filepath, metadata=None, colors=None):
 
     return filepath
 
+
+def pack_variable_mesh_batch(vertices, faces, colors=None):
+    batch_size = len(vertices)
+    max_vertices = max(v.shape[0] for v in vertices)
+    max_faces = max(f.shape[0] for f in faces)
+
+    packed_vertices = vertices[0].new_zeros((batch_size, max_vertices, vertices[0].shape[1]))
+    packed_faces = faces[0].new_zeros((batch_size, max_faces, faces[0].shape[1]))
+    vertex_counts = torch.tensor([v.shape[0] for v in vertices], device=vertices[0].device, dtype=torch.int64)
+    face_counts = torch.tensor([f.shape[0] for f in faces], device=faces[0].device, dtype=torch.int64)
+
+    for i, (v, f) in enumerate(zip(vertices, faces)):
+        packed_vertices[i, :v.shape[0]] = v
+        packed_faces[i, :f.shape[0]] = f
+
+    mesh = Types.MESH(packed_vertices, packed_faces)
+    mesh.vertex_counts = vertex_counts
+    mesh.face_counts = face_counts
+
+    if colors is not None:
+        max_colors = max(c.shape[0] for c in colors)
+        packed_colors = colors[0].new_zeros((batch_size, max_colors, colors[0].shape[1]))
+        color_counts = torch.tensor([c.shape[0] for c in colors], device=colors[0].device, dtype=torch.int64)
+        for i, c in enumerate(colors):
+            packed_colors[i, :c.shape[0]] = c
+        mesh.colors = packed_colors
+        mesh.color_counts = color_counts
+
+    return mesh
+
+
+def get_mesh_batch_item(mesh, index):
+    if hasattr(mesh, "vertex_counts"):
+        vertex_count = int(mesh.vertex_counts[index].item())
+        face_count = int(mesh.face_counts[index].item())
+        vertices = mesh.vertices[index, :vertex_count]
+        faces = mesh.faces[index, :face_count]
+        colors = None
+        if hasattr(mesh, "colors") and mesh.colors is not None:
+            if hasattr(mesh, "color_counts"):
+                color_count = int(mesh.color_counts[index].item())
+                colors = mesh.colors[index, :color_count]
+            else:
+                colors = mesh.colors[index, :vertex_count]
+        return vertices, faces, colors
+
+    colors = None
+    if hasattr(mesh, "colors") and mesh.colors is not None:
+        colors = mesh.colors[index]
+    return mesh.vertices[index], mesh.faces[index], colors
+
 class SaveGLB(IO.ComfyNode):
     @classmethod
     def define_schema(cls):

comfy_extras/nodes_trellis2.py

@@ -1,7 +1,6 @@
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, IO, Types
 from comfy.ldm.trellis2.vae import SparseTensor
-from comfy_extras.mesh_batch_utils import pack_variable_mesh_batch, get_mesh_batch_item
 import comfy.model_management
 from PIL import Image
 import numpy as np
@@ -9,6 +8,57 @@ import torch
 import scipy
 import copy
 
+
+def pack_variable_mesh_batch(vertices, faces, colors=None):
+    batch_size = len(vertices)
+    max_vertices = max(v.shape[0] for v in vertices)
+    max_faces = max(f.shape[0] for f in faces)
+
+    packed_vertices = vertices[0].new_zeros((batch_size, max_vertices, vertices[0].shape[1]))
+    packed_faces = faces[0].new_zeros((batch_size, max_faces, faces[0].shape[1]))
+    vertex_counts = torch.tensor([v.shape[0] for v in vertices], device=vertices[0].device, dtype=torch.int64)
+    face_counts = torch.tensor([f.shape[0] for f in faces], device=faces[0].device, dtype=torch.int64)
+
+    for i, (v, f) in enumerate(zip(vertices, faces)):
+        packed_vertices[i, :v.shape[0]] = v
+        packed_faces[i, :f.shape[0]] = f
+
+    mesh = Types.MESH(packed_vertices, packed_faces)
+    mesh.vertex_counts = vertex_counts
+    mesh.face_counts = face_counts
+
+    if colors is not None:
+        max_colors = max(c.shape[0] for c in colors)
+        packed_colors = colors[0].new_zeros((batch_size, max_colors, colors[0].shape[1]))
+        color_counts = torch.tensor([c.shape[0] for c in colors], device=colors[0].device, dtype=torch.int64)
+        for i, c in enumerate(colors):
+            packed_colors[i, :c.shape[0]] = c
+        mesh.colors = packed_colors
+        mesh.color_counts = color_counts
+
+    return mesh
+
+
+def get_mesh_batch_item(mesh, index):
+    if hasattr(mesh, "vertex_counts"):
+        vertex_count = int(mesh.vertex_counts[index].item())
+        face_count = int(mesh.face_counts[index].item())
+        vertices = mesh.vertices[index, :vertex_count]
+        faces = mesh.faces[index, :face_count]
+        colors = None
+        if hasattr(mesh, "colors") and mesh.colors is not None:
+            if hasattr(mesh, "color_counts"):
+                color_count = int(mesh.color_counts[index].item())
+                colors = mesh.colors[index, :color_count]
+            else:
+                colors = mesh.colors[index, :vertex_count]
+        return vertices, faces, colors
+
+    colors = None
+    if hasattr(mesh, "colors") and mesh.colors is not None:
+        colors = mesh.colors[index]
+    return mesh.vertices[index], mesh.faces[index], colors
+
 shape_slat_normalization = {
     "mean": torch.tensor([
         0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,