File size: 4,056 Bytes
f7fabd1 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # SPDX-License-Identifier: Apache-2.0
"""Check that a downloaded so101-pi05-base checkpoint is complete and loadable.
Reads the safetensors header and config.json directly, so it needs no torch, no
lerobot and no vlash - only the Python standard library. It does not run the
model; it establishes that the files are intact and that the policy config is
the one this release documents.
Usage:
python verify_checkpoint.py /path/to/checkpoint
"""
import argparse
import json
import math
import os
import struct
REQUIRED_FILES = ["config.json", "model.safetensors"]
EXPECTED_SLOTS = [
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
]
def read_safetensors_header(path):
"""Return (header dict, header byte length). The first 8 bytes are its size."""
with open(path, "rb") as f:
(n,) = struct.unpack("<Q", f.read(8))
return json.loads(f.read(n)), n
def line(text):
print(f" {text}")
def status(n, label, ok):
print(f"[{n}/4] {label:<38} -- {'PASS' if ok else 'FAIL'}")
def main():
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("checkpoint", help="directory holding config.json and model.safetensors")
args = ap.parse_args()
ckpt = args.checkpoint
print(f"checkpoint: {ckpt}\n")
failed = 0
missing = [n for n in REQUIRED_FILES if not os.path.isfile(os.path.join(ckpt, n))]
ok = not missing
status(1, "required files", ok)
if ok:
for n in REQUIRED_FILES:
line(f"{n:<20} {os.path.getsize(os.path.join(ckpt, n)):,} bytes")
else:
line(f"missing: {', '.join(missing)}")
print("\n0/4 checks passed")
return 1
weights = os.path.join(ckpt, "model.safetensors")
err = None
header_len, tensors, params, dtypes = 0, {}, 0, []
try:
header, header_len = read_safetensors_header(weights)
tensors = {k: v for k, v in header.items() if k != "__metadata__"}
params = sum(math.prod(v["shape"]) for v in tensors.values())
dtypes = sorted({v["dtype"] for v in tensors.values()})
except Exception as exc:
err = exc
ok = err is None and params > 0
status(2, "safetensors header", ok)
if ok:
line(f"{'tensors':<20} {len(tensors):,}")
line(f"{'parameters':<20} {params:,}")
line(f"{'dtypes':<20} {', '.join(dtypes)}")
else:
failed += 1
line(f"header unreadable: {err}")
with open(os.path.join(ckpt, "config.json")) as f:
cfg = json.load(f)
slots = [k for k in cfg.get("input_features", {}) if k.startswith("observation.images.")]
state = cfg.get("input_features", {}).get("observation.state", {}).get("shape", [None])[0]
action = cfg.get("output_features", {}).get("action", {}).get("shape", [None])[0]
ok = cfg.get("type") == "pi05" and slots == EXPECTED_SLOTS and state == 6 and action == 6
status(3, "policy config", ok)
line(f"{'type':<20} {cfg.get('type')}")
for s in slots:
line(f"{'camera slot':<20} {s.rsplit('.', 1)[-1]}")
line(f"{'state / action dim':<20} {state} / {action}")
line(f"{'chunk size':<20} {cfg.get('chunk_size')}")
line(f"{'image resolution':<20} {cfg.get('image_resolution')}")
line(f"{'normalization':<20} {cfg.get('normalization_mapping')}")
line(f"{'empty_cameras':<20} {cfg.get('empty_cameras')} (override per camera count at run time)")
if not ok:
failed += 1
end = max((v["data_offsets"][1] for v in tensors.values()), default=0)
expected = 8 + header_len + end
actual = os.path.getsize(weights)
ok = expected == actual
status(4, "header offsets vs file size", ok)
line(f"{'declared':<20} {expected:,} bytes")
line(f"{'on disk':<20} {actual:,} bytes")
if not ok:
failed += 1
line("truncated or partially downloaded")
print(f"\n{4 - failed}/4 checks passed")
return 1 if failed else 0
if __name__ == "__main__":
raise SystemExit(main())
|