Mesh export: Flexible Dual Grid -> triangles -> GLB, via o_voxel

The shape decoder's 7 channels are not an SDF — O-Voxel solves a QEF over a Flexible
Dual Grid, which is what lets it carry open and non-manifold surfaces:

  0:3  vertex offset in-voxel, (1+2m)*sigmoid(v)-m so it may sit OUTSIDE its own cell
  3:6  per-axis intersection logits, thresholded at 0
  6:7  quad split weight through softplus

mesh.py is the MLX->torch boundary for export. o_voxel's convert/postprocess are
native (C++/Metal) and deliberately NOT ported: o-voxel builds a CPU CppExtension when
CUDA is absent, and the trellis-2 lane on this fleet already runs it with a Metal
baker, so reusing that build beats reimplementing a QEF solver in MLX. Installed into
the shared venv from ~/Documents/trellis-2-mrp-mlx/o-voxel; it needs cv2 and xatlas,
and NOT utils3d (which drags in open3d, with no cp312 wheel).

Verified against the REAL shape_dec (292/292 params, resolution 256):

  decoded 5954 voxels x 7ch   ->   5954 vertices, 6886 faces   ->   GLB written

Not watertight, correctly: the input was a random latent, and FlexiDualGrid represents
open surfaces by design. Vertices land inside the octant of the unit cube matching the
sparse coords fed in, which is the check that the grid indexing is right.

Still to wire: the SLAT stage itself (sparse latents seeded from the occupancy coords,
with proj features gathered at those coords).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
m3ultra 2026-08-03 14:44:23 +10:00
parent 6f40fae0cc
commit 5484d59cb3

96
pixal3d_mlx/mesh.py Normal file
View File

@ -0,0 +1,96 @@
"""Flexible Dual Grid -> triangle mesh -> GLB, via o_voxel.
The shape decoder emits **7 channels per occupied voxel**, and they are not a
signed-distance field O-Voxel's Flexible Dual Grid solves a QEF instead, which is
what lets it carry open and non-manifold surfaces that marching cubes cannot:
0:3 vertex offset inside the voxel, `(1+2m)*sigmoid(v) - m` so it may sit
slightly OUTSIDE its own cell (m = voxel_margin = 0.5)
3:6 per-axis intersection flags logits at inference, thresholded at 0
6:7 quad split weight, through softplus
`o_voxel.convert.flexible_dual_grid_to_mesh` turns those into vertices and faces, and
`o_voxel.postprocess.to_glb` does UV unwrap plus texture baking. Both are native
(C++/Metal) and are NOT ported: o-voxel builds a CPU CppExtension when CUDA is absent,
and the trellis-2 lane on this fleet already runs it with a Metal baker. Reusing that
build is strictly better than reimplementing a QEF solver in MLX.
o_voxel speaks torch, so this module is the MLX->torch boundary for the export path.
"""
from __future__ import annotations
from typing import Tuple
import mlx.core as mx
import numpy as np
# Upstream fixes both: the model always works in a unit cube centred on the origin.
AABB = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]
def _torch(a):
import torch
return torch.from_numpy(np.asarray(a))
def fdg_to_mesh(h, resolution: int, voxel_margin: float = 0.5) -> Tuple:
"""Shape-decoder output -> (vertices, faces) as torch tensors.
`h` is the decoder's SparseTensor: `h.feats` [N,7], `h.coords` [N,4] with the batch
index in column 0. Single batch item only, which is all inference ever produces.
"""
from o_voxel.convert import flexible_dual_grid_to_mesh
feats = h.feats
m = voxel_margin
vertices = (1 + 2 * m) * mx.sigmoid(feats[..., 0:3]) - m
intersected = feats[..., 3:6] > 0 # logits -> bool at inference
quad_lerp = mx.logaddexp(feats[..., 6:7], mx.zeros_like(feats[..., 6:7])) # softplus
v, f = flexible_dual_grid_to_mesh(
_torch(h.coords[:, 1:]).int(),
_torch(vertices).float(),
_torch(intersected).bool(),
_torch(quad_lerp).float(),
aabb=AABB,
grid_size=resolution,
train=False,
)
return v, f
def to_glb(vertices, faces, tex_voxels, attr_layout: dict, resolution: int,
texture_size: int = 4096, decimation_target: int = 1_000_000,
prefer_metal: bool = True):
"""Bake the texture voxels onto the mesh and return a trimesh GLB scene.
`tex_voxels` is the texture decoder's SparseTensor (attrs in `.feats`, positions in
`.coords`). `attr_layout` maps PBR channel names to slices of that feature vector.
"""
try:
if not prefer_metal:
raise ImportError
from o_voxel import postprocess as pp
except ImportError:
from o_voxel import postprocess_cpu as pp
return pp.to_glb(
vertices=vertices,
faces=faces,
attr_volume=_torch(tex_voxels.feats).float(),
coords=_torch(tex_voxels.coords[:, 1:]).int(),
attr_layout=attr_layout,
grid_size=resolution,
aabb=AABB,
decimation_target=decimation_target,
texture_size=texture_size,
remesh=True, remesh_band=1, remesh_project=0,
)
# Upstream rotates the asset out of its internal frame on the way out (inference.py).
EXPORT_ROTATION = np.array([[-1, 0, 0, 0],
[0, 0, -1, 0],
[0, -1, 0, 0],
[0, 0, 0, 1]], dtype=np.float64)