# Any4D + Wan 2.1 Fun 1.3B InP

This document describes the Wan-based backbone integrated into Any4D as an
alternative to the default Cosmos-Predict2 2B DiT. The base
[Any4D README](../../README.md) covers the data pipeline, configuration system,
and Cosmos-based recipes; this file only documents the Wan-specific pieces.

## Setup

### 1. Default text embeddings

The pipeline loads a precomputed UMT5-XXL embedding for a fixed prompt 
``a high quality video of driving, realistic, cinematic" 
from `custom/wan/default_text_wan.pkl`. To regenerate (e.g. with a different
default prompt), run:

```
python -m custom.wan.build_default_text_wan \
    --prompt "a high quality video of driving, realistic, cinematic" \
    --umt5_tokenizer_path /path/to/Wan2.1-Fun-1.3B-InP/google/umt5-xxl \
    --umt5_weights_path   /path/to/Wan2.1-Fun-1.3B-InP/models_t5_umt5-xxl-enc-bf16.pth \
    --out custom/wan/default_text_wan.pkl
```

### 2. (Optional) Run inference with an init checkpoint

Save a checkpoint with Any4D's surgery-added params to zero. 
Use this as the checkpoint in inference 
script to check video quality before training. 

```
python -m custom.eval.save_init_checkpoint_wan \
    --dit_path  /path/to/Wan2.1-Fun-1.3B-InP \
    --vae_path  /path/to/Wan2.1_VAE.pth \
    --out       /path/to/init_any4d_wan.pt
```

## Training

`custom/experiment/examples_anydata/rwm5_drive_wan.py` is the
driving-world-model counterpart of `rwm5_drive.py`. 
Launch (inside the docker container described in the main README):

```
torchrun --nproc_per_node=1 --master_port=22045 -m scripts.train \
    -- experiment=any4d_rwm5_drive_wan \
    model.config.fsdp_shard_size=1 \
    dataloader_train.batch_size=1 \
    job.group=example job.name=wm5_drive_wan \
    job.local_root=/raid/$USER/any4d \
    job.s3_root=s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d
```

For robotics, `examples_anydata/rwm3_robot_wan.py` is the counterpart of
`rwm3_robot.py`.


## Inference

An example inference input can be found in `s3://gaia-e2e-driving-components/any4d/examples/anydrive_val`. 
Checkpoint for this example input is in `s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/a4d2/example/05-18-18-25_wm5_drive_wan_4gpu_6cam_perblock/model`

```
python -m custom.eval.infer_wan \
    --ckpt_path /path/to/iter_NNNNN.pt \
    --vae_path  /path/to/Wan2.1-Fun-1.3B-InP/Wan2.1_VAE.pth \
    --out_dir   /tmp/wan_pred \
    --video_dir /path/with/CAMERA_01.mp4 CAMERA_05.mp4 ... \
    --num_views 6 --num_frames 21 \
    --target_h 192 --target_w 320 \
    --traj_npy /path/to/traj.npy
```

## Mask convention and difference

Wan's `patch_embedding` and `head` are shared across views (single shared
module, no per-view duplicates). This is different from the Cosmos
`Any4DDiT`, which instantiates a dedicated `x_embedder_newviews[i]` /
`final_layer_newviews[i]` per extra view. Per-view embedders can give different views different input channel
counts (e.g. RGB-only vs RGB+depth). 
The Wan integration reshapes every view to a uniform 36-channel layout via `_to_wan_36ch`, so the pretrained
`patch_embedding` and `head` can be applied as-is to all views. This keeps the network closer to the pretrained model, hence non-garbage videos at initialization.


Mask inputs are also slightly different from the Cosmos backbone but provide similar functionality. 
The goal is to be consistent with the pretrained model and retain pretrained capability. 

Cosmos: uses a 1-channel mask. Any4D streams carry `rgb` (16-ch latent) +
`rgb_input_mask` (1-ch) + `rgb_output_mask` (1-ch). At assemble time, the
input_mask gates clean-x0-pinning at given frames; the rgb slot is the
denoising target everywhere else, i.e. input and output shares the 16-ch latent

Wan 2.1 Fun InP: uses a 4-channel mask plus a 16-channel reference latent,
giving a 36-channel DiT input layout:

```
ch [ 0:16]  noisy_slot   — xt (denoised everywhere; what the DiT predicts)
ch [16:20]  mask_slot    — 4-ch broadcast of the input mask
ch [20:36]  ref_slot     — reference latent (VAE-encoded pixel video with
                          real pixels at given frames, gray padding elsewhere)
```

Conditioning is delivered through the ref slot, not through clean-frame
replacement in the noisy slot. To make Any4D's stream-and-mask machinery
produce this layout, the integration uses 
`Any4DWanDiT._to_wan_36ch` to reshape to Wan's 36-channel input:

```
Any4D stream (34 ch)                  →  Wan input (36 ch)
ch [ 0:16] rgb latent (noisy)         →  ch [ 0:16] noisy_slot
ch [16:17] input_mask (1 ch)          →  ch [16:20] mask_slot (1→4 broadcast)
ch [17:18] output_mask (kept; unused) → (dropped)
ch [18:34] rgb_ref (16 ch)            →  ch [20:36] ref_slot
```

## Supported Any4D features and TODOs

Status

| Feature | Status | Notes |
| --- | --- | --- |
| Multi-view (`num_views≥2`, `video_concat_mode='view'`) | ✅ | Tested with 2/6/11 views. Joint self-attention over per-view token sequences. |
| Surgery-added params (view embs, adaln, sattn, xattn) | ✅ | All zero-init so the network starts as identity-of-Wan. |
| FSDP | ✅ | `Any4DWanDiT.fully_shard` shards `self.wan.blocks`. Tested at 4-GPU. |
| CP | ❌ | TODO. |
| Variable cross-modal (rgb, depth, flow) tasks | ❌ | `task_probs.cross_modal>0` is not implemented. |
| Plucker embedding and DVS | ❌ | Not implemented. For world model task, cameras poses are given by input frames|


## File layout

```
custom/wan/
├── __init__.py
├── README.md                # this file
├── a4d_config_wan.py            # Any4DWanConfig
├── a4d_model_wan.py             # Any4DWanModel
├── a4d_network_wan.py           # Any4DWanDiT
├── a4d_pipe_wan.py              # Any4DWanPipeline (denoise / generate / load_checkpoint)
├── a4d_vae_wan.py               # VAE wrapper
├── build_default_text_wan.py    # Precompute default UMT5 embedding
├── save_init_checkpoint_wan.py  # save a initialized checkpoint
└── wan_vae_adapter.py           # TokenizerInterface adapter around cosmos_predict2.tokenizers.tokenizer.VAE

cosmos_predict2/schedulers/wan_flow_match_scheduler.py   # flow matching scheduler

custom/eval/
└── infer_wan.py                 # minimal Any4D + Wan inference

custom/experiment/examples_anydata/
├── rwm3_robot_wan.py         # robotics
└── rwm5_drive_wan.py         # driving
```

`use_wan_backbone=True` in `any4d_config` swaps the entire video diffusion
backbone for Wan 2.1 Fun 1.3B InP while keeping the rest of Any4D (logistics, 
surgery, dataloader, visualizer, multi-view streams) unchanged.

* DiT: `custom/wan/a4d_network_wan.py` — `Any4DWanDiT`, embeds the Wan 2.1 DiT.
* Pipeline: `custom/wan/a4d_pipe_wan.py` — `Any4DWanPipeline`, subclass of
  `Any4DPipeline` with flow-matching `denoise()` (no EDM preconditioning).
* VAE wrapper: `custom/wan/a4d_vae_wan.py` — Wan2.1 VAE (16-ch latent, 4×8×8
  compression). The underlying VAE model is `cosmos_predict2.tokenizers.tokenizer.VAE` 
  (same as Wan 2.1 VAE); the adapter lives at `custom/wan/wan_vae_adapter.py`.
* Scheduler: `cosmos_predict2/schedulers/wan_flow_match_scheduler.py`.
* Config subclass: `custom/wan/a4d_config_wan.py` — `Any4DWanConfig`, 
  overrides inference defaults (`val_sigma_max=1.0`, `val_sigma_min=0.003`,
  `val_cfg_scale=5.0`).
* Training model class: `custom/wan/a4d_model_wan.py` — `Any4DWanModel`,
  overrides sigma sampling / forward process / loss weighting to flow-match
  conventions while inheriting masks, streams, conditioner, viz, etc.

The experiment-level entry point is `template_any4d_wan_1_3b` in
`custom/experiment/template.py`, which routes through the
`any4d_wan_1_3b` model registered in
`cosmos_predict2/configs/base/defaults/model.py`.