[feat]: add bidirectional and causal WanTrack support

2026-07-23 02:23:10 -07:00
parent 41dcd065ae
commit ce4fa8ef2d
27 changed files with 2868 additions and 44 deletions
@@ -0,0 +1,188 @@
+# WanTrack training
+
+WanTrack extends Wan I2V with sparse point tracks. The same model wrapper and
+conditioning path are used for preprocessing, training, validation, and
+standalone inference. Both bidirectional and causal checkpoints are supported.
+
+## Build an initialization checkpoint
+
+WanTrack uses the pretrained control slot from Wan2.1-Fun Control as its track
+slot. Convert a VideoX-Fun control checkpoint together with a diffusers-format
+Wan2.1-Fun InP base:
+
+```bash
+python scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py \
+    --inp-base models/Wan2.1-Fun-1.3B-InP-Diffusers \
+    --control-ckpt models/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors \
+    --out models/wantrack-control-init
+```
+
+For the causal transformer, add `--causal` and use a separate output directory:
+
+```bash
+python scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py \
+    --inp-base models/Wan2.1-Fun-1.3B-InP-Diffusers \
+    --control-ckpt models/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors \
+    --out models/wantrack-control-causal-init \
+    --causal
+```
+
+`--causal` selects `CausalTrackWanTransformer3DModel` in the generated
+transformer config. Without it, the converter selects
+`TrackWanTransformer3DModel`.
+
+The converted patch input has 52 channels:
+
+| Channels | Meaning |
+|----------|---------|
+| `0:16` | Noisy video latent |
+| `16:20` | I2V mask |
+| `20:36` | First-frame latent |
+| `36:52` | Track map |
+
+The first three entries form the usual 36-channel Wan I2V input. A
+`TrackEncoder` rasterizes sparse points on the VAE grid and appends the
+16-channel track map.
+
+## Prepare point-track data
+
+Each video needs an `.npz` sidecar referenced by `points_path` in its metadata.
+The archive must contain:
+
+- `tracks`: floating-point source-pixel coordinates with shape `[T, N, 2]`.
+  The final dimension is `(x, y)`.
+- `visibility`: a boolean or numeric visibility mask with shape `[T, N]`.
+
+`T` is the source-video timeline and `N` is the number of point tracks. The
+preprocessor applies the same temporal sample and center-crop/resize transform
+to the sidecar as it applies to the video. Coordinates should therefore be in
+the original video's pixel space, not normalized or pre-cropped. Points outside
+the retained crop are marked invisible.
+
+For example, a merged-dataset annotation can contain:
+
+```json
+[
+  {
+    "path": "videos/clip_0001.mp4",
+    "points_path": "tracks/clip_0001.npz",
+    "cap": ["A cyclist follows a winding road."],
+    "resolution": {"width": 1920, "height": 1080},
+    "fps": 24.0,
+    "duration": 5.0,
+    "num_frames": 120
+  }
+]
+```
+
+The paths are relative to the dataset root named in the merge file:
+
+```text
+data/wantrack/raw,data/wantrack/metadata.json
+```
+
+All examples combined into one training batch must have a stackable point
+dimension. Keeping `N` fixed across the dataset is the simplest option.
+
+Run the I2V-track preprocessor with matching pixel and latent lengths. Wan's
+temporal compression is four, so 81 pixel frames produce 21 latent frames:
+
+```bash
+torchrun --nproc_per_node=1 fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path models/wantrack-control-init \
+    --data_merge_path data/wantrack/data_merge.txt \
+    --output_dir data/wantrack/preprocessed \
+    --preprocess_task i2v_track \
+    --num_frames 81 \
+    --num_latent_t 21 \
+    --max_height 480 \
+    --max_width 832 \
+    --train_fps 16 \
+    --preprocess_video_batch_size 1
+```
+
+The trainer reads the resulting
+`data/wantrack/preprocessed/combined_parquet_dataset` directory with
+`preprocessed_data_type: i2v_track`.
+
+## Train the bidirectional model
+
+The bidirectional training wrapper is
+`fastvideo.train.models.wantrack.WanTrackModel`; it loads
+`TrackWanTransformer3DModel`.
+
+```bash
+torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \
+    --config examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml
+```
+
+The example optionally subsamples points and applies temporal track masking.
+Track IDs are sampled once for a training sample and then reused by every
+denoising call and its conditional/unconditional branches. Re-sampling IDs per
+call would change the point embeddings even when the coordinates are
+unchanged.
+
+## Train the causal model
+
+The causal wrapper is
+`fastvideo.train.models.wantrack.WanTrackCausalModel`; it loads
+`CausalTrackWanTransformer3DModel`. The example uses
+`TeacherForcingSFTMethod`, so clean history and the noisy current chunk receive
+the same I2V and track conditioning:
+
+```bash
+torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \
+    --config examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml
+```
+
+Track encoding is causal at the VAE temporal boundary. The implementation
+left-pads the complete source track sequence at its global beginning, encodes
+it, and then slices the latent track map at the current latent `start_frame`.
+It must not independently pad and encode each later chunk: doing so would
+reset temporal context and produce a different feature for the same point
+history.
+
+The same stable `track_ids` are reused across all chunks, denoising steps, and
+CFG branches. Causal sampling follows the RobotWM streaming contract:
+`predict_noise_streaming()` owns the KV caches, each denoised block is committed
+once as clean context, and caches are cleared at the sample boundary.
+
+## Validation and inference
+
+Both example configs use
+`fastvideo.train.callbacks.track_validation.TrackValidationCallback`. It loads
+fixed samples from the preprocessed WanTrack parquet, builds conditions through
+the student's normal `prepare_batch()` path, generates videos, overlays the
+active tracks, and logs them to the configured tracker. Set
+`callbacks.track_validation.val_data_path` to use a dedicated validation
+parquet; otherwise it samples from the training parquet.
+
+The callback and standalone callers share two helpers:
+
+```python
+from fastvideo.train.models.wantrack.inference import (
+    prepare_wantrack_batch,
+    sample_wantrack,
+)
+
+batch = prepare_wantrack_batch(
+    model,
+    raw_batch,
+    seed=1000,
+    latents_source="zeros",
+)
+latents = sample_wantrack(
+    model,
+    batch,
+    num_inference_steps=30,
+    seed=1000,
+    text_guidance_scale=3.0,
+    motion_guidance_scale=1.5,
+)
+video = model.decode_latents(latents)
+```
+
+`sample_wantrack()` denoises the complete clip for `WanTrackModel`. For
+`WanTrackCausalModel`, it uses the existing `CausalModelBase` streaming API and
+the transformer's configured block size; it does not modify or fork the common
+Wan causal denoising stage.
@@ -0,0 +1,82 @@
+# WanTrack bidirectional I2V fine-tuning.
+#
+# Build models/wantrack-control-init with
+# scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py, then preprocess
+# the dataset with --preprocess_task i2v_track.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.WanTrackModel
+    init_from: models/wantrack-control-init
+    trainable: true
+    track_augmentation:
+      enabled: true
+      min_points: 1000
+      max_points: 2500
+      temporal_mask_probability: 0.2
+      temporal_mask_chunk_size: 8
+      motion_dropout_probability: 0.0
+      text_dropout_probability: 0.0
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/wantrack/preprocessed/combined_parquet_dataset
+    preprocessed_data_type: i2v_track
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wantrack_bidirectional_i2v
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: null
+
+  tracker:
+    project_name: fastvideo
+    run_name: wantrack_bidirectional_i2v
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 2
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    validate_at_start: false
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,88 @@
+# WanTrack causal I2V teacher-forcing fine-tuning.
+#
+# Build models/wantrack-control-causal-init with the checkpoint converter's
+# --causal flag. Track IDs and the full encoded track map are shared by every
+# latent chunk.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.WanTrackCausalModel
+    init_from: models/wantrack-control-causal-init
+    trainable: true
+    track_augmentation:
+      enabled: true
+      min_points: 1000
+      max_points: 2500
+      temporal_mask_probability: 0.2
+      temporal_mask_chunk_size: 8
+      motion_dropout_probability: 0.0
+      text_dropout_probability: 0.0
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
+  chunk_size: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/wantrack/preprocessed/combined_parquet_dataset
+    preprocessed_data_type: i2v_track
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 21
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wantrack_causal_i2v
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: null
+
+  tracker:
+    project_name: fastvideo
+    run_name: wantrack_causal_i2v
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 250
+    num_val_samples: 2
+    num_inference_steps: 30
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    validate_at_start: false
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1
+    rope_cache_policy: relativistic
+    causal_train_attention: triton
@@ -11,6 +11,7 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
+from fastvideo.configs.models.dits.trackwan import CausalTrackWanVideoConfig, TrackWanVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
@@ -23,5 +24,5 @@ __all__ = [
    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
    "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig",
-    "ZImageDiTConfig"
+    "ZImageDiTConfig", "TrackWanVideoConfig", "CausalTrackWanVideoConfig"
 ]
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for track-conditioned Wan transformers."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig
+from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
+
+
+def _default_track_config() -> dict[str, int | bool]:
+    return {
+        "id_dim": 64,
+        "track_channels": 16,
+        "vae_spatial_compression": 8,
+        "vae_temporal_compression": 4,
+        "max_track_id": 100_000,
+        "zero_init_head": False,
+        "use_bias": False,
+    }
+
+
+@dataclass
+class TrackWanVideoArchConfig(WanVideoArchConfig):
+    """Wan I2V input widened with a latent-aligned point-track map.
+
+    Channel order is fixed and checkpoint-visible:
+    noisy latent (16), I2V mask (4), first-frame latent (16), track map (16).
+    """
+
+    in_channels: int = 52
+    out_channels: int = 16
+    image_dim: int = 1280
+    added_kv_proj_dim: int | None = 5120
+    track_config: dict[str, int | bool] = field(default_factory=_default_track_config)
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        track_channels = int(self.track_config.get("track_channels", 0))
+        expected_in_channels = self.num_channels_latents + 20 + track_channels
+        if track_channels <= 0:
+            raise ValueError("track_config.track_channels must be positive")
+        if self.in_channels != expected_in_channels:
+            raise ValueError("TrackWan in_channels must equal latent channels + 20 I2V "
+                             f"channels + track channels; got {self.in_channels}, expected "
+                             f"{expected_in_channels}")
+
+
+@dataclass
+class TrackWanVideoConfig(WanVideoConfig):
+    arch_config: DiTArchConfig = field(default_factory=TrackWanVideoArchConfig)
+    prefix: str = "Wan"
+
+
+@dataclass
+class CausalTrackWanVideoArchConfig(TrackWanVideoArchConfig):
+    """Explicit config type for the causal TrackWan architecture."""
+
+
+@dataclass
+class CausalTrackWanVideoConfig(TrackWanVideoConfig):
+    arch_config: DiTArchConfig = field(default_factory=CausalTrackWanVideoArchConfig)
@@ -50,6 +50,46 @@ pyarrow_schema_i2v = pa.schema([
    pa.field("fps", pa.float64()),
 ])

+pyarrow_schema_i2v_track = pa.schema([
+    pa.field("id", pa.string()),
+    pa.field("vae_latent_bytes", pa.binary()),
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    pa.field("vae_latent_dtype", pa.string()),
+    pa.field("text_embedding_bytes", pa.binary()),
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("first_frame_latent_bytes", pa.binary()),
+    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
+    pa.field("first_frame_latent_dtype", pa.string()),
+    pa.field("clip_feature_bytes", pa.binary()),
+    pa.field("clip_feature_shape", pa.list_(pa.int64())),
+    pa.field("clip_feature_dtype", pa.string()),
+    # MotionStream point tracks, normalized to [0, 1].
+    pa.field("track_points_bytes", pa.binary()),
+    pa.field("track_points_shape", pa.list_(pa.int64())),  # [T, N, 2]
+    pa.field("track_points_dtype", pa.string()),
+    pa.field("track_visibility_bytes", pa.binary()),
+    pa.field("track_visibility_shape", pa.list_(pa.int64())),  # [T, N]
+    pa.field("track_visibility_dtype", pa.string()),
+    # Optional preprocessing signals used by object-covering track sampling.
+    # New preprocessing writes dense fallback values when segmentation or
+    # informativeness weights are unavailable.
+    pa.field("object_ids_bytes", pa.binary()),
+    pa.field("object_ids_shape", pa.list_(pa.int64())),  # [N]
+    pa.field("object_ids_dtype", pa.string()),
+    pa.field("track_weights_bytes", pa.binary()),
+    pa.field("track_weights_shape", pa.list_(pa.int64())),  # [N]
+    pa.field("track_weights_dtype", pa.string()),
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])
+

 pyarrow_schema_t2v = pa.schema([
    pa.field("id", pa.string()),
@@ -189,4 +229,4 @@ pyarrow_schema_matrixgame2_ode_trajectory = pa.schema([
    pa.field("file_name", pa.string()),
    pa.field("caption", pa.string()),
    pa.field("media_type", pa.string()),  # 'image' or 'video'
-])
+])
@@ -41,6 +41,7 @@ class PreprocessBatch:
    sample_frame_index: list[int] | None = None
    sample_num_frames: int | None = None
    action_path: str | None = None
+    points_path: str | None = None

    # Processed data
    pixel_values: torch.Tensor | None = None
@@ -471,6 +472,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
            item["path"] = opj(folder, item["path"])
            if "action_path" in item and item["action_path"]:
                item["action_path"] = opj(folder, item["action_path"])
+            if "points_path" in item and item["points_path"]:
+                item["points_path"] = opj(folder, item["points_path"])

        return data_items

@@ -492,7 +495,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                                    resolution=item.get("resolution"),
                                    fps=item.get("fps"),
                                    duration=item.get("duration"),
-                                    action_path=item.get("action_path"))
+                                    action_path=item.get("action_path"),
+                                    points_path=item.get("points_path"))

            # Apply filtering stages
            if not self._apply_filter_stages(batch, filter_counts):
@@ -570,6 +574,16 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        # Add action_path
        if batch.action_path:
            result["action_path"] = batch.action_path
+        if batch.points_path:
+            result["points_path"] = batch.points_path
+        if batch.sample_frame_index is not None:
+            result["sample_frame_index"] = torch.tensor(batch.sample_frame_index,
+                                                        dtype=torch.long)
+        if batch.resolution is not None:
+            if "width" in batch.resolution:
+                result["source_width"] = int(batch.resolution["width"])
+            if "height" in batch.resolution:
+                result["source_height"] = int(batch.resolution["height"])

        return result

@@ -40,7 +40,7 @@ from fastvideo.models.dits._causal_train_attention import (
 )
 from fastvideo.models.dits._relative_rope import relativistic_window_offsets
 from fastvideo.models.dits.base import BaseDiT
-from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding
+from fastvideo.models.dits.wanvideo import WanI2VCrossAttention, WanT2VCrossAttention, WanTimeTextImageEmbedding
 from fastvideo.platforms import AttentionBackendEnum, current_platform

 logger = init_logger(__name__)
@@ -311,12 +311,25 @@ class CausalWanTransformerBlock(nn.Module):
            elementwise_affine=True,
            dtype=torch.float32)

-        # 2. Cross-attention
-        # Only T2V for now
-        self.attn2 = WanT2VCrossAttention(dim,
-                                            num_heads,
-                                            qk_norm=qk_norm,
-                                            eps=eps)
+        # 2. Cross-attention. I2V checkpoints prepend CLIP image tokens and
+        # carry a second pair of projections, which must be preserved during
+        # causal rollouts.
+        if added_kv_proj_dim is not None:
+            self.attn2 = WanI2VCrossAttention(
+                dim,
+                num_heads,
+                qk_norm=qk_norm,
+                eps=eps,
+                prefix=f"{prefix}.attn2",
+            )
+        else:
+            self.attn2 = WanT2VCrossAttention(
+                dim,
+                num_heads,
+                qk_norm=qk_norm,
+                eps=eps,
+                prefix=f"{prefix}.attn2",
+            )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
            norm_type="layer",
@@ -757,11 +770,10 @@ class CausalWanTransformer3DModel(BaseDiT):
        orig_dtype = hidden_states.dtype
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (
+                encoder_hidden_states_image[0]
+                if encoder_hidden_states_image else None)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
@@ -876,11 +888,10 @@ class CausalWanTransformer3DModel(BaseDiT):
        teacher_forcing = clean_x is not None
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (
+                encoder_hidden_states_image[0]
+                if encoder_hidden_states_image else None)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
@@ -0,0 +1,17 @@
+from fastvideo.models.dits.trackwan.model import (
+    CausalTrackWanTransformer3DModel,
+    TrackWanTransformer3DModel,
+)
+from fastvideo.models.dits.trackwan.track_encoder import TrackEncoder
+
+__all__ = [
+    "TrackEncoder",
+    "TrackWanTransformer3DModel",
+    "CausalTrackWanTransformer3DModel",
+]
+
+# Entry points for model registry discovery.
+EntryClass = [
+    TrackWanTransformer3DModel,
+    CausalTrackWanTransformer3DModel,
+]
@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bidirectional and causal track-conditioned Wan transformers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.configs.models.dits.trackwan import TrackWanVideoConfig
+from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel
+from fastvideo.models.dits.trackwan.track_encoder import TrackEncoder
+from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+
+
+class _TrackConditioningMixin:
+    supports_track_input = True
+    track_channels: int
+    track_encoder: TrackEncoder
+
+    def _init_track_conditioning(
+        self,
+        config: TrackWanVideoConfig,
+        hf_config: dict[str, Any],
+    ) -> None:
+        track_config = dict(getattr(config, "track_config", None) or {})
+        # Checkpoint config is authoritative for checkpoint-visible TrackEncoder
+        # shapes and legacy bias tensors. Defaults only fill missing fields.
+        if isinstance(hf_config, dict):
+            track_config.update(dict(hf_config.get("track_config", {}) or {}))
+
+        self.track_channels = int(
+            track_config.get("track_channels", self.in_channels - 36))
+        base_channels = self.in_channels - self.track_channels
+        expected_base_channels = self.num_channels_latents + 20
+        if self.track_channels <= 0:
+            raise ValueError("TrackWan requires positive track_channels")
+        if base_channels != expected_base_channels:
+            raise ValueError(
+                "TrackWan expects noisy latent + 20 I2V channels before "
+                f"track conditioning; got {base_channels}, expected "
+                f"{expected_base_channels}")
+
+        self.track_encoder = TrackEncoder(
+            id_dim=int(track_config.get("id_dim", 64)),
+            track_channels=self.track_channels,
+            vae_spatial_compression=int(
+                track_config.get("vae_spatial_compression", 8)),
+            vae_temporal_compression=int(
+                track_config.get("vae_temporal_compression", 4)),
+            max_track_id=int(track_config.get("max_track_id", 100_000)),
+            zero_init=bool(track_config.get("zero_init_head", False)),
+            use_bias=bool(track_config.get("use_bias", False)),
+        )
+
+    def _append_track_conditioning(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        track_points: torch.Tensor | None,
+        track_visibility: torch.Tensor | None,
+        track_ids: torch.Tensor | None,
+        start_frame: int,
+    ) -> torch.Tensor:
+        expected_channels = self.in_channels - self.track_channels
+        if hidden_states.ndim != 5:
+            raise ValueError(
+                "TrackWan hidden_states must be [B, C, T, H, W], got "
+                f"{tuple(hidden_states.shape)}")
+        if hidden_states.shape[1] != expected_channels:
+            raise ValueError(
+                "TrackWan received an unexpected pre-track channel count: "
+                f"{hidden_states.shape[1]} (expected {expected_channels})")
+        if (track_points is None) != (track_visibility is None):
+            raise ValueError(
+                "track_points and track_visibility must be provided together")
+
+        batch_size, _, latent_t, latent_h, latent_w = hidden_states.shape
+        if track_points is None:
+            track_map = hidden_states.new_zeros(
+                batch_size,
+                self.track_channels,
+                latent_t,
+                latent_h,
+                latent_w,
+            )
+        else:
+            if track_points.shape[0] != batch_size:
+                raise ValueError(
+                    "track batch size must match hidden_states batch size")
+            temporal_ratio = self.track_encoder.vae_temporal_compression
+            full_latent_t = ((track_points.shape[1] - 1) //
+                             temporal_ratio + 1)
+            end_frame = int(start_frame) + latent_t
+            if start_frame < 0 or end_frame > full_latent_t:
+                raise ValueError(
+                    "track sequence does not cover the requested latent "
+                    f"window [{start_frame}, {end_frame}); available "
+                    f"latent frames: {full_latent_t}")
+
+            # Encode from the global beginning so the causal temporal conv is
+            # left-padded exactly once, then take the requested latent chunk.
+            full_track_map = self.track_encoder(
+                track_points,
+                track_visibility,
+                full_latent_t,
+                latent_h,
+                latent_w,
+                track_ids=track_ids,
+            )
+            track_map = full_track_map[:, :, start_frame:end_frame]
+            track_map = track_map.to(dtype=hidden_states.dtype)
+
+        return torch.cat([hidden_states, track_map], dim=1)
+
+
+class TrackWanTransformer3DModel(
+        _TrackConditioningMixin, WanTransformer3DModel):
+    """Bidirectional Wan I2V transformer conditioned on sparse point tracks."""
+
+    def __init__(
+        self,
+        config: TrackWanVideoConfig,
+        hf_config: dict[str, Any],
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        self._init_track_conditioning(config, hf_config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+        | None = None,
+        guidance: torch.Tensor | None = None,
+        r_timestep: torch.Tensor | None = None,
+        track_points: torch.Tensor | None = None,
+        track_visibility: torch.Tensor | None = None,
+        track_ids: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        hidden_states = self._append_track_conditioning(
+            hidden_states,
+            track_points=track_points,
+            track_visibility=track_visibility,
+            track_ids=track_ids,
+            start_frame=0,
+        )
+        return super().forward(
+            hidden_states,
+            encoder_hidden_states,
+            timestep,
+            encoder_hidden_states_image=encoder_hidden_states_image,
+            guidance=guidance,
+            r_timestep=r_timestep,
+            **kwargs,
+        )
+
+
+class CausalTrackWanTransformer3DModel(
+        _TrackConditioningMixin, CausalWanTransformer3DModel):
+    """Causal Wan I2V transformer with globally aligned track chunks."""
+
+    def __init__(
+        self,
+        config: TrackWanVideoConfig,
+        hf_config: dict[str, Any],
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        self._init_track_conditioning(config, hf_config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
+        | None = None,
+        start_frame: int = 0,
+        clean_x: torch.Tensor | None = None,
+        aug_t: torch.Tensor | None = None,
+        track_points: torch.Tensor | None = None,
+        track_visibility: torch.Tensor | None = None,
+        track_ids: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        hidden_states = self._append_track_conditioning(
+            hidden_states,
+            track_points=track_points,
+            track_visibility=track_visibility,
+            track_ids=track_ids,
+            start_frame=start_frame,
+        )
+        if clean_x is not None:
+            clean_x = self._append_track_conditioning(
+                clean_x,
+                track_points=track_points,
+                track_visibility=track_visibility,
+                track_ids=track_ids,
+                start_frame=start_frame,
+            )
+        return super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            encoder_hidden_states_image=encoder_hidden_states_image,
+            start_frame=start_frame,
+            clean_x=clean_x,
+            aug_t=aug_t,
+            **kwargs,
+        )
@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sparse point tracks to a dense, latent-aligned conditioning map."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def sinusoidal_embedding(
+    ids: torch.Tensor,
+    dim: int,
+    max_period: int = 10_000,
+) -> torch.Tensor:
+    """Map integer identity labels to sinusoidal embeddings."""
+    half = dim // 2
+    frequencies = torch.exp(
+        -math.log(max_period) *
+        torch.arange(half, device=ids.device, dtype=torch.float32) /
+        max(half, 1))
+    angles = ids.float().unsqueeze(-1) * frequencies
+    embedding = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1)
+    if dim % 2:
+        embedding = F.pad(embedding, (0, 1))
+    return embedding
+
+
+class TrackEncoder(nn.Module):
+    """Encode ``[B, T, N, 2]`` normalized tracks on the Wan latent grid."""
+
+    def __init__(
+        self,
+        id_dim: int,
+        track_channels: int,
+        vae_spatial_compression: int = 8,
+        vae_temporal_compression: int = 4,
+        max_track_id: int = 100_000,
+        zero_init: bool = False,
+        use_bias: bool = False,
+    ) -> None:
+        super().__init__()
+        if id_dim <= 0 or track_channels <= 0:
+            raise ValueError("id_dim and track_channels must be positive")
+        if vae_spatial_compression <= 0 or vae_temporal_compression <= 0:
+            raise ValueError("VAE compression ratios must be positive")
+        if max_track_id <= 0:
+            raise ValueError("max_track_id must be positive")
+
+        self.id_dim = int(id_dim)
+        self.track_channels = int(track_channels)
+        self.vae_spatial_compression = int(vae_spatial_compression)
+        self.vae_temporal_compression = int(vae_temporal_compression)
+        self.max_track_id = int(max_track_id)
+        self.use_bias = bool(use_bias)
+
+        # Bias-free convolutions keep empty cells exactly zero. Left-padding
+        # below gives the same temporal mapping as the causal Wan VAE:
+        # T_latent = (T_pixel - 1) // ratio + 1.
+        self.temporal_conv = nn.Conv3d(
+            self.id_dim,
+            self.track_channels,
+            kernel_size=(self.vae_temporal_compression, 1, 1),
+            stride=(self.vae_temporal_compression, 1, 1),
+            bias=self.use_bias,
+        )
+        self.proj = nn.Conv3d(
+            self.track_channels,
+            self.track_channels,
+            kernel_size=1,
+            bias=self.use_bias,
+        )
+        if self.temporal_conv.bias is not None:
+            nn.init.zeros_(self.temporal_conv.bias)
+        if self.proj.bias is not None:
+            nn.init.zeros_(self.proj.bias)
+        if zero_init:
+            nn.init.zeros_(self.proj.weight)
+
+    def sample_ids(
+        self,
+        batch_size: int,
+        num_tracks: int,
+        device: torch.device,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        """Sample IDs once at a batch/session boundary."""
+        if num_tracks > self.max_track_id:
+            raise ValueError(
+                f"num_tracks ({num_tracks}) exceeds max_track_id "
+                f"({self.max_track_id}); unique IDs are unavailable")
+        return torch.stack([
+            torch.randperm(
+                self.max_track_id,
+                device=device,
+                generator=generator,
+            )[:num_tracks] for _ in range(batch_size)
+        ])
+
+    def forward(
+        self,
+        coords: torch.Tensor,
+        visibility: torch.Tensor,
+        latent_t: int,
+        latent_h: int,
+        latent_w: int,
+        track_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if coords.ndim != 4 or coords.shape[-1] != 2:
+            raise ValueError(
+                "track coords must have shape [B, T, N, 2], got "
+                f"{tuple(coords.shape)}")
+        if visibility.shape != coords.shape[:-1]:
+            raise ValueError(
+                "track visibility must have shape [B, T, N] matching "
+                f"coords, got {tuple(visibility.shape)}")
+        if latent_t <= 0 or latent_h <= 0 or latent_w <= 0:
+            raise ValueError("latent dimensions must be positive")
+
+        batch_size, _, num_tracks, _ = coords.shape
+        device = coords.device
+        if num_tracks > self.max_track_id:
+            raise ValueError(
+                f"num_tracks ({num_tracks}) exceeds max_track_id "
+                f"({self.max_track_id})")
+        if track_ids is None:
+            # A deterministic fallback is required for repeated denoising and
+            # causal chunks. Training samples random IDs once in its wrapper.
+            track_ids = torch.arange(
+                num_tracks,
+                device=device,
+            ).unsqueeze(0).expand(batch_size, num_tracks)
+        elif track_ids.shape != (batch_size, num_tracks):
+            raise ValueError(
+                "track_ids must have shape [B, N], got "
+                f"{tuple(track_ids.shape)}")
+        else:
+            track_ids = track_ids.to(device=device)
+
+        identities = sinusoidal_embedding(track_ids, self.id_dim)
+        coords_float = coords.float()
+        x = (coords_float[..., 0].clamp(0.0, 1.0) *
+             (latent_w - 1)).round().long()
+        y = (coords_float[..., 1].clamp(0.0, 1.0) *
+             (latent_h - 1)).round().long()
+        visible = visibility.to(device=device, dtype=torch.float32)
+        contribution = visible.unsqueeze(-1) * identities.unsqueeze(1)
+
+        cell_index = y * latent_w + x
+        dense = torch.zeros(
+            (
+                batch_size,
+                coords.shape[1],
+                latent_h * latent_w,
+                self.id_dim,
+            ),
+            device=device,
+            dtype=torch.float32,
+        )
+        dense.scatter_add_(
+            2,
+            cell_index.unsqueeze(-1).expand(-1, -1, -1, self.id_dim),
+            contribution,
+        )
+        dense = dense.view(
+            batch_size,
+            coords.shape[1],
+            latent_h,
+            latent_w,
+            self.id_dim,
+        ).permute(0, 4, 1, 2, 3).contiguous()
+        dense = dense.to(dtype=self.temporal_conv.weight.dtype)
+        dense = F.pad(
+            dense,
+            (0, 0, 0, 0, self.vae_temporal_compression - 1, 0),
+        )
+        dense = self.temporal_conv(dense)
+        if dense.shape[2:] != (latent_t, latent_h, latent_w):
+            dense = F.interpolate(
+                dense,
+                size=(latent_t, latent_h, latent_w),
+                mode="nearest",
+            )
+        return self.proj(dense)
@@ -258,7 +258,7 @@ class WanI2VCrossAttention(WanSelfAttention):
        self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()

-    def forward(self, x, context, context_lens):
+    def forward(self, x, context, context_lens, crossattn_cache=None):
        r"""
        Args:
            x(Tensor): Shape [B, L1, C]
@@ -269,13 +269,29 @@ class WanI2VCrossAttention(WanSelfAttention):
        context = context[:, 257:]
        b, n, d = x.size(0), self.num_heads, self.head_dim

-        # compute query, key, value
+        # Query depends on the current latent tokens. Text/image keys and
+        # values only depend on conditioning and can be reused by causal
+        # streaming rollouts.
        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
-        v = self.to_v(context)[0].view(b, -1, n, d)
-        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(
-            b, -1, n, d)
-        v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
+        if crossattn_cache is not None and crossattn_cache.get("is_init", False):
+            k = crossattn_cache["k"]
+            v = crossattn_cache["v"]
+            k_img = crossattn_cache["k_img"]
+            v_img = crossattn_cache["v_img"]
+        else:
+            k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+            v = self.to_v(context)[0].view(b, -1, n, d)
+            k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(
+                b, -1, n, d)
+            v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
+            if crossattn_cache is not None:
+                crossattn_cache.update({
+                    "is_init": True,
+                    "k": k,
+                    "v": v,
+                    "k_img": k_img,
+                    "v_img": v_img,
+                })
        img_x = self.attn(q, k_img, v_img)
        # compute attention
        x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q)
@@ -694,11 +710,10 @@ class WanTransformer3DModel(BaseDiT):
        orig_dtype = hidden_states.dtype
        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (
+                encoder_hidden_states_image[0]
+                if encoder_hidden_states_image else None)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
@@ -32,9 +32,13 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
    "HYWorldTransformer3DModel":
    ("dits", "hyworld", "HYWorldTransformer3DModel"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "TrackWanTransformer3DModel":
+    ("dits", "trackwan", "TrackWanTransformer3DModel"),
    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
+    "CausalTrackWanTransformer3DModel":
+    ("dits", "trackwan", "CausalTrackWanTransformer3DModel"),
    "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
    "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
@@ -55,9 +59,13 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
 _IMAGE_TO_VIDEO_DIT_MODELS = {
    # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "TrackWanTransformer3DModel":
+    ("dits", "trackwan", "TrackWanTransformer3DModel"),
    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
+    "CausalTrackWanTransformer3DModel":
+    ("dits", "trackwan", "CausalTrackWanTransformer3DModel"),
    "LingBotWorld2CausalFastTransformer3DModel": (
        "dits",
        "lingbotworld2",
@@ -301,6 +301,11 @@ class TrainingBatch:
    image_latents: torch.Tensor | None = None
    infos: list[dict[str, Any]] | None = None
    mask_lat_size: torch.Tensor | None = None
+    # WanTrack point trajectories. IDs are sampled once per batch/session
+    # and reused by all denoising steps and causal chunks.
+    track_points: torch.Tensor | None = None
+    track_visibility: torch.Tensor | None = None
+    track_ids: torch.Tensor | None = None

    # ODE trajectory supervision
    trajectory_latents: torch.Tensor | None = None
@@ -286,6 +286,14 @@ class BasePreprocessPipeline(ComposedPipelineBase):

                if "action_path" in data:
                    valid_data["action_path"] = [data["action_path"][i] for i in valid_indices]
+                if "points_path" in data:
+                    valid_data["points_path"] = [data["points_path"][i] for i in valid_indices]
+                if "sample_frame_index" in data:
+                    valid_data["sample_frame_index"] = [data["sample_frame_index"][i] for i in valid_indices]
+                if "source_width" in data:
+                    valid_data["source_width"] = [data["source_width"][i] for i in valid_indices]
+                if "source_height" in data:
+                    valid_data["source_height"] = [data["source_height"][i] for i in valid_indices]

                # VAE
                with torch.autocast("cuda", dtype=torch.float32):
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+"""I2V preprocessing with point-track conditioning."""
+
+from typing import Any
+
+import numpy as np
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (
+    PreprocessPipeline_I2V, )
+
+
+class PreprocessPipeline_I2V_Track(PreprocessPipeline_I2V):
+    """Store I2V features and spatially aligned MotionStream point tracks."""
+
+    def get_pyarrow_schema(self):
+        return pyarrow_schema_i2v_track
+
+    @staticmethod
+    def _get_source_size(
+        sidecar: Any,
+        valid_data: dict[str, Any],
+        idx: int,
+        points_path: str,
+    ) -> tuple[float, float]:
+        if "width" in sidecar and "height" in sidecar:
+            width = float(np.asarray(sidecar["width"]).item())
+            height = float(np.asarray(sidecar["height"]).item())
+        elif "source_width" in valid_data and "source_height" in valid_data:
+            width = float(valid_data["source_width"][idx])
+            height = float(valid_data["source_height"][idx])
+        else:
+            raise ValueError(f"{points_path}: source width and height are required to align "
+                             "pixel-space tracks. Store them in the sidecar or manifest resolution.")
+
+        if width <= 0 or height <= 0:
+            raise ValueError(f"{points_path}: invalid source size ({width}, {height})")
+        return width, height
+
+    @staticmethod
+    def _normalize_after_center_crop(
+        tracks: np.ndarray,
+        visibility: np.ndarray,
+        *,
+        source_width: float,
+        source_height: float,
+        target_width: int,
+        target_height: int,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Apply ``CenterCropResizeVideo`` geometry and normalize coordinates."""
+        if target_width <= 0 or target_height <= 0:
+            raise ValueError(f"Invalid target size ({target_width}, {target_height})")
+        target_ratio = target_height / target_width
+        if source_height / source_width > target_ratio:
+            crop_height = int(source_width * target_ratio)
+            crop_width = int(source_width)
+        else:
+            crop_height = int(source_height)
+            crop_width = int(source_height / target_ratio)
+        if crop_width <= 0 or crop_height <= 0:
+            raise ValueError(f"Center crop is empty for source size "
+                             f"({source_width}, {source_height}) and target size "
+                             f"({target_width}, {target_height})")
+
+        crop_top = int(round((source_height - crop_height) / 2.0))
+        crop_left = int(round((source_width - crop_width) / 2.0))
+
+        x = tracks[..., 0]
+        y = tracks[..., 1]
+        finite = np.isfinite(x) & np.isfinite(y)
+        in_crop = (finite
+                   & (x >= crop_left)
+                   & (x < crop_left + crop_width)
+                   & (y >= crop_top)
+                   & (y < crop_top + crop_height))
+
+        normalized = np.empty_like(tracks, dtype=np.float32)
+        normalized[..., 0] = (x - crop_left) / crop_width
+        normalized[..., 1] = (y - crop_top) / crop_height
+        normalized = np.nan_to_num(normalized, nan=0.0, posinf=1.0, neginf=0.0)
+        np.clip(normalized, 0.0, 1.0, out=normalized)
+
+        aligned_visibility = np.asarray(visibility, dtype=np.float32) * in_crop
+        aligned_visibility = np.nan_to_num(aligned_visibility, nan=0.0, posinf=1.0, neginf=0.0)
+        np.clip(aligned_visibility, 0.0, 1.0, out=aligned_visibility)
+        return (
+            np.ascontiguousarray(normalized, dtype=np.float32),
+            np.ascontiguousarray(aligned_visibility, dtype=np.float32),
+        )
+
+    def get_extra_features(
+        self,
+        valid_data: dict[str, Any],
+        fastvideo_args: FastVideoArgs,
+    ) -> dict[str, Any]:
+        features = super().get_extra_features(valid_data, fastvideo_args)
+
+        points_paths = valid_data.get("points_path")
+        frame_indices_batch = valid_data.get("sample_frame_index")
+        if not points_paths or frame_indices_batch is None:
+            raise ValueError("i2v_track preprocessing requires points_path and sampled frame "
+                             "indices for every video.")
+        non_video_paths = [path for path in valid_data["path"] if not str(path).lower().endswith(".mp4")]
+        if non_video_paths:
+            raise ValueError("i2v_track preprocessing currently supports videos only; "
+                             f"received {non_video_paths[0]!r}")
+
+        target_height = int(valid_data["pixel_values"].shape[-2])
+        target_width = int(valid_data["pixel_values"].shape[-1])
+        expected_frames = int(valid_data["pixel_values"].shape[2])
+        track_points: list[np.ndarray] = []
+        track_visibility: list[np.ndarray] = []
+        object_ids: list[np.ndarray] = []
+        track_weights: list[np.ndarray] = []
+
+        samples = zip(points_paths, frame_indices_batch, strict=True)
+        for idx, (points_path, frame_indices) in enumerate(samples):
+            indices = np.asarray(frame_indices, dtype=np.int64)
+            if indices.ndim != 1 or indices.size != expected_frames:
+                raise ValueError(f"{points_path}: expected {expected_frames} sampled frame "
+                                 f"indices, got shape {indices.shape}")
+            if indices.size == 0 or indices.min() < 0:
+                raise ValueError(f"{points_path}: sampled frame indices must be non-negative")
+
+            with np.load(points_path) as sidecar:
+                if "tracks" not in sidecar or "visibility" not in sidecar:
+                    raise ValueError(f"{points_path}: sidecar must contain 'tracks' and 'visibility'")
+                tracks = np.asarray(sidecar["tracks"], dtype=np.float32)
+                visibility = np.asarray(sidecar["visibility"], dtype=np.float32)
+                if tracks.ndim != 3 or tracks.shape[-1] != 2:
+                    raise ValueError(f"{points_path}: tracks must have shape [T, N, 2], got "
+                                     f"{tracks.shape}")
+                if visibility.shape != tracks.shape[:2]:
+                    raise ValueError(f"{points_path}: visibility shape {visibility.shape} does "
+                                     f"not match tracks {tracks.shape[:2]}")
+                if indices.max() >= tracks.shape[0]:
+                    raise ValueError(f"{points_path}: sampled frame {indices.max()} exceeds "
+                                     f"{tracks.shape[0]} track frames")
+
+                source_width, source_height = self._get_source_size(sidecar, valid_data, idx, points_path)
+                num_tracks = tracks.shape[1]
+                sample_object_ids = np.asarray(
+                    sidecar["object_ids"] if "object_ids" in sidecar else np.full(num_tracks, -1, dtype=np.float32),
+                    dtype=np.float32,
+                )
+                sample_track_weights = np.asarray(
+                    sidecar["track_weights"] if "track_weights" in sidecar else np.zeros(num_tracks, dtype=np.float32),
+                    dtype=np.float32,
+                )
+                if sample_object_ids.shape != (num_tracks, ):
+                    raise ValueError(f"{points_path}: object_ids must have shape "
+                                     f"[{num_tracks}], got {sample_object_ids.shape}")
+                if sample_track_weights.shape != (num_tracks, ):
+                    raise ValueError(f"{points_path}: track_weights must have shape "
+                                     f"[{num_tracks}], got {sample_track_weights.shape}")
+                tracks = tracks[indices]
+                visibility = visibility[indices]
+
+            tracks, visibility = self._normalize_after_center_crop(
+                tracks,
+                visibility,
+                source_width=source_width,
+                source_height=source_height,
+                target_width=target_width,
+                target_height=target_height,
+            )
+            track_points.append(tracks)
+            track_visibility.append(visibility)
+            object_ids.append(np.ascontiguousarray(sample_object_ids, dtype=np.float32))
+            track_weights.append(np.ascontiguousarray(sample_track_weights, dtype=np.float32))
+
+        features["track_points"] = track_points
+        features["track_visibility"] = track_visibility
+        features["object_ids"] = object_ids
+        features["track_weights"] = track_weights
+        return features
+
+    def create_record(
+        self,
+        video_name: str,
+        vae_latent: np.ndarray,
+        text_embedding: np.ndarray,
+        valid_data: dict[str, Any],
+        idx: int,
+        extra_features: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        record = super().create_record(
+            video_name=video_name,
+            vae_latent=vae_latent,
+            text_embedding=text_embedding,
+            valid_data=valid_data,
+            idx=idx,
+            extra_features=extra_features,
+        )
+
+        for name in (
+                "track_points",
+                "track_visibility",
+                "object_ids",
+                "track_weights",
+        ):
+            if extra_features is None or name not in extra_features:
+                raise ValueError(f"Missing required WanTrack feature {name!r} for {video_name}")
+            array = np.ascontiguousarray(extra_features[name], dtype=np.float32)
+            record[f"{name}_bytes"] = array.tobytes()
+            record[f"{name}_shape"] = list(array.shape)
+            record[f"{name}_dtype"] = str(array.dtype)
+            record.pop(name, None)
+
+        return record
+
+
+EntryClass = PreprocessPipeline_I2V_Track
@@ -8,6 +8,7 @@ from fastvideo.distributed import (get_world_size, maybe_init_distributed_enviro
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (PreprocessPipeline_I2V)
+from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v_track import (PreprocessPipeline_I2V_Track)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (PreprocessPipeline_ODE_Trajectory)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (PreprocessPipeline_T2V)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_text import (PreprocessPipeline_Text)
@@ -52,6 +53,8 @@ def main(args) -> None:
        PreprocessPipeline = PreprocessPipeline_T2V
    elif args.preprocess_task == "i2v":
        PreprocessPipeline = PreprocessPipeline_I2V
+    elif args.preprocess_task == "i2v_track":
+        PreprocessPipeline = PreprocessPipeline_I2V_Track
    elif args.preprocess_task == "text_only":
        PreprocessPipeline = PreprocessPipeline_Text
    elif args.preprocess_task == "ode_trajectory":
@@ -72,7 +75,8 @@ def main(args) -> None:
        PreprocessPipeline = PreprocessPipeline_MatrixGame2_ODE_Trajectory
    else:
        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
-                         f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame2, matrixgame2_ode_trajectory")
+                         f"Valid options: t2v, i2v, i2v_track, ode_trajectory, text_only, "
+                         f"matrixgame2, matrixgame2_ode_trajectory")

    logger.info("Preprocess task: %s using %s", args.preprocess_task, PreprocessPipeline.__name__)

@@ -115,6 +119,7 @@ if __name__ == "__main__":
        choices=[
            "t2v",
            "i2v",
+            "i2v_track",
            "text_only",
            "ode_trajectory",
            "matrixgame2",
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Track-conditioned validation shared by bidirectional and causal WanTrack."""
+
+from __future__ import annotations
+
+import colorsys
+import glob
+import os
+from typing import Any
+
+import imageio.v2 as imageio
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+from PIL import Image, ImageDraw
+
+from fastvideo.distributed import get_world_group
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.train.models.wantrack.inference import (
+    prepare_wantrack_batch,
+    sample_wantrack,
+)
+from fastvideo.training.trackers import DummyTracker
+
+logger = init_logger(__name__)
+
+
+def _track_colors(count: int) -> np.ndarray:
+    colors = np.empty((count, 3), dtype=np.uint8)
+    for index in range(count):
+        hue = index / max(count, 1)
+        red, green, blue = colorsys.hsv_to_rgb(hue, 0.9, 1.0)
+        colors[index] = (
+            int(red * 255),
+            int(green * 255),
+            int(blue * 255),
+        )
+    return colors
+
+
+def _overlay_tracks(
+    frames: np.ndarray,
+    track_points: torch.Tensor,
+    track_visibility: torch.Tensor,
+    *,
+    stride: int,
+    tail: int,
+    radius: int,
+    visibility_threshold: float,
+) -> list[np.ndarray]:
+    points = track_points[0].float().cpu().numpy()
+    visibility = track_visibility[0].float().cpu().numpy()
+    frame_count = min(
+        int(frames.shape[0]),
+        int(points.shape[0]),
+        int(visibility.shape[0]),
+    )
+    frames = frames[:frame_count]
+    points = points[:frame_count, ::stride]
+    visibility = visibility[:frame_count, ::stride]
+    colors = _track_colors(int(points.shape[1]))
+
+    height, width = int(frames.shape[1]), int(frames.shape[2])
+    points = points.copy()
+    points[..., 0] *= width
+    points[..., 1] *= height
+
+    output: list[np.ndarray] = []
+    for frame_index, frame in enumerate(frames):
+        image = Image.fromarray(frame).convert("RGB")
+        draw = ImageDraw.Draw(image)
+        start = max(0, frame_index - tail)
+        for track_index, color_value in enumerate(colors):
+            color = tuple(int(value) for value in color_value)
+            trail: list[tuple[float, float]] = []
+            for history_index in range(start, frame_index + 1):
+                if visibility[history_index, track_index] < visibility_threshold:
+                    trail = []
+                    continue
+                x, y = points[history_index, track_index]
+                if 0 <= x < width and 0 <= y < height:
+                    trail.append((float(x), float(y)))
+            if len(trail) >= 2:
+                draw.line(trail, fill=color, width=1)
+            if visibility[frame_index, track_index] >= visibility_threshold:
+                x, y = points[frame_index, track_index]
+                if 0 <= x < width and 0 <= y < height:
+                    draw.ellipse(
+                        (
+                            float(x - radius),
+                            float(y - radius),
+                            float(x + radius),
+                            float(y + radius),
+                        ),
+                        fill=color,
+                    )
+        output.append(np.asarray(image))
+    return output
+
+
+class TrackValidationCallback(Callback):
+    """Generate fixed track-conditioned samples through the live student."""
+
+    def __init__(
+        self,
+        *,
+        every_steps: int = 250,
+        val_data_path: str | None = None,
+        num_val_samples: int = 1,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 1.0,
+        motion_guidance_scale: float = 1.0,
+        output_dir: str | None = None,
+        fps: int = 24,
+        grid_stride: int = 3,
+        tail: int = 12,
+        radius: int = 2,
+        visibility_threshold: float = 0.5,
+        seed: int = 1000,
+        validate_at_start: bool = False,
+    ) -> None:
+        self.every_steps = int(every_steps)
+        self.val_data_path = (str(val_data_path) if val_data_path is not None else None)
+        self.num_val_samples = int(num_val_samples)
+        self.num_inference_steps = int(num_inference_steps)
+        self.guidance_scale = float(guidance_scale)
+        self.motion_guidance_scale = float(motion_guidance_scale)
+        self.output_dir = (str(output_dir) if output_dir is not None else None)
+        self.fps = int(fps)
+        self.grid_stride = int(grid_stride)
+        self.tail = int(tail)
+        self.radius = int(radius)
+        self.visibility_threshold = float(visibility_threshold)
+        self.seed = int(seed)
+        self.validate_at_start = bool(validate_at_start)
+
+        if self.every_steps <= 0:
+            raise ValueError("every_steps must be positive")
+        if self.num_val_samples <= 0:
+            raise ValueError("num_val_samples must be positive")
+        if self.num_inference_steps <= 0:
+            raise ValueError("num_inference_steps must be positive")
+        if self.grid_stride <= 0:
+            raise ValueError("grid_stride must be positive")
+
+        self.tracker: Any = DummyTracker()
+        self._samples: list[dict[str, Any]] = []
+        self._is_main = False
+        self._did_start_validation = False
+
+    def on_train_start(self, method: Any, iteration: int = 0) -> None:
+        del iteration
+        self.tracker = getattr(method, "tracker", None) or DummyTracker()
+        self._is_main = int(get_world_group().rank) == 0
+        try:
+            self._samples = self._load_samples()
+            logger.info(
+                "WanTrack validation loaded %d fixed sample(s)",
+                len(self._samples),
+            )
+        except Exception as error:  # noqa: BLE001
+            logger.warning(
+                "WanTrack validation setup failed; disabling callback: %s",
+                error,
+            )
+            self._samples = []
+
+    def on_validation_begin(self, method: Any, iteration: int = 0) -> None:
+        if not self._samples:
+            return
+        run_at_start = (self.validate_at_start and not self._did_start_validation and iteration == 0)
+        run_periodic = iteration > 0 and iteration % self.every_steps == 0
+        if not run_at_start and not run_periodic:
+            return
+        self._did_start_validation = True
+        try:
+            self._run(method, iteration)
+        except Exception as error:  # noqa: BLE001
+            logger.warning(
+                "WanTrack validation failed at step %d: %s",
+                iteration,
+                error,
+            )
+
+    def _load_samples(self) -> list[dict[str, Any]]:
+        from fastvideo.dataset.dataloader.schema import (
+            pyarrow_schema_i2v_track, )
+        from fastvideo.dataset.utils import (
+            collate_rows_from_parquet_schema, )
+
+        data_path = self.val_data_path or str(self.training_config.data.data_path)
+        files = sorted(glob.glob(
+            os.path.join(data_path, "**", "*.parquet"),
+            recursive=True,
+        ))
+        if not files:
+            raise FileNotFoundError(f"No WanTrack validation parquet under {data_path}")
+
+        rows: list[dict[str, Any]] = []
+        for path in files:
+            table = pq.read_table(
+                path,
+                columns=pyarrow_schema_i2v_track.names,
+            )
+            rows.extend(table.to_pylist())
+            if len(rows) >= self.num_val_samples:
+                break
+
+        text_len = int(self.training_config.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+        collated = collate_rows_from_parquet_schema(
+            rows[:self.num_val_samples],
+            pyarrow_schema_i2v_track,
+            text_padding_length=text_len,
+            cfg_rate=0.0,
+            seed=self.seed,
+        )
+
+        infos = collated.get("info_list") or [{} for _ in rows]
+        samples: list[dict[str, Any]] = []
+        batch_size = int(collated["text_embedding"].shape[0])
+        for index in range(batch_size):
+            sample = {
+                name: value[index:index + 1].clone()
+                for name, value in collated.items() if torch.is_tensor(value)
+            }
+            sample["info_list"] = [infos[index]]
+            samples.append(sample)
+        return samples
+
+    @torch.no_grad()
+    def _run(self, method: Any, iteration: int) -> None:
+        student = method.student
+        transformer = student.transformer
+        was_training = bool(transformer.training)
+        transformer.eval()
+        try:
+            artifacts: list[Any] = []
+            for sample_index, raw_sample in enumerate(self._samples):
+                seed = self.seed + sample_index
+                batch = prepare_wantrack_batch(
+                    student,
+                    raw_sample,
+                    seed=seed,
+                    latents_source="zeros",
+                )
+                sampled = sample_wantrack(
+                    student,
+                    batch,
+                    num_inference_steps=self.num_inference_steps,
+                    seed=seed,
+                    text_guidance_scale=self.guidance_scale,
+                    motion_guidance_scale=self.motion_guidance_scale,
+                )
+                decoded = student.decode_latents(sampled)
+
+                if not self._is_main:
+                    continue
+                output_dir = self.output_dir or os.path.join(
+                    self.training_config.checkpoint.output_dir,
+                    "track_validation",
+                )
+                os.makedirs(output_dir, exist_ok=True)
+                output_path = os.path.join(
+                    output_dir,
+                    f"step{iteration:06d}_sample{sample_index}.mp4",
+                )
+                video = (decoded[0].clamp(0, 1).float().cpu().numpy() * 255.0).astype(np.uint8)
+                frames = video.transpose(1, 2, 3, 0)
+                conditional = batch.conditional_dict
+                if conditional is None:
+                    raise RuntimeError("WanTrack validation batch has no conditioning")
+                frames_with_tracks = _overlay_tracks(
+                    frames,
+                    conditional["track_points"],
+                    conditional["track_visibility"],
+                    stride=self.grid_stride,
+                    tail=self.tail,
+                    radius=self.radius,
+                    visibility_threshold=self.visibility_threshold,
+                )
+                imageio.mimsave(
+                    output_path,
+                    frames_with_tracks,
+                    fps=self.fps,
+                    macro_block_size=1,
+                )
+                info = raw_sample.get("info_list") or [{}]
+                caption = str(info[0].get("caption", ""))
+                artifact = self.tracker.video(
+                    output_path,
+                    caption=(f"WanTrack step {iteration}: "
+                             f"{caption[:120]}"),
+                )
+                if artifact is not None:
+                    artifacts.append(artifact)
+
+            if self._is_main and artifacts:
+                self.tracker.log_artifacts(
+                    {"track_validation/generated": artifacts},
+                    iteration,
+                )
+        finally:
+            if was_training:
+                transformer.train()
@@ -592,7 +592,9 @@ class WanModel(ModelBase):
        text_dict: dict[str, torch.Tensor] | None,
        clean_x: torch.Tensor | None = None,
        aug_t: torch.Tensor | None = None,
+        start_frame: int = 0,
    ) -> dict[str, Any]:
+        del start_frame
        if text_dict is None:
            raise ValueError("text_dict cannot be None for "
                             "Wan distillation")
@@ -172,6 +172,7 @@ class WanCausalModel(WanModel, CausalModelBase):
                noisy_latents,
                timestep_full,
                text_dict,
+                start_frame=cur_start_frame,
            ))
            input_kwargs["timestep"] = (timestep_full.to(
                device=self.device,
@@ -0,0 +1,11 @@
+from fastvideo.train.models.wantrack.wantrack import (
+    TrackAugmentationConfig,
+    WanTrackModel,
+)
+from fastvideo.train.models.wantrack.wantrack_causal import WanTrackCausalModel
+
+__all__ = [
+    "TrackAugmentationConfig",
+    "WanTrackModel",
+    "WanTrackCausalModel",
+]
@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared bidirectional and causal WanTrack sampling helpers."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any, Literal
+
+import torch
+
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.train.models.base import CausalModelBase, ModelBase
+
+_Branch = Literal["full", "no_text", "no_motion"]
+
+
+def prepare_wantrack_batch(
+    model: ModelBase,
+    raw_batch: dict[str, Any],
+    *,
+    seed: int,
+    latents_source: Literal["data", "zeros"] = "zeros",
+) -> TrainingBatch:
+    """Build inference conditions through the same path used by training."""
+    augmentation = getattr(model, "track_augmentation", None)
+    if augmentation is None:
+        raise TypeError("WanTrack inference requires a WanTrack training model")
+
+    generator = torch.Generator(device=model.device).manual_seed(int(seed))
+    model.track_augmentation = replace(
+        augmentation,
+        track_dropout_probability=0.0,
+        temporal_mask_probability=0.0,
+        motion_dropout_probability=0.0,
+        text_dropout_probability=0.0,
+    )
+    try:
+        batch = model.prepare_batch(
+            raw_batch,
+            generator=generator,
+            latents_source=latents_source,
+        )
+    finally:
+        model.track_augmentation = augmentation
+
+    # Streaming inference owns its cache/mask geometry. Training attention
+    # metadata must not leak into validation or standalone sampling.
+    batch.attn_metadata = None
+    batch.attn_metadata_vsa = None
+    return batch
+
+
+def _branch_args(branch: _Branch) -> tuple[bool, dict[str, Any] | None]:
+    if branch == "full":
+        return True, None
+    if branch == "no_text":
+        return False, {
+            "text": "zero",
+            "track": "keep",
+            "on_missing": "ignore",
+        }
+    if branch == "no_motion":
+        return False, {
+            "text": "keep",
+            "track": "drop",
+            "on_missing": "ignore",
+        }
+    raise ValueError(f"Unknown WanTrack CFG branch: {branch!r}")
+
+
+def _predict(
+    model: ModelBase,
+    latents: torch.Tensor,
+    timestep: torch.Tensor,
+    batch: TrainingBatch,
+    *,
+    branch: _Branch,
+    cache_tag: str,
+    start_frame: int,
+    store_kv: bool,
+) -> torch.Tensor | None:
+    conditional, cfg_uncond = _branch_args(branch)
+    if isinstance(model, CausalModelBase):
+        return model.predict_noise_streaming(
+            latents,
+            timestep,
+            batch,
+            conditional=conditional,
+            cache_tag=cache_tag,
+            store_kv=store_kv,
+            cur_start_frame=start_frame,
+            cfg_uncond=cfg_uncond,
+            attn_kind="dense",
+        )
+    if store_kv:
+        return None
+    return model.predict_noise(
+        latents,
+        timestep,
+        batch,
+        conditional=conditional,
+        cfg_uncond=cfg_uncond,
+        attn_kind="dense",
+    )
+
+
+def _guided_prediction(
+    model: ModelBase,
+    latents: torch.Tensor,
+    timestep: torch.Tensor,
+    batch: TrainingBatch,
+    *,
+    start_frame: int,
+    text_guidance_scale: float,
+    motion_guidance_scale: float,
+    motion_cfg: bool,
+) -> tuple[torch.Tensor, tuple[_Branch, ...]]:
+    full = _predict(
+        model,
+        latents,
+        timestep,
+        batch,
+        branch="full",
+        cache_tag="wantrack_full",
+        start_frame=start_frame,
+        store_kv=False,
+    )
+    if full is None:
+        raise RuntimeError("WanTrack prediction unexpectedly returned None")
+
+    text_scale = float(text_guidance_scale)
+    motion_scale = float(motion_guidance_scale)
+    if text_scale == 1.0 and motion_scale == 1.0:
+        return full, ("full", )
+
+    no_text = _predict(
+        model,
+        latents,
+        timestep,
+        batch,
+        branch="no_text",
+        cache_tag="wantrack_no_text",
+        start_frame=start_frame,
+        store_kv=False,
+    )
+    if no_text is None:
+        raise RuntimeError("WanTrack no-text prediction returned None")
+
+    if not motion_cfg:
+        return no_text + text_scale * (full - no_text), (
+            "full",
+            "no_text",
+        )
+
+    no_motion = _predict(
+        model,
+        latents,
+        timestep,
+        batch,
+        branch="no_motion",
+        cache_tag="wantrack_no_motion",
+        start_frame=start_frame,
+        store_kv=False,
+    )
+    if no_motion is None:
+        raise RuntimeError("WanTrack no-motion prediction returned None")
+
+    denominator = text_scale + motion_scale
+    alpha = text_scale / denominator if denominator > 0 else 0.5
+    base = alpha * no_text + (1.0 - alpha) * no_motion
+    guided = (base + text_scale * (full - no_text) + motion_scale * (full - no_motion))
+    return guided, ("full", "no_text", "no_motion")
+
+
+def _store_causal_context(
+    model: CausalModelBase,
+    latents: torch.Tensor,
+    batch: TrainingBatch,
+    *,
+    start_frame: int,
+    branches: tuple[_Branch, ...],
+) -> None:
+    timestep = torch.zeros(
+        latents.shape[:2],
+        device=latents.device,
+        dtype=torch.float32,
+    )
+    batch.timesteps = timestep
+    for branch in branches:
+        _predict(
+            model,
+            latents,
+            timestep,
+            batch,
+            branch=branch,
+            cache_tag=f"wantrack_{branch}",
+            start_frame=start_frame,
+            store_kv=True,
+        )
+
+
+def _sample_block(
+    model: ModelBase,
+    latents: torch.Tensor,
+    batch: TrainingBatch,
+    *,
+    scheduler: FlowMatchEulerDiscreteScheduler,
+    num_inference_steps: int,
+    start_frame: int,
+    text_guidance_scale: float,
+    motion_guidance_scale: float,
+    motion_cfg: bool,
+) -> tuple[torch.Tensor, tuple[_Branch, ...]]:
+    scheduler.set_timesteps(int(num_inference_steps), device=latents.device)
+    branches: tuple[_Branch, ...] = ("full", )
+    for current_timestep in scheduler.timesteps:
+        timestep = torch.full(
+            latents.shape[:2],
+            float(current_timestep.item()),
+            device=latents.device,
+            dtype=torch.float32,
+        )
+        batch.timesteps = timestep
+        prediction, branches = _guided_prediction(
+            model,
+            latents,
+            timestep,
+            batch,
+            start_frame=start_frame,
+            text_guidance_scale=text_guidance_scale,
+            motion_guidance_scale=motion_guidance_scale,
+            motion_cfg=motion_cfg,
+        )
+        latents = scheduler.step(
+            prediction.float(),
+            current_timestep,
+            latents.float(),
+            return_dict=False,
+        )[0]
+    return latents, branches
+
+
+@torch.no_grad()
+def sample_wantrack(
+    model: ModelBase,
+    batch: TrainingBatch,
+    *,
+    num_inference_steps: int = 30,
+    seed: int = 0,
+    text_guidance_scale: float = 1.0,
+    motion_guidance_scale: float = 1.0,
+    motion_cfg: bool = True,
+    chunk_size: int | None = None,
+) -> torch.Tensor:
+    """Generate normalized latents in ``[B, T, C, H, W]`` layout.
+
+    Bidirectional models denoise the complete clip. Causal models reuse the
+    same streaming prediction and cache API as RobotWM.
+    """
+    if batch.latents is None or batch.latents.ndim != 5:
+        raise ValueError("WanTrack inference requires [B, T, C, H, W] latents")
+    if num_inference_steps <= 0:
+        raise ValueError("num_inference_steps must be positive")
+
+    generator = torch.Generator(device="cpu").manual_seed(int(seed))
+    latents = torch.randn(
+        tuple(batch.latents.shape),
+        generator=generator,
+        dtype=torch.float32,
+    ).to(device=batch.latents.device, dtype=batch.latents.dtype)
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=float(getattr(model, "timestep_shift", 5.0)), )
+
+    if not isinstance(model, CausalModelBase):
+        sampled, _ = _sample_block(
+            model,
+            latents,
+            batch,
+            scheduler=scheduler,
+            num_inference_steps=num_inference_steps,
+            start_frame=0,
+            text_guidance_scale=text_guidance_scale,
+            motion_guidance_scale=motion_guidance_scale,
+            motion_cfg=motion_cfg,
+        )
+        return sampled
+
+    if chunk_size is None:
+        transformer = getattr(model, "transformer", None)
+        chunk_size = int(
+            getattr(
+                transformer,
+                "num_frame_per_block",
+                getattr(transformer.config.arch_config, "num_frames_per_block", 3),
+            ))
+    chunk_size = int(chunk_size)
+    if chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+    num_frames = int(latents.shape[1])
+    if num_frames % chunk_size != 0:
+        raise ValueError("Causal WanTrack inference requires latent frames "
+                         f"divisible by chunk_size; got {num_frames} and "
+                         f"{chunk_size}")
+
+    for branch in ("full", "no_text", "no_motion"):
+        model.clear_caches(cache_tag=f"wantrack_{branch}")
+
+    sampled_blocks: list[torch.Tensor] = []
+    try:
+        for start_frame in range(0, num_frames, chunk_size):
+            block = latents[:, start_frame:start_frame + chunk_size]
+            block, branches = _sample_block(
+                model,
+                block,
+                batch,
+                scheduler=scheduler,
+                num_inference_steps=num_inference_steps,
+                start_frame=start_frame,
+                text_guidance_scale=text_guidance_scale,
+                motion_guidance_scale=motion_guidance_scale,
+                motion_cfg=motion_cfg,
+            )
+            sampled_blocks.append(block)
+            _store_causal_context(
+                model,
+                block,
+                batch,
+                start_frame=start_frame,
+                branches=branches,
+            )
+    finally:
+        for branch in ("full", "no_text", "no_motion"):
+            model.clear_caches(cache_tag=f"wantrack_{branch}")
+
+    return torch.cat(sampled_blocks, dim=1)
@@ -0,0 +1,595 @@
+# SPDX-License-Identifier: Apache-2.0
+"""WanTrack training integration shared by bidirectional and causal models."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Mapping
+from dataclasses import dataclass, fields
+from typing import Any, Literal
+
+import torch
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.train.models.wan.wan import WanModel
+from fastvideo.train.utils.dataloader import build_parquet_t2v_train_dataloader
+from fastvideo.train.utils.moduleloader import load_module_from_path
+from fastvideo.training.training_utils import normalize_dit_input
+
+
+@dataclass(slots=True)
+class TrackAugmentationConfig:
+    """Configuration-driven MotionStream-style training augmentation."""
+
+    enabled: bool = True
+    min_points: int = 1000
+    max_points: int = 2500
+    sparse_object_sampling: bool = False
+    extra_points: int = 0
+    extra_point_sampling: Literal["random", "weighted"] = "random"
+    track_dropout_probability: float = 0.0
+    temporal_mask_probability: float = 0.2
+    temporal_mask_chunk_size: int = 8
+    motion_dropout_probability: float = 0.0
+    text_dropout_probability: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.min_points < 0 or self.max_points < 0:
+            raise ValueError("min_points and max_points must be non-negative")
+        if self.max_points and self.min_points > self.max_points:
+            raise ValueError("min_points cannot exceed max_points")
+        if self.extra_points < 0:
+            raise ValueError("extra_points must be non-negative")
+        if self.extra_point_sampling not in {"random", "weighted"}:
+            raise ValueError("extra_point_sampling must be 'random' or "
+                             "'weighted'")
+        if self.temporal_mask_chunk_size <= 0:
+            raise ValueError("temporal_mask_chunk_size must be positive")
+        for name in (
+                "track_dropout_probability",
+                "temporal_mask_probability",
+                "motion_dropout_probability",
+                "text_dropout_probability",
+        ):
+            value = float(getattr(self, name))
+            if not 0.0 <= value <= 1.0:
+                raise ValueError(f"{name} must be in [0, 1], got {value}")
+
+    @classmethod
+    def from_value(
+        cls,
+        value: TrackAugmentationConfig | Mapping[str, Any] | None,
+    ) -> TrackAugmentationConfig:
+        if value is None:
+            return cls()
+        if isinstance(value, cls):
+            return value
+        if not isinstance(value, Mapping):
+            raise TypeError("track_augmentation must be a mapping or "
+                            "TrackAugmentationConfig")
+        valid_fields = {item.name for item in fields(cls)}
+        unknown = set(value) - valid_fields
+        if unknown:
+            raise ValueError("Unknown track_augmentation fields: "
+                             f"{sorted(unknown)}")
+        return cls(**dict(value))
+
+
+class WanTrackModel(WanModel):
+    """Bidirectional WanTrack model and shared causal conditioning."""
+
+    _transformer_cls_name = "TrackWanTransformer3DModel"
+
+    def __init__(
+        self,
+        *,
+        track_augmentation: TrackAugmentationConfig
+        | Mapping[str, Any]
+        | None = None,
+        freeze_track_encoder: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.track_augmentation = TrackAugmentationConfig.from_value(track_augmentation)
+        self.freeze_track_encoder = bool(freeze_track_encoder)
+        super().__init__(**kwargs)
+        if self.freeze_track_encoder:
+            track_encoder = getattr(self.transformer, "track_encoder", None)
+            if track_encoder is None:
+                raise ValueError("freeze_track_encoder=true requires a "
+                                 "transformer.track_encoder")
+            track_encoder.requires_grad_(False)
+
+    def init_preprocessors(self, training_config: Any) -> None:
+        self.vae = load_module_from_path(
+            model_path=str(training_config.model_path),
+            module_type="vae",
+            training_config=training_config,
+        )
+        self.world_group = get_world_group()
+        self.sp_group = get_sp_group()
+        self._init_timestep_mechanics()
+
+        preprocessed_data_type = str(getattr(
+            training_config.data,
+            "preprocessed_data_type",
+            "",
+        )).strip().lower()
+        if preprocessed_data_type != "i2v_track":
+            raise ValueError("WanTrack requires data.preprocessed_data_type='i2v_track', "
+                             f"got {preprocessed_data_type!r}")
+
+        text_len = (training_config.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+        self.dataloader = build_parquet_t2v_train_dataloader(
+            training_config.data,
+            text_len=int(text_len),
+            parquet_schema=pyarrow_schema_i2v_track,
+        )
+        self.start_step = 0
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        if self._requires_negative_conditioning:
+            self.ensure_negative_conditioning()
+        assert self.training_config is not None
+        training_config = self.training_config
+        dtype = self._get_training_dtype()
+        device = self.device
+
+        required = (
+            "first_frame_latent",
+            "clip_feature",
+            "track_points",
+            "track_visibility",
+            "text_embedding",
+            "text_attention_mask",
+        )
+        missing = [name for name in required if name not in raw_batch]
+        if missing:
+            raise ValueError(f"WanTrack batch is missing required fields: {missing}")
+
+        latent_t = int(training_config.data.num_latent_t)
+        if latents_source == "data":
+            if "vae_latent" not in raw_batch:
+                raise ValueError("WanTrack requires vae_latent when "
+                                 "latents_source='data'")
+            latents = raw_batch["vae_latent"][:, :, :latent_t].to(device=device, dtype=dtype)
+        elif latents_source == "zeros":
+            batch_size = int(raw_batch["text_embedding"].shape[0])
+            vae_config = (training_config.pipeline_config.vae_config.arch_config)
+            latents = torch.zeros(
+                batch_size,
+                int(vae_config.z_dim),
+                latent_t,
+                int(training_config.data.num_height) // int(vae_config.spatial_compression_ratio),
+                int(training_config.data.num_width) // int(vae_config.spatial_compression_ratio),
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            raise ValueError(f"Unknown latents_source: {latents_source!r}")
+        first_frame_latent = raw_batch["first_frame_latent"][:, :, :latent_t].to(device=device, dtype=dtype)
+        image_embeds = raw_batch["clip_feature"]
+        if not torch.is_tensor(image_embeds) or image_embeds.numel() == 0:
+            raise ValueError("WanTrack requires non-empty clip_feature")
+        image_embeds = image_embeds.to(device=device, dtype=dtype)
+
+        expected_track_frames = ((latent_t - 1) * self._temporal_compression_ratio() + 1)
+        track_points = raw_batch["track_points"]
+        track_visibility = raw_batch["track_visibility"]
+        if track_points.ndim != 4 or track_points.shape[-1] != 2:
+            raise ValueError("track_points must have shape [B, T, N, 2], got "
+                             f"{tuple(track_points.shape)}")
+        if track_visibility.shape != track_points.shape[:-1]:
+            raise ValueError("track_visibility must have shape [B, T, N], got "
+                             f"{tuple(track_visibility.shape)}")
+        if track_points.shape[1] < expected_track_frames:
+            raise ValueError("track sequence is shorter than the requested latent clip: "
+                             f"{track_points.shape[1]} < {expected_track_frames}")
+        track_points = track_points[:, :expected_track_frames].to(device=device, dtype=torch.float32)
+        track_visibility = track_visibility[:, :expected_track_frames].to(device=device, dtype=torch.float32)
+        num_tracks = int(track_points.shape[2])
+
+        object_ids = raw_batch.get("object_ids")
+        if (torch.is_tensor(object_ids) and object_ids.numel() > 0
+                and object_ids.shape == (track_points.shape[0], num_tracks)):
+            object_ids = object_ids.to(device=device, dtype=torch.long)
+        else:
+            object_ids = None
+
+        track_weights = raw_batch.get("track_weights")
+        if (torch.is_tensor(track_weights) and track_weights.numel() > 0
+                and track_weights.shape == (track_points.shape[0], num_tracks)):
+            track_weights = track_weights.to(device=device, dtype=torch.float32)
+        else:
+            track_weights = None
+
+        encoder_hidden_states = raw_batch["text_embedding"].to(device=device, dtype=dtype)
+        encoder_hidden_states, track_visibility = self._augment_conditions(
+            encoder_hidden_states,
+            track_visibility,
+            object_ids=object_ids,
+            track_weights=track_weights,
+            generator=generator,
+        )
+        batch_size = int(track_points.shape[0])
+        max_track_id = int(getattr(
+            getattr(self.transformer, "track_encoder", None),
+            "max_track_id",
+            100_000,
+        ))
+        if num_tracks > max_track_id:
+            raise ValueError(f"num_tracks ({num_tracks}) exceeds max_track_id "
+                             f"({max_track_id})")
+        track_ids = torch.stack([
+            torch.randperm(
+                max_track_id,
+                generator=generator,
+                device=device,
+            )[:num_tracks] for _ in range(batch_size)
+        ])
+
+        training_batch = TrainingBatch(
+            latents=latents,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=raw_batch["text_attention_mask"].to(device=device, dtype=dtype),
+            image_embeds=image_embeds,
+            image_latents=first_frame_latent,
+            track_points=track_points,
+            track_visibility=track_visibility,
+            track_ids=track_ids,
+            infos=raw_batch.get("info_list"),
+        )
+        training_batch.latents = normalize_dit_input(
+            "wan",
+            training_batch.latents,
+            self.vae,
+        )
+        training_batch = self._prepare_dit_inputs(training_batch, generator)
+        training_batch = self._build_attention_metadata(training_batch)
+        training_batch.attn_metadata_vsa = copy.copy(training_batch.attn_metadata)
+        if training_batch.attn_metadata is not None:
+            training_batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore[attr-defined]
+        return training_batch
+
+    def _augment_conditions(
+        self,
+        encoder_hidden_states: torch.Tensor,
+        track_visibility: torch.Tensor,
+        *,
+        object_ids: torch.Tensor | None,
+        track_weights: torch.Tensor | None,
+        generator: torch.Generator,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        config = self.track_augmentation
+        if not config.enabled:
+            return encoder_hidden_states, track_visibility
+
+        visibility = track_visibility.clone()
+        batch_size, num_frames, num_tracks = visibility.shape
+        device = visibility.device
+
+        # Keep tensors rectangular for the collator; sampling drops tracks by
+        # zeroing their visibility. Sparse mode preserves one visible query per
+        # segmented object, then adds background/context queries.
+        if config.sparse_object_sampling:
+            if object_ids is None:
+                raise ValueError("sparse_object_sampling requires object_ids "
+                                 "in the WanTrack parquet")
+            for batch_index in range(batch_size):
+                valid = visibility[batch_index, 0] > 0.5
+                keep = torch.zeros(
+                    num_tracks,
+                    device=device,
+                    dtype=torch.bool,
+                )
+                sample_object_ids = object_ids[batch_index]
+                for object_id in torch.unique(sample_object_ids):
+                    if int(object_id.item()) < 0:
+                        continue
+                    candidates = torch.nonzero(
+                        (sample_object_ids == object_id) & valid,
+                        as_tuple=True,
+                    )[0]
+                    if candidates.numel() == 0:
+                        candidates = torch.nonzero(
+                            sample_object_ids == object_id,
+                            as_tuple=True,
+                        )[0]
+                    if candidates.numel() > 0:
+                        choice = torch.randint(
+                            int(candidates.numel()),
+                            (1, ),
+                            generator=generator,
+                            device=device,
+                        )
+                        keep[candidates[choice]] = True
+
+                pool = valid & ~keep
+                extra_count = min(
+                    config.extra_points,
+                    int(pool.sum().item()),
+                )
+                if extra_count > 0:
+                    weights = pool.to(dtype=torch.float32)
+                    if (config.extra_point_sampling == "weighted" and track_weights is not None):
+                        weights *= (0.05 + track_weights[batch_index].clamp_min(0))
+                    selected = torch.multinomial(
+                        weights,
+                        extra_count,
+                        replacement=False,
+                        generator=generator,
+                    )
+                    keep[selected] = True
+                visibility[batch_index] *= keep.to(dtype=visibility.dtype).unsqueeze(0)
+        elif config.max_points > 0 and config.min_points < num_tracks:
+            high = min(config.max_points, num_tracks)
+            low = min(config.min_points, high)
+            for batch_index in range(batch_size):
+                keep_count = int(torch.randint(
+                    low,
+                    high + 1,
+                    (1, ),
+                    generator=generator,
+                    device=device,
+                ).item())
+                permutation = torch.randperm(
+                    num_tracks,
+                    generator=generator,
+                    device=device,
+                )
+                keep = torch.zeros(
+                    num_tracks,
+                    device=device,
+                    dtype=visibility.dtype,
+                )
+                keep[permutation[:keep_count]] = 1
+                visibility[batch_index] *= keep.unsqueeze(0)
+
+        if config.track_dropout_probability > 0:
+            for batch_index in range(batch_size):
+                kept = visibility[batch_index].amax(dim=0) > 0.5
+                kept_indices = torch.nonzero(kept, as_tuple=True)[0]
+                if kept_indices.numel() <= 1:
+                    continue
+                drop = (torch.rand(
+                    num_tracks,
+                    generator=generator,
+                    device=device,
+                ) < config.track_dropout_probability) & kept
+                if int((kept & ~drop).sum().item()) == 0:
+                    spare = kept_indices[torch.randint(
+                        int(kept_indices.numel()),
+                        (1, ),
+                        generator=generator,
+                        device=device,
+                    )]
+                    drop[spare] = False
+                visibility[batch_index, :, drop] = 0
+
+        chunk_size = config.temporal_mask_chunk_size
+        num_chunks = (num_frames + chunk_size - 1) // chunk_size
+        if config.temporal_mask_probability > 0:
+            drop_chunks = torch.rand(
+                batch_size,
+                num_chunks,
+                generator=generator,
+                device=device,
+            ) < config.temporal_mask_probability
+            for chunk_index in range(num_chunks):
+                start = chunk_index * chunk_size
+                end = min(start + chunk_size, num_frames)
+                visibility[:, start:end] *= (~drop_chunks[:, chunk_index]).to(dtype=visibility.dtype).view(-1, 1, 1)
+
+        if config.motion_dropout_probability > 0:
+            drop_motion = torch.rand(
+                batch_size,
+                generator=generator,
+                device=device,
+            ) < config.motion_dropout_probability
+            visibility[drop_motion] = 0
+
+        if config.text_dropout_probability > 0:
+            drop_text = torch.rand(
+                batch_size,
+                generator=generator,
+                device=device,
+            ) < config.text_dropout_probability
+            encoder_hidden_states = encoder_hidden_states.clone()
+            encoder_hidden_states[drop_text] = 0
+
+        return encoder_hidden_states, visibility
+
+    def _prepare_dit_inputs(
+        self,
+        training_batch: TrainingBatch,
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        training_batch = super()._prepare_dit_inputs(
+            training_batch,
+            generator,
+        )
+        if training_batch.image_latents is None:
+            raise RuntimeError("WanTrack requires first_frame_latent")
+        if training_batch.image_embeds is None:
+            raise RuntimeError("WanTrack requires clip_feature")
+
+        image_condition = self._build_i2v_condition(training_batch.image_latents)
+        training_batch.image_latents = image_condition
+        assert training_batch.noisy_model_input is not None
+        training_batch.noisy_model_input = torch.cat(
+            [training_batch.noisy_model_input, image_condition],
+            dim=1,
+        )
+        assert training_batch.conditional_dict is not None
+        extras = {
+            "image_latents": image_condition,
+            "encoder_hidden_states_image": training_batch.image_embeds,
+            "track_points": training_batch.track_points,
+            "track_visibility": training_batch.track_visibility,
+            "track_ids": training_batch.track_ids,
+        }
+        training_batch.conditional_dict.update(extras)
+        if training_batch.unconditional_dict is None:
+            training_batch.unconditional_dict = dict(training_batch.conditional_dict)
+        else:
+            training_batch.unconditional_dict.update(extras)
+        return training_batch
+
+    def _build_distill_input_kwargs(
+        self,
+        noise_input: torch.Tensor,
+        timestep: torch.Tensor,
+        text_dict: dict[str, Any] | None,
+        clean_x: torch.Tensor | None = None,
+        aug_t: torch.Tensor | None = None,
+        start_frame: int = 0,
+    ) -> dict[str, Any]:
+        if text_dict is None:
+            raise ValueError("text_dict cannot be None for WanTrack")
+        hidden_states = self._append_i2v_for_window(
+            noise_input,
+            text_dict,
+            start_frame=start_frame,
+        )
+        kwargs: dict[str, Any] = {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": text_dict["encoder_hidden_states"],
+            "encoder_attention_mask": text_dict["encoder_attention_mask"],
+            "encoder_hidden_states_image": text_dict.get("encoder_hidden_states_image"),
+            "track_points": text_dict.get("track_points"),
+            "track_visibility": text_dict.get("track_visibility"),
+            "track_ids": text_dict.get("track_ids"),
+            "timestep": timestep,
+            "return_dict": False,
+        }
+        if clean_x is not None:
+            kwargs["clean_x"] = self._append_i2v_for_window(
+                clean_x,
+                text_dict,
+                start_frame=start_frame,
+            )
+            kwargs["aug_t"] = aug_t
+        return kwargs
+
+    def _append_i2v_for_window(
+        self,
+        latents: torch.Tensor,
+        text_dict: dict[str, Any],
+        *,
+        start_frame: int,
+    ) -> torch.Tensor:
+        hidden_states = latents.permute(0, 2, 1, 3, 4)
+        if hidden_states.shape[1] == 36:
+            return hidden_states
+        if hidden_states.shape[1] != 16:
+            raise ValueError("WanTrack expects a 16-channel latent or a 36-channel "
+                             f"latent+I2V input, got {hidden_states.shape[1]}")
+        image_condition = text_dict.get("image_latents")
+        if image_condition is None:
+            raise RuntimeError("WanTrack requires image_latents for a 16-channel input")
+        start = int(start_frame)
+        end = start + hidden_states.shape[2]
+        if start < 0 or end > image_condition.shape[2]:
+            raise ValueError("I2V condition does not cover the requested latent window "
+                             f"[{start}, {end})")
+        return torch.cat(
+            [hidden_states, image_condition[:, :, start:end]],
+            dim=1,
+        )
+
+    def _get_uncond_text_dict(
+        self,
+        batch: TrainingBatch,
+        *,
+        cfg_uncond: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        sanitized_cfg = None
+        track_policy = "keep"
+        if cfg_uncond is not None:
+            sanitized_cfg = dict(cfg_uncond)
+            policies = []
+            for channel in ("track", "motion"):
+                value = sanitized_cfg.pop(channel, None)
+                if value is not None:
+                    if not isinstance(value, str):
+                        raise ValueError(f"cfg_uncond.{channel} must be a string")
+                    policies.append(value.strip().lower())
+            if policies:
+                if len(set(policies)) != 1:
+                    raise ValueError("cfg_uncond.track and cfg_uncond.motion conflict")
+                track_policy = policies[0]
+
+        result = dict(super()._get_uncond_text_dict(
+            batch,
+            cfg_uncond=sanitized_cfg,
+        ))
+        conditional = batch.conditional_dict
+        if conditional is None:
+            raise RuntimeError("Missing conditional_dict in TrainingBatch")
+        for key in (
+                "image_latents",
+                "encoder_hidden_states_image",
+                "track_points",
+                "track_visibility",
+                "track_ids",
+        ):
+            result.setdefault(key, conditional.get(key))
+
+        if track_policy not in {"keep", "zero", "drop"}:
+            raise ValueError("cfg_uncond.track/motion must be one of {keep, zero, drop}, "
+                             f"got {track_policy!r}")
+        if track_policy in {"zero", "drop"}:
+            result["track_points"] = None
+            result["track_visibility"] = None
+            result["track_ids"] = None
+        return result
+
+    def _build_i2v_condition(
+        self,
+        image_latents: torch.Tensor,
+    ) -> torch.Tensor:
+        if image_latents.ndim != 5:
+            raise ValueError("first_frame_latent must be [B, C, T, H, W], got "
+                             f"{tuple(image_latents.shape)}")
+        if image_latents.shape[1] == 20:
+            return image_latents
+        if image_latents.shape[1] != 16:
+            raise ValueError("first_frame_latent must have 16 or 20 channels, got "
+                             f"{image_latents.shape[1]}")
+
+        ratio = self._temporal_compression_ratio()
+        batch_size, _, latent_t, height, width = image_latents.shape
+        pixel_frames = (latent_t - 1) * ratio + 1
+        mask = torch.ones(
+            batch_size,
+            1,
+            pixel_frames,
+            height,
+            width,
+            device=image_latents.device,
+            dtype=image_latents.dtype,
+        )
+        mask[:, :, 1:] = 0
+        mask = torch.cat([
+            mask[:, :, :1].repeat_interleave(ratio, dim=2),
+            mask[:, :, 1:],
+        ], dim=2)
+        mask = mask.view(
+            batch_size,
+            -1,
+            ratio,
+            height,
+            width,
+        ).transpose(1, 2)
+        return torch.cat([mask, image_latents], dim=1)
+
+    def _temporal_compression_ratio(self) -> int:
+        assert self.training_config is not None
+        return int(self.training_config.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Causal WanTrack training model."""
+
+from fastvideo.train.models.wan.wan_causal import WanCausalModel
+from fastvideo.train.models.wantrack.wantrack import WanTrackModel
+
+
+class WanTrackCausalModel(WanTrackModel, WanCausalModel):
+    """WanTrack with block-causal attention and streaming cache support."""
+
+    _transformer_cls_name = "CausalTrackWanTransformer3DModel"
@@ -54,8 +54,7 @@ class RunConfig:
            if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
                return {
                    f.name: _safe_asdict(getattr(obj, f.name))
-                    for f in dataclasses.fields(obj)
-                    if not callable(getattr(obj, f.name))
+                    for f in dataclasses.fields(obj) if not callable(getattr(obj, f.name))
                }
            if isinstance(obj, dict):
                return {k: _safe_asdict(v) for k, v in obj.items()}
@@ -349,12 +348,10 @@ def _build_training_config(
        data_path = str(raw_data_path)

    preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower()
-    if preprocessed_data_type not in {"t2v", "text_only"}:
-        raise ValueError(
-            "training.data.preprocessed_data_type must be one of "
-            "{'t2v', 'text_only'}, got "
-            f"{preprocessed_data_type!r}"
-        )
+    if preprocessed_data_type not in {"t2v", "i2v_track", "text_only"}:
+        raise ValueError("training.data.preprocessed_data_type must be one of "
+                         "{'t2v', 'i2v_track', 'text_only'}, got "
+                         f"{preprocessed_data_type!r}")

    return TrainingConfig(
        distributed=DistributedConfig(
@@ -424,9 +421,7 @@ def _build_training_config(
    )


-def _parse_cli_overrides(
-    overrides: list[str],
-) -> dict[str, Any]:
+def _parse_cli_overrides(overrides: list[str], ) -> dict[str, Any]:
    """Parse ``--dotted.key value`` CLI overrides.

    Returns a flat dict mapping dotted keys to parsed
@@ -172,6 +172,7 @@ nav:
    - Experiment Tracking: training/trackers.md
    - Data Preprocessing: training/data_preprocess.md
    - Fine-tuning: training/finetune.md
+    - WanTrack Training: training/wantrack.md
    - Attn-QAT Training: training/attn_qat.md
    - Examples:
      - Training Index: training/examples/examples_training_index.md
@@ -0,0 +1,412 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Build a WanTrack initialization checkpoint from Wan2.1-Fun Control weights.
+
+The VideoX-Fun Control DiT receives 48 input channels in this order:
+
+    noisy latent (16) | spatial control (16) | reference latent (16)
+
+WanTrack receives 52 input channels:
+
+    noisy latent (16) | I2V mask (4) | first-frame latent (16) | track map (16)
+
+This converter preserves the pretrained control pathway by moving the Fun
+Control channels into WanTrack's track slot. The four new I2V mask channels are
+initialized to zero. The remainder of the model is renamed to the Diffusers
+key layout and checked against a compatible Fun-InP Diffusers checkpoint.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+
+SOURCE_IN_CHANNELS = 48
+TARGET_IN_CHANNELS = 52
+LATENT_CHANNELS = 16
+TRACK_CHANNELS = 16
+VAE_SPATIAL_COMPRESSION = 8
+VAE_TEMPORAL_COMPRESSION = 4
+
+# VideoX-Fun Wan names to the Diffusers/FastVideo Wan state-dict surface.
+# patch_embedding is handled separately because its input channels are remapped.
+PARAM_NAME_MAP: tuple[tuple[re.Pattern[str], str], ...] = tuple(
+    (re.compile(pattern), replacement)
+    for pattern, replacement in {
+        r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1",
+        r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1",
+        r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1",
+        r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1",
+        r"^time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1",
+        r"^img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1",
+        r"^img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1",
+        r"^img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1",
+        r"^img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1",
+        r"^head\.modulation$": r"scale_shift_table",
+        r"^head\.head\.(.*)$": r"proj_out.\1",
+        r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
+        r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
+        r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
+        r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
+        r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2",
+        r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2",
+        r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+        r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+        r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2",
+        r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+        r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2",
+        r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
+        r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
+        r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
+        r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2",
+        r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2",
+        r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2",
+        r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
+        r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2",
+    }.items()
+)
+
+
+def load_sharded_safetensors(directory: Path) -> dict[str, torch.Tensor]:
+    """Load every safetensors shard in a directory and reject duplicate keys."""
+    shard_paths = sorted(directory.glob("*.safetensors"))
+    if not shard_paths:
+        raise FileNotFoundError(f"No safetensors files found under {directory}")
+
+    state_dict: dict[str, torch.Tensor] = {}
+    for shard_path in shard_paths:
+        shard = load_file(str(shard_path), device="cpu")
+        duplicate_keys = state_dict.keys() & shard.keys()
+        if duplicate_keys:
+            examples = ", ".join(sorted(duplicate_keys)[:8])
+            raise RuntimeError(f"Duplicate parameter keys across base shards: {examples}")
+        state_dict.update(shard)
+    return state_dict
+
+
+def rename_control_state_dict(source: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Rename all VideoX-Fun parameters and reject unknown or colliding keys."""
+    converted: dict[str, torch.Tensor] = {}
+    for source_key, value in source.items():
+        if source_key.startswith("patch_embedding."):
+            target_key = source_key
+        else:
+            target_key = ""
+            for pattern, replacement in PARAM_NAME_MAP:
+                if pattern.match(source_key):
+                    target_key = pattern.sub(replacement, source_key)
+                    break
+            if not target_key:
+                raise ValueError(f"No target mapping for source parameter: {source_key}")
+
+        if target_key in converted:
+            raise RuntimeError(f"Multiple source parameters map to target parameter: {target_key}")
+        converted[target_key] = value
+    return converted
+
+
+def synthesize_image_query_norms(
+    converted: dict[str, torch.Tensor],
+    reference: dict[str, torch.Tensor],
+) -> None:
+    """Create Diffusers image-query norm tensors absent from VideoX-Fun."""
+    expected_keys = sorted(key for key in reference if key.endswith("attn2.norm_added_q.weight"))
+    for query_key in expected_keys:
+        key_norm_key = query_key.replace("norm_added_q", "norm_added_k")
+        if key_norm_key not in converted:
+            raise RuntimeError(f"Cannot synthesize {query_key}: converted state lacks {key_norm_key}")
+        if query_key in converted:
+            raise RuntimeError(f"Refusing to overwrite converted parameter: {query_key}")
+        converted[query_key] = torch.zeros_like(converted[key_norm_key])
+
+
+def remap_patch_embedding(source_weight: torch.Tensor) -> torch.Tensor:
+    """Remap Fun Control's 48-channel patch projection into WanTrack's 52 channels."""
+    if source_weight.ndim != 5:
+        raise ValueError(f"Expected a 5D patch embedding weight, got shape {tuple(source_weight.shape)}")
+    if source_weight.shape[1] != SOURCE_IN_CHANNELS:
+        raise ValueError(
+            f"Expected patch_embedding.weight input channels={SOURCE_IN_CHANNELS}, "
+            f"got {source_weight.shape[1]}"
+        )
+
+    target_shape = (source_weight.shape[0], TARGET_IN_CHANNELS, *source_weight.shape[2:])
+    target_weight = source_weight.new_zeros(target_shape)
+    target_weight[:, 0:16] = source_weight[:, 0:16]
+    target_weight[:, 20:36] = source_weight[:, 32:48]
+    target_weight[:, 36:52] = source_weight[:, 16:32]
+
+    if torch.count_nonzero(target_weight[:, 16:20]).item() != 0:
+        raise RuntimeError("WanTrack I2V mask channels must be initialized to zero")
+    return target_weight
+
+
+def validate_converted_body(
+    converted: dict[str, torch.Tensor],
+    reference: dict[str, torch.Tensor],
+) -> None:
+    """Require the converted Wan body to match the Fun-InP key and shape surface."""
+    converted_keys = set(converted)
+    reference_keys = set(reference)
+    missing = sorted(reference_keys - converted_keys)
+    unexpected = sorted(converted_keys - reference_keys)
+    if missing or unexpected:
+        raise RuntimeError(
+            "Converted transformer keys do not match the Fun-InP reference. "
+            f"missing={missing[:8]} ({len(missing)} total), "
+            f"unexpected={unexpected[:8]} ({len(unexpected)} total)"
+        )
+
+    patch_key = "patch_embedding.weight"
+    if patch_key not in converted:
+        raise RuntimeError(f"Converted transformer lacks required parameter: {patch_key}")
+    reference_patch = reference[patch_key]
+    converted_patch = converted[patch_key]
+    if reference_patch.ndim != converted_patch.ndim:
+        raise RuntimeError(
+            f"{patch_key} rank mismatch: reference={reference_patch.ndim}, converted={converted_patch.ndim}"
+        )
+    if reference_patch.shape[0] != converted_patch.shape[0] or reference_patch.shape[2:] != converted_patch.shape[2:]:
+        raise RuntimeError(
+            f"{patch_key} non-input dimensions differ: "
+            f"reference={tuple(reference_patch.shape)}, converted={tuple(converted_patch.shape)}"
+        )
+    if converted_patch.shape[1] != TARGET_IN_CHANNELS:
+        raise RuntimeError(f"{patch_key} must have {TARGET_IN_CHANNELS} input channels")
+
+    shape_mismatches = []
+    for key in sorted(reference_keys - {patch_key}):
+        if converted[key].shape != reference[key].shape:
+            shape_mismatches.append((key, tuple(reference[key].shape), tuple(converted[key].shape)))
+    if shape_mismatches:
+        raise RuntimeError(
+            "Converted transformer shapes do not match the Fun-InP reference: "
+            f"{shape_mismatches[:8]} ({len(shape_mismatches)} total)"
+        )
+
+
+def build_track_encoder_state(
+    id_dim: int,
+    dtype: torch.dtype,
+    seed: int,
+    zero_init_head: bool,
+) -> dict[str, torch.Tensor]:
+    """Initialize the two bias-free TrackEncoder convolution weights reproducibly."""
+    if id_dim <= 0:
+        raise ValueError(f"--id-dim must be positive, got {id_dim}")
+
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(seed)
+        temporal_conv = nn.Conv3d(
+            id_dim,
+            TRACK_CHANNELS,
+            kernel_size=(VAE_TEMPORAL_COMPRESSION, 1, 1),
+            stride=(VAE_TEMPORAL_COMPRESSION, 1, 1),
+            bias=False,
+        )
+        projection = nn.Conv3d(TRACK_CHANNELS, TRACK_CHANNELS, kernel_size=1, bias=False)
+        if zero_init_head:
+            nn.init.zeros_(projection.weight)
+
+    state = {
+        "track_encoder.temporal_conv.weight": temporal_conv.weight.detach().to(dtype=dtype).contiguous(),
+        "track_encoder.proj.weight": projection.weight.detach().to(dtype=dtype).contiguous(),
+    }
+    expected_shapes = {
+        "track_encoder.temporal_conv.weight": (
+            TRACK_CHANNELS,
+            id_dim,
+            VAE_TEMPORAL_COMPRESSION,
+            1,
+            1,
+        ),
+        "track_encoder.proj.weight": (TRACK_CHANNELS, TRACK_CHANNELS, 1, 1, 1),
+    }
+    for key, expected_shape in expected_shapes.items():
+        if state[key].shape != expected_shape:
+            raise RuntimeError(f"Unexpected initialized shape for {key}: {tuple(state[key].shape)}")
+    return state
+
+
+def copy_passthrough_components(source_dir: Path, output_dir: Path) -> None:
+    """Copy every base bundle entry except the transformer directory."""
+    for source_path in sorted(source_dir.iterdir(), key=lambda path: path.name):
+        if source_path.name == "transformer":
+            continue
+        target_path = output_dir / source_path.name
+        if source_path.is_dir():
+            shutil.copytree(source_path, target_path, symlinks=True)
+        else:
+            shutil.copy2(source_path, target_path, follow_symlinks=False)
+
+
+def update_model_index(output_dir: Path, transformer_class: str) -> None:
+    """Point the copied bundle index at the selected FastVideo transformer."""
+    model_index_path = output_dir / "model_index.json"
+    if not model_index_path.is_file():
+        raise FileNotFoundError(f"Fun-InP base does not provide {model_index_path.name}")
+
+    model_index = json.loads(model_index_path.read_text(encoding="utf-8"))
+    transformer_entry = model_index.get("transformer")
+    if not isinstance(transformer_entry, list) or len(transformer_entry) != 2:
+        raise ValueError(f"Unsupported transformer entry in {model_index_path}: {transformer_entry!r}")
+    transformer_entry[1] = transformer_class
+    model_index_path.write_text(json.dumps(model_index, indent=2) + "\n", encoding="utf-8")
+
+
+def write_transformer_config(
+    base_dir: Path,
+    output_dir: Path,
+    transformer_class: str,
+    id_dim: int,
+    max_track_id: int,
+    zero_init_head: bool,
+) -> None:
+    """Write the widened, loader-compatible WanTrack transformer config."""
+    source_config_path = base_dir / "transformer" / "config.json"
+    if not source_config_path.is_file():
+        raise FileNotFoundError(f"Missing Fun-InP transformer config: {source_config_path}")
+
+    config = json.loads(source_config_path.read_text(encoding="utf-8"))
+    config.update(
+        {
+            "_class_name": transformer_class,
+            "in_channels": TARGET_IN_CHANNELS,
+            "out_channels": LATENT_CHANNELS,
+            "track_config": {
+                "id_dim": id_dim,
+                "track_channels": TRACK_CHANNELS,
+                "vae_spatial_compression": VAE_SPATIAL_COMPRESSION,
+                "vae_temporal_compression": VAE_TEMPORAL_COMPRESSION,
+                "max_track_id": max_track_id,
+                "zero_init_head": zero_init_head,
+                "use_bias": False,
+            },
+        }
+    )
+    config_path = output_dir / "transformer" / "config.json"
+    config_path.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
+
+
+def validate_paths(base_dir: Path, control_checkpoint: Path, output_dir: Path) -> None:
+    """Validate input paths and refuse ambiguous or destructive output layouts."""
+    if not base_dir.is_dir():
+        raise NotADirectoryError(f"Fun-InP base directory does not exist: {base_dir}")
+    if not control_checkpoint.is_file():
+        raise FileNotFoundError(f"Fun Control checkpoint does not exist: {control_checkpoint}")
+    if control_checkpoint.suffix != ".safetensors":
+        raise ValueError(f"Fun Control checkpoint must be a .safetensors file: {control_checkpoint}")
+
+    base_resolved = base_dir.resolve()
+    output_resolved = output_dir.resolve()
+    if output_resolved == base_resolved:
+        raise ValueError("Output directory must differ from the Fun-InP base directory")
+    if output_resolved.is_relative_to(base_resolved) or base_resolved.is_relative_to(output_resolved):
+        raise ValueError("Output directory and Fun-InP base directory must not contain one another")
+    if output_dir.exists() and not output_dir.is_dir():
+        raise NotADirectoryError(f"Output path exists and is not a directory: {output_dir}")
+    if output_dir.exists() and any(output_dir.iterdir()):
+        raise FileExistsError(f"Refusing to overwrite non-empty output directory: {output_dir}")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Convert VideoX-Fun Wan Control weights into a FastVideo WanTrack bundle.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--inp-base",
+        required=True,
+        type=Path,
+        help="Compatible Fun-InP Diffusers bundle used for config, key validation, and passthrough components.",
+    )
+    parser.add_argument(
+        "--control-ckpt",
+        required=True,
+        type=Path,
+        help="VideoX-Fun Wan Control diffusion_pytorch_model.safetensors.",
+    )
+    parser.add_argument("--out", required=True, type=Path, help="Output directory for the converted WanTrack bundle.")
+    parser.add_argument("--id-dim", type=int, default=128, help="Sinusoidal track-ID embedding dimension.")
+    parser.add_argument(
+        "--max-track-id",
+        type=int,
+        default=100_000,
+        help="Exclusive upper bound for sampled track IDs.",
+    )
+    parser.add_argument("--seed", type=int, default=0, help="Seed used to initialize TrackEncoder weights.")
+    parser.add_argument(
+        "--zero-init-head",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Zero-initialize TrackEncoder's final projection.",
+    )
+    parser.add_argument(
+        "--causal",
+        action="store_true",
+        help="Emit CausalTrackWanTransformer3DModel instead of the bidirectional transformer class.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Convert the checkpoint and create a complete WanTrack bundle."""
+    args = parse_args()
+    if args.max_track_id <= 0:
+        raise ValueError(f"--max-track-id must be positive, got {args.max_track_id}")
+    validate_paths(args.inp_base, args.control_ckpt, args.out)
+
+    reference_state = load_sharded_safetensors(args.inp_base / "transformer")
+    control_state = load_file(str(args.control_ckpt), device="cpu")
+    converted_state = rename_control_state_dict(control_state)
+    synthesize_image_query_norms(converted_state, reference_state)
+
+    patch_key = "patch_embedding.weight"
+    if patch_key not in converted_state:
+        raise RuntimeError(f"Fun Control checkpoint lacks required parameter: {patch_key}")
+    converted_state[patch_key] = remap_patch_embedding(converted_state[patch_key])
+    validate_converted_body(converted_state, reference_state)
+
+    converted_state.update(
+        build_track_encoder_state(
+            id_dim=args.id_dim,
+            dtype=converted_state[patch_key].dtype,
+            seed=args.seed,
+            zero_init_head=args.zero_init_head,
+        )
+    )
+
+    transformer_class = "CausalTrackWanTransformer3DModel" if args.causal else "TrackWanTransformer3DModel"
+    args.out.mkdir(parents=True, exist_ok=True)
+    copy_passthrough_components(args.inp_base, args.out)
+    (args.out / "transformer").mkdir()
+    update_model_index(args.out, transformer_class)
+    write_transformer_config(
+        base_dir=args.inp_base,
+        output_dir=args.out,
+        transformer_class=transformer_class,
+        id_dim=args.id_dim,
+        max_track_id=args.max_track_id,
+        zero_init_head=args.zero_init_head,
+    )
+
+    output_weights = args.out / "transformer" / "diffusion_pytorch_model.safetensors"
+    contiguous_state = {key: value.contiguous() for key, value in converted_state.items()}
+    save_file(contiguous_state, str(output_weights), metadata={"format": "pt"})
+    print(
+        f"Converted {len(reference_state)} Wan parameters and 2 TrackEncoder parameters "
+        f"to {output_weights} ({transformer_class})."
+    )
+
+
+if __name__ == "__main__":
+    main()