From 22dcd81fb3b8360f4d602dda97f62187921e97f3 Mon Sep 17 00:00:00 2001
From: Yousef Rafat <81116377+yousef-rafat@users.noreply.github.com>
Date: Wed, 25 Mar 2026 23:51:54 +0200
Subject: [PATCH] fixed color addition

---
 comfy/ldm/trellis2/model.py     | 3 +--
 comfy_extras/nodes_hunyuan3d.py | 3 +++
 comfy_extras/nodes_trellis2.py  | 1 +
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py
index 34aeba3e1..613f7ef50 100644
--- a/comfy/ldm/trellis2/model.py
+++ b/comfy/ldm/trellis2/model.py
@@ -782,7 +782,7 @@ class Trellis2(nn.Module):
             mode = "structure_generation"
             not_struct_mode = False
 
-        if is_1024 and mode == "shape_generation" and not is_512_run:
+        if is_1024 and not_struct_mode and not is_512_run:
             context = embeds
 
         sigmas = transformer_options.get("sigmas")[0].item()
@@ -825,7 +825,6 @@ class Trellis2(nn.Module):
             x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32))
 
         if mode == "shape_generation":
-            # TODO
             if is_512_run:
                 out = self.img2shape_512(x_st, t_eval, c_eval)
             else:
diff --git a/comfy_extras/nodes_hunyuan3d.py b/comfy_extras/nodes_hunyuan3d.py
index 692834c2b..ac91fe0a7 100644
--- a/comfy_extras/nodes_hunyuan3d.py
+++ b/comfy_extras/nodes_hunyuan3d.py
@@ -591,6 +591,9 @@ def save_glb(vertices, faces, filepath, metadata=None, colors=None):
         gltf["bufferViews"].append({"buffer": 0, "byteOffset": colors_byte_offset, "byteLength": colors_byte_length, "target": 34962})
         gltf["accessors"].append({"bufferView": 2, "byteOffset": 0, "componentType": 5126, "count": len(colors_np), "type": "VEC3"})
         gltf["meshes"][0]["primitives"][0]["attributes"]["COLOR_0"] = 2
+        # Define a base material so Three.js actually activates vertex coloring
+        gltf["materials"] =[{"pbrMetallicRoughness": {"baseColorFactor": [1.0, 1.0, 1.0, 1.0]}}]
+        gltf["meshes"][0]["primitives"][0]["material"] = 0
 
     if metadata is not None:
         gltf["asset"]["extras"] = metadata
diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py
index 469b460eb..ff95e8332 100644
--- a/comfy_extras/nodes_trellis2.py
+++ b/comfy_extras/nodes_trellis2.py
@@ -53,6 +53,7 @@ def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution):
     device = comfy.model_management.vae_offload_device()
 
     origin = torch.tensor([-0.5, -0.5, -0.5], device=device)
+    # TODO: generic independent node? if so: figure how pass the resolution parameter
     voxel_size = 1.0 / resolution
 
     # map voxels