| 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)) |
|
|
