File size: 3,070 Bytes
301b0eb 750aacd 301b0eb 750aacd 301b0eb 750aacd 301b0eb | 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 37 38 39 40 41 42 | """Inspect local inference assets; never load pickle metadata, model weights, or a robot SDK."""
from pathlib import Path
import argparse,hashlib,json
def check(model,checkpoint,resources=None,cell_data=None,verify_hashes=False,checkpoint_only=False):
root=Path(checkpoint);required=['model.safetensors'];result={'model':model,'checkpoint_name':root.name,'hashes_verified':False,'resources_checked':False}
if model=='flux':required+=['config.json','manifest.json']
else:required+=['assets/umi_v3/norm_stats.json']
for name in required:
p=root/name
if not p.is_file() or p.stat().st_size==0:raise ValueError('Missing/empty checkpoint asset: '+name)
if model=='flux':
c=json.loads((root/'config.json').read_text());m=json.loads((root/'manifest.json').read_text())
if c.get('action_dim')!=20 or c.get('chunk_size')!=32 or c.get('torch_dtype')!='bfloat16':raise ValueError('Expected 20-D / 32-step BF16 UMI export')
if m.get('kind')!='policy_export' or m.get('format_version')!=1:raise ValueError('Not a policy-export manifest')
if not checkpoint_only:
if resources is None:raise ValueError('FLUX needs --resources for frozen encoder assets')
if not (Path(resources)/'video_vae.safetensors').is_file() or not (Path(resources)/'text_encoder').is_dir():raise ValueError('Missing FLUX VAE or text encoder')
result['resources_checked']=True
if verify_hashes:
for name in ['config.json','model.safetensors']:
h=hashlib.sha256()
with (root/name).open('rb') as f:
while block:=f.read(8*1024*1024):h.update(block)
if h.hexdigest()!=m.get('sha256',{}).get(name):raise ValueError('Manifest hash mismatch: '+name)
result['hashes_verified']=True
else:
stats=json.loads((root/'assets/umi_v3/norm_stats.json').read_text())
if not isinstance(stats,dict) or not stats:raise ValueError('Empty/invalid normalization statistics')
result['optimizer_required_for_inference']=False
if cell_data:
for name in ['E_left.npy','E_right.npy','E_from_urdf.json']:
if not (Path(cell_data)/name).is_file():raise ValueError('Missing cell extrinsic: '+name)
result.update(status='passed_inventory_checks',model_bytes=(root/'model.safetensors').stat().st_size,weights_loaded=False)
return result
def main():
p=argparse.ArgumentParser(description=__doc__);p.add_argument('--model',choices=['flux','pi05'],required=True);p.add_argument('--checkpoint',required=True);p.add_argument('--resources');p.add_argument('--cell-data');p.add_argument('--verify-hashes',action='store_true');p.add_argument('--checkpoint-only',action='store_true',help='Check checkpoint without external encoders');a=p.parse_args()
try:print(json.dumps(check(a.model,a.checkpoint,a.resources,a.cell_data,a.verify_hashes,a.checkpoint_only),indent=2))
except (ValueError,OSError,json.JSONDecodeError) as exc:p.exit(2,str(exc)+'\n')
if __name__=='__main__':main()
|