File size: 2,509 Bytes
41ecabc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | 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()
|