fix train/infer flow shift discrepency issue

Feat/kaiqin/mg overlay validation

examples/training/finetune/MatrixGame2.0/finetune_df.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers"
+DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "matrixgame_finetune"
+  --output_dir "checkpoints/matrixgame_finetune"
+  --override_transformer_cls_name "CausalMatrixGameWanModel"
+  --max_train_steps 1500
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 4
+  --num_latent_t 21
+  --num_height 352
+  --num_width 640
+  --num_frames 81
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 2
+  --tp_size 1
+  --hsdp_replicate_dim 4
+  --hsdp_shard_dim 2
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 100
+  --training_state_checkpointing_steps 100
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/matrixgame_ar_diffusion_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

fastvideo/dataset/dataloader/record_schema.py

@@ -188,3 +188,110 @@ def text_only_record_creator(text_name: str, text_embedding: np.ndarray,
         "caption": caption,
     }
     return record
+
+
+def matrixgame_ode_record_creator(
+        video_name: str,
+        clip_feature: np.ndarray,
+        first_frame_latent: np.ndarray,
+        trajectory_latents: np.ndarray,
+        trajectory_timesteps: np.ndarray,
+        pil_image: np.ndarray | None = None,
+        keyboard_cond: np.ndarray | None = None,
+        mouse_cond: np.ndarray | None = None,
+        caption: str = "") -> dict[str, Any]:
+    """Create a ODE trajectory record matching pyarrow_schema_matrixgame_ode_trajectory.
+
+    Args:
+        video_name: Base name/id for the sample.
+        clip_feature: CLIP image features.
+        first_frame_latent: First frame VAE latent.
+        trajectory_latents: Collected trajectory latents array.
+        trajectory_timesteps: Collected timesteps array.
+        pil_image: Optional PIL image array for validation.
+        keyboard_cond: Optional keyboard action array.
+        mouse_cond: Optional mouse action array.
+        caption: Optional caption.
+
+    Returns:
+        dict suitable for records_to_table(…, pyarrow_schema_matrixgame_ode_trajectory)
+    """
+    assert trajectory_latents is not None, "trajectory_latents is required"
+    assert trajectory_timesteps is not None, "trajectory_timesteps is required"
+    assert clip_feature is not None, "clip_feature is required"
+    assert first_frame_latent is not None, "first_frame_latent is required"
+
+    record = {
+        "id": video_name,
+        "file_name": video_name,
+        "caption": caption,
+        "media_type": "video",
+    }
+
+    # I2V features
+    record.update({
+        "clip_feature_bytes": clip_feature.tobytes(),
+        "clip_feature_shape": list(clip_feature.shape),
+        "clip_feature_dtype": str(clip_feature.dtype),
+    })
+
+    record.update({
+        "first_frame_latent_bytes": first_frame_latent.tobytes(),
+        "first_frame_latent_shape": list(first_frame_latent.shape),
+        "first_frame_latent_dtype": str(first_frame_latent.dtype),
+    })
+
+    # Optional PIL Image
+    if pil_image is not None:
+        record.update({
+            "pil_image_bytes": pil_image.tobytes(),
+            "pil_image_shape": list(pil_image.shape),
+            "pil_image_dtype": str(pil_image.dtype),
+        })
+    else:
+        record.update({
+            "pil_image_bytes": b"",
+            "pil_image_shape": [],
+            "pil_image_dtype": "",
+        })
+
+    # Actions
+    if keyboard_cond is not None:
+        record.update({
+            "keyboard_cond_bytes": keyboard_cond.tobytes(),
+            "keyboard_cond_shape": list(keyboard_cond.shape),
+            "keyboard_cond_dtype": str(keyboard_cond.dtype),
+        })
+    else:
+        record.update({
+            "keyboard_cond_bytes": b"",
+            "keyboard_cond_shape": [],
+            "keyboard_cond_dtype": "",
+        })
+
+    if mouse_cond is not None:
+        record.update({
+            "mouse_cond_bytes": mouse_cond.tobytes(),
+            "mouse_cond_shape": list(mouse_cond.shape),
+            "mouse_cond_dtype": str(mouse_cond.dtype),
+        })
+    else:
+        record.update({
+            "mouse_cond_bytes": b"",
+            "mouse_cond_shape": [],
+            "mouse_cond_dtype": "",
+        })
+
+    record.update({
+        "trajectory_latents_bytes": trajectory_latents.tobytes(),
+        "trajectory_latents_shape": list(trajectory_latents.shape),
+        "trajectory_latents_dtype": str(trajectory_latents.dtype),
+    })
+
+    record.update({
+        "trajectory_timesteps_bytes": trajectory_timesteps.tobytes(),
+        "trajectory_timesteps_shape": list(trajectory_timesteps.shape),
+        "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
+    })
+
+    return record

fastvideo/dataset/dataloader/schema.py

@@ -157,3 +157,36 @@ pyarrow_schema_matrixgame = pa.schema([
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
 ])
+
+pyarrow_schema_matrixgame_ode_trajectory = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Action ---
+    pa.field("mouse_cond_bytes", pa.binary()),
+    pa.field("mouse_cond_shape", pa.list_(pa.int64())),  # [T, 2]
+    pa.field("mouse_cond_dtype", pa.string()),
+    pa.field("keyboard_cond_bytes", pa.binary()),
+    pa.field("keyboard_cond_shape", pa.list_(pa.int64())),  # [T, 4]
+    pa.field("keyboard_cond_dtype", pa.string()),
+    #I2V
+    pa.field("clip_feature_bytes", pa.binary()),
+    pa.field("clip_feature_shape", pa.list_(pa.int64())),
+    pa.field("clip_feature_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()),
+    # I2V Validation
+    pa.field("pil_image_bytes", pa.binary()),
+    pa.field("pil_image_shape", pa.list_(pa.int64())),
+    pa.field("pil_image_dtype", pa.string()),
+    # --- ODE Trajectory ---
+    pa.field("trajectory_latents_bytes", pa.binary()),
+    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_latents_dtype", pa.string()),
+    pa.field("trajectory_timesteps_bytes", pa.binary()),
+    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_timesteps_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+])

fastvideo/dataset/validation_dataset.py

@@ -139,6 +139,16 @@ class ValidationDataset(IterableDataset):
                 else:
                     sample["video"] = load_video(video_path)
 
+            if sample.get("ref_video", None) is not None:
+                ref_video_path = sample["ref_video"]
+                ref_video_path = os.path.join(self.dir, ref_video_path)
+                if not pathlib.Path(ref_video_path).is_file():
+                    logger.warning("Reference video file %s does not exist.",
+                                   ref_video_path)
+                    sample.pop("ref_video", None)
+                else:
+                    sample["ref_video"] = ref_video_path
+
             if sample.get("control_image_path", None) is not None:
                 control_image_path = sample["control_image_path"]
                 control_image_path = os.path.join(self.dir, control_image_path)

fastvideo/models/registry.py

@@ -93,6 +93,9 @@ _SCHEDULERS = {
     "FlowMatchEulerDiscreteScheduler":
     ("schedulers", "scheduling_flow_match_euler_discrete",
      "FlowMatchEulerDiscreteScheduler"),
+    "DiffusionForcingScheduler":
+    ("schedulers", "scheduling_diffusion_forcing",
+     "DiffusionForcingScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
     "FlowUniPCMultistepScheduler":

fastvideo/models/schedulers/scheduling_diffusion_forcing.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+import torch
+
+from fastvideo.models.schedulers.base import BaseScheduler
+
+
+class DiffusionForcingSchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+
+
+class DiffusionForcingScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
+
+    config_name = "scheduler_config.json"
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_inference_steps: int = 100,
+        num_train_timesteps: int = 1000,
+        shift: float = 5.0,
+        sigma_max: float = 1.0,
+        sigma_min: float = 0.0,
+        extra_one_step: bool = True,
+        training: bool = False,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.extra_one_step = extra_one_step
+        self.set_timesteps(num_inference_steps, training=training)
+
+    def _apply_shift(self, sigmas: torch.Tensor) -> torch.Tensor:
+        """Apply time-shift to sigmas, matching FlowMatchEulerDiscreteScheduler."""
+        if self.shift != 1.0:
+            sigmas = (self.shift * sigmas
+                      / (1 + (self.shift - 1) * sigmas))
+        return sigmas
+
+    def _lookup_sigma(self, timestep: torch.Tensor) -> torch.Tensor:
+        """Map timestep values to sigmas via table lookup."""
+        timestep = timestep.to(torch.float32)
+        self.timesteps = self.timesteps.to(timestep.device)
+        self.sigmas = self.sigmas.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0)
+             - timestep.unsqueeze(1)).abs(),
+            dim=1,
+        )
+        return self.sigmas[timestep_id]
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = 100,
+        denoising_strength: float = 1.0,
+        training: bool = False,
+        return_dict: bool = False,
+        **kwargs,
+    ):
+        if training:
+            # Linear sigmas from 1.0 to 0.0 (inclusive),
+            # then apply the same shift used by the student model.
+            raw_sigmas = torch.arange(
+                self.num_train_timesteps, -1, -1,
+                dtype=torch.float32,
+            ) / float(self.num_train_timesteps)
+            sigmas = self._apply_shift(raw_sigmas)
+            self.sigmas = sigmas
+            self.timesteps = (
+                sigmas * float(self.num_train_timesteps)
+            ).to(torch.float32)
+            self.linear_timesteps_weights = torch.ones_like(
+                self.timesteps)
+            return
+
+        sigma_start = self.sigma_min + (
+            self.sigma_max - self.sigma_min
+        ) * denoising_strength
+        if self.extra_one_step:
+            sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1
+            )[:-1]
+        else:
+            sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps
+            )
+        sigmas = self._apply_shift(sigmas)
+        self.sigmas = sigmas.to(torch.float32)
+        self.timesteps = (
+            self.sigmas * float(self.num_train_timesteps)
+        ).to(torch.float32)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: torch.FloatTensor,
+        sample: torch.FloatTensor,
+        to_final: bool = False,
+        return_dict: bool = False,
+        **kwargs,
+    ):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        elif timestep.ndim == 0:
+            timestep = timestep.unsqueeze(0)
+
+        timestep = timestep.to(model_output.device, dtype=torch.float32)
+        self.timesteps = self.timesteps.to(model_output.device)
+        self.sigmas = self.sigmas.to(model_output.device)
+
+        sigma = self._lookup_sigma(timestep).reshape(-1, 1, 1, 1)
+        if to_final:
+            sigma_next = torch.zeros_like(sigma)
+        else:
+            timestep_id = torch.argmin(
+                (self.timesteps.unsqueeze(0)
+                 - timestep.unsqueeze(1)).abs(),
+                dim=1,
+            )
+            sigma_next = torch.zeros_like(sigma)
+            valid = timestep_id + 1 < len(self.timesteps)
+            if valid.any():
+                sigma_next[valid] = self.sigmas[
+                    timestep_id[valid] + 1
+                ].reshape(-1, 1, 1, 1)
+
+        prev_sample = sample + model_output * (sigma_next - sigma)
+        if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
+            return (prev_sample,)
+        return DiffusionForcingSchedulerOutput(prev_sample=prev_sample)
+
+    @staticmethod
+    def calculate_alpha_beta_high(sigma, sigma_bound):
+        alpha = (1 - sigma) / (1 - sigma_bound)
+        beta = torch.sqrt(sigma**2 - (alpha * sigma_bound) ** 2)
+        return alpha, beta
+
+    def add_noise(self, original_samples, noise, timestep):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        sigma = self._lookup_sigma(timestep).to(noise.device).reshape(
+            -1, 1, 1, 1
+        )
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def add_noise_high(
+        self, original_samples, noise, timestep, boundary_timestep
+    ):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        if boundary_timestep.ndim == 2:
+            boundary_timestep = boundary_timestep.flatten(0, 1)
+        sigma = self._lookup_sigma(timestep).to(noise.device).reshape(
+            -1, 1, 1, 1
+        )
+        sigma_boundary = self._lookup_sigma(
+            boundary_timestep
+        ).to(noise.device).reshape(-1, 1, 1, 1)
+        alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary)
+        sample = alpha * original_samples + beta * noise
+        return sample.type_as(noise)
+
+    def sigma_from_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
+        """Public API used by ``pred_noise_to_pred_video``."""
+        return self._lookup_sigma(timestep)
+
+    def training_target(self, sample, noise, timestep):
+        return noise - sample
+
+    def training_weight(self, timestep):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        return torch.ones_like(timestep, dtype=torch.float32)
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        return sample
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = shift
+
+
+EntryClass = DiffusionForcingScheduler

fastvideo/models/utils.py

@@ -173,11 +173,15 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     device = pred_noise.device
     pred_noise = pred_noise.double().to(device)
     noise_input_latent = noise_input_latent.double().to(device)
-    sigmas = scheduler.sigmas.double().to(device)
-    timesteps = scheduler.timesteps.double().to(device)
-    timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    if hasattr(scheduler, "sigma_from_timestep"):
+        sigma_t = scheduler.sigma_from_timestep(timestep).double().to(
+            device).reshape(-1, 1, 1, 1)
+    else:
+        sigmas = scheduler.sigmas.double().to(device)
+        timesteps = scheduler.timesteps.double().to(device)
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
     pred_video = noise_input_latent - sigma_t * pred_noise
     return pred_video.to(dtype)
 
@@ -217,14 +221,20 @@ def pred_noise_to_x_bound(pred_noise: torch.Tensor,
     device = pred_noise.device
     pred_noise = pred_noise.double().to(device)
     noise_input_latent = noise_input_latent.double().to(device)
-    sigmas = scheduler.sigmas.double().to(device)
-    timesteps = scheduler.timesteps.double().to(device)
-    timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    if hasattr(scheduler, "sigma_from_timestep"):
+        sigma_t = scheduler.sigma_from_timestep(timestep).double().to(
+            device).reshape(-1, 1, 1, 1)
+        sigma_t_boundary = scheduler.sigma_from_timestep(
+            boundary_timestep).double().to(device).reshape(-1, 1, 1, 1)
+    else:
+        sigmas = scheduler.sigmas.double().to(device)
+        timesteps = scheduler.timesteps.double().to(device)
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
 
-    boundary_timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
+        boundary_timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
     pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise
     return pred_video.to(dtype)

fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py

@@ -60,7 +60,7 @@ class PreprocessPipeline_MatrixGame(BasePreprocessPipeline):
             clip_features = clip_features.last_hidden_state
 
         features["clip_feature"] = clip_features
-        """Get VAE features from the first frame of each video"""
+        # Get CLIP features from the first frame of each video.
         video_conditions = []
         for frame in first_frame:
             processed_img = frame.to(device="cpu", dtype=torch.float32)

fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline_ode_trajectory.py

@@ -0,0 +1,391 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+ODE Trajectory Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
+using the modular pipeline architecture.
+
+Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import getdataset
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, records_to_table)
+from fastvideo.dataset.dataloader.record_schema import (matrixgame_ode_record_creator)
+from fastvideo.dataset.dataloader.schema import (pyarrow_schema_matrixgame_ode_trajectory)
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (SelfForcingFlowMatchScheduler)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, LatentPreparationStage,
+                                        MatrixGameImageEncodingStage, TimestepPreparationStage)
+from fastvideo.pipelines.stages.matrixgame_denoising import (MatrixGameCausalDenoisingStage)
+from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline_MatrixGame_ODE_Trajectory(BasePreprocessPipeline):
+    """ODE Trajectory preprocessing pipeline implementation."""
+
+    _required_config_modules = ["vae", "image_encoder", "image_processor", "transformer", "scheduler"]
+
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+    pbar: Any
+    num_processed_samples: int
+
+    def get_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for ODE Trajectory pipeline."""
+        return pyarrow_schema_matrixgame_ode_trajectory
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        assert fastvideo_args.pipeline_config.flow_shift == 5
+        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
+                                                                  sigma_min=0.0,
+                                                                  extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=48, denoising_strength=1.0)
+
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=MatrixGameImageEncodingStage(
+                           image_encoder=self.get_module("image_encoder"),
+                           image_processor=self.get_module("image_processor"),
+                       ))
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
+                                                    transformer=self.get_module("transformer", None)))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=MatrixGameCausalDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                           vae=self.get_module("vae"),
+                       ))
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
+
+    def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+
+        # TODO(will): move these to cpu at some point
+        self.get_module("image_encoder").to(get_local_torch_device())
+        self.get_module("vae").to(get_local_torch_device())
+
+        features = {}
+        """Get CLIP features from the first frame of each video."""
+        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
+        _, _, num_frames, height, width = valid_data["pixel_values"].shape
+        # latent_height = height // self.get_module(
+        #     "vae").spatial_compression_ratio
+        # latent_width = width // self.get_module("vae").spatial_compression_ratio
+
+        processed_images = []
+        # Frame has values between -1 and 1
+        for frame in first_frame:
+            frame = (frame + 1) * 127.5
+            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
+            processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt")
+            processed_images.append(processed_img)
+
+        # Get CLIP features
+        pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device())
+        with torch.no_grad():
+            image_inputs = {'pixel_values': pixel_values}
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                clip_features = self.get_module("image_encoder")(**image_inputs)
+            clip_features = clip_features.last_hidden_state
+
+        features["clip_feature"] = clip_features
+        features["pil_image"] = first_frame
+        # Get CLIP features from the first frame of each video.
+        video_conditions = []
+        for frame in first_frame:
+            processed_img = frame.to(device="cpu", dtype=torch.float32)
+            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
+            # (B, H, W, C) -> (B, C, 1, H, W)
+            video_condition = torch.cat([
+                processed_img,
+                processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1], num_frames - 1, height, width)
+            ],
+                                        dim=2)
+            video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32)
+            video_conditions.append(video_condition)
+
+        video_conditions = torch.cat(video_conditions, dim=0)
+
+        with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
+            encoder_outputs = self.get_module("vae").encode(video_conditions)
+
+        # Use mode() instead of mean
+        latent_condition = encoder_outputs.mode()
+
+        # Use latents_mean/latents_std normalization to match
+        vae = self.get_module("vae")
+        if (hasattr(vae.config, 'latents_mean') and hasattr(vae.config, 'latents_std')):
+            latents_mean = torch.tensor(vae.config.latents_mean,
+                                        device=latent_condition.device,
+                                        dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
+            latents_std = torch.tensor(vae.config.latents_std,
+                                       device=latent_condition.device,
+                                       dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
+            latent_condition = (latent_condition - latents_mean) / latents_std
+        elif (hasattr(vae, "shift_factor") and vae.shift_factor is not None):
+            if isinstance(vae.shift_factor, torch.Tensor):
+                latent_condition -= vae.shift_factor.to(latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= vae.shift_factor
+
+            if isinstance(vae.scaling_factor, torch.Tensor):
+                latent_condition = latent_condition * vae.scaling_factor.to(latent_condition.device,
+                                                                            latent_condition.dtype)
+            else:
+                latent_condition = latent_condition * vae.scaling_factor
+
+        # Create mask_cond: ones for first frame, zeros for rest
+        # Shape: (B, 16, latent_frames, latent_height, latent_width)
+        mask_cond = torch.ones_like(latent_condition)
+        mask_cond[:, :, 1:] = 0  # Set all frames except first to 0
+        # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond
+        # Shape: (B, 20, latent_frames, latent_height, latent_width)
+        cond_concat = torch.cat([mask_cond[:, :4], latent_condition], dim=1)
+        features["first_frame_latent"] = cond_concat
+
+        if "action_path" in valid_data and valid_data["action_path"]:
+            keyboard_cond_list = []
+            mouse_cond_list = []
+            keyboard_dim = self.get_module("transformer").config.arch_config.action_config["keyboard_dim_in"]
+            for action_path in valid_data["action_path"]:
+                if action_path:
+                    action_data = np.load(action_path, allow_pickle=True)
+                    if isinstance(action_data, np.ndarray) and action_data.dtype == np.dtype('O'):
+                        action_dict = action_data.item()
+                        if "keyboard" in action_dict:
+                            keyboard_cond_list.append(action_dict["keyboard"][:, :keyboard_dim].astype(np.float32))
+                        if "mouse" in action_dict:
+                            mouse_cond_list.append(action_dict["mouse"])
+                    else:
+                        keyboard_cond_list.append(action_data[:, :keyboard_dim].astype(np.float32))
+            if keyboard_cond_list:
+                features["keyboard_cond"] = keyboard_cond_list
+            if mouse_cond_list:
+                features["mouse_cond"] = mouse_cond_list
+        return features
+
+    def preprocess_action_and_trajectory(self, fastvideo_args: FastVideoArgs, args):
+        """Preprocess data and generate trajectory information."""
+
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # Filter out invalid samples (those with all zeros)
+                valid_indices = []
+                for i, pixel_values in enumerate(data["pixel_values"]):
+                    if not torch.all(pixel_values == 0):  # Check if all values are zero
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples
+                valid_data = {
+                    "pixel_values": torch.stack([data["pixel_values"][i] for i in valid_indices]),
+                    "path": [data["path"][i] for i in valid_indices],
+                }
+
+                if "fps" in data:
+                    valid_data["fps"] = [data["fps"][i] for i in valid_indices]
+                if "duration" in data:
+                    valid_data["duration"] = [data["duration"][i] for i in valid_indices]
+                if "action_path" in data:
+                    valid_data["action_path"] = [data["action_path"][i] for i in valid_indices]
+
+                pixel_values = valid_data["pixel_values"]
+                if pixel_values.shape[2] == 1 and args.num_frames is not None:
+                    pixel_values = pixel_values.repeat(1, 1, args.num_frames, 1, 1)
+                    valid_data["pixel_values"] = pixel_values
+
+                # Get extra features if needed
+                extra_features = self.get_extra_features(valid_data, fastvideo_args)
+
+                clip_features = extra_features['clip_feature']
+                image_latents = extra_features['first_frame_latent']
+                image_latents = image_latents[:, :, :args.num_latent_t]
+                pil_image = extra_features['pil_image']
+                if "keyboard_cond" in extra_features:
+                    keyboard_cond = extra_features['keyboard_cond']
+                else:
+                    keyboard_cond = None
+                if "mouse_cond" in extra_features:
+                    mouse_cond = extra_features['mouse_cond']
+                else:
+                    mouse_cond = None
+
+                sampling_params = SamplingParam.from_pretrained(args.model_path)
+
+                trajectory_latents = []
+                trajectory_timesteps = []
+                trajectory_decoded = []
+
+                device = get_local_torch_device()
+                for i in range(len(valid_indices)):
+                    # Collect the trajectory data
+                    batch = ForwardBatch(**shallow_asdict(sampling_params), )
+                    batch.image_embeds = [clip_features[i].unsqueeze(0)]
+                    batch.image_latent = image_latents[i].unsqueeze(0)
+                    batch.keyboard_cond = (torch.from_numpy(keyboard_cond[i]).unsqueeze(0).to(device)
+                                           if keyboard_cond is not None else None)
+                    batch.mouse_cond = (torch.from_numpy(mouse_cond[i]).unsqueeze(0).to(device)
+                                        if mouse_cond is not None else None)
+                    batch.num_inference_steps = 48
+                    batch.return_trajectory_latents = True
+                    # Enabling this will save the decoded trajectory videos.
+                    # Used for debugging.
+                    batch.return_trajectory_decoded = False
+                    batch.height = args.max_height
+                    batch.width = args.max_width
+                    batch.fps = args.train_fps
+                    batch.num_frames = valid_data["pixel_values"].shape[2]
+                    batch.guidance_scale = 6.0
+                    batch.do_classifier_free_guidance = False
+                    batch.prompt = ""
+
+                    result_batch = self.input_validation_stage(batch, fastvideo_args)
+                    result_batch = self.timestep_preparation_stage(batch, fastvideo_args)
+                    result_batch.timesteps = result_batch.timesteps.to(device)
+                    result_batch = self.latent_preparation_stage(result_batch, fastvideo_args)
+                    result_batch = self.denoising_stage(result_batch, fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch, fastvideo_args)
+
+                    trajectory_latents.append(result_batch.trajectory_latents.cpu())
+                    trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
+
+                # Prepare extra features
+                extra_features = {
+                    "trajectory_latents": trajectory_latents,
+                    "trajectory_timesteps": trajectory_timesteps
+                }
+
+                if batch.return_trajectory_decoded:
+                    for i, decoded_frames in enumerate(trajectory_decoded):
+                        for j, decoded_frame in enumerate(decoded_frames):
+                            save_decoded_latents_as_video(decoded_frame,
+                                                          f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
+                                                          args.train_fps)
+
+                # Prepare batch data for Parquet dataset
+                batch_data: list[dict[str, Any]] = []
+
+                # Add progress bar for saving outputs
+                save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False)
+
+                for idx, video_path in save_pbar:
+                    video_name = os.path.basename(video_path).split(".")[0]
+
+                    clip_feature_np = clip_features[idx].cpu().numpy()
+                    first_frame_latent_np = image_latents[idx].cpu().numpy()
+                    pil_image_np = pil_image[idx].cpu().numpy()
+                    keyboard_cond_np = keyboard_cond[idx] if keyboard_cond is not None else None
+                    mouse_cond_np = mouse_cond[idx] if mouse_cond is not None else None
+
+                    # Get trajectory features for this sample
+                    traj_latents = extra_features["trajectory_latents"][idx]
+                    traj_timesteps = extra_features["trajectory_timesteps"][idx]
+                    if isinstance(traj_latents, torch.Tensor):
+                        traj_latents = traj_latents.cpu().float().numpy()
+                    if isinstance(traj_timesteps, torch.Tensor):
+                        traj_timesteps = traj_timesteps.cpu().float().numpy()
+
+                    # Create record for Parquet dataset
+                    record: dict[str, Any] = matrixgame_ode_record_creator(video_name=video_name,
+                                                                           clip_feature=clip_feature_np,
+                                                                           first_frame_latent=first_frame_latent_np,
+                                                                           trajectory_latents=traj_latents,
+                                                                           trajectory_timesteps=traj_timesteps,
+                                                                           pil_image=pil_image_np,
+                                                                           keyboard_cond=keyboard_cond_np,
+                                                                           mouse_cond=mouse_cond_np,
+                                                                           caption="")
+                    batch_data.append(record)
+
+                if batch_data:
+                    write_pbar = tqdm(total=1, desc="Writing to Parquet dataset", unit="batch")
+                    table = records_to_table(batch_data, self.get_pyarrow_schema())
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    if not hasattr(self, 'dataset_writer'):
+                        self.dataset_writer = ParquetDatasetWriter(
+                            out_dir=self.combined_parquet_dir,
+                            samples_per_file=args.samples_per_file,
+                        )
+                    self.dataset_writer.append_table(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    written = self.dataset_writer.flush()
+                    logger.info("Flushed %s samples to parquet", written)
+                    self.num_processed_samples = 0
+
+        # Final flush for any remaining samples
+        if hasattr(self, 'dataset_writer'):
+            written = self.dataset_writer.flush(write_remainder=True)
+            if written:
+                logger.info("Final flush wrote %s samples", written)
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading dataset
+        train_dataset = getdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for video preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing videos",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: dict[str, Any] = {}  # Store latent tensors
+        self.preprocess_action_and_trajectory(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_MatrixGame_ODE_Trajectory

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -12,6 +12,8 @@ from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (P
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (PreprocessPipeline_T2V)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_text import (PreprocessPipeline_Text)
 from fastvideo.pipelines.preprocess.matrixgame.matrixgame_preprocess_pipeline import (PreprocessPipeline_MatrixGame)
+from fastvideo.pipelines.preprocess.matrixgame.matrixgame_preprocess_pipeline_ode_trajectory import (
+    PreprocessPipeline_MatrixGame_ODE_Trajectory)
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -58,9 +60,11 @@ def main(args) -> None:
         PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
     elif args.preprocess_task == "matrixgame":
         PreprocessPipeline = PreprocessPipeline_MatrixGame
+    elif args.preprocess_task == "matrixgame_ode_trajectory":
+        PreprocessPipeline = PreprocessPipeline_MatrixGame_ODE_Trajectory
     else:
         raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
-                         f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame")
+                         f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame, matrixgame_ode_trajectory")
 
     logger.info("Preprocess task: %s using %s", args.preprocess_task, PreprocessPipeline.__name__)
 
@@ -96,11 +100,12 @@ if __name__ == "__main__":
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
     parser.add_argument("--flow_shift", type=float, default=None)
-    parser.add_argument("--preprocess_task",
-                        type=str,
-                        default="t2v",
-                        choices=["t2v", "i2v", "text_only", "ode_trajectory", "matrixgame"],
-                        help="Type of preprocessing task to run")
+    parser.add_argument(
+        "--preprocess_task",
+        type=str,
+        default="t2v",
+        choices=["t2v", "i2v", "text_only", "ode_trajectory", "matrixgame", "matrixgame_ode_trajectory"],
+        help="Type of preprocessing task to run")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)

fastvideo/training/matrixgame_ar_diffusion_pipeline.py

@@ -0,0 +1,476 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from copy import deepcopy
+from typing import Any, cast
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_diffusion_forcing import (
+    DiffusionForcingScheduler, )
+from fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline import (
+    MatrixGameCausalDMDPipeline, )
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases, )
+from fastvideo.utils import shallow_asdict
+
+logger = init_logger(__name__)
+
+
+class MatrixGameARDiffusionPipeline(TrainingPipeline):
+
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = DiffusionForcingScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True,
+        )
+        self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)
+
+    def set_schemas(self):
+        self.train_dataset_schema = pyarrow_schema_matrixgame
+
+    def _get_temporal_compression_ratio(self) -> int:
+        assert self.training_args is not None
+        return int(self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+
+    def _resolve_num_frame_per_block(self, training_args: TrainingArgs) -> int:
+        transformer_num_frame_per_block = getattr(self.transformer, "num_frame_per_block", None)
+        requested_num_frame_per_block = getattr(training_args, "num_frame_per_block", None)
+
+        if (transformer_num_frame_per_block is not None and requested_num_frame_per_block is not None
+                and transformer_num_frame_per_block != requested_num_frame_per_block):
+            raise ValueError("num_frame_per_block mismatch between loaded transformer and "
+                             "training args: "
+                             f"transformer={transformer_num_frame_per_block}, "
+                             f"training_args={requested_num_frame_per_block}")
+
+        if transformer_num_frame_per_block is not None:
+            return int(transformer_num_frame_per_block)
+        if requested_num_frame_per_block is not None:
+            return int(requested_num_frame_per_block)
+        return 3
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        super().initialize_training_pipeline(training_args)
+
+        self.vae = self.get_module("vae")
+        self.vae.requires_grad_(False)
+
+        self.num_frame_per_block = self._resolve_num_frame_per_block(training_args)
+        self.chunkwise_timestep = bool(getattr(training_args, "chunkwise_timestep", False))
+        self.timestep_shift = training_args.pipeline_config.flow_shift
+        self.ar_noise_scheduler = DiffusionForcingScheduler(shift=self.timestep_shift,
+                                                            sigma_min=0.0,
+                                                            extra_one_step=True)
+        self.ar_noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)
+
+        logger.info(
+            "MatrixGame AR diffusion pipeline initialized with "
+            "num_frame_per_block=%d, diffusion_forcing_shift=%.1f, "
+            "chunkwise_timestep=%s",
+            self.num_frame_per_block,
+            self.timestep_shift,
+            self.chunkwise_timestep,
+        )
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing MatrixGame AR validation pipeline...")
+        args_copy = deepcopy(training_args)
+        args_copy.inference_mode = True
+
+        validation_scheduler = DiffusionForcingScheduler(
+            shift=args_copy.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True,
+        )
+        validation_scheduler.set_timesteps(num_inference_steps=1000, training=False)
+
+        num_val_steps = int(training_args.validation_sampling_steps.split(",")[0])
+        step_size = 1000 // num_val_steps
+        args_copy.pipeline_config.dmd_denoising_steps = list(range(1000, 0, -step_size))
+        args_copy.pipeline_config.warp_denoising_step = True
+        training_args.pipeline_config.dmd_denoising_steps = (args_copy.pipeline_config.dmd_denoising_steps)
+        training_args.pipeline_config.warp_denoising_step = True
+
+        logger.info(
+            "Validation: %d-step MatrixGame causal denoising, "
+            "dmd_denoising_steps has %d entries",
+            num_val_steps,
+            len(args_copy.pipeline_config.dmd_denoising_steps),
+        )
+
+        loaded_modules = {
+            "transformer": self.get_module("transformer"),
+            "vae": self.get_module("vae"),
+            "scheduler": validation_scheduler,
+        }
+        image_encoder = self.get_module("image_encoder")
+        image_processor = self.get_module("image_processor")
+        if image_encoder is not None:
+            loaded_modules["image_encoder"] = image_encoder
+        if image_processor is not None:
+            loaded_modules["image_processor"] = image_processor
+
+        self.validation_pipeline = MatrixGameCausalDMDPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,
+            inference_mode=True,
+            loaded_modules=loaded_modules,
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True,
+        )
+
+    def _get_timestep(
+        self,
+        batch_size: int,
+        num_frame: int,
+        num_frame_per_block: int,
+        uniform_timestep: bool = False,
+    ) -> torch.Tensor:
+        """Sample timestep indices and return the scheduler's
+        (shifted) timestep values.
+        """
+        device = get_local_torch_device()
+        num_schedule_steps = len(self.ar_noise_scheduler.timesteps)
+
+        if uniform_timestep:
+            indices = torch.randint(
+                0,
+                num_schedule_steps,
+                [batch_size, 1],
+                device=device,
+                dtype=torch.long,
+            ).repeat(1, num_frame)
+        else:
+            indices = torch.randint(
+                0,
+                num_schedule_steps,
+                [batch_size, num_frame],
+                device=device,
+                dtype=torch.long,
+            )
+
+        if getattr(self, "chunkwise_timestep", False):
+            chunk_size = int(num_frame_per_block)
+            if chunk_size <= 0:
+                raise ValueError(
+                    f"num_frame_per_block must be > 0, "
+                    f"got {chunk_size}")
+            num_chunks = ((num_frame + chunk_size - 1)
+                          // chunk_size)
+            padded = num_chunks * chunk_size
+            if padded != num_frame:
+                pad = indices[:, -1:].expand(
+                    batch_size, padded - num_frame)
+                indices = torch.cat([indices, pad], dim=1)
+            indices = indices.reshape(
+                batch_size, num_chunks, chunk_size)
+            indices[:, :, 1:] = indices[:, :, :1]
+            indices = indices.reshape(
+                batch_size, padded)[:, :num_frame]
+
+        schedule_t = self.ar_noise_scheduler.timesteps.to(device)
+        return schedule_t[indices]
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            self.train_loader_iter = iter(self.train_dataloader)
+            batch = next(self.train_loader_iter)
+
+        latents = batch["vae_latent"]
+        latents = latents[:, :, :self.training_args.num_latent_t]
+        clip_features = batch["clip_feature"]
+        image_latents = batch["first_frame_latent"]
+        image_latents = image_latents[:, :, :self.training_args.num_latent_t]
+        pil_image = batch["pil_image"]
+        infos = batch["info_list"]
+
+        training_batch.latents = latents.to(get_local_torch_device(), dtype=torch.bfloat16)
+        training_batch.encoder_hidden_states = None
+        training_batch.encoder_attention_mask = None
+        training_batch.preprocessed_image = pil_image.to(get_local_torch_device())
+        training_batch.image_embeds = clip_features.to(get_local_torch_device())
+        training_batch.image_latents = image_latents.to(get_local_torch_device())
+        training_batch.infos = infos
+
+        if "mouse_cond" in batch and batch["mouse_cond"].numel() > 0:
+            training_batch.mouse_cond = batch["mouse_cond"].to(get_local_torch_device(), dtype=torch.bfloat16)
+        else:
+            training_batch.mouse_cond = None
+
+        if "keyboard_cond" in batch and batch["keyboard_cond"].numel() > 0:
+            training_batch.keyboard_cond = batch["keyboard_cond"].to(get_local_torch_device(), dtype=torch.bfloat16)
+        else:
+            training_batch.keyboard_cond = None
+
+        temporal_compression_ratio = self._get_temporal_compression_ratio()
+        expected_num_frames = (self.training_args.num_latent_t - 1) * temporal_compression_ratio + 1
+        if training_batch.keyboard_cond is not None:
+            assert training_batch.keyboard_cond.shape[1] >= expected_num_frames, (
+                f"keyboard_cond has {training_batch.keyboard_cond.shape[1]} "
+                f"frames but need at least {expected_num_frames}")
+            training_batch.keyboard_cond = training_batch.keyboard_cond[:, :expected_num_frames]
+        if training_batch.mouse_cond is not None:
+            assert training_batch.mouse_cond.shape[1] >= expected_num_frames, (
+                f"mouse_cond has {training_batch.mouse_cond.shape[1]} frames "
+                f"but need at least {expected_num_frames}")
+            training_batch.mouse_cond = training_batch.mouse_cond[:, :expected_num_frames]
+
+        return training_batch
+
+    def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch:
+        """Prepare diffusion-forcing inputs and MatrixGame I2V concat."""
+        assert self.training_args is not None
+        latents = training_batch.latents
+        assert latents is not None
+        batch_size = latents.shape[0]
+        num_latent_t = latents.shape[2]
+
+        latents_btchw = latents.permute(0, 2, 1, 3, 4)
+
+        timesteps = self._get_timestep(
+            batch_size=batch_size,
+            num_frame=num_latent_t,
+            num_frame_per_block=self.num_frame_per_block,
+            uniform_timestep=False,
+        )
+
+        noise_generator = getattr(self, "noise_gen_cuda", None)
+        noise = torch.randn(
+            latents_btchw.shape,
+            generator=noise_generator,
+            device=latents_btchw.device,
+            dtype=latents_btchw.dtype,
+        )
+        if (self.training_args.sp_size > 1 and self.sp_group is not None and hasattr(self.sp_group, "broadcast")):
+            self.sp_group.broadcast(timesteps, src=0)
+            self.sp_group.broadcast(noise, src=0)
+
+        noisy_latents = self.ar_noise_scheduler.add_noise(
+            latents_btchw.flatten(0, 1),
+            noise.flatten(0, 1),
+            timesteps.flatten(0, 1),
+        ).unflatten(0, (batch_size, num_latent_t))
+
+        noisy_model_input = noisy_latents.permute(0, 2, 1, 3, 4)
+
+        assert isinstance(training_batch.image_latents, torch.Tensor)
+        image_latents = training_batch.image_latents.to(get_local_torch_device(), dtype=torch.bfloat16)
+
+        temporal_compression_ratio = self._get_temporal_compression_ratio()
+        num_frames = (num_latent_t - 1) * temporal_compression_ratio + 1
+        _, _, _, latent_height, latent_width = image_latents.shape
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        mask_lat_size[:, :, 1:] = 0
+
+        first_frame_mask = mask_lat_size[:, :, :1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask,
+            dim=2,
+            repeats=temporal_compression_ratio,
+        )
+        mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], dim=2)
+        mask_lat_size = mask_lat_size.view(
+            batch_size,
+            -1,
+            temporal_compression_ratio,
+            latent_height,
+            latent_width,
+        )
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(image_latents.device, dtype=torch.bfloat16)
+
+        noisy_model_input = torch.cat([noisy_model_input, mask_lat_size, image_latents], dim=1)
+
+        training_target = self.ar_noise_scheduler.training_target(
+            latents_btchw.flatten(0, 1),
+            noise.flatten(0, 1),
+            timesteps.flatten(0, 1),
+        ).unflatten(0, (batch_size, num_latent_t))
+
+        training_batch.noisy_model_input = noisy_model_input
+        training_batch.timesteps = timesteps
+        training_batch.noise = noise.permute(0, 2, 1, 3, 4)
+        training_batch.raw_latent_shape = latents.shape
+        training_batch._ar_training_target = training_target
+
+        return training_batch
+
+    def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:
+        image_embeds = training_batch.image_embeds
+        assert isinstance(image_embeds, torch.Tensor)
+        assert torch.isnan(image_embeds).sum() == 0
+        image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16)
+
+        timesteps = training_batch.timesteps
+        assert isinstance(timesteps, torch.Tensor)
+        assert timesteps.ndim == 2, (f"Expected per-frame timesteps [B, T], got shape {timesteps.shape}")
+
+        training_batch.input_kwargs = {
+            "hidden_states": training_batch.noisy_model_input,
+            "encoder_hidden_states": None,
+            "timestep": timesteps.to(get_local_torch_device()),
+            "encoder_hidden_states_image": image_embeds,
+            "mouse_cond": training_batch.mouse_cond,
+            "keyboard_cond": training_batch.keyboard_cond,
+            "num_frame_per_block": self.num_frame_per_block,
+            "return_dict": False,
+        }
+        return training_batch
+
+    def _transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -> TrainingBatch:
+        """
+        Run transformer forward pass and compute diffusion-forcing loss.
+        """
+        input_kwargs = training_batch.input_kwargs
+
+        with set_forward_context(
+                current_timestep=training_batch.timesteps,
+                attn_metadata=None,
+                forward_batch=None,
+        ):
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                model_pred = self.transformer(**input_kwargs)
+            model_pred_btchw = model_pred.permute(0, 2, 1, 3, 4)
+
+            training_target = training_batch._ar_training_target.to(model_pred_btchw.device,
+                                                                    dtype=model_pred_btchw.dtype)
+
+            loss = F.mse_loss(
+                model_pred_btchw.float(),
+                training_target.float(),
+                reduction="mean",
+            )
+            loss = loss / self.training_args.gradient_accumulation_steps
+            loss.backward()
+
+        avg_loss = loss.detach().clone()
+        training_batch.total_loss += avg_loss.item()
+
+        return training_batch
+
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        self.transformer.train()
+        self.optimizer.zero_grad()
+        training_batch.total_loss = 0.0
+        args = cast(TrainingArgs, self.training_args)
+
+        for _ in range(args.gradient_accumulation_steps):
+            training_batch = self._get_next_batch(training_batch)
+            training_batch = self._normalize_dit_input(training_batch)
+            training_batch = self._prepare_dit_inputs(training_batch)
+            training_batch = self._build_input_kwargs(training_batch)
+            training_batch = self._transformer_forward_and_compute_loss(training_batch)
+
+        grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+            [p for p in self.transformer.parameters() if p.requires_grad],
+            args.max_grad_norm if args.max_grad_norm is not None else 0.0,
+        )
+
+        self.optimizer.step()
+        self.lr_scheduler.step()
+
+        if grad_norm is None:
+            grad_value = 0.0
+        else:
+            try:
+                if isinstance(grad_norm, torch.Tensor):
+                    grad_value = float(grad_norm.detach().float().item())
+                else:
+                    grad_value = float(grad_norm)
+            except Exception:
+                grad_value = 0.0
+        training_batch.grad_norm = grad_value
+        assert training_batch.latents is not None
+        training_batch.raw_latent_shape = training_batch.latents.shape
+        return training_batch
+
+    def _prepare_validation_batch(
+        self,
+        sampling_param: SamplingParam,
+        training_args: TrainingArgs,
+        validation_batch: dict[str, Any],
+        num_inference_steps: int,
+    ) -> ForwardBatch:
+        sampling_param.prompt = validation_batch["prompt"]
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.image_path = validation_batch.get("image_path") or validation_batch.get("video_path")
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [
+            (sampling_param.num_frames - 1) // 4 + 1,
+            sampling_param.height // 8,
+            sampling_param.width // 8,
+        ]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = (training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+        num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=torch.Generator(device="cpu").manual_seed(self.seed),
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+        if "image" in validation_batch and validation_batch["image"] is not None:
+            batch.pil_image = validation_batch["image"]
+
+        if ("keyboard_cond" in validation_batch and validation_batch["keyboard_cond"] is not None):
+            keyboard_cond = validation_batch["keyboard_cond"]
+            keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
+            keyboard_cond = keyboard_cond.unsqueeze(0)
+            batch.keyboard_cond = keyboard_cond
+
+        if ("mouse_cond" in validation_batch and validation_batch["mouse_cond"] is not None):
+            mouse_cond = validation_batch["mouse_cond"]
+            mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
+            mouse_cond = mouse_cond.unsqueeze(0)
+            batch.mouse_cond = mouse_cond
+
+        return batch
+
+
+def main(args) -> None:
+    logger.info("Starting MatrixGame AR diffusion training pipeline...")
+
+    pipeline = MatrixGameARDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("MatrixGame AR diffusion training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/matrixgame_ode_causal_pipeline.py

@@ -0,0 +1,467 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+from typing import Any, cast
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset.dataloader.schema import (pyarrow_schema_matrixgame_ode_trajectory)
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (SelfForcingFlowMatchScheduler)
+from fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline import (MatrixGameCausalDMDPipeline)
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases)
+from fastvideo.utils import shallow_asdict
+
+logger = init_logger(__name__)
+
+
+class MatrixGameODEInitTrainingPipeline(TrainingPipeline):
+    """
+    Training pipeline for ODE-init using precomputed denoising trajectories.
+
+    Supervision: predict the next latent in the stored trajectory by
+    - feeding current latent at timestep t into the transformer to predict noise
+    - stepping the scheduler with the predicted noise
+    - minimizing MSE to the stored next latent at timestep t_next
+    """
+
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # Match the preprocess/generation scheduler for consistent stepping
+        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
+                                                                  sigma_min=0.0,
+                                                                  extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)
+
+    def set_schemas(self):
+        self.train_dataset_schema = pyarrow_schema_matrixgame_ode_trajectory
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        super().initialize_training_pipeline(training_args)
+
+        self.noise_scheduler = self.get_module("scheduler")
+        self.vae = self.get_module("vae")
+        self.vae.requires_grad_(False)
+
+        self.timestep_shift = self.training_args.pipeline_config.flow_shift
+        self.noise_scheduler = SelfForcingFlowMatchScheduler(shift=self.timestep_shift,
+                                                             sigma_min=0.0,
+                                                             extra_one_step=True)
+        self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)
+
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
+
+        logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps)
+        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250, 0],
+                                                dtype=torch.long,
+                                                device=get_local_torch_device())
+        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
+            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
+            logger.info("timesteps: %s", timesteps)
+            self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
+            logger.info("warped self.dmd_denoising_steps: %s", self.dmd_denoising_steps)
+        else:
+            raise ValueError("warp_denoising_step must be true")
+
+        self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())
+
+        logger.info("denoising_step_list: %s", self.dmd_denoising_steps)
+
+        logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
+        # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
+        self._cached_closest_idx_per_dmd = None
+        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
+        self.manual_idx = 0
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+        args_copy.inference_mode = True
+        # Use the same flow-matching scheduler as training for consistent validation.
+        validation_scheduler = SelfForcingFlowMatchScheduler(shift=args_copy.pipeline_config.flow_shift,
+                                                             sigma_min=0.0,
+                                                             extra_one_step=True)
+        validation_scheduler.set_timesteps(num_inference_steps=1000, training=True)
+        # Warm start validation with current transformer
+        self.validation_pipeline = MatrixGameCausalDMDPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+                "vae": self.get_module("vae"),
+                "scheduler": validation_scheduler,
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+    def _get_next_batch(self, training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]:
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            self.train_loader_iter = iter(self.train_dataloader)
+            batch = next(self.train_loader_iter)
+
+        # Required fields from parquet (ODE trajectory schema)
+        clip_feature = batch['clip_feature']
+        first_frame_latent = batch['first_frame_latent']
+        keyboard_cond = batch.get('keyboard_cond', None)
+        mouse_cond = batch.get('mouse_cond', None)
+        infos = batch['info_list']
+
+        # Trajectory tensors may include a leading singleton batch dim per row
+        trajectory_latents = batch['trajectory_latents']
+        if trajectory_latents.dim() == 7:
+            # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
+            trajectory_latents = trajectory_latents[:, 0]
+        elif trajectory_latents.dim() == 6:
+            # already [B, S, C, T, H, W]
+            pass
+        else:
+            raise ValueError(f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}")
+
+        trajectory_timesteps = batch['trajectory_timesteps']
+        if trajectory_timesteps.dim() == 3:
+            # [B, 1, S] -> [B, S]
+            trajectory_timesteps = trajectory_timesteps[:, 0]
+        elif trajectory_timesteps.dim() == 2:
+            # [B, S]
+            pass
+        else:
+            raise ValueError(f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}")
+        # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
+        trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
+
+        # Move to device
+        device = get_local_torch_device()
+        training_batch.image_embeds = clip_feature.to(device, dtype=torch.bfloat16)
+        training_batch.image_latents = first_frame_latent.to(device, dtype=torch.bfloat16)
+        if keyboard_cond is not None and keyboard_cond.numel() > 0:
+            training_batch.keyboard_cond = keyboard_cond.to(device, dtype=torch.bfloat16)
+        else:
+            training_batch.keyboard_cond = None
+        if mouse_cond is not None and mouse_cond.numel() > 0:
+            training_batch.mouse_cond = mouse_cond.to(device, dtype=torch.bfloat16)
+        else:
+            training_batch.mouse_cond = None
+        training_batch.infos = infos
+
+        return training_batch, trajectory_latents[:, :, :self.training_args.num_latent_t].to(
+            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
+
+    def _get_timestep(self,
+                      min_timestep: int,
+                      max_timestep: int,
+                      batch_size: int,
+                      num_frame: int,
+                      num_frame_per_block: int,
+                      uniform_timestep: bool = False) -> torch.Tensor:
+        if uniform_timestep:
+            timestep = torch.randint(min_timestep, max_timestep, [batch_size, 1], device=self.device,
+                                     dtype=torch.long).repeat(1, num_frame)
+            return timestep
+        else:
+            timestep = torch.randint(min_timestep,
+                                     max_timestep, [batch_size, num_frame],
+                                     device=self.device,
+                                     dtype=torch.long)
+            # logger.info(f"individual timestep: {timestep}")
+            # make the noise level the same within every block
+            timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block)
+            timestep[:, :, 1:] = timestep[:, :, 0:1]
+            timestep = timestep.reshape(timestep.shape[0], -1)
+            return timestep
+
+    def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch:
+        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""
+
+        # First, call parent method to prepare noise, timesteps, etc. for video latents
+        training_batch = super()._prepare_dit_inputs(training_batch)
+
+        assert isinstance(training_batch.image_latents, torch.Tensor)
+        image_latents = training_batch.image_latents.to(get_local_torch_device(), dtype=torch.bfloat16)
+
+        temporal_compression_ratio = 4
+        num_frames = (self.training_args.num_latent_t - 1) * temporal_compression_ratio + 1
+        batch_size, num_channels, _, latent_height, latent_width = image_latents.shape
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        mask_lat_size[:, :, 1:] = 0
+
+        first_frame_mask = mask_lat_size[:, :, :1]
+        first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=temporal_compression_ratio)
+        mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], dim=2)
+        mask_lat_size = mask_lat_size.view(batch_size, -1, temporal_compression_ratio, latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(image_latents.device).to(dtype=torch.bfloat16)
+
+        training_batch.noisy_model_input = torch.cat([training_batch.noisy_model_input, mask_lat_size, image_latents],
+                                                     dim=1)
+
+        return training_batch
+
+    def _step_predict_next_latent(
+            self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, image_embeds: torch.Tensor,
+            image_latents: torch.Tensor, keyboard_cond: torch.Tensor | None, mouse_cond: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
+        latent_vis_dict: dict[str, torch.Tensor] = {}
+        device = get_local_torch_device()
+        target_latent = traj_latents[:, -1]
+
+        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
+        B, S, num_frames, num_channels, height, width = traj_latents.shape
+        expected_action_frames = 1 + 4 * (num_frames - 1)
+        if keyboard_cond is None or mouse_cond is None:
+            raise ValueError("keyboard_cond/mouse_cond must both be provided for action-follow training. "
+                             f"keyboard_cond={None if keyboard_cond is None else tuple(keyboard_cond.shape)}, "
+                             f"mouse_cond={None if mouse_cond is None else tuple(mouse_cond.shape)}")
+        if keyboard_cond.shape[1] < expected_action_frames:
+            raise ValueError("keyboard_cond length is shorter than required for latent frames. "
+                             f"got={keyboard_cond.shape[1]}, required>={expected_action_frames}, "
+                             f"num_latent_frames={num_frames}")
+        if mouse_cond.shape[1] < expected_action_frames:
+            raise ValueError("mouse_cond length is shorter than required for latent frames. "
+                             f"got={mouse_cond.shape[1]}, required>={expected_action_frames}, "
+                             f"num_latent_frames={num_frames}")
+
+        # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
+        if self._cached_closest_idx_per_dmd is None:
+            traj_ts = traj_timesteps[0].float().cpu()
+            dmd_steps = self.dmd_denoising_steps.float().cpu()
+            closest_idx = torch.argmin(torch.abs(traj_ts.unsqueeze(0) - dmd_steps.unsqueeze(1)), dim=1)
+            self._cached_closest_idx_per_dmd = closest_idx.to(torch.long).cpu()
+            logger.info("self._cached_closest_idx_per_dmd: %s", self._cached_closest_idx_per_dmd)
+            logger.info("corresponding timesteps: %s", traj_ts[self._cached_closest_idx_per_dmd])
+
+        # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
+        # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
+        # Output: [B, K, C, T, H, W]
+        assert self._cached_closest_idx_per_dmd is not None
+        relevant_traj_latents = torch.index_select(traj_latents,
+                                                   dim=1,
+                                                   index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
+        logger.info("relevant_traj_latents: %s", relevant_traj_latents.shape)
+
+        indexes = self._get_timestep(  # [B, num_frames]
+            0, len(self.dmd_denoising_steps), B, num_frames, 3, uniform_timestep=False)
+        logger.info("indexes: %s", indexes.shape)
+        logger.info("indexes: %s", indexes)
+        noisy_input = torch.gather(relevant_traj_latents,
+                                   dim=1,
+                                   index=indexes.reshape(B, 1, num_frames, 1, 1,
+                                                         1).expand(-1, -1, -1, num_channels, height,
+                                                                   width).to(self.device)).squeeze(1)
+        latent_model_input = noisy_input.permute(0, 2, 1, 3, 4)
+        if image_latents is not None:
+            latent_model_input = torch.cat([
+                latent_model_input,
+                image_latents.to(latent_model_input.device, latent_model_input.dtype),
+            ],
+                                           dim=1)
+        timestep = self.dmd_denoising_steps[indexes]
+        logger.info("selected timestep for rank %s: %s", self.global_rank, timestep, local_main_process_only=False)
+
+        # Prepare inputs for transformer
+        latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu()
+        latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu()
+
+        latent_model_input = latent_model_input.to(device, dtype=torch.bfloat16)
+        timestep = timestep.to(device, dtype=torch.bfloat16)
+
+        logger.info("========== Transformer Input ==========")
+        logger.info("hidden_states (latent_model_input) shape: %s, dtype: %s", latent_model_input.shape,
+                    latent_model_input.dtype)
+        logger.info("hidden_states min/max/mean: %.4f / %.4f / %.4f",
+                    latent_model_input.min().item(),
+                    latent_model_input.max().item(),
+                    latent_model_input.mean().item())
+        logger.info("encoder_hidden_states_image (image_embeds) shape: %s",
+                    image_embeds.shape if image_embeds is not None else None)
+        logger.info("timestep shape: %s, dtype: %s", timestep.shape, timestep.dtype)
+        logger.info("keyboard_cond: %s", keyboard_cond.shape if keyboard_cond is not None else None)
+        logger.info("mouse_cond: %s", mouse_cond.shape if mouse_cond is not None else None)
+
+        input_kwargs = {
+            "hidden_states": latent_model_input,
+            "encoder_hidden_states": None,
+            "encoder_hidden_states_image": image_embeds,
+            "timestep": timestep,
+            "mouse_cond": mouse_cond,
+            "keyboard_cond": keyboard_cond,
+            "return_dict": False,
+        }
+        # Predict noise and step the scheduler to obtain next latent
+        with set_forward_context(current_timestep=timestep, attn_metadata=None, forward_batch=None):
+            noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
+
+        logger.info("========== Transformer Output ==========")
+        logger.info("noise_pred shape: %s", noise_pred.shape)
+        logger.info("noise_pred min/max/mean: %.4f / %.4f / %.4f",
+                    noise_pred.min().item(),
+                    noise_pred.max().item(),
+                    noise_pred.mean().item())
+
+        from fastvideo.models.utils import pred_noise_to_pred_video
+        pred_video = pred_noise_to_pred_video(pred_noise=noise_pred.flatten(0, 1),
+                                              noise_input_latent=noisy_input.flatten(0, 1),
+                                              timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1),
+                                              scheduler=self.modules["scheduler"]).unflatten(0, noise_pred.shape[:2])
+        latent_vis_dict["pred_video"] = pred_video.permute(0, 2, 1, 3, 4).detach().clone().cpu()
+
+        return pred_video, target_latent, timestep, latent_vis_dict
+
+    def train_one_step(self, training_batch):  # type: ignore[override]
+        self.transformer.train()
+        self.optimizer.zero_grad()
+        training_batch.total_loss = 0.0
+        args = cast(TrainingArgs, self.training_args)
+
+        # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
+
+        for _ in range(args.gradient_accumulation_steps):
+            training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
+            image_embeds = training_batch.image_embeds
+            image_latents = training_batch.image_latents
+            keyboard_cond = training_batch.keyboard_cond
+            mouse_cond = training_batch.mouse_cond
+            assert traj_latents.shape[0] == 1
+
+            # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
+            _, S = traj_latents.shape[0], traj_latents.shape[1]
+            if S < 2:
+                raise ValueError("Trajectory must contain at least 2 steps")
+
+            # Forward to predict next latent by stepping scheduler with predicted noise
+            noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
+                traj_latents, traj_timesteps, image_embeds, image_latents, keyboard_cond, mouse_cond)
+
+            training_batch.latent_vis_dict.update(latent_vis_dict)
+
+            mask = t != 0
+
+            # Compute loss
+            loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
+            loss = loss / args.gradient_accumulation_steps
+
+            with set_forward_context(current_timestep=t, attn_metadata=None, forward_batch=None):
+                loss.backward()
+            avg_loss = loss.detach().clone()
+            training_batch.total_loss += avg_loss.item()
+
+        # Clip grad and step optimizers
+        grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+            [p for p in self.transformer.parameters() if p.requires_grad],
+            args.max_grad_norm if args.max_grad_norm is not None else 0.0)
+
+        self.optimizer.step()
+        self.lr_scheduler.step()
+
+        if grad_norm is None:
+            grad_value = 0.0
+        else:
+            try:
+                if isinstance(grad_norm, torch.Tensor):
+                    grad_value = float(grad_norm.detach().float().item())
+                else:
+                    grad_value = float(grad_norm)
+            except Exception:
+                grad_value = 0.0
+        training_batch.grad_norm = grad_value
+        B, S, T, C, H, W = traj_latents.shape
+        training_batch.raw_latent_shape = (B, C, T, H, W)
+        return training_batch
+
+    def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args: TrainingArgs,
+                                  validation_batch: dict[str, Any], num_inference_steps: int) -> ForwardBatch:
+        sampling_param.prompt = validation_batch['prompt']
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.image_path = validation_batch.get('image_path') or validation_batch.get('video_path')
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=torch.Generator(device="cpu").manual_seed(self.seed),
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+        if "image" in validation_batch and validation_batch["image"] is not None:
+            batch.pil_image = validation_batch["image"]
+
+        if "keyboard_cond" in validation_batch and validation_batch["keyboard_cond"] is not None:
+            keyboard_cond = validation_batch["keyboard_cond"]
+            keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
+            keyboard_cond = keyboard_cond.unsqueeze(0)
+            batch.keyboard_cond = keyboard_cond
+
+        if "mouse_cond" in validation_batch and validation_batch["mouse_cond"] is not None:
+            mouse_cond = validation_batch["mouse_cond"]
+            mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
+            mouse_cond = mouse_cond.unsqueeze(0)
+            batch.mouse_cond = mouse_cond
+
+        return batch
+
+    def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
+        tracker_loss_dict: dict[str, Any] = {}
+        latents_vis_dict = training_batch.latent_vis_dict
+        latent_log_keys = ['noisy_input', 'x0', 'pred_video']
+        for latent_key in latent_log_keys:
+            assert latent_key in latents_vis_dict and latents_vis_dict[latent_key] is not None
+            latent = latents_vis_dict[latent_key]
+            pixel_latent = self.decoding_stage.decode(latent, training_args)
+
+            video = pixel_latent.cpu().float()
+            video = video.permute(0, 2, 1, 3, 4)
+            video = (video * 255).numpy().astype(np.uint8)
+            video_artifact = self.tracker.video(video, fps=16, format="mp4")  # change to 16 for Wan2.1
+            if video_artifact is not None:
+                tracker_loss_dict[latent_key] = video_artifact
+            # Clean up references
+            del video, pixel_latent, latent
+
+        if self.global_rank == 0 and tracker_loss_dict:
+            self.tracker.log_artifacts(tracker_loss_dict, step)
+
+
+def main(args) -> None:
+    logger.info("Starting ODE-init training pipeline...")
+    pipeline = MatrixGameODEInitTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("ODE-init training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/trackers.py

@@ -154,6 +154,16 @@ class WandbTracker(BaseTracker):
             self._run.log(metrics, step=step)
         self._timed_metrics = {}
 
+    def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
+        if not artifacts:
+            return
+
+        normalized_artifacts = dict(artifacts)
+        if "validation_ref_videos" in normalized_artifacts:
+            normalized_artifacts["reference_video/videos"] = (normalized_artifacts.pop("validation_ref_videos"))
+
+        self.log(normalized_artifacts, step)
+
     def log_file(
         self,
         file_path: str,

fastvideo/training/training_pipeline.py

@@ -80,6 +80,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         set_random_seed(fastvideo_args.seed)  # for lora param init
         super().__init__(model_path, fastvideo_args, required_config_modules, loaded_modules)  # type: ignore
         self.tracker = DummyTracker()
+        self.validation_ref_videos_logged = False
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         raise RuntimeError("create_pipeline_stages should not be called for training pipeline")
@@ -727,6 +728,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                         local_main_process_only=False)
             step_videos: list[np.ndarray] = []
             step_captions: list[str] = []
+            step_ref_videos: list[str | None] = []
 
             step_audio: list[np.ndarray | None] = []
             step_sample_rates: list[int | None] = []
@@ -742,6 +744,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 assert batch.prompt is not None and isinstance(batch.prompt, str)
                 step_captions.append(batch.prompt)
+                step_ref_videos.append(validation_batch.get("ref_video"))
 
                 # Run validation inference
                 output_batch = self.validation_pipeline.forward(batch, training_args)
@@ -775,6 +778,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 # Global rank 0 collects results from all sp_group leaders
                 all_videos = step_videos  # Start with own results
                 all_captions = step_captions
+                all_ref_videos = step_ref_videos
                 all_audios = step_audio
                 all_sample_rates = step_sample_rates
 
@@ -783,11 +787,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
                     src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
                     recv_videos = world_group.recv_object(src=src_rank)
                     recv_captions = world_group.recv_object(src=src_rank)
+                    recv_ref_videos = world_group.recv_object(src=src_rank)
                     recv_audios = world_group.recv_object(src=src_rank)
                     recv_sample_rates = world_group.recv_object(src=src_rank)
 
                     all_videos.extend(recv_videos)
                     all_captions.extend(recv_captions)
+                    all_ref_videos.extend(recv_ref_videos)
                     all_audios.extend(recv_audios)
                     all_sample_rates.extend(recv_sample_rates)
 
@@ -816,10 +822,24 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 if artifacts:
                     logs = {f"validation_videos_{num_inference_steps}_steps": artifacts}
                     self.tracker.log_artifacts(logs, global_step)
+                if not self.validation_ref_videos_logged:
+                    ref_artifacts = []
+                    for filename, caption in zip(all_ref_videos, all_captions, strict=True):
+                        if filename is None:
+                            continue
+                        ref_frames = np.stack([np.asarray(frame) for frame in load_video(filename)], axis=0)
+                        ref_frames = np.ascontiguousarray(ref_frames.transpose(0, 3, 1, 2))
+                        video_artifact = self.tracker.video(ref_frames, caption=caption, fps=sampling_param.fps)
+                        if video_artifact is not None:
+                            ref_artifacts.append(video_artifact)
+                    if ref_artifacts:
+                        self.tracker.log_artifacts({"validation_ref_videos": ref_artifacts}, global_step)
+                        self.validation_ref_videos_logged = True
             elif self.rank_in_sp_group == 0:
                 # Other sp_group leaders send their results to global rank 0
                 world_group.send_object(step_videos, dst=0)
                 world_group.send_object(step_captions, dst=0)
+                world_group.send_object(step_ref_videos, dst=0)
                 world_group.send_object(step_audio, dst=0)
                 world_group.send_object(step_sample_rates, dst=0)