import json,os,sys,time
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-slat-val',allow_patterns=['engines/l4-sm89/slat-dae-custom/*.plan','runtime/trt_dds.py','plugins/tensorrt/companion_sparse_trt.py']))
sys.path.insert(0,str(root/'runtime'));sys.path.insert(0,str(root/'plugins/tensorrt'));import companion_sparse_trt
from trt_dds import TensorRTDDS
# AniGen checkpoint/reference.
model_root=Path('/tmp/anigen-model');snapshot_download('VAST-AI/AniGen',token=tok,local_dir=model_root,allow_patterns=['ckpts/anigen/slat_dae/config.json','ckpts/anigen/slat_dae/ckpts/decoder_final.pt']);os.chdir(model_root)
from anigen.utils.model_utils import load_decoder
from anigen.modules.sparse import SparseTensor
from anigen.models.structured_latent_vae.anigen_base import AniGenSparseTransformerBase
D=load_decoder('ckpts/anigen/slat_dae','final','cuda').eval()
E={n:TensorRTDDS(root/'engines/l4-sm89/slat-dae-custom'/f'{n}.plan') for n in ['geo','skin','skl','geo-head','skin-head','skl-head']}

def diff(a,b):
 a=a.float();b=b.float();z=(a-b).abs();return {'max_abs':float(z.max()),'mean_abs':float(z.mean()),'cos':float(torch.nn.functional.cosine_similarity(a.reshape(1,-1),b.reshape(1,-1)).item())}

def unique_coords(n):
 # Deterministic unique points within 64^3.
 idx=torch.arange(n,device='cuda',dtype=torch.int32);x=idx//(64*64);r=idx%(64*64);y=r//64;z=r%64;return torch.stack([torch.zeros_like(idx),x,y,z],1)

torch.manual_seed(2026);n=512;ns=128;c=unique_coords(n);cs=unique_coords(ns)
xg=torch.randn(n,D.latent_channels,device='cuda',dtype=torch.float16);xs=torch.randn(n,D.latent_channels_vertskin,device='cuda',dtype=torch.float16);xj=torch.randn(ns,D.latent_channels_skl,device='cuda',dtype=torch.float16)
x=SparseTensor(xg,c);skin=SparseTensor(xs,c.clone());skl=SparseTensor(xj,cs)
with torch.inference_mode():
 with torch.autocast('cuda',dtype=torch.float16):
  t=time.time();hg,hj,hs=AniGenSparseTransformerBase.forward(D,x,skl,skin);torch.cuda.synchronize();native_branch_ms=(time.time()-t)*1000
 t=time.time();og=E['geo'].run({'feats':xg,'coords':c})['out_feats'];oskin=E['skin'].run({'feats':xs,'coords':c})['out_feats'];oj=E['skl'].run({'feats':xj,'coords':cs})['out_feats'];torch.cuda.synchronize();trt_branch_ms=(time.time()-t)*1000
 branch={'geo':diff(og,hg.feats),'skin':diff(oskin,hs.feats),'skl':diff(oj,hj.feats),'native_ms':native_branch_ms,'trt_ms':trt_branch_ms}
 # Native heads from native branch outputs.
 with torch.autocast('cuda',dtype=torch.float16):
  ng=hg
  for b in D.upsample:ng=b(ng)
  ng=D.out_layer(ng.type(x.dtype))
  nskin=hs
  for b in D.upsample_skin_net:nskin=b(nskin)
  nskin=D.out_layer_skin(nskin.type(x.dtype))
  nj=D.out_layer_skl(hj.type(skl.dtype));njs=D.out_layer_skl_skin(hj.type(skl.dtype));nj=nj.replace(torch.cat([nj.feats,njs.feats],-1))
 # TRT heads, fed TRT branch output to validate real composition.
 tg=E['geo-head'].run({'feats':og,'coords':c});ts=E['skin-head'].run({'feats':oskin,'coords':c});tj=E['skl-head'].run({'feats':oj})
 heads={'geo':diff(tg['out_feats'],ng.feats),'geo_coords':bool(torch.equal(tg['out_coords'].to(torch.int32),ng.coords.to(torch.int32))),'skin':diff(ts['out_feats'],nskin.feats),'skin_coords':bool(torch.equal(ts['out_coords'].to(torch.int32),nskin.coords.to(torch.int32))),'skl':diff(tj['out_feats'],nj.feats)}
report={'tokens':{'geo_skin':n,'skl':ns},'branches':branch,'heads':heads,'gpu':torch.cuda.get_device_name(0)}
print('SLAT_DAE_VALIDATION',json.dumps(report,indent=2),flush=True);Path('/tmp/slat_dae_custom_validation.json').write_text(json.dumps(report,indent=2))