WAM_DIT4DIT — context pooling with a dedicated cross-attention pooler (xattn), RoboCasa kitchen
Wan2.2-TI2V-5B video DiT (vendored) fine-tuned on RoboCasa, base recipe, effective batch 64, 4-latin history (nin=25/nout=41, fdf=2). Before block L=3 the 3 PAST cond latent frames are pooled into ONE motion frame = mean(past) + ContextPooler(q = current-frame tokens, k/v = past tokens, DiT-sized 24 heads, 3D rope on q/k, zero-init output) -> step 0 == plain avg pooling.
| folder | pooler |
|---|---|
xattn_L3/checkpoint-<step> |
cross-attn only |
xattn_ffn_L3/checkpoint-<step> |
cross-attn + FFN refine (4x hidden, zero-init) |
Every 20k steps. Weights + configs only. Code: https://github.com/HEMMO0208/wam
(gr00t/model/wam_dit4dit/wan22/wan_video_dit.py ContextPooler / _ctx_pool_apply).
- Downloads last month
- -