commit 827e5cf6fe92194842be1b1a6dfd68afa9812280 Author: Dodzilla Date: Sun Feb 22 17:24:22 2026 -0800 Add Trellis mesh normals postprocess ComfyUI node diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..43ae0e2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.py[cod] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..2b4c350 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Dodzilla + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..9f365ca --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# ComfyUI-TrellisMeshPostprocess + +ComfyUI custom node for Trellis mesh normal postprocessing. + +## Node + +- `Trellis2 - Mesh Postprocess Normals` +- Class: `PixelArtistryTrellisMeshPostprocessNormals` +- Input: `TRIMESH` +- Output: `TRIMESH` + +This node recalculates normals by clustering face normals in quantized position space. It is designed for Trellis outputs where duplicated/split vertices can produce hard lighting seams. + +## Inputs + +- `trimesh`: mesh from Trellis texturing/refinement stages. +- `position_epsilon` (default `1e-5`): weld tolerance used only for normal clustering. +- `normal_crease_deg` (default `55`): smoothing angle threshold. + +## Install + +Clone into `ComfyUI/custom_nodes`: + +```bash +git clone https://github.com/Dodzilla/ComfyUI-TrellisMeshPostprocess.git +``` + +Restart ComfyUI. + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..0fbfb5d --- /dev/null +++ b/nodes.py @@ -0,0 +1,187 @@ +"""ComfyUI custom nodes for Trellis mesh postprocessing.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from typing import Dict, List, Tuple + +import numpy as np + + +def _normalize(v: np.ndarray) -> np.ndarray: + length = float(np.linalg.norm(v)) + if not np.isfinite(length) or length <= 1e-12: + return np.array([0.0, 1.0, 0.0], dtype=np.float64) + return v / length + + +def _position_key(vertices: np.ndarray, index: int, epsilon: float) -> str: + p = vertices[index] + qx = int(round(float(p[0]) / epsilon)) + qy = int(round(float(p[1]) / epsilon)) + qz = int(round(float(p[2]) / epsilon)) + return f"{qx},{qy},{qz}" + + +@dataclass +class _Cluster: + avg: np.ndarray + count: int + + +class PixelArtistryTrellisMeshPostprocessNormals: + """ + Recalculate normals for duplicated/split Trellis vertices by clustering + face normals in quantized position space. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "trimesh": ("TRIMESH",), + "position_epsilon": ( + "FLOAT", + { + "default": 0.00001, + "min": 0.00000001, + "max": 0.01, + "step": 0.000001, + }, + ), + "normal_crease_deg": ( + "FLOAT", + { + "default": 55.0, + "min": 1.0, + "max": 180.0, + "step": 0.5, + }, + ), + } + } + + RETURN_TYPES = ("TRIMESH",) + RETURN_NAMES = ("trimesh",) + FUNCTION = "process" + CATEGORY = "Trellis2/Mesh" + + def _assign_face_cluster( + self, + groups: Dict[str, List[_Cluster]], + pos_key: str, + face_normal: np.ndarray, + min_dot: float, + ) -> int: + cluster_list = groups.get(pos_key) + if cluster_list is None: + cluster_list = [] + groups[pos_key] = cluster_list + + for idx, cluster in enumerate(cluster_list): + if cluster.count <= 0: + continue + + dot = float(np.dot(cluster.avg, face_normal)) + if dot >= min_dot: + next_count = cluster.count + 1 + cluster.avg = _normalize((cluster.avg * cluster.count + face_normal) / next_count) + cluster.count = next_count + return idx + + cluster_list.append(_Cluster(avg=face_normal.copy(), count=1)) + return len(cluster_list) - 1 + + def process(self, trimesh, position_epsilon: float, normal_crease_deg: float): + mesh = trimesh.copy() + vertices = np.asarray(mesh.vertices, dtype=np.float64) + faces = np.asarray(mesh.faces, dtype=np.int64) + + if vertices.ndim != 2 or vertices.shape[1] != 3: + return (mesh,) + + if faces.ndim != 2 or faces.shape[1] != 3 or faces.shape[0] == 0: + return (mesh,) + + epsilon = max(1e-8, min(1e-2, float(position_epsilon))) + crease = max(1.0, min(180.0, float(normal_crease_deg))) + min_dot = float(np.cos(np.deg2rad(crease))) + + vertex_count = int(vertices.shape[0]) + position_keys = [_position_key(vertices, i, epsilon) for i in range(vertex_count)] + groups_by_position: Dict[str, List[_Cluster]] = {} + preferred_cluster_by_vertex: List[Dict[int, float]] = [defaultdict(float) for _ in range(vertex_count)] + accum: Dict[str, np.ndarray] = {} + + for face in faces: + i0, i1, i2 = (int(face[0]), int(face[1]), int(face[2])) + if i0 < 0 or i1 < 0 or i2 < 0: + continue + if i0 >= vertex_count or i1 >= vertex_count or i2 >= vertex_count: + continue + + p0 = vertices[i0] + p1 = vertices[i1] + p2 = vertices[i2] + + face_vec = np.cross(p1 - p0, p2 - p0) + magnitude = float(np.linalg.norm(face_vec)) + if not np.isfinite(magnitude) or magnitude <= 1e-12: + continue + + face_normal = face_vec / magnitude + + for vidx in (i0, i1, i2): + pos_key = position_keys[vidx] + cluster_index = self._assign_face_cluster(groups_by_position, pos_key, face_normal, min_dot) + key = f"{pos_key}|{cluster_index}" + + if key in accum: + accum[key] += face_vec + else: + accum[key] = face_vec.copy() + + preferred_cluster_by_vertex[vidx][cluster_index] += magnitude + + fallback_normals = None + try: + current_normals = np.asarray(mesh.vertex_normals, dtype=np.float64) + if current_normals.shape == vertices.shape: + fallback_normals = current_normals + except Exception: + fallback_normals = None + + normals = np.zeros_like(vertices, dtype=np.float64) + up = np.array([0.0, 1.0, 0.0], dtype=np.float64) + + for vidx in range(vertex_count): + pos_key = position_keys[vidx] + preferred = preferred_cluster_by_vertex[vidx] + cluster_index = 0 + + if preferred: + cluster_index = max(preferred.items(), key=lambda item: item[1])[0] + + key = f"{pos_key}|{cluster_index}" + summed = accum.get(key) + + if summed is None: + if fallback_normals is not None: + summed = fallback_normals[vidx] + else: + summed = up + + normals[vidx] = _normalize(summed) + + mesh.vertex_normals = normals.astype(np.float32) + return (mesh,) + + +NODE_CLASS_MAPPINGS = { + "PixelArtistryTrellisMeshPostprocessNormals": PixelArtistryTrellisMeshPostprocessNormals, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PixelArtistryTrellisMeshPostprocessNormals": "Trellis2 - Mesh Postprocess Normals", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..24ce15a --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +numpy