from __future__ import annotations
import argparse, json, os, sys, time, gc
from pathlib import Path
os.environ.setdefault("ATTN_BACKEND", "sdpa")
import torch

MODEL_ID="VAST-AI/AniGen"

class SkinExport(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m=m
    def forward(self, vertex_features, joint_features, parents):
        return self.m(vertex_features, joint_features, parents)


def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--out',default=os.environ.get('CF_ONNX_OUT','/tmp/cf-skin')); ap.add_argument('--model-root',default='/tmp/anigen-model'); ap.add_argument('--app-root',default='/home/user/app'); a=ap.parse_args()
    out=Path(a.out); root=Path(a.model_root); app=Path(a.app_root); od=out/'onnx/anigen/skin-decoder'; od.mkdir(parents=True,exist_ok=True); sys.path.insert(0,str(app))
    from huggingface_hub import snapshot_download
    snapshot_download(MODEL_ID,token=os.environ.get('HF_TOKEN'),local_dir=root,allow_patterns=['ckpts/anigen/slat_dae/config.json','ckpts/anigen/slat_dae/ckpts/decoder_final.pt'])
    os.chdir(root)
    from anigen.utils.model_utils import load_decoder
    dec=load_decoder(str(root/'ckpts/anigen/slat_dae'),'final','cuda').eval()
    model=SkinExport(dec.skin_decoder).cuda().eval()
    vf=torch.zeros((1,256,4),device='cuda',dtype=torch.float32)
    jf=torch.zeros((1,24,4),device='cuda',dtype=torch.float32)
    parents=torch.arange(24,device='cuda',dtype=torch.int64).unsqueeze(0)-1
    p=od/'model.onnx'; t=time.time()
    with torch.inference_mode():
        torch.onnx.export(model,(vf,jf,parents),str(p),input_names=['vertex_features','joint_features','parents'],output_names=['skin_weights'],dynamic_axes={'vertex_features':{1:'n_vertices'},'joint_features':{1:'n_joints'},'parents':{1:'n_joints'},'skin_weights':{1:'n_vertices',2:'n_joints'}},opset_version=18,do_constant_folding=True,external_data=True,dynamo=False)
    import onnx; onnx.checker.check_model(str(p))
    meta={'component':'anigen-skin-decoder','source':MODEL_ID,'checkpoint':'slat_dae/decoder_final.pt::skin_decoder','opset':18,'precision':'mixed-fp16','dynamic':{'n_vertices':[128,10000,100000],'n_joints':[2,24,64]},'export_seconds':round(time.time()-t,3),'gpu':torch.cuda.get_device_name(0),'torch':torch.__version__}
    (od/'export_meta.json').write_text(json.dumps(meta,indent=2)); print('SKIN_DECODER_EXPORTED',json.dumps(meta),flush=True)
    del model,dec;gc.collect();torch.cuda.empty_cache()
if __name__=='__main__':main()