OmniWM + Wan2.2-TI2V-5B backbone

Every additive branch is gated so the SD/NanoWM/Omni path is never entered for Wan and vice-versa. The seams below were confirmed by code investigation.

The registry (src/models/__init__.py:11) merges NanoWM_models + Omni_models; dispatch guards on the substring 'NanoWM'/'Omni' in args.model.arch. New arch Omni-Wan-5B (registered in new src/models/omni_wan.py → OmniWan_models) deliberately keeps the Omni- prefix so the joint-knob plumbing still passes action_expert_depth/action_loss_dim.

# src/models/__init__.py — additive Wan-only knobs
if "Omni-Wan" in arch:
    common["video_backbone"] = "wan2.2-ti2v-5b"
    common["wan_repo"] = args.model.get("wan_repo", "Wan-AI/Wan2.2-TI2V-5B-Diffusers")
    common["peft"] = OmegaConf.to_container(args.model.get("peft", {}), resolve=True)
    common["latent_channels"] = args.model.get("latent_channels", 48)

New module kind omni_wan: an Omni-Wan arch with any other module raises. OmniWanTrainingModule lives in a new src/experiments/train_wan_experiment.py, imported lazily inside the omni_wan branch so diffusers' Wan classes / peft are not required for NanoWM/Omni runs.

ActionInEmbedder, ActionBlock, ActionHead, TimestepEmbedder, CrossAttention are reused by import from omni.py/nanowm.py — backbone-agnostic; only the cross-attn K/V width changes. The empty-prompt umT5 embedding [1, 512, 4096] is precomputed once and registered as a buffer (world model is null-text).

f0 token surgery

Wan has no output_hidden_states; we fork the block loop (copying WanTransformer3DModel.forward) and after each of the 30 blocks slice the first-frame latent tokens. Token order is frame-major (patch_embedding(Conv3d).flatten(2).transpose(1,2)), so f0 = the first n_f0 = (H_lat/2)·(W_lat/2) tokens.

Cross-attn realization

Clean-action conditioning into Wan

Wan TI2V conditions on the initial image by latent inpainting + per-token 2D timestep masking (expand_timesteps), not a separate cross-attn stream. We adopt the same mechanism for f0 = current obs: place the clean current-obs latent at latent-frame 0, set its per-token timestep to 0, and run latent_model_input = (1-mask)·cond + mask·noisy. Action conditioning of the video goes additively into Wan's timestep_proj/AdaLN temb per latent-frame via a zero-init action_to_temb hook, time-shifted so latent-frame 0 receives zero action — f0 stays a leak-free, policy-readable observation.

The core break: the repo's contract is per-frame 2D SD latents [T, V, 4, 32, 32] sliced by frame index. Wan's VAE is 3D-causal: T_lat = 1 + (T-1)//4, z_dim=48, spatial 16×. On-disk latents become [T_lat, V, 48, 16, 16]; the naive per-frame slice is wrong for Wan.

Precompute over on-the-fly: new tool src/tools/precompute_latents_wan.py (fork; SD tool untouched) encodes each view independently in 3D over the whole episode and stores normalized fp16 [T_lat, V, 48, 16, 16]. meta.json records z_dim, scale factors, latents_mean/std, and the frame_to_latent map.

Parallel cache key + slice logic: a new mixture knob latents_backbone (default sdvae) switches the latent dir suffix (<name>_wanvae_<size>) and dispatches precompute. Data sources gain a temporal_compression field (default 1) mapping a raw-frame window to its latent-frame window; with 1 the behavior is bit-identical to today.

Wan is a single video stream; the repo carries V views (e.g. 2 wrist cams). Encode each view independently through the Wan VAE, then token-concat the views' patch tokens along the sequence with a learned per-view embedding — matching the existing view_embed + (v p) contract. f0 tokens become V·n_f0 (both views' frame-0 tokens). Do not concat views along time (mixes views temporally, breaks f0). First deliverable runs V=1 to de-risk; lift to V=2 once V=1 trains.

Wan is flow-matching (UniPC flow_prediction, flow_shift=5.0); the Gaussian compute_joint_loss stays untouched. A parallel compute_joint_loss_flow lands in train_wan_experiment.py:

A new wrapper src/diffusion/flow_match.py holds add_noise, training_target, Euler step, and σ shifting — gaussian_diffusion.py untouched.

The existing dfot_action_sample is coupled to Gaussian DDIM (alphas_cumprod); a parallel dfot_action_sample_flow (Euler ODE over sigmas) is added — df_sample.py untouched. This is what makes the 5B policy tractable, the FastWAM prefill_video_cache analogue and identical in spirit to the current Omni prefill. OmniWanTrainingModule.policy_sample mirrors the Omni version and reuses _denorm_action unchanged.

Two new config files, siblings of model/omni_s2.yaml / experiment/train_omni.yaml (which stay untouched):

# src/configs/model/omni_wan_5b.yaml
arch: Omni-Wan-5B
image_size: 256
latent_channels: 48
num_frames: 33          # -> T_lat = 9
action_expert_hidden: 1024
peft: {type: lora, r: 32, alpha: 64}
wan_repo: Wan-AI/Wan2.2-TI2V-5B-Diffusers

# src/configs/experiment/train_omni_wan.yaml
module: omni_wan
mixture: {enabled: true, latents: auto, latents_backbone: wanvae}
diffusion: {objective: flow_match, flow_shift: 5.0, action_loss_weight: 1.0}
infra: {gradient_checkpointing: true, vae_precision: fp32}