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