diff --git a/pixal3d_mlx/mesh.py b/pixal3d_mlx/mesh.py index fc51200..483994e 100644 --- a/pixal3d_mlx/mesh.py +++ b/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, ) diff --git a/pixal3d_mlx/models.py b/pixal3d_mlx/models.py index 5936c9c..40f3750 100644 --- a/pixal3d_mlx/models.py +++ b/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 diff --git a/pixal3d_mlx/pipeline.py b/pixal3d_mlx/pipeline.py index d6420ff..50a8abe 100644 --- a/pixal3d_mlx/pipeline.py +++ b/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( diff --git a/pixal3d_mlx/sampler.py b/pixal3d_mlx/sampler.py index a9f9fd5..9f3ac48 100644 --- a/pixal3d_mlx/sampler.py +++ b/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( diff --git a/pixal3d_mlx/slat_flow.py b/pixal3d_mlx/slat_flow.py index 9e85391..dc76078 100644 --- a/pixal3d_mlx/slat_flow.py +++ b/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) diff --git a/scripts/image_to_mesh.py b/scripts/image_to_mesh.py index b02aebd..3afd297 100644 --- a/scripts/image_to_mesh.py +++ b/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)