dexjoco_bimanual_multitask_videoonly_wan22ti2v5b_L18_txtcache-60k

dexjoco_bimanual_multitask_videoonly_wan22ti2v5b_L18_txtcache @ step 60,000

Video-only WAM-DiT4DiT checkpoint: a Wan2.2-TI2V-5B backbone finetuned on the DexJoCo bimanual multitask dataset (seungkukim/dexjoco_lerobot_v20), with cached text embeddings (txtcache) and feature extraction at layer 18.

What is in here

Inference files only -- 12 files, 11.41 GB. The 1,021 tensors in the shards are:

There is no action head (this is a video-only run) and no text encoder -- the umT5 encoder and tokenizer are never saved into the checkpoint.

The training-only artifacts are deliberately excluded: the DeepSpeed global_step60000/ optimizer state (~68 GB), rng_state_*.pth, scheduler.pt, trainer_state.json, training_args.bin, zero_to_fp32.py, latest, wandb_config.json.

Requires the base Wan2.2 snapshot

This repo is not self-contained. WAMDiT4DiT.from_pretrained() first builds the backbone from the base snapshot named by config.json:wan_model_path, then overlays these shards onto it. The text encoder and tokenizer always come from that base snapshot, so you need it available:

huggingface-cli download Wan-AI/Wan2.2-TI2V-5B-Diffusers --local-dir /path/to/Wan2.2-TI2V-5B-Diffusers

Currently wan_model_path = /data/seungku/hf_cache/Wan2.2-TI2V-5B-Diffusers.

config.json still points at the local training path /data/seungku/hf_cache/Wan2.2-TI2V-5B-Diffusers with wan_local_files_only=true. Override wan_model_path to Wan-AI/Wan2.2-TI2V-5B-Diffusers (and set wan_local_files_only=false) before loading this on another machine.

Load

from gr00t.model.wam_dit4dit import WAMDiT4DiT

model = WAMDiT4DiT.from_pretrained("seungkukim/dexjoco_bimanual_multitask_videoonly_wan22ti2v5b_L18_txtcache-60k", torch_dtype="bfloat16")