200 lines
10 KiB
Python
200 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Reference dumps for the PBR-texture stages, for validating the C++/ggml port.
|
|
|
|
Produces dumps/reference_texture.gguf with per-stage golden tensors so each C++
|
|
stage validates in isolation (feed identical inputs, compare outputs):
|
|
|
|
cond [T, Cc] DINOv3 conditioning at texture resolution R
|
|
enc_vert [N, 3] shape-encoder input: dual-vertex offsets (QEF)
|
|
enc_inter [N, 3] shape-encoder input: intersection flags (QEF)
|
|
enc_coords [N, 4] active voxels at R (batch-idx 0 + xyz)
|
|
shape_slat [Nl, 32] shape encoder output (mean), the concat_cond
|
|
shape_coords [Nl, 4] latent voxels at R/16
|
|
tex_noise [Nl, 32] fixed sampling noise (seed)
|
|
tex_flow_t500 [Nl, 32] tex-flow forward at t=500 (concat_cond, f32)
|
|
tex_slat [Nl, 32] full sampler output, denormalized
|
|
pbr [M, 6] decoded PBR voxels (base_color,metal,rough,alpha), *0.5+0.5
|
|
pbr_coords [M, 4] decoded voxels at R (should equal enc_coords set)
|
|
shape_slat_mean/std [32] shape_slat_normalization (concat_cond in-norm)
|
|
tex_slat_mean/std [32] tex_slat_normalization (output de-norm)
|
|
|
|
This dump validates the standalone arbitrary-mesh texturing path, whose encoder
|
|
is fed a reproducible QEF dual grid. Integrated image-to-3D generation instead
|
|
retains the generated shape SLat and replays the shape decoder's subdivisions;
|
|
that wiring is covered by test_slat plus the sparse PBR sampling regression.
|
|
|
|
Run inside the reference container (real o-voxel), e.g.:
|
|
docker exec t2tex bash -lc \
|
|
'cd /work && TRELLIS2_PY=/trellis2 python scripts/dump_texture_reference.py \
|
|
--mesh /s/mesh_shipped.bin --image dumps/fixture_rgba.png --resolution 512'
|
|
"""
|
|
import argparse, json, os, struct, sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import o_voxel # noqa: E402 real o-voxel BEFORE ref_common (skips its stub)
|
|
import o_voxel.convert # noqa: E402
|
|
import ref_common # noqa: E402
|
|
ref_common.setup() # sdpa attention + pure-torch sparse conv
|
|
|
|
import numpy as np # noqa: E402
|
|
import torch # noqa: E402
|
|
import torch.nn.functional as F # 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
|
|
from trellis2.modules import sparse as sp # 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 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("--t", type=float, default=500.0)
|
|
ap.add_argument("--out", default=os.path.join(ref_common.DUMPS, "reference_texture.gguf"))
|
|
args = ap.parse_args()
|
|
dev = "cuda"
|
|
R = args.resolution
|
|
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, fp32_flags=False, **extra):
|
|
cfg = json.load(open(os.path.join(CK, stem + ".json")))["args"]
|
|
if fp32_flags:
|
|
cfg.pop("use_fp16", None); cfg["use_fp16"] = False
|
|
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)
|
|
|
|
print(f"loading models (res {R}) ...", flush=True)
|
|
shape_enc = load_model(FlexiDualGridVaeEncoder, "shape_enc_next_dc_f16c32_fp16", fp32_flags=True)
|
|
tex_dec = load_model(SparseUnetVaeDecoder, "tex_dec_next_dc_f16c32_fp16", fp32_flags=True)
|
|
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, Ftri = load_t2mesh(args.mesh)
|
|
mesh = trimesh.Trimesh(vertices=V, faces=Ftri, process=False)
|
|
img = Image.open(args.image).convert("RGBA")
|
|
print(f"mesh {len(V)} verts / {len(Ftri)} tris; image {img.size}", flush=True)
|
|
|
|
caps = {}
|
|
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)
|
|
caps["cond"] = cond["cond"][0] # [T, Cc]
|
|
|
|
# ── shape encoder (QEF dual grid input) ──────────────────────────────
|
|
vertices = torch.from_numpy(mesh.vertices).float()
|
|
faces = torch.from_numpy(mesh.faces).long()
|
|
voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
|
|
vertices.cpu(), faces.cpu(), grid_size=R,
|
|
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
face_weight=1.0, boundary_weight=0.2, regularization_weight=1e-2, timing=True)
|
|
enc_vert_feats = dual_vertices * R - voxel_indices # offset in voxel
|
|
vtx = sp.SparseTensor(
|
|
feats=enc_vert_feats,
|
|
coords=torch.cat([torch.zeros_like(voxel_indices[:, 0:1]), voxel_indices], dim=-1)).to(dev)
|
|
inter = vtx.replace(intersected).to(dev)
|
|
caps["enc_vert"] = vtx.feats
|
|
caps["enc_inter"] = inter.feats.float()
|
|
caps["enc_coords"] = vtx.coords.float()
|
|
print(f"enc input: {vtx.feats.shape[0]} voxels at {R}^3", flush=True)
|
|
|
|
shape_slat = shape_enc(vtx, inter)
|
|
caps["shape_slat"] = shape_slat.feats
|
|
caps["shape_coords"] = shape_slat.coords.float()
|
|
print(f"shape_slat: {shape_slat.feats.shape} coords {shape_slat.coords.shape}", flush=True)
|
|
|
|
# ── tex flow: forward parity point + full sampler ────────────────────
|
|
s_std = torch.tensor(tpa["shape_slat_normalization"]["std"], device=dev)[None]
|
|
s_mean = torch.tensor(tpa["shape_slat_normalization"]["mean"], device=dev)[None]
|
|
shape_slat_n = shape_slat.replace((shape_slat.feats - s_mean) / s_std)
|
|
caps["shape_slat_mean"] = s_mean[0]
|
|
caps["shape_slat_std"] = s_std[0]
|
|
|
|
Nl = shape_slat.coords.shape[0]
|
|
g = torch.Generator(device=dev).manual_seed(args.seed)
|
|
noise_feats = torch.randn(Nl, tex_flow.out_channels, generator=g, device=dev)
|
|
caps["tex_noise"] = noise_feats
|
|
x0 = shape_slat.replace(noise_feats.clone())
|
|
|
|
out = tex_flow(x0, torch.tensor([args.t], device=dev), cond["cond"], concat_cond=shape_slat_n)
|
|
caps["tex_flow_t500"] = out.feats
|
|
print(f"tex_flow t={args.t}: mean={out.feats.mean().item():.6f} l2={out.feats.norm().item():.4f}", flush=True)
|
|
|
|
slat = sampler.sample(
|
|
tex_flow, x0, concat_cond=shape_slat_n,
|
|
cond=cond["cond"], neg_cond=cond["neg_cond"],
|
|
**tpa["tex_slat_sampler"]["params"], verbose=True).samples
|
|
t_std = torch.tensor(tpa["tex_slat_normalization"]["std"], device=dev)[None]
|
|
t_mean = torch.tensor(tpa["tex_slat_normalization"]["mean"], device=dev)[None]
|
|
slat = slat.replace(slat.feats * t_std + t_mean)
|
|
caps["tex_slat"] = slat.feats
|
|
caps["tex_slat_mean"] = t_mean[0]
|
|
caps["tex_slat_std"] = t_std[0]
|
|
print(f"tex_slat: mean={slat.feats.mean().item():.5f} std={slat.feats.std().item():.5f}", flush=True)
|
|
|
|
# ── tex decoder → PBR voxels ─────────────────────────────────────────
|
|
pbr = tex_dec(slat) * 0.5 + 0.5
|
|
caps["pbr"] = pbr.feats
|
|
caps["pbr_coords"] = pbr.coords.float()
|
|
print(f"pbr: {pbr.feats.shape} range [{pbr.feats.min():.3f},{pbr.feats.max():.3f}]", flush=True)
|
|
|
|
# sanity: the tex decoder should reconstruct exactly the encoder's voxel set
|
|
enc_set = set(map(tuple, vtx.coords[:, 1:].cpu().numpy().tolist()))
|
|
pbr_set = set(map(tuple, pbr.coords[:, 1:].cpu().numpy().tolist()))
|
|
print(f"voxel-set match: pbr=={len(pbr_set)} enc=={len(enc_set)} "
|
|
f"symdiff={len(enc_set ^ pbr_set)} ({100.0*len(enc_set & pbr_set)/max(1,len(enc_set)):.2f}% common)",
|
|
flush=True)
|
|
|
|
import gguf
|
|
writer = gguf.GGUFWriter(args.out, "reference")
|
|
manifest = {"shapes": {}, "atol": 2e-3, "rtol": 2e-3, "resolution": R}
|
|
for name, t in caps.items():
|
|
a = t.detach().cpu().float().numpy()
|
|
manifest["shapes"][name] = list(a.shape)
|
|
writer.add_tensor(name, np.ascontiguousarray(a.reshape(-1), dtype=np.float32))
|
|
writer.write_header_to_file(); writer.write_kv_data_to_file(); writer.write_tensors_to_file(); writer.close()
|
|
with open(os.path.join(ref_common.DUMPS, "manifest_texture.json"), "w") as f:
|
|
json.dump(manifest, f, indent=1)
|
|
print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes), {len(caps)} tensors", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|