bench/validate_sparse_mesh_plugin.py

import json,os,sys
from pathlib import Path
import torch
from huggingface_hub import snapshot_download

repo='patdev/Companion-Forge-L4-ONNX';tok=os.environ.get('HF_TOKEN')
root=Path(snapshot_download(repo,repo_type='model',token=tok,local_dir='/tmp/cf-smesh-val',allow_patterns=['engines/l4-sm89/custom-ops/SparseMeshTopologyExtract.plan','plugins/tensorrt/companion_sparse_trt.py','runtime/trt_dds.py']))
sys.path.insert(0,str(root/'plugins/tensorrt'));sys.path.insert(0,str(root/'runtime'));import companion_sparse_trt
from trt_dds import TensorRTDDS
from anigen.modules.sparse import SparseTensor
from anigen.representations.mesh.cube2mesh_skeleton import AniGenSparseFeatures2Mesh

# 16^3 compact sparse cube field centered around 128 in the production res256 grid.
r=16;base=120
xyz=torch.cartesian_prod(torch.arange(base,base+r),torch.arange(base,base+r),torch.arange(base,base+r)).to(torch.int32).cuda();n=xyz.shape[0]
coords=torch.cat([torch.zeros((n,1),device='cuda',dtype=torch.int32),xyz],1)
# Per-cube 8 corner values. Corner order matches sparse_cube2verts upstream.
off=torch.tensor([[x,y,z] for x in (0,1) for y in (0,1) for z in (0,1)],device='cuda',dtype=torch.float32)
vp=xyz.float()[:,None,:]+off[None,:,:]
center=torch.tensor([128.,128.,128.],device='cuda');sdf=(torch.linalg.vector_norm(vp-center,dim=-1)-6.0).half()
deform=torch.zeros((n,8,3),device='cuda',dtype=torch.float16)
weights=torch.zeros((n,21),device='cuda',dtype=torch.float16)
# 6-channel color/normal values per corner; keep deterministic nonzero values.
color=torch.zeros((n,8,6),device='cuda',dtype=torch.float16);color[...,0]=(vp[...,0]-base)/float(r);color[...,1]=(vp[...,1]-base)/float(r);color[...,2]=(vp[...,2]-base)/float(r)
skin=torch.zeros((n,8,4),device='cuda',dtype=torch.float16);skin[...,0]=1.0
geo=torch.cat([sdf.reshape(n,-1),deform.reshape(n,-1),weights,color.reshape(n,-1)],-1)
feats=torch.cat([geo,skin.reshape(n,-1)],-1).contiguous();assert feats.shape[1]==133
native_ext=AniGenSparseFeatures2Mesh(res=256,use_color=True,skin_feat_channels=4,predict_skin=True,device='cuda')
with torch.inference_mode():native=native_ext(SparseTensor(feats,coords),training=False)
print('NATIVE',n,native.vertices.shape,native.faces.shape,flush=True)
eng=TensorRTDDS(root/'engines/l4-sm89/custom-ops/SparseMeshTopologyExtract.plan');o=eng.run({'cube_feats':feats,'cube_coords':coords});nv=int(o['vertex_count'].item());nf=int(o['face_count'].item())
verts=o['vertices'][:nv];faces=o['faces'][:nf].long();attrs=o['vertex_attrs'][:nv];skin_o=o['vertex_skin_feats'][:nv]
def diff(a,b):
 if a is None and b is None:return {'both_none':True}
 z=(a.float()-b.float()).abs();return {'max_abs':float(z.max()) if z.numel() else 0.,'mean_abs':float(z.mean()) if z.numel() else 0.}
rep={'sparse_cubes':int(n),'vertex_count':nv,'ref_vertex_count':int(native.vertices.shape[0]),'face_count':nf,'ref_face_count':int(native.faces.shape[0]),'vertices':diff(verts,native.vertices),'faces_equal':bool(torch.equal(faces,native.faces.long())),'attrs':diff(attrs,native.vertex_attrs),'skin':diff(skin_o,native.vertex_skin_feats),'gpu':torch.cuda.get_device_name(0)}
print('SPARSE_MESH_VALIDATION',json.dumps(rep,indent=2),flush=True);Path('/tmp/sparse_mesh_validation.json').write_text(json.dumps(rep,indent=2))