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