| from __future__ import annotations |
| import json,os |
| from pathlib import Path |
| import torch |
| from safetensors.torch import save_file |
| from huggingface_hub import snapshot_download |
|
|
| MODEL='VAST-AI/AniGen' |
| root=Path('/tmp/anigen-model');out=Path('/tmp/slat-shell');out.mkdir(parents=True,exist_ok=True) |
| snapshot_download(MODEL,token=os.environ.get('HF_TOKEN'),local_dir=root,allow_patterns=['ckpts/anigen/slat_flow_auto/config.json','ckpts/anigen/slat_flow_auto/ckpts/**']) |
| config=json.loads((root/'ckpts/anigen/slat_flow_auto/config.json').read_text()) |
| ckpts=sorted((root/'ckpts/anigen/slat_flow_auto/ckpts').glob('*.pt')) |
| ckpt=[p for p in ckpts if 'ema' not in p.name and 'misc' not in p.name][-1] |
| sd=torch.load(ckpt,map_location='cpu') |
| if next(iter(sd)).startswith('module.'): sd={k[7:]:v for k,v in sd.items()} |
| CORE_PREFIXES=('blocks.','blocks_vert_skin.','blocks_skl.','adapter_geo_to_skin.') |
| shell={k:v.contiguous() for k,v in sd.items() if not k.startswith(CORE_PREFIXES)} |
| core={k:v.contiguous() for k,v in sd.items() if k.startswith(CORE_PREFIXES)} |
| save_file(shell,str(out/'shell.safetensors')) |
| (out/'config.json').write_text(json.dumps(config,indent=2)) |
| meta={'source_checkpoint':ckpt.name,'shell_tensors':len(shell),'core_tensors_removed':len(core),'shell_bytes':(out/'shell.safetensors').stat().st_size,'original_bytes':ckpt.stat().st_size,'removed_prefixes':list(CORE_PREFIXES)} |
| (out/'meta.json').write_text(json.dumps(meta,indent=2));print(json.dumps(meta,indent=2)) |
|
|
