Enhance color assignment in generate3d function by utilizing original vertex data

- Added functionality to retrieve colors for original vertices and assign them to new vertices based on nearest neighbor search using cKDTree.
- Updated mesh creation to include vertex colors, improving the visual fidelity of the generated mesh.
- Streamlined the vertex handling process by directly using original vertex data for color assignment.

inference.py

@@ -62,14 +62,27 @@ def generate3d(model, rgb, ccm, device):
         data_config['faces'] = faces
 
     from kiui.mesh_utils import clean_mesh
+    orig_verts = data_config['verts'].squeeze().cpu().numpy()
+    orig_colors = ... # get the color for each original vertex, shape [N, 3]
     verts, faces = clean_mesh(
-        data_config['verts'].squeeze().cpu().numpy().astype(np.float32),
+        orig_verts.astype(np.float32),
         data_config['faces'].squeeze().cpu().numpy().astype(np.int32),
         repair=False, remesh=True, remesh_size=0.01, remesh_iters=2
     )
     data_config['verts'] = torch.from_numpy(verts).to(device).contiguous()
     data_config['faces'] = torch.from_numpy(faces).to(device).contiguous()
 
+    # Build KDTree from original verts
+    tree = cKDTree(orig_verts)
+
+    # For each new vertex, find the nearest old vertex and copy its color
+    k = 3
+    dists, idxs = tree.query(verts, k=k)
+    new_colors = np.mean(orig_colors[idxs], axis=1)
+
+    # Create the new mesh with colors
+    mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_colors=new_colors)
+
     # === Export OBJ/MTL/PNG ===
     obj_path = tempfile.NamedTemporaryFile(suffix=".obj", delete=False).name
     base_path = obj_path[:-4]  # remove .obj
@@ -79,15 +92,4 @@ def generate3d(model, rgb, ccm, device):
 
     model.export_mesh(data_config, base_path, tri_fea_2=triplane_feature2)  # writes .obj
 
-    # Build KDTree for nearest neighbor search
-    tree = cKDTree(data_config['verts'].squeeze().cpu().numpy())
-
-    # For each new vertex, find the nearest old vertex and copy its color
-    k = 3  # or 5, etc.
-    dists, idxs = tree.query(verts, k=k)
-    new_colors = np.mean(color_tri[idxs], axis=1)
-
-    # Create the new mesh with colors
-    mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_colors=new_colors)
-
     return obj_path