Refactor generate3d function in inference.py to streamline noise addition and improve mesh processing. Enhanced color generation for vertices and updated GLB export logic, ensuring efficient handling of mesh data and improved readability.

inference.py

@@ -2,12 +2,9 @@ import numpy as np
 import torch
 import time
 import tempfile
-import os
-from PIL import Image
 import trimesh
 
 from util.utils import get_tri
-from mesh import Mesh
 from util.renderer import Renderer
 
 def generate3d(model, rgb, ccm, device):
@@ -42,61 +39,36 @@ def generate3d(model, rgb, ccm, device):
 
     model.eval()
 
-    if model.denoising:
-        tnew = torch.randint(20, 21, [triplane.shape[0]], dtype=torch.long, device=triplane.device)
-        noise_new = torch.randn_like(triplane) * 0.5 + 0.5
-        triplane = model.scheduler.add_noise(triplane, noise_new, tnew)
-        with torch.no_grad():
+    with torch.no_grad():
+        if model.denoising:
+            tnew = torch.randint(20, 21, [triplane.shape[0]], dtype=torch.long, device=triplane.device)
+            noise_new = torch.randn_like(triplane) * 0.5 + 0.5
+            triplane = model.scheduler.add_noise(triplane, noise_new, tnew)
             triplane_feature2 = model.unet2(triplane, tnew)
-    else:
-        with torch.no_grad():
+        else:
             triplane_feature2 = model.unet2(triplane)
 
-    data_config = {
-        'resolution': [1024, 1024],
-        "triview_color": triplane_color.to(device),
-    }
-
     with torch.no_grad():
+        data_config = {
+            'resolution': [1024, 1024],
+            "triview_color": triplane_color.to(device),
+        }
         verts, faces = model.decode(data_config, triplane_feature2)
-        data_config['verts'] = verts[0]
-        data_config['faces'] = faces
 
     from kiui.mesh_utils import clean_mesh
-    verts, faces = clean_mesh(
-        data_config['verts'].squeeze().cpu().numpy().astype(np.float32),
-        data_config['faces'].squeeze().cpu().numpy().astype(np.int32),
+    verts_np, faces_np = clean_mesh(
+        verts[0].squeeze().cpu().numpy().astype(np.float32),
+        faces.squeeze().cpu().numpy().astype(np.int32),
         repair=False, remesh=True, remesh_size=0.005, remesh_iters=1
     )
-    data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
-    data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
-
-    # Export to .obj with UV and texture image
-    mesh_path_obj = tempfile.NamedTemporaryFile(suffix="", delete=False).name
-    model.export_mesh_wt_uv(
-        None,  # GL context skipped (no dr.RasterizeGLContext)
-        data_config,
-        mesh_path_obj,
-        "",
-        device,
-        res=(1024, 1024),
-        tri_fea_2=triplane_feature2
-    )
-
-    # Convert to GLB with embedded texture
-    obj_path = mesh_path_obj + ".obj"
-    texture_path = mesh_path_obj + ".png"
-
-    if not os.path.exists(obj_path) or not os.path.exists(texture_path):
-        raise RuntimeError("OBJ or texture file missing — cannot convert to GLB.")
 
-    mesh_scene = trimesh.load(obj_path, force='scene')
-    texture_image = np.array(Image.open(texture_path))
+    # === Generate per-vertex color (approximate)
+    colors = np.tile(np.mean(rgb.reshape(-1, 3), axis=0, keepdims=True), (verts_np.shape[0], 1)) / 255.0
+    # Optionally, use more sophisticated color mapping logic if you have UVs
 
-    for name, geom in mesh_scene.geometry.items():
-        geom.visual = trimesh.visual.texture.TextureVisuals(image=texture_image)
+    mesh = trimesh.Trimesh(vertices=verts_np, faces=faces_np, vertex_colors=colors, process=False)
 
-    glb_path = mesh_path_obj + "_final.glb"
-    mesh_scene.export(glb_path)
+    glb_path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False).name
+    mesh.export(glb_path)
 
     return glb_path