Modelpublic

seungkukim/dexjoco_bimanual_multitask_handaction_joint_wan22ti2v5b_L18_txtcache-60k

sourceHugging Faceotherupdated 1mo agoView on Hugging Face

0likes11downloads

Model Card

dexjocobimanualmultitaskhandactionjointwan22ti2v5bL18_txtcache @ step 60,000

DiT4DiT joint checkpoint -- a Wan2.2-TI2V-5B video backbone plus a DiT-B action head with a hand-geometry stream (66-D) -- finetuned on the DexJoCo bimanual multitask dataset (`seungkukim/dexjoco_lerobot_v20`), with cached umT5 text embeddings (txtcache) and feature extraction at layer 18.

What is in here

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

prefix	tensors	what
`backbone.extractor.transformer.dit.*`	825	finetuned Wan2.2 DiT
`backbone.extractor.vae.*`	196	Wan2.2 VAE
`action_head.*`	289	DiT-B action head, incl. `hand_encoder` / `hand_decoder`

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

Arm configuration

key	value
`training_mode`	`joint`
`hand_action`	`True`
`hand_target_kind`	`dexjoco_bimanual`
`hand_pose_mode`	`absolute`
`hand_action_weight`	`0.5`
`hand_action_stride` / `shift`	`2` / `1`
`hand_dn`	`false` -- plain hand-stream arm, no DN group
`wan_feature_extraction_layer`	`18` of 30

Hand target is 66-D: [right 0:33 | left 33:66]; per hand [flange pos 0:3 | flange rot6d 3:9 | 4 palm-local fingertips 9:21 | 4 palm-local medial (PIP) origins 21:33].

Hand whitening stats are shipped WITH this repo

action_head.py registers the de-whitening constants as non-persistent buffers, so they are not in the shards, and load_hand_norm_stats() raises rather than falling back to identity. Both config.json and processor_config.json name them by the repo-relative path

scripts/dexjoco/dexjoco_bimanual_hand_action_stats_absolute.json

which resolves against the current working directory. That file is included in this repo at exactly that relative path (66-D, mode=absolute, kind=dexjoco_bimanual), so loading works from the snapshot directory. If you load from elsewhere, either cd into the snapshot or set hand_action_norm_stats_path to an absolute path.

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 umT5 text encoder and tokenizer are never saved into the checkpoint and always come from that base snapshot:

bash

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

python

from gr00t.model.wam_dit4dit import WAMDiT4DiT

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