| """Load the IWM inference-only safetensors release; never loads private data. |
| |
| Use with the source checkout pinned in manifest.json. No training is dispatched. |
| """ |
| import hashlib |
| import json |
| from pathlib import Path |
| import sys |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| SOURCE_COMMIT = "516a71c252ff8f895dcd036bb343c996254c7a76" |
| BANK = ("feature_mean", "feature_std", "coefficient", "action_mean", |
| "expert_weight", "expert_mean", "residual_std") |
| GUIDE = ("probe_mean", "probe_std", "probe_weight", "probe_bias", |
| "class_centers", "within_std", "relation_keys", "relation_labels", "relation_ids") |
| EXCLUDED = ("prototypes", "amplitude_limit", "variance") + tuple("guide_" + k for k in GUIDE) |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as f: |
| for block in iter(lambda: f.read(8 * 1024**2), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def source_paths(source_root): |
| root = Path(source_root).resolve() |
| archive = root / "experiments/videorae3d" |
| paths = [archive / "cg2_dualcontext", archive / "cg2_ng", |
| archive / "eeg", archive / "teacher", archive, root / "src", root / "scripts"] |
| if not all(p.is_dir() for p in paths): |
| raise ValueError("Use the complete pinned intention-world-model source checkout") |
| for path in reversed(paths): |
| if str(path) not in sys.path: |
| sys.path.insert(0, str(path)) |
| return archive |
|
|
|
|
| def prepared_from_state(state, prefix="", include_training=False): |
| """Constructor adapter, not a replacement for private training preparation.""" |
| normal = {"mean": state[prefix + "normal_mean"], "std": state[prefix + "normal_std"]} |
| bank = {k: state[prefix + k] for k in BANK} |
| if include_training: |
| bank.update({k: state[prefix + k] for k in ("prototypes", "amplitude_limit", "variance")}) |
| guide = {k: state[prefix + "guide_" + k] for k in GUIDE} |
| else: |
| |
| |
| guide = dict(probe_mean=torch.zeros(256), probe_std=torch.ones(256), |
| probe_weight=torch.zeros(5, 256), probe_bias=torch.zeros(5), |
| class_centers=torch.zeros(5, 256), within_std=torch.ones(256), |
| relation_keys=torch.empty(0, 256), |
| relation_labels=torch.empty(0, dtype=torch.long), |
| relation_ids=torch.empty(0, dtype=torch.long)) |
| return dict(normalizers={"target": normal}, bank=bank, classguide=guide) |
|
|
|
|
| def build_model(config, state, source_root, original=False): |
| |
| |
| state = dict(state) |
| for name, strides in config.get("tensor_layouts", {}).items(): |
| t = state[name] |
| if len(strides) != t.ndim or any(type(x) is not int or x <= 0 for x in strides): |
| raise ValueError("Invalid tensor layout") |
| axes = sorted(zip(t.shape, strides), key=lambda pair: pair[1]) |
| span = 1 |
| for size, stride in axes: |
| if size > 1: |
| if stride != span: |
| raise ValueError("Only dense, non-overlapping layouts are supported") |
| span *= size |
| if tuple(t.stride()) != tuple(strides): |
| state[name] = torch.empty_strided(t.shape, strides, dtype=t.dtype).copy_(t) |
| source_paths(source_root) |
| kind = config["architecture"] |
| if kind == "Teacher3D": |
| from model import Teacher3D |
| m = Teacher3D() |
| else: |
| if kind not in ("CG2-NG", "NG-D"): |
| raise ValueError("Unsupported released architecture") |
| cfg = dict(config["model_config"]) |
| if not (cfg["timing"] == "target" and cfg["method"] == "CG2" |
| and cfg["variant"] == "CG2-NG" and cfg["representation"] == "tokens" |
| and cfg["head"] == "continuous"): |
| raise ValueError("Unexpected model protocol") |
| prefix = "base." if kind == "NG-D" else "" |
| prepared = prepared_from_state(state, prefix, include_training=original) |
| if kind == "NG-D": |
| assert not cfg["past_aux"] and cfg["separate_motion"] |
| from dual_models import Controller |
| else: |
| from noise_models import Controller |
| m = Controller(cfg, prepared) |
| if not original: |
| base = m.base if kind == "NG-D" else m |
| for name in EXCLUDED: |
| if name in base._buffers: |
| delattr(base, name) |
| m.load_state_dict(state, strict=True) |
| return m.eval().requires_grad_(False) |
|
|
|
|
| def load_checkpoint(release_root, model_id, source_root, device="cpu", verify=True): |
| """Return an eval-only Teacher3D or EEG controller. |
| |
| Controllers take (z_current[B,1024,64], eeg_raw[B,32,500]) and return |
| (class_logits[B,5], action[B,1,256]). Zero is raw zero, before normalization. |
| Use teacher.fdm(z_current, action) for the predicted H1 latent. |
| Training losses/probes and optimizer-resume state are intentionally absent. |
| """ |
| root = Path(release_root).resolve() |
| manifest = json.loads((root / "manifest.json").read_text()) |
| entry = next((r for r in manifest["models"] if r["id"] == model_id), None) |
| if entry is None: |
| raise ValueError("Unknown model_id; see manifest.json") |
| weights, config = root / entry["weights"], root / entry["config"] |
| if verify: |
| for p in (weights, config): |
| rel = p.relative_to(root).as_posix() |
| if digest(p) != manifest["files"][rel]["sha256"]: |
| raise ValueError("Release checksum mismatch: " + rel) |
| cfg = json.loads(config.read_text()) |
| state = load_file(str(weights), device="cpu") |
| return build_model(cfg, state, source_root).to(device) |
|
|
