File size: 2,056 Bytes
872cf4d | 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 | """Load the frozen LeWM world model for PushT.
The published `quentinll/lewm-pusht` checkpoint was saved with transformers 4.x
ViT parameter names. transformers >= 5 renamed them, so the state dict needs a
1:1 key remap before it will load (shapes are unchanged).
"""
import re
import torch
from hydra.utils import instantiate
from stable_worldmodel.data import get_cache_dir
from stable_worldmodel.wm.utils import _resolve
# transformers 4.x ViT name -> transformers 5.x name. Shapes are identical.
_VIT_RENAMES = (
(r'^encoder\.encoder\.layer\.', 'encoder.layers.'),
(r'\.attention\.attention\.query\.', '.attention.q_proj.'),
(r'\.attention\.attention\.key\.', '.attention.k_proj.'),
(r'\.attention\.attention\.value\.', '.attention.v_proj.'),
(r'\.attention\.output\.dense\.', '.attention.o_proj.'),
(r'\.intermediate\.dense\.', '.mlp.fc1.'),
(r'(\.layers\.\d+)\.output\.dense\.', r'\1.mlp.fc2.'),
)
def _remap_vit_keys(state_dict: dict) -> dict:
out = {}
for key, value in state_dict.items():
for pattern, repl in _VIT_RENAMES:
key = re.sub(pattern, repl, key)
out[key] = value
return out
def load_lewm(
name: str = 'quentinll/lewm-pusht',
device: str = 'cuda',
cache_dir: str | None = None,
):
"""Instantiate LeWM and load the pretrained weights, frozen and in eval."""
cache_dir = get_cache_dir(cache_dir, sub_folder='checkpoints')
ckpt_path, config = _resolve(name, cache_dir)
model = instantiate(config)
state_dict = torch.load(ckpt_path, map_location='cpu')
missing, unexpected = model.load_state_dict(
_remap_vit_keys(state_dict), strict=False
)
if missing or unexpected:
raise RuntimeError(
f'LeWM checkpoint mismatch after remap.\n'
f' missing: {sorted(missing)[:8]}\n'
f' unexpected: {sorted(unexpected)[:8]}'
)
model = model.to(device).eval()
model.requires_grad_(False)
return model
|