WAM_DIT4DIT — Stage-1 Context Compressor (RoboCasa)
Stage-1 of the context compressor for WAM_DIT4DIT (Wan2.2-TI2V-5B video DiT). It compresses T context latent frames into N motion tokens via an inverse-dynamics (IDM) head + skip-reconstruction losses. The 5B video DiT is frozen; only the ~compressor trains. Trained on RoboCasa (robocasa_mg_gr00t_300), T=4, 25k steps.
Weights only — model.safetensors + config. The deepspeed optimizer state
(global_step*, 28 GB) and resume artifacts are omitted; these are not needed to
load the compressor.
Variants (each in its own subfolder)
| subfolder | config | train metrics @ 25k (normalized, <1 = working) |
|---|---|---|
orig/ |
motion_dim=192, attn pool (baseline) | recon_rel 0.045 / idm_rel 0.158 |
mdimnone/ |
motion_dim=None (384, identity) | recon_rel 0.012 / idm_rel 0.268 |
encspatial/ |
enc_spatial_mix on (mdim192, attn) | recon_rel 0.039 / idm_rel 0.061 |
Metrics are training values (not held-out eval).
Load
from huggingface_hub import snapshot_download
path = snapshot_download("hmkang/wam_ctxcomp_stage1", allow_patterns="encspatial/*")
- Downloads last month
- 32