Initial release

scripts/ref_texture.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+"""
+Reference PBR texturing, end-to-end, on OUR generated mesh -- the golden target
+the C++/ggml texture port validates against, and a quick eyeball of the texture
+NN stages.
+
+Runs the real Trellis2TexturingPipeline NN stages (shape encoder -> tex SLAT
+flow -> tex decoder) with sparse ops monkeypatched to pure torch (ref_common),
+so it needs NO custom CUDA kernels except o-voxel's CPU mesh->dual-grid.
+
+For the eyeball we skip the CUDA-only UV bake (nvdiffrast/cumesh/flexgemm) and
+instead trilinear-sample the decoded 6-channel PBR voxels at each mesh vertex
+(base_color / metallic / roughness), then dump a coloured mesh to render.
+
+  python scripts/ref_texture.py --mesh <T2MESH01.bin> --image <rgba.png> \
+      --resolution 512 --out dumps/tex_vcolor.bin
+"""
+import argparse, json, os, struct, sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import o_voxel  # noqa: E402  (import the REAL o-voxel BEFORE ref_common so its
+import o_voxel.convert          # stub guard skips it; we need mesh_to_flexible_dual_grid)
+import ref_common  # noqa: E402
+ref_common.setup()
+
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from PIL import Image  # noqa: E402
+import trimesh  # noqa: E402
+
+from safetensors.torch import load_file  # noqa: E402
+from trellis2.models.sc_vaes.fdg_vae import FlexiDualGridVaeEncoder  # noqa: E402
+from trellis2.models.sc_vaes.sparse_unet_vae import SparseUnetVaeDecoder  # noqa: E402
+from trellis2.models.structured_latent_flow import SLatFlowModel  # noqa: E402
+from trellis2.pipelines.samplers import FlowEulerGuidanceIntervalSampler  # noqa: E402
+from trellis2.pipelines import Trellis2TexturingPipeline  # noqa: E402
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor  # noqa: E402
+
+
+def load_t2mesh(path):
+    b = open(path, "rb").read()
+    assert b[:8] == b"T2MESH01", b[:8]
+    nv, nt = struct.unpack("<II", b[8:16]); o = 16
+    V = np.frombuffer(b, "<f4", 3 * nv, o).reshape(-1, 3).copy(); o += 12 * nv
+    o += 12 * nv  # skip normals
+    F = np.frombuffer(b, "<i4", 3 * nt, o).reshape(-1, 3).copy()
+    return V, F
+
+
+def trilinear_sample_sparse(feats, coords_xyz, query_vox):
+    """feats [M,C], coords_xyz [M,3] int voxel coords, query_vox [Q,3] float.
+    Returns (sampled [Q,C], weight [Q]) trilinear over present voxels."""
+    dev = feats.device
+    MULT = 4096
+    enc = lambda c: (c[:, 0].long() * MULT + c[:, 1].long()) * MULT + c[:, 2].long()
+    ck = enc(coords_xyz)
+    order = torch.argsort(ck)
+    ck_s = ck[order]; f_s = feats[order]
+    base = torch.floor(query_vox)
+    frac = query_vox - base
+    out = torch.zeros(query_vox.shape[0], feats.shape[1], device=dev)
+    wsum = torch.zeros(query_vox.shape[0], device=dev)
+    for dx in (0, 1):
+        for dy in (0, 1):
+            for dz in (0, 1):
+                corner = base + torch.tensor([dx, dy, dz], device=dev, dtype=base.dtype)
+                w = (frac[:, 0] if dx else 1 - frac[:, 0]) * \
+                    (frac[:, 1] if dy else 1 - frac[:, 1]) * \
+                    (frac[:, 2] if dz else 1 - frac[:, 2])
+                cck = enc(corner)
+                pos = torch.searchsorted(ck_s, cck).clamp(max=ck_s.shape[0] - 1)
+                found = (ck_s[pos] == cck)
+                out += (w * found)[:, None] * f_s[pos]
+                wsum += w * found
+    return out / wsum.clamp(min=1e-6)[:, None], wsum
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--mesh", required=True)
+    ap.add_argument("--image", required=True)
+    ap.add_argument("--resolution", type=int, default=512, choices=[512, 1024])
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--out", default=os.path.join(ref_common.DUMPS, "tex_vcolor.bin"))
+    args = ap.parse_args()
+    dev = "cuda"
+    CK = os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "ckpts")
+    tpa = json.load(open(os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "texturing_pipeline.json")))["args"]
+
+    def load_model(cls, stem, **extra):
+        cfg = json.load(open(os.path.join(CK, stem + ".json")))["args"]
+        m = cls(**{**cfg, **extra})
+        sd = {k: v.float() for k, v in load_file(os.path.join(CK, stem + ".safetensors")).items()}
+        m.load_state_dict(sd)
+        return m.eval().to(dev)
+
+    R = args.resolution
+    print(f"loading models (res {R}) ...", flush=True)
+    shape_enc = load_model(FlexiDualGridVaeEncoder, "shape_enc_next_dc_f16c32_fp16")
+    tex_dec = load_model(SparseUnetVaeDecoder, "tex_dec_next_dc_f16c32_fp16")
+    tex_flow = load_model(SLatFlowModel, f"slat_flow_imgshape2tex_dit_1_3B_{R}_bf16", dtype="float32")
+
+    dino = DinoV3FeatureExtractor(os.path.join(ref_common.MODELS, "dinov3-vitl16"), image_size=R)
+    dino.model = dino.model.to(dev).eval()
+
+    sampler = FlowEulerGuidanceIntervalSampler(**tpa["tex_slat_sampler"]["args"])
+    pipe = Trellis2TexturingPipeline(
+        models={"shape_slat_encoder": shape_enc, "tex_slat_decoder": tex_dec,
+                f"tex_slat_flow_model_{R}": tex_flow},
+        tex_slat_sampler=sampler,
+        tex_slat_sampler_params=tpa["tex_slat_sampler"]["params"],
+        shape_slat_normalization=tpa["shape_slat_normalization"],
+        tex_slat_normalization=tpa["tex_slat_normalization"],
+        image_cond_model=dino, rembg_model=None, low_vram=False,
+    )
+    pipe._device = dev
+
+    V, F = load_t2mesh(args.mesh)
+    mesh = trimesh.Trimesh(vertices=V, faces=F, process=False)
+    img = Image.open(args.image).convert("RGBA")
+    print(f"mesh {len(V)} verts / {len(F)} tris; image {img.size}", flush=True)
+
+    torch.manual_seed(args.seed)
+    with torch.no_grad():
+        image = pipe.preprocess_image(img)
+        mesh = pipe.preprocess_mesh(mesh)
+        cond = pipe.get_cond([image], R)
+        print("encoding shape SLAT (o-voxel dual grid + shape encoder) ...", flush=True)
+        shape_slat = pipe.encode_shape_slat(mesh, R)
+        print(f"shape_slat: {shape_slat.feats.shape} coords {shape_slat.coords.shape}", flush=True)
+        print("sampling tex SLAT (flow, concat_cond) ...", flush=True)
+        tex_slat = pipe.sample_tex_slat(cond, tex_flow, shape_slat)
+        print("decoding tex SLAT -> PBR voxels ...", flush=True)
+        pbr = pipe.decode_tex_slat(tex_slat)   # SparseTensor, 6ch, already *0.5+0.5
+        print(f"pbr voxels: {pbr.feats.shape} coords {pbr.coords.shape} "
+              f"range [{pbr.feats.min():.3f},{pbr.feats.max():.3f}]", flush=True)
+
+        # per-vertex trilinear sample of the PBR voxels (mesh is now normalized to [-.5,.5])
+        Vt = torch.from_numpy(mesh.vertices).float().to(dev)
+        qvox = (Vt + 0.5) * R
+        feats = pbr.feats.float()
+        coords_xyz = pbr.coords[:, 1:].to(dev)
+        vals, w = trilinear_sample_sparse(feats, coords_xyz, qvox)
+        vals = vals.clamp(0, 1).cpu().numpy()
+        hit = (w > 1e-4).float().mean().item()
+        print(f"per-vertex sample hit-rate {hit*100:.1f}%", flush=True)
+
+    base_color = vals[:, 0:3]
+    metallic = vals[:, 3:4]; roughness = vals[:, 4:5]
+    Vout = mesh.vertices.astype("<f4")
+    Fout = mesh.faces.astype("<i4")
+    with open(args.out, "wb") as f:
+        f.write(b"T2VCOL01")
+        f.write(struct.pack("<II", len(Vout), len(Fout)))
+        f.write(Vout.tobytes())
+        f.write(base_color.astype("<f4").tobytes())
+        f.write(np.concatenate([metallic, roughness], 1).astype("<f4").tobytes())
+        f.write(Fout.tobytes())
+    print(f"wrote {args.out}  (verts+rgb+metal/rough+tris)", flush=True)
+    print(f"base_color mean {base_color.mean(0)}  metallic {metallic.mean():.3f}  rough {roughness.mean():.3f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()