256 lines
9.3 KiB
Python
256 lines
9.3 KiB
Python
"""
|
|
Mesh export + texture baking for the TrellisMac ComfyUI nodes.
|
|
|
|
This is the entire bake+export branch of trellis-mac/generate.py:156-297,
|
|
extracted verbatim into a reusable function. It picks one of three paths:
|
|
|
|
1. Metal-accelerated bake via o_voxel.postprocess.to_glb + mtldiffrast,
|
|
when the Metal backend is available and has a working nvdiffrast port.
|
|
2. KDTree fallback via trellis-mac's backends/texture_baker.py
|
|
(xatlas UV unwrap + scipy cKDTree nearest-voxel lookup).
|
|
3. Vertex-color fallback when there's no voxel attribute data.
|
|
|
|
Path selection and the `_grid_sample_3d` monkey-patch are load-bearing
|
|
mac-specific details from generate.py; do not "simplify" them without
|
|
understanding why they exist (see inline comments).
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def _detect_metal_backend():
|
|
"""
|
|
Detect whether o_voxel has a working Metal bake backend.
|
|
|
|
Mirrors generate.py:164-196. Returns True if o_voxel.postprocess is
|
|
importable and reports _BACKEND == 'metal' with a working dr (mtldiffrast).
|
|
|
|
Also installs the _grid_sample_3d patch when flex_gemm is absent, because
|
|
o_voxel's fallback grid sampler returns the wrong layout ([B*C, M] vs the
|
|
[M, C] the bake expects) and we must transpose before reshape.
|
|
"""
|
|
try:
|
|
import o_voxel.postprocess # noqa: F401
|
|
except (ImportError, AttributeError):
|
|
# ImportError: o_voxel not installed at all.
|
|
# AttributeError: the stubs/o_voxel/ stub package has no .postprocess.
|
|
return False
|
|
|
|
backend = getattr(o_voxel.postprocess, "_BACKEND", None)
|
|
has_dr = getattr(o_voxel.postprocess, "_HAS_DR", False)
|
|
use_metal = backend == "metal" and has_dr
|
|
|
|
if use_metal and not getattr(o_voxel.postprocess, "_HAS_FLEX_GEMM", False):
|
|
# o_voxel's _grid_sample_3d fallback returns [B*C, M] but the bake
|
|
# consumes it as [M, C]. Patch it to transpose before the reshape.
|
|
# We avoid installing flex_gemm itself because its import slows the
|
|
# diffusion hot path ~10x on MPS. (generate.py:173-194)
|
|
import torch.nn.functional as _F_gs
|
|
import torch
|
|
|
|
def _gs3d_fix(feats, coords, shape, grid, mode="trilinear"):
|
|
B, C = shape[0], shape[1]
|
|
D, H, W = shape[2], shape[3], shape[4]
|
|
device = feats.device
|
|
dense_vol = torch.zeros(B, C, D, H, W, dtype=feats.dtype, device=device)
|
|
batch_idx = coords[:, 0].long()
|
|
cx = coords[:, 1].long()
|
|
cy = coords[:, 2].long()
|
|
cz = coords[:, 3].long()
|
|
dense_vol[batch_idx, :, cx, cy, cz] = feats
|
|
grid_norm = torch.stack(
|
|
[
|
|
grid[..., 2] / (W - 1) * 2 - 1,
|
|
grid[..., 1] / (H - 1) * 2 - 1,
|
|
grid[..., 0] / (D - 1) * 2 - 1,
|
|
],
|
|
dim=-1,
|
|
).reshape(B, 1, 1, -1, 3)
|
|
sampled = _F_gs.grid_sample(
|
|
dense_vol,
|
|
grid_norm,
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
padding_mode="border",
|
|
)
|
|
M = grid.shape[1]
|
|
return sampled.reshape(B, C, M).permute(0, 2, 1).reshape(B * M, C)
|
|
|
|
o_voxel.postprocess._grid_sample_3d = _gs3d_fix
|
|
|
|
return use_metal
|
|
|
|
|
|
def _bake_metal(mesh_out, glb_path, texture_size, target_faces=200000):
|
|
"""
|
|
Metal-accelerated bake via o_voxel.postprocess.to_glb.
|
|
|
|
Mirrors generate.py:201-237. Raises RuntimeError on failure so the caller
|
|
can fall back to KDTree.
|
|
"""
|
|
import torch
|
|
import fast_simplification
|
|
import o_voxel
|
|
|
|
verts_np = mesh_out.vertices.cpu().numpy()
|
|
faces_np = mesh_out.faces.cpu().numpy()
|
|
target_faces = min(target_faces, len(faces_np))
|
|
|
|
if len(faces_np) > target_faces:
|
|
# Pre-simplify to avoid mtlbvh crash on large meshes. Keeps detail,
|
|
# avoids Metal BVH issues. (generate.py:208-217)
|
|
ratio = 1.0 - (target_faces / len(faces_np))
|
|
simp_verts, simp_faces = fast_simplification.simplify(verts_np, faces_np, ratio)
|
|
simp_verts_t = torch.from_numpy(simp_verts).float().to(mesh_out.vertices.device)
|
|
simp_faces_t = torch.from_numpy(simp_faces.astype("int32")).to(mesh_out.faces.device)
|
|
else:
|
|
simp_verts_t = mesh_out.vertices
|
|
simp_faces_t = mesh_out.faces
|
|
|
|
# Move all mesh tensors to CPU — o_voxel.to_glb mixes a device-neutral
|
|
# AABB tensor with mesh tensors; keep everything on CPU to avoid mismatch.
|
|
# (generate.py:222-235)
|
|
glb = o_voxel.postprocess.to_glb(
|
|
vertices=simp_verts_t.cpu(),
|
|
faces=simp_faces_t.cpu(),
|
|
attr_volume=mesh_out.attrs.cpu(),
|
|
coords=mesh_out.coords.cpu(),
|
|
attr_layout=mesh_out.layout,
|
|
voxel_size=mesh_out.voxel_size,
|
|
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
decimation_target=target_faces,
|
|
texture_size=texture_size,
|
|
verbose=True,
|
|
)
|
|
glb.export(glb_path)
|
|
return glb_path
|
|
|
|
|
|
def _bake_kdtree(mesh_out, glb_path, texture_size, target_faces=200000):
|
|
"""
|
|
KDTree CPU fallback bake via trellis-mac's backends/texture_baker.py.
|
|
|
|
Mirrors generate.py:243-277. xatlas UV unwrap + scipy cKDTree nearest-voxel
|
|
lookup + inverse-distance weighting.
|
|
"""
|
|
from PIL import Image
|
|
from backends.texture_baker import (
|
|
uv_unwrap,
|
|
bake_texture,
|
|
export_glb_with_texture,
|
|
)
|
|
|
|
verts = mesh_out.vertices.cpu().numpy()
|
|
faces = mesh_out.faces.cpu().numpy()
|
|
|
|
voxel_coords = mesh_out.coords.cpu().float()
|
|
voxel_attrs = mesh_out.attrs.cpu().float()
|
|
origin = mesh_out.origin.cpu().float()
|
|
vs = mesh_out.voxel_size
|
|
|
|
# Simplify before UV unwrap — xatlas is very slow on 800K+ vertex meshes.
|
|
# (generate.py:252-262)
|
|
bake_verts, bake_faces = verts, faces
|
|
target_faces = min(target_faces, len(faces))
|
|
if len(faces) > target_faces:
|
|
try:
|
|
import fast_simplification
|
|
ratio = 1.0 - (target_faces / len(faces))
|
|
bake_verts, bake_faces = fast_simplification.simplify(verts, faces, ratio)
|
|
except ImportError:
|
|
pass # xatlas will be slow but still works
|
|
|
|
new_verts, new_faces, uvs, _vmapping = uv_unwrap(bake_verts, bake_faces)
|
|
|
|
base_color_img, mr_img, _mask = bake_texture(
|
|
new_verts,
|
|
new_faces,
|
|
uvs,
|
|
voxel_coords.numpy(),
|
|
voxel_attrs.numpy(),
|
|
origin.numpy(),
|
|
vs,
|
|
texture_size=texture_size,
|
|
)
|
|
|
|
basecolor_path = glb_path.rsplit(".", 1)[0] + "_basecolor.png"
|
|
Image.fromarray(base_color_img).save(basecolor_path)
|
|
export_glb_with_texture(new_verts, new_faces, uvs, base_color_img, mr_img, glb_path)
|
|
return glb_path
|
|
|
|
|
|
def _export_vertex_color(mesh_out, glb_path):
|
|
"""Vertex-color-only fallback. Mirrors generate.py:282-288."""
|
|
import trimesh
|
|
|
|
verts = mesh_out.vertices.cpu().numpy()
|
|
faces = mesh_out.faces.cpu().numpy()
|
|
tm = trimesh.Trimesh(vertices=verts, faces=faces)
|
|
tm.export(glb_path)
|
|
return glb_path
|
|
|
|
|
|
def _export_obj(mesh_out, obj_path):
|
|
"""Plain OBJ export. Mirrors generate.py:290-297."""
|
|
verts = mesh_out.vertices.cpu().numpy()
|
|
faces = mesh_out.faces.cpu().numpy()
|
|
with open(obj_path, "w") as f:
|
|
for v in verts:
|
|
f.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
|
|
for face in faces:
|
|
f.write(f"f {face[0] + 1} {face[1] + 1} {face[2] + 1}\n")
|
|
return obj_path
|
|
|
|
|
|
def export_mesh(mesh_out, glb_path, texture_size=1024, bake_texture=True,
|
|
also_export_obj=False, verbose=print):
|
|
"""
|
|
Bake + export a MeshWithVoxel to disk.
|
|
|
|
This is the public entry point for the Export node. It replicates the
|
|
full decision tree of generate.py:156-297:
|
|
|
|
- If voxel attrs present and bake_texture requested:
|
|
try Metal → fall back to KDTree on RuntimeError.
|
|
- Else: vertex-color GLB.
|
|
- Optionally also write a plain OBJ alongside.
|
|
|
|
`mesh_out` is the MeshWithVoxel object from pipeline.run() (has
|
|
.vertices, .faces, .attrs, .coords, .origin, .voxel_size, .layout).
|
|
`glb_path` is the absolute output path (already resolved by the node).
|
|
Returns the glb_path.
|
|
"""
|
|
try:
|
|
texture_size = int(texture_size)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(f"texture_size must be 512, 1024, or 2048; got {texture_size!r}")
|
|
if texture_size not in (512, 1024, 2048):
|
|
raise ValueError(f"texture_size must be 512, 1024, or 2048; got {texture_size}")
|
|
|
|
has_voxels = hasattr(mesh_out, "attrs") and mesh_out.attrs is not None
|
|
|
|
if has_voxels and bake_texture:
|
|
use_metal = _detect_metal_backend()
|
|
|
|
if use_metal:
|
|
try:
|
|
verbose(f"Baking PBR textures via Metal ({texture_size}x{texture_size})...")
|
|
_bake_metal(mesh_out, glb_path, texture_size)
|
|
except RuntimeError as e:
|
|
verbose(f"Metal bake failed: {e}")
|
|
verbose("Falling back to KDTree texture baker...")
|
|
use_metal = False
|
|
|
|
if not use_metal:
|
|
verbose(f"Baking PBR textures via KDTree ({texture_size}x{texture_size})...")
|
|
_bake_kdtree(mesh_out, glb_path, texture_size)
|
|
else:
|
|
verbose("Exporting with vertex colors (no PBR bake)...")
|
|
_export_vertex_color(mesh_out, glb_path)
|
|
|
|
if also_export_obj:
|
|
obj_path = glb_path.rsplit(".", 1)[0] + ".obj"
|
|
_export_obj(mesh_out, obj_path)
|
|
|
|
return glb_path
|