plugins/tests/validate_mesh_only.py

import json,os,sys
from pathlib import Path
import numpy as np
import torch,tensorrt as trt
from huggingface_hub import snapshot_download
ROOT=Path(snapshot_download('patdev/Companion-Forge-L4-ONNX',repo_type='model',token=os.environ.get('HF_TOKEN'),local_dir='/tmp/cf-mesh',allow_patterns=['engines/l4-sm89/custom-ops/MeshTopologyExtract.plan','plugins/tensorrt/companion_sparse_trt.py']))
sys.path.insert(0,str(ROOT/'plugins/tensorrt'));import companion_sparse_trt
from anigen.representations.mesh.utils_cube import construct_dense_grid,get_defomed_verts
from anigen.representations.mesh.flexicubes.flexicubes import FlexiCubes
LOGGER=trt.Logger(trt.Logger.ERROR);rt=trt.Runtime(LOGGER);eng=rt.deserialize_cuda_engine((ROOT/'engines/l4-sm89/custom-ops/MeshTopologyExtract.plan').read_bytes());ctx=eng.create_execution_context()
DT={trt.float16:torch.float16,trt.float32:torch.float32,trt.int32:torch.int32,trt.int64:torch.int64}
class OA(trt.IOutputAllocator):
 def __init__(self,dt):super().__init__();self.dt=dt;self.t=None;self.shape=None
 def reallocate_output(self,n,mem,size,align):
  es=torch.empty((),dtype=self.dt).element_size();self.t=torch.empty(max(1,(int(size)+es-1)//es),device='cuda',dtype=self.dt);print('ALLOC',n,int(size)//(1024*1024),'MiB',flush=True);return int(self.t.data_ptr())
 def reallocate_output_async(self,n,mem,size,align,stream):return self.reallocate_output(n,mem,size,align)
 def notify_shape(self,n,d):self.shape=tuple(int(x) for x in d);print('SHAPE',n,self.shape,flush=True)
res=64;grid,cubes=construct_dense_grid(res,'cuda');deform=torch.zeros((grid.shape[0],3),device='cuda');vg=get_defomed_verts(grid,deform,res).float();sdf=(torch.linalg.vector_norm(vg,dim=1)-0.30).float();beta=torch.zeros((cubes.shape[0],12),device='cuda');alpha=torch.zeros((cubes.shape[0],8),device='cuda');gamma=torch.zeros((cubes.shape[0],),device='cuda');colors=torch.cat([(vg+0.5).clamp(0,1),torch.zeros_like(vg)],dim=1).float();cubes64=cubes.to(torch.int64)
fc=FlexiCubes(device='cuda',use_color=True);rv,rf,_ld,rc=fc(voxelgrid_vertices=vg,scalar_field=sdf,cube_idx=cubes64,resolution=res,beta=beta,alpha=alpha,gamma_f=gamma,voxelgrid_colors=colors,training=False,no_sigmoid=True);torch.cuda.synchronize();print('REF',rv.shape,rf.shape,rc.shape,flush=True)
feeds={'verts_grid':vg,'sdf':sdf,'cube_idx':cubes64,'beta':beta,'alpha':alpha,'gamma':gamma,'colors_grid':colors};keep=[];alloc={};outs={}
for n,x in feeds.items():ctx.set_input_shape(n,tuple(x.shape));ctx.set_tensor_address(n,int(x.data_ptr()));keep.append(x)
for i in range(eng.num_io_tensors):
 n=eng.get_tensor_name(i)
 if eng.get_tensor_mode(n)!=trt.TensorIOMode.OUTPUT:continue
 dt=DT[eng.get_tensor_dtype(n)];sh=tuple(ctx.get_tensor_shape(n));print('OUTDESC',n,sh,dt,flush=True)
 if any(int(q)<0 for q in sh):a=OA(dt);alloc[n]=a;ctx.set_output_allocator(n,a)
 else:y=torch.empty(sh if sh else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);ctx.set_tensor_address(n,int(y.data_ptr()))
ok=ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream);print('EXEC',ok,flush=True);assert ok;torch.cuda.synchronize()
for n,a in alloc.items():
 num=int(np.prod(a.shape)) if a.shape else 1;outs[n]=a.t[:num].view(a.shape)
nv=int(outs['vertex_count'].item());nf=int(outs['face_count'].item());v=outs['vertices'][:nv];f=outs['faces'][:nf];col=outs['colors'][:nv]
def d(a,b):
 z=(a.float()-b.float()).abs();return {'max_abs':float(z.max()),'mean_abs':float(z.mean())}
rep={'vertex_count':nv,'face_count':nf,'ref_vertex_count':int(rv.shape[0]),'ref_face_count':int(rf.shape[0]),'vertices':d(v,rv),'faces_equal':bool(torch.equal(f.to(torch.long),rf.to(torch.long))),'colors':d(col,rc)}
print('MESH_REPORT',json.dumps(rep,indent=2),flush=True);Path('/tmp/mesh_validation.json').write_text(json.dumps(rep,indent=2))