Texture stage: tex SLAT + PBR decode, and the bake order that makes it tractable

The texture flow is imgshape2tex - it denoises 32 PBR channels while SEEING the shape
latent, so in_channels is 64 against out_channels 32. Upstream feeds the shape latent
as concat_cond and the model does sparse_cat([x, concat_cond], dim=-1); both share
coords, so it reduces to a channel concat. Added to slat_flow and carried on the
sampler (it is fixed for the whole trajectory and must reach BOTH CFG branches).

Verified running on the real checkpoints: 3,988,052 PBR voxels x 6 channels in 65.8s
(base_color 0:3, metallic 3:4, roughness 4:5, alpha 5:6).

Two things here fail SILENTLY rather than loudly, so both are asserted in comments:

1. shape_slat arrives DENORMALISED - the shape stage un-standardises it for the
   decoder - but the texture flow was trained against the standardised form. It is
   re-normalised before use as concat_cond. Skipping that gives a plausible mesh with
   wrong colours, not an error.
2. tex_dec has pred_subdiv=False: it cannot invent subdivisions and must be handed the
   shape decoder's subs as guides, so texture voxels land on the geometry that was
   actually built.

The decoder's output is mapped * 0.5 + 0.5 into [0,1], the range o_voxel expects.

BAKE ORDER. Handing o_voxel the raw ~8M-face mesh hangs - the same wall the standalone
remesh test hit (killed at 20min), and the trellis2 lane's own operator note says the
uncapped bake peaks at 75GB. So the mesh is welded, stripped of floaters and decimated
BEFORE baking; the baker samples the attribute VOLUME at mesh positions, so a
decimated mesh still gets correct colours. Measured on the way through:

  welded      3,988,052 -> 3,983,672 verts
  floaters           12 components -> 1 kept, 6,332 faces dropped
  decimated   7,996,876 -> 214,322 faces
  pre-bake    34.6s

That floater count is worth noting: 12 components, not the 52,855 the first health
pass reported. Welding first is what makes the difference.

remesh now defaults OFF in to_glb, unlike upstream. Upstream runs on CUDA; this is the
CPU/Metal build and its remesher took >20 minutes on a 214k-face mesh. It is also
handed an already-clean mesh, so there is far less for it to fix.

Operator gains texture + texture_size params; geometry-only stays the default because
it is ~3min against the textured path's extra flow and bake.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

pixal3d_mlx/mesh.py

@@ -93,7 +93,7 @@ def fdg_to_mesh(h, resolution: int, voxel_margin: float = 0.5) -> Tuple:
 
 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):
+           prefer_metal: bool = True, remesh: bool = False):
     """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
@@ -116,7 +116,11 @@ def to_glb(vertices, faces, tex_voxels, attr_layout: dict, resolution: int,
         aabb=AABB,
         decimation_target=decimation_target,
         texture_size=texture_size,
-        remesh=True, remesh_band=1, remesh_project=0,
+        # `remesh` defaults OFF here, unlike upstream. Upstream runs on CUDA; this
+        # build is the CPU/Metal one and its remesher took >20 minutes on a 214k-face
+        # mesh before being killed. We also hand it an already-welded, floater-free,
+        # decimated mesh, so the remesh has much less to fix than it would upstream.
+        remesh=remesh, remesh_band=1, remesh_project=0,
     )
 
 

pixal3d_mlx/models.py

@@ -21,6 +21,7 @@ FILES = {
     "ss_dec": "ss_dec_conv3d_16l8_fp16",
     "slat_512": "slat_flow_img2shape_dit_1_3B_512_bf16",
     "slat_1024": "slat_flow_img2shape_dit_1_3B_1024_bf16",
+    "slat_tex": "slat_flow_imgshape2tex_dit_1_3B_1024_bf16",
     "shape_dec": "shape_dec_next_dc_f16c32_fp16",
     "tex_dec": "tex_dec_next_dc_f16c32_fp16",
 }
@@ -46,9 +47,13 @@ def load_all(device: str | None = None, with_texture: bool = False) -> dict:
     def cfg(stem):
         return CKPTS / f"{stem}.json"
 
+    flows = [("ss_flow", ss_flow.load), ("ss_dec", ss_dec.load),
+             ("slat_512", slat_flow.load), ("slat_1024", slat_flow.load)]
+    if with_texture:
+        flows.append(("slat_tex", slat_flow.load))
+
     models = {}
-    for key, loader in (("ss_flow", ss_flow.load), ("ss_dec", ss_dec.load),
-                        ("slat_512", slat_flow.load), ("slat_1024", slat_flow.load)):
+    for key, loader in flows:
         stem = FILES[key]
         model, rep = loader(_weights(stem), cfg(stem))
         if rep["missing"] or rep["unmapped"]:
@@ -65,4 +70,8 @@ def load_all(device: str | None = None, with_texture: bool = False) -> dict:
 
     models["cond_512"] = ProjConditioner("shape_512", device=device)
     models["cond_1024"] = ProjConditioner("shape_1024", device=device)
+    if with_texture:
+        # tex_1024 differs from shape_1024 only in naf_target_size (1024 vs 512), but
+        # that changes the high-res branch it samples, so it needs its own conditioner
+        models["cond_tex"] = ProjConditioner("tex_1024", device=device)
     return models

pixal3d_mlx/pipeline.py

@@ -153,6 +153,63 @@ def run_slat_stage(
     return slat
 
 
+def run_tex_stage(
+    flow_model,
+    tex_decoder,
+    cond: dict,
+    shape_slat,
+    subs,
+    neg_cond=None,
+    shape_normalization: dict | None = None,
+    tex_normalization: dict | None = None,
+    seed: int = 0,
+    progress: Optional[Callable[[int, int], None]] = None,
+    **overrides,
+):
+    """Texture SLAT + decode -> PBR voxels (base_color / metallic / roughness / alpha).
+
+    Two things here are easy to get wrong and both fail silently rather than loudly:
+
+    1. `shape_slat` arrives DENORMALISED (the shape stage un-standardises it for the
+       decoder), but the texture flow was trained against the standardised form, so it
+       is re-normalised before being used as `concat_cond`. Feeding the denormalised
+       latent gives a plausible mesh with wrong colours.
+    2. The texture decoder has `pred_subdiv=False` — it cannot invent its own
+       subdivisions and must be handed the shape decoder's `subs` as guides, so the
+       texture voxels land on exactly the geometry that was built.
+
+    The decoder's tanh-ish output is mapped `* 0.5 + 0.5` into [0,1], which is the
+    range o_voxel's baker expects.
+    """
+    from trellis_sparse_mlx import SparseTensor
+
+    if shape_normalization:
+        mean = mx.array(np.asarray(shape_normalization["mean"], dtype=np.float32))[None]
+        std = mx.array(np.asarray(shape_normalization["std"], dtype=np.float32))[None]
+        shape_slat = shape_slat.replace((shape_slat.feats - mean) / std)
+
+    params = {**TEX_SLAT_PARAMS, **overrides}
+    mx.random.seed(seed)
+    noise_channels = flow_model.in_channels - shape_slat.feats.shape[1]
+    noise = shape_slat.replace(
+        mx.random.normal((shape_slat.coords.shape[0], noise_channels))
+    )
+
+    sampler = FlowEulerSampler(concat_cond=shape_slat)
+    slat = sampler.sample(flow_model, noise, cond=cond, neg_cond=neg_cond,
+                          progress=progress, **params)
+
+    if tex_normalization:
+        mean = mx.array(np.asarray(tex_normalization["mean"], dtype=np.float32))[None]
+        std = mx.array(np.asarray(tex_normalization["std"], dtype=np.float32))[None]
+        slat = slat.replace(slat.feats * std + mean)
+
+    voxels = tex_decoder(slat, guide_subs=subs)
+    voxels = voxels.replace(voxels.feats * 0.5 + 0.5)
+    mx.eval(voxels.feats)
+    return voxels
+
+
 def refine_coords(shape_decoder, lr_slat, hr_resolution: int = 1024,
                   lr_resolution: int = 512, max_num_tokens: int = 49152):
     """Low-res SLAT -> refined high-res coordinate set. Stage 3a of the cascade.
@@ -190,6 +247,8 @@ def image_to_mesh(
     mesh_scale: float = 1.0,
     seed: int = 0,
     normalization: dict | None = None,
+    tex_normalization: dict | None = None,
+    texture: bool = False,
     log=print,
 ):
     """The full shipped geometry cascade: image -> (vertices, faces).
@@ -258,9 +317,29 @@ def image_to_mesh(
 
     info = {"hr_resolution": hr_res, "grid_resolution": grid_res,
             "output_resolution": res_out, "num_tokens": int(hr_coords.shape[0]),
-            "seconds": round(time.time() - t0, 1),
             "peak_gb": round(mx.get_peak_memory() / 2 ** 30, 1)}
-    return v, f, {**info, "subs": subs, "hr_slat": hr_slat}
+    extra = {"subs": subs, "hr_slat": hr_slat}
+
+    if texture:
+        t = time.time()
+        ctex = models["cond_tex"]
+        img = load_image(image_path, ctex.image_size)
+        cond, uncond = ctex(img, camera_angle_x, mesh_scale=mesh_scale,
+                            grid_resolution=grid_res)
+        voxels = run_tex_stage(
+            models["slat_tex"], models["tex_dec"],
+            gather_proj_at_coords(cond, hr_coords, grid_res), hr_slat, subs,
+            neg_cond=gather_proj_at_coords(uncond, hr_coords, grid_res),
+            shape_normalization=normalization, tex_normalization=tex_normalization,
+            seed=seed)
+        log(f"  texture    {voxels.feats.shape[0]:6} PBR voxels x "
+            f"{voxels.feats.shape[1]}      {time.time() - t:6.1f}s")
+        extra["tex_voxels"] = voxels
+        info["textured"] = True
+
+    info["seconds"] = round(time.time() - t0, 1)
+    info["peak_gb"] = round(mx.get_peak_memory() / 2 ** 30, 1)
+    return v, f, {**info, **extra}
 
 
 def image_to_occupancy(

pixal3d_mlx/sampler.py

@@ -41,8 +41,12 @@ def _like(ref, feats):
 
 
 class FlowEulerSampler:
-    def __init__(self, sigma_min: float = 1e-5):
+    def __init__(self, sigma_min: float = 1e-5, concat_cond: SparseTensor | None = None):
         self.sigma_min = sigma_min
+        # Carried on the sampler rather than threaded through every call: it is fixed
+        # for the whole trajectory (the shape latent never changes while the texture
+        # denoises) and it must reach BOTH the positive and negative CFG branches.
+        self.concat_cond = concat_cond
 
     # -- conversions between the model's velocity and x0/eps -------------------
     def v_to_xstart_eps(self, x_t, t: float, v):
@@ -68,6 +72,8 @@ class FlowEulerSampler:
     def _call(self, model: Callable, x_t, t: float, cond) -> Any:
         b = _feats(x_t).shape[0] if not isinstance(x_t, SparseTensor) else len(x_t.layout)
         tt = mx.full((b,), 1000.0 * t, dtype=mx.float32)  # upstream scales t by 1000
+        if self.concat_cond is not None:
+            return model(x_t, tt, cond, concat_cond=self.concat_cond)
         return model(x_t, tt, cond)
 
     def _inference(

pixal3d_mlx/slat_flow.py

@@ -82,7 +82,15 @@ class SLatFlowModel(nn.Module):
         ]
         self.out_layer = SparseLinear(model_channels, out_channels)
 
-    def __call__(self, x: SparseTensor, t: mx.array, cond) -> SparseTensor:
+    def __call__(self, x: SparseTensor, t: mx.array, cond,
+                 concat_cond: SparseTensor | None = None) -> SparseTensor:
+        # The texture flow is `imgshape2tex`: it denoises 32 PBR channels while SEEING
+        # the shape latent, so in_channels is 64 against out_channels 32. Upstream does
+        # `sparse_cat([x, concat_cond], dim=-1)`; both share coords, so it reduces to a
+        # plain channel concat. Without it the input layer gets half its expected width.
+        if concat_cond is not None:
+            x = x.replace(mx.concatenate([x.feats, concat_cond.feats], axis=-1))
+
         h = self.input_layer(x)
 
         t_emb = self.t_embedder(t)

scripts/image_to_mesh.py

@@ -20,7 +20,7 @@ import trimesh
 REPO = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(REPO))
 
-from pixal3d_mlx.mesh import to_camera_frame                       # noqa: E402
+from pixal3d_mlx.mesh import PBR_ATTR_LAYOUT, to_camera_frame     # noqa: E402
 from pixal3d_mlx.models import load_all, normalization             # noqa: E402
 from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_mesh        # noqa: E402
 
@@ -70,16 +70,53 @@ def main():
     ap.add_argument("--min-iou", type=float, default=0.85,
                     help="fail the run below this silhouette IoU; 0 disables the gate")
     ap.add_argument("--json", help="write run metadata here")
+    ap.add_argument("--texture", action="store_true",
+                    help="run the texture stage and bake PBR maps through o_voxel")
+    ap.add_argument("--texture-size", type=int, default=2048)
     a = ap.parse_args()
 
     t = time.time()
-    models = load_all()
+    models = load_all(with_texture=a.texture)
     print(f"loaded models ({time.time() - t:.1f}s, lazy — weights fault in on first use)")
 
     v, f, info = image_to_mesh(a.image, models, camera_angle_x=a.fov, seed=a.seed,
-                               normalization=normalization("shape"))
-    info.pop("subs", None)
+                               normalization=normalization("shape"),
+                               tex_normalization=normalization("tex") if a.texture else None,
+                               texture=a.texture)
+    subs = info.pop("subs", None)
     info.pop("hr_slat", None)
+    tex_voxels = info.pop("tex_voxels", None)
+
+    if a.texture and tex_voxels is not None:
+        # CLEAN BEFORE BAKING. o_voxel's remesh+unwrap on the raw ~8M-face mesh hangs
+        # (killed at 20min); the trellis2 lane hit the same wall and its operator note
+        # says the uncapped bake peaks at 75GB. Welding and stripping floaters first
+        # makes it tractable, and the baker samples the attribute VOLUME at mesh
+        # positions, so a decimated mesh still gets correct colours.
+        from pixal3d_mlx.cleanup import clean
+        from pixal3d_mlx.mesh import to_glb
+
+        t = time.time()
+        pre, _ = clean(v.cpu().numpy(), f.cpu().numpy(),
+                       target_faces=a.target_faces or 500_000)
+        print(f"  pre-bake   {len(pre.faces):,} faces   {time.time() - t:6.1f}s")
+
+        import torch
+        t = time.time()
+        scene = to_glb(torch.from_numpy(np.asarray(pre.vertices, np.float32)),
+                       torch.from_numpy(np.asarray(pre.faces, np.int32)),
+                       tex_voxels, PBR_ATTR_LAYOUT, info["output_resolution"],
+                       texture_size=a.texture_size,
+                       decimation_target=a.target_faces or 500_000)
+        mesh = scene if isinstance(scene, trimesh.Trimesh) else scene.dump(concatenate=True)
+        print(f"  bake       {len(mesh.faces):,} faces, {a.texture_size}px  "
+              f"{time.time() - t:6.1f}s")
+        mesh.export(a.output)
+        info["faces"], info["vertices"] = len(mesh.faces), len(mesh.vertices)
+        print(f"\nTOTAL {info['seconds']}s   peak {info['peak_gb']} GB   -> {a.output}")
+        if a.json:
+            Path(a.json).write_text(json.dumps(info, indent=2))
+        return
 
     if a.raw:
         mesh = trimesh.Trimesh(v.cpu().numpy(), f.cpu().numpy(), process=False)