132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
#!/usr/bin/env python
|
|
"""Convert LeWM HuggingFace weights into eval-compatible object checkpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import stable_pretraining as spt
|
|
import torch
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from jepa import JEPA
|
|
from module import ARPredictor, Embedder, MLP
|
|
|
|
|
|
def _load_json(path: Path) -> dict:
|
|
with path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _strip_target(config: dict) -> dict:
|
|
return {key: value for key, value in config.items() if key != "_target_"}
|
|
|
|
|
|
def infer_config_from_state_dict(state_dict: dict) -> dict:
|
|
action_dim = state_dict["action_encoder.patch_embed.weight"].shape[1]
|
|
return {
|
|
"encoder": {
|
|
"size": "tiny",
|
|
"patch_size": 14,
|
|
"image_size": 224,
|
|
"pretrained": False,
|
|
"use_mask_token": False,
|
|
},
|
|
"predictor": {
|
|
"num_frames": 3,
|
|
"input_dim": 192,
|
|
"hidden_dim": 192,
|
|
"output_dim": 192,
|
|
"depth": 6,
|
|
"heads": 16,
|
|
"mlp_dim": 2048,
|
|
"dim_head": 64,
|
|
"dropout": 0.1,
|
|
"emb_dropout": 0.0,
|
|
},
|
|
"action_encoder": {
|
|
"input_dim": action_dim,
|
|
"emb_dim": 192,
|
|
},
|
|
"projector": {
|
|
"input_dim": 192,
|
|
"output_dim": 192,
|
|
"hidden_dim": 2048,
|
|
},
|
|
"pred_proj": {
|
|
"input_dim": 192,
|
|
"output_dim": 192,
|
|
"hidden_dim": 2048,
|
|
},
|
|
}
|
|
|
|
|
|
def build_model(config: dict) -> JEPA:
|
|
encoder = spt.backbone.utils.vit_hf(**_strip_target(config["encoder"]))
|
|
predictor = ARPredictor(**_strip_target(config["predictor"]))
|
|
action_encoder = Embedder(**_strip_target(config["action_encoder"]))
|
|
|
|
projector_cfg = _strip_target(config["projector"])
|
|
projector_cfg["norm_fn"] = torch.nn.BatchNorm1d
|
|
projector = MLP(**projector_cfg)
|
|
|
|
pred_proj_cfg = _strip_target(config["pred_proj"])
|
|
pred_proj_cfg["norm_fn"] = torch.nn.BatchNorm1d
|
|
pred_proj = MLP(**pred_proj_cfg)
|
|
|
|
return JEPA(
|
|
encoder=encoder,
|
|
predictor=predictor,
|
|
action_encoder=action_encoder,
|
|
projector=projector,
|
|
pred_proj=pred_proj,
|
|
)
|
|
|
|
|
|
def convert_checkpoint(input_dir: Path, output_name: str) -> tuple[Path, Path]:
|
|
config_path = input_dir / "config.json"
|
|
weights_path = input_dir / "weights.pt"
|
|
if not weights_path.exists():
|
|
raise FileNotFoundError(f"Missing weights file: {weights_path}")
|
|
|
|
state_dict = torch.load(weights_path, map_location="cpu")
|
|
config = _load_json(config_path) if config_path.exists() else infer_config_from_state_dict(state_dict)
|
|
model = build_model(config)
|
|
missing, unexpected = model.load_state_dict(state_dict, strict=True)
|
|
if missing or unexpected:
|
|
raise RuntimeError(
|
|
f"State dict mismatch: missing={missing}, unexpected={unexpected}"
|
|
)
|
|
model.eval()
|
|
|
|
object_path = input_dir / f"{output_name}_object.ckpt"
|
|
weight_path = input_dir / f"{output_name}_weight.ckpt"
|
|
torch.save(model, object_path)
|
|
torch.save(model.state_dict(), weight_path)
|
|
return object_path, weight_path
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"input_dir",
|
|
type=Path,
|
|
help="Directory containing weights.pt and optionally config.json.",
|
|
)
|
|
parser.add_argument("--output-name", default="lewm")
|
|
args = parser.parse_args()
|
|
|
|
object_path, weight_path = convert_checkpoint(args.input_dir, args.output_name)
|
|
print(f"wrote {object_path}")
|
|
print(f"wrote {weight_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|