From d6dfe954669d64af45dea6476ac5ca1c552b511c Mon Sep 17 00:00:00 2001 From: Kaiqin Kong Date: Sun, 10 May 2026 22:13:12 -0700 Subject: [PATCH] [feat] FastVideo World Model Training (#1179) Co-authored-by: mignonjia Co-authored-by: alexzms <3036648523@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .agents/memory/codebase-map/README.md | 3 +- .../onboarding/worldmodel-training/README.md | 3 + examples/distill/MatrixGame2.0/distill_dmd.sh | 126 + .../distill/MatrixGame2.0/distill_dmd.slurm | 168 + .../distill/MatrixGame2.0/validation.json | 356 ++ .../fine_tuning/wan/dfsft_causal_t2v.yaml | 10 +- examples/train/scenario/worldmodel/README.md | 32 + .../scenario/worldmodel/dfsft_causal_i2v.yaml | 74 + .../scenario/worldmodel/finetune_i2v.yaml | 72 + .../worldmodel/self_forcing_causal_i2v.yaml | 100 + .../finetune/MatrixGame2.0/finetune_df.sh | 94 + .../finetune/MatrixGame2.0/validation.json | 5344 +---------------- fastvideo/dataset/dataloader/record_schema.py | 107 + fastvideo/dataset/dataloader/schema.py | 33 + fastvideo/dataset/validation_dataset.py | 10 + .../models/dits/matrixgame/action_module.py | 7 + .../models/dits/matrixgame/causal_model.py | 188 +- .../matrixgame_preprocess_pipeline.py | 2 +- ...game_preprocess_pipeline_ode_trajectory.py | 385 ++ .../pipelines/preprocess/v1_preprocess.py | 17 +- fastvideo/train/methods/fine_tuning/dfsft.py | 43 +- .../train/models/matrixgame/matrixgame.py | 280 + .../models/matrixgame/matrixgame_causal.py | 436 ++ fastvideo/train/utils/dataloader.py | 23 + .../matrixgame_ar_diffusion_pipeline.py | 451 ++ .../matrixgame_ode_causal_pipeline.py | 416 ++ ...game_self_forcing_distillation_pipeline.py | 873 +++ fastvideo/training/trackers.py | 10 + fastvideo/training/training_pipeline.py | 21 + 29 files changed, 4541 insertions(+), 5143 deletions(-) create mode 100644 examples/distill/MatrixGame2.0/distill_dmd.sh create mode 100644 examples/distill/MatrixGame2.0/distill_dmd.slurm create mode 100644 examples/distill/MatrixGame2.0/validation.json create mode 100644 examples/train/scenario/worldmodel/README.md create mode 100644 examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml create mode 100644 examples/train/scenario/worldmodel/finetune_i2v.yaml create mode 100644 examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml create mode 100644 examples/training/finetune/MatrixGame2.0/finetune_df.sh create mode 100644 fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline_ode_trajectory.py create mode 100644 fastvideo/train/models/matrixgame/matrixgame.py create mode 100644 fastvideo/train/models/matrixgame/matrixgame_causal.py create mode 100644 fastvideo/training/matrixgame_ar_diffusion_pipeline.py create mode 100644 fastvideo/training/matrixgame_ode_causal_pipeline.py create mode 100644 fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py diff --git a/.agents/memory/codebase-map/README.md b/.agents/memory/codebase-map/README.md index 3800ce7a2..64474c2ba 100644 --- a/.agents/memory/codebase-map/README.md +++ b/.agents/memory/codebase-map/README.md @@ -30,7 +30,8 @@ FastVideo-WorldModel/ │ │ │ ├── fine_tuning/ # FineTuneMethod, DiffusionForcingSFTMethod │ │ │ └── distribution_matching/ # DMD2Method, SelfForcingMethod │ │ ├── models/ # Per-role model wrappers (ModelBase, CausalModelBase) -│ │ │ └── wan/ # WanModel, WanCausalModel +│ │ │ ├── wan/ # WanModel, WanCausalModel +│ │ │ └── matrixgame/ # MatrixGameModel, MatrixGameCausalModel │ │ ├── callbacks/ # Composable hooks (grad_clip, ema, validation) │ │ └── utils/ # Config, builder, checkpoint, optimizer, tracking │ ├── training/ # Legacy training infrastructure (being phased out) diff --git a/.agents/onboarding/worldmodel-training/README.md b/.agents/onboarding/worldmodel-training/README.md index be1d9a61d..3ac445357 100644 --- a/.agents/onboarding/worldmodel-training/README.md +++ b/.agents/onboarding/worldmodel-training/README.md @@ -226,6 +226,9 @@ callbacks: | Wan T2V finetune | `fastvideo/training/wan_training_pipeline.py` | Standard text-to-video finetune / LoRA | | Wan I2V finetune | `fastvideo/training/wan_i2v_training_pipeline.py` | Image-to-video (first frame conditioned) | | MatrixGame finetune | `fastvideo/training/matrixgame_training_pipeline.py` | Action-conditioned world model | +| MatrixGame AR diffusion | `fastvideo/training/matrixgame_ar_diffusion_pipeline.py` | AR diffusion-forcing training | +| MatrixGame ODE-init | `fastvideo/training/matrixgame_ode_causal_pipeline.py` | ODE-trajectory init | +| MatrixGame self-forcing distill | `fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py` | Self-forcing distillation | | LTX-2 finetune | `fastvideo/training/ltx2_training_pipeline.py` | LTX-2 architecture finetuning | | Wan DMD distillation | `fastvideo/training/wan_distillation_pipeline.py` | Few-step distillation via DMD | | Self-Forcing distill | `fastvideo/training/wan_self_forcing_distillation_pipeline.py` | Causal streaming distillation | diff --git a/examples/distill/MatrixGame2.0/distill_dmd.sh b/examples/distill/MatrixGame2.0/distill_dmd.sh new file mode 100644 index 000000000..f7f7fa447 --- /dev/null +++ b/examples/distill/MatrixGame2.0/distill_dmd.sh @@ -0,0 +1,126 @@ +#!/bin/bash + +export WANDB_API_KEY="${WANDB_API_KEY:-}" +export WANDB_MODE="offline" +export TOKENIZERS_PARALLELISM=false + +RUN_NAME=$(date +"%m%d_%H%M") +echo "RUN_NAME: $RUN_NAME" + +# Model paths for Self-Forcing DMD distillation. +GENERATOR_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" +REAL_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers" # Teacher +FAKE_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers" # Critic + +DATA_DIR="data/matrixgame" +VALIDATION_DATASET_FILE="examples/distill/MatrixGame2.0/validation.json" +NUM_GPUS=1 + +# Training arguments +training_args=( + --tracker_project_name "matrixgame_sf" + --output_dir "checkpoints/matrixgame_sf_${RUN_NAME}" + --wandb_run_name "${RUN_NAME}_test" + --max_train_steps 5 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 352 + --num_width 640 + # --enable_gradient_checkpointing_type "full" + --simulate_generator_forward + --num_frames 81 + --num_frame_per_block 3 # Frame generation block size for self-forcing + # --enable_gradient_masking + # --gradient_mask_last_n_frames 21 + # --init_weights_from_safetensors "path/to/generator_ema.safetensors" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim $NUM_GPUS +) + +model_args=( + --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --pretrained_model_name_or_path $GENERATOR_MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_MODEL_PATH +) + +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Validation arguments +validation_args=( + --log_validation + --log_visualization + --visualization-steps 100 + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "4" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 3e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 400 + --training_state_checkpointing_steps 400 + --weight_decay 0 + --betas "0.9,0.95" + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --flow_shift 5 + --seed 1000 + --use_ema True + --ema_decay 0.99 + --ema_start_step 200 +) + +dmd_args=( + --dmd_denoising_steps '1000,750,500,250' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --dfake_gen_update_ratio 5 + --real_score_guidance_scale 3.0 + --fake_score_learning_rate 3e-7 + --fake_score_betas "0.9,0.95" + --warp_denoising_step +) + +self_forcing_args=( + --independent_first_frame False # Whether to treat first frame independently + --same_step_across_blocks True # Whether to use same denoising step across all blocks + --last_step_only False # Whether to only use the last denoising step + --context_noise 0 # Amount of noise to add during context caching (0 = no noise) +) + +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ + "${self_forcing_args[@]}" \ No newline at end of file diff --git a/examples/distill/MatrixGame2.0/distill_dmd.slurm b/examples/distill/MatrixGame2.0/distill_dmd.slurm new file mode 100644 index 000000000..b64919770 --- /dev/null +++ b/examples/distill/MatrixGame2.0/distill_dmd.slurm @@ -0,0 +1,168 @@ +#!/bin/bash +#SBATCH --job-name=matrixgame_sf_dmd +#SBATCH --partition=main +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:4 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=logs/slurm/matrixgame_sf_dmd_%j.out +#SBATCH --error=logs/slurm/matrixgame_sf_dmd_%j.err +#SBATCH --exclusive + +set -euo pipefail +set -x + +CONDA_ROOT="${CONDA_ROOT:-$HOME/miniforge3}" +CONDA_BIN="${CONDA_ROOT}/bin/conda" +ENV_NAME="${ENV_NAME:-FastVideo_kaiqin}" +MASTER_PORT="${MASTER_PORT:-29500}" +WANDB_API_KEY="${WANDB_API_KEY:-}" +WANDB_MODE="${WANDB_MODE:-online}" +GPUS_PER_NODE="${SLURM_GPUS_ON_NODE:-4}" +TOTAL_GPUS="$(( SLURM_JOB_NUM_NODES * GPUS_PER_NODE ))" + +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export TRITON_CACHE_DIR="/tmp/triton_cache_${SLURM_PROCID}" +export TOKENIZERS_PARALLELISM=false +export WANDB_API_KEY +export WANDB_MODE + +if [ -x "${CONDA_BIN}" ]; then + eval "$("${CONDA_BIN}" shell.bash hook)" + conda activate "${ENV_NAME}" +else + echo "conda not found at ${CONDA_BIN}" >&2 + exit 1 +fi + +cd "${SLURM_SUBMIT_DIR}" +mkdir -p logs/slurm +export PYTHONPATH="${SLURM_SUBMIT_DIR}${PYTHONPATH:+:${PYTHONPATH}}" + +nodes=( $(scontrol show hostnames "${SLURM_JOB_NODELIST}") ) +export MASTER_ADDR="${nodes[0]}" +export MASTER_PORT +export NODE_RANK="${SLURM_PROCID}" + +echo "MASTER_ADDR: ${MASTER_ADDR}" +echo "MASTER_PORT: ${MASTER_PORT}" +echo "NODE_RANK: ${NODE_RANK}" +echo "GPUS_PER_NODE: ${GPUS_PER_NODE}" +echo "TOTAL_GPUS: ${TOTAL_GPUS}" + +RUN_NAME="$(date +"%m%d_%H%M")" +echo "RUN_NAME: ${RUN_NAME}" + +GENERATOR_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" +REAL_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers" +FAKE_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers" + +DATA_DIR="data/matrixgame" +VALIDATION_DATASET_FILE="examples/distill/MatrixGame2.0/validation.json" + +training_args=( + --tracker_project_name "matrixgame_sf" + --output_dir "checkpoints/matrixgame_sf_${RUN_NAME}" + --wandb_run_name "${RUN_NAME}_test" + --max_train_steps 1200 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 2 + --num_latent_t 21 + --num_height 352 + --num_width 640 + --simulate_generator_forward + --num_frames 81 + --num_frame_per_block 3 + # --init_weights_from_safetensors "path/to/generator_ema.safetensors" +) + +parallel_args=( + --num_gpus "${TOTAL_GPUS}" + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim "${TOTAL_GPUS}" +) + +model_args=( + --model_path "${GENERATOR_MODEL_PATH}" + --pretrained_model_name_or_path "${GENERATOR_MODEL_PATH}" + --real_score_model_path "${REAL_SCORE_MODEL_PATH}" + --fake_score_model_path "${FAKE_SCORE_MODEL_PATH}" +) + +dataset_args=( + --data_path "${DATA_DIR}" + --dataloader_num_workers 4 +) + +validation_args=( + --log_validation + --log_visualization + --visualization-steps 100 + --validation_dataset_file "${VALIDATION_DATASET_FILE}" + --validation_steps 100 + --validation_sampling_steps "4" + --validation_guidance_scale "6.0" +) + +optimizer_args=( + --learning_rate 3e-6 + --mixed_precision "bf16" + --weight_only_checkpointing_steps 400 + --training_state_checkpointing_steps 400 + --weight_decay 0 + --betas "0.9,0.95" + --max_grad_norm 1.0 +) + +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --flow_shift 5 + --seed 1000 + --use_ema True + --ema_decay 0.99 + --ema_start_step 200 +) + +dmd_args=( + --dmd_denoising_steps "1000,750,500,250" + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --dfake_gen_update_ratio 5 + --real_score_guidance_scale 3.0 + --fake_score_learning_rate 3e-7 + --fake_score_betas "0.9,0.95" + --warp_denoising_step +) + +self_forcing_args=( + --independent_first_frame False + --same_step_across_blocks True + --last_step_only False + --context_noise 0 +) + +srun python -m torch.distributed.run \ + --nnodes "${SLURM_JOB_NUM_NODES}" \ + --nproc_per_node "${GPUS_PER_NODE}" \ + --node_rank "${SLURM_PROCID}" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ + "${self_forcing_args[@]}" diff --git a/examples/distill/MatrixGame2.0/validation.json b/examples/distill/MatrixGame2.0/validation.json new file mode 100644 index 000000000..2d714cbdb --- /dev/null +++ b/examples/distill/MatrixGame2.0/validation.json @@ -0,0 +1,356 @@ +{ + "data": [ + { + "caption": "00", + "image_path": "path_to_validation/00.png", + "ref_video": "path_to_validation/00.mp4", + "action_path": "path_to_validation/00.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "01", + "image_path": "path_to_validation/01.png", + "ref_video": "path_to_validation/01.mp4", + "action_path": "path_to_validation/01.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "02", + "image_path": "path_to_validation/02.png", + "ref_video": "path_to_validation/02.mp4", + "action_path": "path_to_validation/02.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "03", + "image_path": "path_to_validation/03.png", + "ref_video": "path_to_validation/03.mp4", + "action_path": "path_to_validation/03.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "04", + "image_path": "path_to_validation/04.png", + "ref_video": "path_to_validation/04.mp4", + "action_path": "path_to_validation/04.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "05", + "image_path": "path_to_validation/05.png", + "ref_video": "path_to_validation/05.mp4", + "action_path": "path_to_validation/05.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "06", + "image_path": "path_to_validation/06.png", + "ref_video": "path_to_validation/06.mp4", + "action_path": "path_to_validation/06.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "07", + "image_path": "path_to_validation/07.png", + "ref_video": "path_to_validation/07.mp4", + "action_path": "path_to_validation/07.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "08", + "image_path": "path_to_validation/08.png", + "ref_video": "path_to_validation/08.mp4", + "action_path": "path_to_validation/08.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "09", + "image_path": "path_to_validation/09.png", + "ref_video": "path_to_validation/09.mp4", + "action_path": "path_to_validation/09.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "10", + "image_path": "path_to_validation/10.png", + "ref_video": "path_to_validation/10.mp4", + "action_path": "path_to_validation/10.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "11", + "image_path": "path_to_validation/11.png", + "ref_video": "path_to_validation/11.mp4", + "action_path": "path_to_validation/11.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "12", + "image_path": "path_to_validation/12.png", + "ref_video": "path_to_validation/12.mp4", + "action_path": "path_to_validation/12.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "13", + "image_path": "path_to_validation/13.png", + "ref_video": "path_to_validation/13.mp4", + "action_path": "path_to_validation/13.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "14", + "image_path": "path_to_validation/14.png", + "ref_video": "path_to_validation/14.mp4", + "action_path": "path_to_validation/14.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "15", + "image_path": "path_to_validation/15.png", + "ref_video": "path_to_validation/15.mp4", + "action_path": "path_to_validation/15.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "16", + "image_path": "path_to_validation/16.png", + "ref_video": "path_to_validation/16.mp4", + "action_path": "path_to_validation/16.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "17", + "image_path": "path_to_validation/17.png", + "ref_video": "path_to_validation/17.mp4", + "action_path": "path_to_validation/17.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "18", + "image_path": "path_to_validation/18.png", + "ref_video": "path_to_validation/18.mp4", + "action_path": "path_to_validation/18.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "19", + "image_path": "path_to_validation/19.png", + "ref_video": "path_to_validation/19.mp4", + "action_path": "path_to_validation/19.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "20", + "image_path": "path_to_validation/20.png", + "ref_video": "path_to_validation/20.mp4", + "action_path": "path_to_validation/20.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "21", + "image_path": "path_to_validation/21.png", + "ref_video": "path_to_validation/21.mp4", + "action_path": "path_to_validation/21.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "22", + "image_path": "path_to_validation/22.png", + "ref_video": "path_to_validation/22.mp4", + "action_path": "path_to_validation/22.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "23", + "image_path": "path_to_validation/23.png", + "ref_video": "path_to_validation/23.mp4", + "action_path": "path_to_validation/23.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "24", + "image_path": "path_to_validation/24.png", + "ref_video": "path_to_validation/24.mp4", + "action_path": "path_to_validation/24.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "25", + "image_path": "path_to_validation/25.png", + "ref_video": "path_to_validation/25.mp4", + "action_path": "path_to_validation/25.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "26", + "image_path": "path_to_validation/26.png", + "ref_video": "path_to_validation/26.mp4", + "action_path": "path_to_validation/26.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "27", + "image_path": "path_to_validation/27.png", + "ref_video": "path_to_validation/27.mp4", + "action_path": "path_to_validation/27.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "28", + "image_path": "path_to_validation/28.png", + "ref_video": "path_to_validation/28.mp4", + "action_path": "path_to_validation/28.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "29", + "image_path": "path_to_validation/29.png", + "ref_video": "path_to_validation/29.mp4", + "action_path": "path_to_validation/29.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "30", + "image_path": "path_to_validation/30.png", + "ref_video": "path_to_validation/30.mp4", + "action_path": "path_to_validation/30.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + }, + { + "caption": "31", + "image_path": "path_to_validation/31.png", + "ref_video": "path_to_validation/31.mp4", + "action_path": "path_to_validation/31.npy", + "video_path": null, + "num_inference_steps": 4, + "height": 352, + "width": 640, + "num_frames": 81 + } + ] +} diff --git a/examples/train/configs/fine_tuning/wan/dfsft_causal_t2v.yaml b/examples/train/configs/fine_tuning/wan/dfsft_causal_t2v.yaml index 59b15172c..50f6309da 100644 --- a/examples/train/configs/fine_tuning/wan/dfsft_causal_t2v.yaml +++ b/examples/train/configs/fine_tuning/wan/dfsft_causal_t2v.yaml @@ -27,10 +27,10 @@ training: train_batch_size: 1 training_cfg_rate: 0.0 seed: 1000 - num_latent_t: 20 + num_latent_t: 18 num_height: 448 num_width: 832 - num_frames: 77 + num_frames: 69 optimizer: learning_rate: 2.0e-6 @@ -50,7 +50,7 @@ training: tracker: project_name: distillation_wan_r - run_name: wan2.1_causal_dfsft + run_name: wan2.1_causal_dfsft_shift5_gauss_weight model: enable_gradient_checkpointing_type: full @@ -60,11 +60,11 @@ callbacks: max_grad_norm: 1.0 validation: pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline - dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json + dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json every_steps: 50 sampling_steps: [40] guidance_scale: 6.0 num_frames: 69 pipeline: - flow_shift: 8 + flow_shift: 5 diff --git a/examples/train/scenario/worldmodel/README.md b/examples/train/scenario/worldmodel/README.md new file mode 100644 index 000000000..6a26e3c6d --- /dev/null +++ b/examples/train/scenario/worldmodel/README.md @@ -0,0 +1,32 @@ +# World-Model: MatrixGame 2.0 I2V + +Three training scenarios for the MatrixGame 2.0 I2V world model on the +new YAML-driven trainer (`fastvideo/train/entrypoint/train.py`). + +| Config | Method | Student | Notes | +|---|---|---|---| +| `finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGameModel` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. | +| `dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGameCausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. | +| `self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGameCausalModel` | DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. | + +## Usage + +```bash +bash examples/train/run.sh \ + examples/train/scenario/worldmodel/finetune_i2v.yaml + +bash examples/train/run.sh \ + examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml + +bash examples/train/run.sh \ + examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml +``` + +Override any field on the command line: + +```bash +bash examples/train/run.sh \ + examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml \ + --training.distributed.num_gpus 8 \ + --training.optimizer.learning_rate 1e-5 +``` diff --git a/examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml b/examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml new file mode 100644 index 000000000..e8ffcb5b7 --- /dev/null +++ b/examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml @@ -0,0 +1,74 @@ +# Diffusion-Forcing SFT: MatrixGame I2V Causal +# +# - Student: trainable causal MatrixGame model +# - Training: inhomogeneous timesteps per chunk (diffusion forcing) + +models: + student: + _target_: fastvideo.train.models.matrixgame.matrixgame_causal.MatrixGameCausalModel + init_from: FastVideo/Matrix-Game-2.0-Foundation-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod + chunk_size: 3 + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/matrixgame + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 352 + num_width: 640 + num_frames: 81 + + optimizer: + learning_rate: 2.0e-5 + betas: [0.9, 0.999] + weight_decay: 1.0e-4 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1500 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/matrixgame_causal_dfsft + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 3 + + tracker: + project_name: dfsft-matrixgame + run_name: matrixgame_causal_dfsft + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline.MatrixGameCausalDMDPipeline + dataset_file: examples/distill/MatrixGame2.0/validation.json + every_steps: 100 + sampling_steps: [40] + sampling_timesteps: [1000, 975, 950, 925, 900, 875, 850, 825, 800, 775, + 750, 725, 700, 675, 650, 625, 600, 575, 550, 525, + 500, 475, 450, 425, 400, 375, 350, 325, 300, 275, + 250, 225, 200, 175, 150, 125, 100, 75, 50, 25] + num_frames: 81 + guidance_scale: 6.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/scenario/worldmodel/finetune_i2v.yaml b/examples/train/scenario/worldmodel/finetune_i2v.yaml new file mode 100644 index 000000000..5642092ea --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_i2v.yaml @@ -0,0 +1,72 @@ +# MatrixGame I2V finetune. + +models: + student: + _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel + init_from: FastVideo/Matrix-Game-2.0-Foundation-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 1 + + data: + data_path: data/matrixgame + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 448 + num_width: 832 + num_frames: 81 + + optimizer: + learning_rate: 1.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/matrixgame_i2v_finetune + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 3 + resume_from_checkpoint: latest + + # tracker: + # project_name: fastvideo + # run_name: matrixgame_i2v_finetune + tracker: + trackers: [none] + project_name: "" + run_name: "" + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + validation: + _target_: fastvideo.train.callbacks.validation.ValidationCallback + pipeline_target: fastvideo.pipelines.basic.matrixgame.matrixgame_i2v_pipeline.MatrixGamePipeline + dataset_file: examples/distill/MatrixGame2.0/validation.json + every_steps: 100 + sampling_steps: [50] + guidance_scale: 5.0 + +pipeline: + flow_shift: 3 diff --git a/examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml b/examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml new file mode 100644 index 000000000..3496143f0 --- /dev/null +++ b/examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml @@ -0,0 +1,100 @@ +# Self-forcing distillation: MatrixGame I2V Causal +# +# - Teacher: frozen pretrained MatrixGame model +# - Student: trainable causal MatrixGame model +# - Critic: trainable MatrixGame model +# - Training: streaming rollout with SDE sampling + +models: + student: + _target_: fastvideo.train.models.matrixgame.matrixgame_causal.MatrixGameCausalModel + init_from: FastVideo/Matrix-Game-2.0-Base-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel + init_from: FastVideo/Matrix-Game-2.0-Foundation-Diffusers + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel + init_from: FastVideo/Matrix-Game-2.0-Foundation-Diffusers + trainable: true + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod + rollout_mode: simulate + generator_update_interval: 5 + real_score_guidance_scale: 4.5 + dmd_denoising_steps: [1000, 750, 500, 250] + + chunk_size: 3 + student_sample_type: sde + context_noise: 0.0 + enable_gradient_in_rollout: true + start_gradient_frame: 0 + + # Critic optimizer + fake_score_learning_rate: 8.0e-6 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/matrixgame + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 352 + num_width: 640 + num_frames: 81 + + optimizer: + learning_rate: 1.0e-5 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/matrixgame_causal_self_forcing + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + + tracker: + project_name: self-forcing-matrixgame + run_name: matrixgame_causal_self_forcing + + model: + enable_gradient_checkpointing_type: full + +callbacks: + ema: + decay: 0.99 + start_iter: 200 + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline.MatrixGameCausalDMDPipeline + dataset_file: examples/distill/MatrixGame2.0/validation.json + every_steps: 50 + sampling_steps: [4] + sampling_timesteps: [1000, 750, 500, 250] + num_frames: 81 + guidance_scale: 6.0 + +pipeline: + flow_shift: 5 diff --git a/examples/training/finetune/MatrixGame2.0/finetune_df.sh b/examples/training/finetune/MatrixGame2.0/finetune_df.sh new file mode 100644 index 000000000..e09bcad01 --- /dev/null +++ b/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[@]}" diff --git a/examples/training/finetune/MatrixGame2.0/validation.json b/examples/training/finetune/MatrixGame2.0/validation.json index b3a1a5326..c903a0c3a 100644 --- a/examples/training/finetune/MatrixGame2.0/validation.json +++ b/examples/training/finetune/MatrixGame2.0/validation.json @@ -1,5020 +1,324 @@ -{ - "data": [ - { - "caption": "all_left", - "image_path": "../../../../footsies-dataset/validate/0.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "all_right", - "image_path": "../../../../footsies-dataset/validate/0.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "all_idle", - "image_path": "../../../../footsies-dataset/validate/0.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "alternating_left_idle_right", - "image_path": "../../../../footsies-dataset/validate/0.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "all_left_with_left_position", - "image_path": "../../../../footsies-dataset/validate/1.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "all_right_with_right_position", - "image_path": "../../../../footsies-dataset/validate/2.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "alternating_left_right_with_1:2 ratio", - "image_path": "../../../../footsies-dataset/validate/1.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ] - ] - }, - { - "caption": "alternating_left_right_with_2:1 ratio", - "image_path": "../../../../footsies-dataset/validate/1.png", - "video_path": null, - "num_inference_steps": 40, - "height": 352, - "width": 640, - "num_frames": 77, - "keyboard_cond": [ - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 1, - 0, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ], - [ - 0, - 0, - 1, - 0, - 0, - 0 - ] - ] - } - ] -} \ No newline at end of file +{ + "data": [ + { + "caption": "00 Val-00: W", + "image_path": "path_to_validation/000002.jpg", + "action_path": "path_to_validation/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01 Val-01: S", + "image_path": "path_to_validation/000003.jpg", + "action_path": "path_to_validation/S.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02 Val-02: A", + "image_path": "path_to_validation/000004.jpg", + "action_path": "path_to_validation/A.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03 Val-03: D", + "image_path": "path_to_validation/000005.jpg", + "action_path": "path_to_validation/D.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "04 Val-04: u", + "image_path": "path_to_validation/000000.jpg", + "action_path": "path_to_validation/u.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "05 Val-05: d", + "image_path": "path_to_validation/000001.jpg", + "action_path": "path_to_validation/d.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "06 Val-06: l", + "image_path": "path_to_validation/000006.jpg", + "action_path": "path_to_validation/l.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "07 Val-07: r", + "image_path": "path_to_validation/000007.jpg", + "action_path": "path_to_validation/r.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "08 Val-00: key rand", + "image_path": "path_to_validation/000002.jpg", + "action_path": "path_to_validation/key_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "09 Val-01: key rand", + "image_path": "path_to_validation/000003.jpg", + "action_path": "path_to_validation/key_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "10 Val-02: camera rand", + "image_path": "path_to_validation/000004.jpg", + "action_path": "path_to_validation/camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "11 Val-03: camera rand", + "image_path": "path_to_validation/000005.jpg", + "action_path": "path_to_validation/camera_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "12 Val-00: key+camera excl rand", + "image_path": "path_to_validation/000002.jpg", + "action_path": "path_to_validation/key_camera_excl_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "13 Val-01: key+camera excl rand", + "image_path": "path_to_validation/000003.jpg", + "action_path": "path_to_validation/key_camera_excl_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "14 Val-02: key+camera rand", + "image_path": "path_to_validation/000004.jpg", + "action_path": "path_to_validation/key_camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "15 Val-03: key+camera rand", + "image_path": "path_to_validation/000005.jpg", + "action_path": "path_to_validation/key_camera_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "16 Val-04: (simultaneous) key rand", + "image_path": "path_to_validation/000000.jpg", + "action_path": "path_to_validation/key_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "17 Val-05: (simultaneous) camera rand", + "image_path": "path_to_validation/000001.jpg", + "action_path": "path_to_validation/camera_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "18 Val-06: (simultaneous) key+camera excl rand", + "image_path": "path_to_validation/000006.jpg", + "action_path": "path_to_validation/key_camera_excl_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "19 Val-07: (simultaneous) key+camera rand", + "image_path": "path_to_validation/000007.jpg", + "action_path": "path_to_validation/key_camera_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "20 Val-08: W+A", + "image_path": "path_to_validation/000005.jpg", + "action_path": "path_to_validation/WA.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "21 Val-09: S+u", + "image_path": "path_to_validation/000013.jpg", + "action_path": "path_to_validation/S_u.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "22 Val-08: Still", + "image_path": "path_to_validation/000005.jpg", + "action_path": "path_to_validation/still.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "23 Val-09: Still", + "image_path": "path_to_validation/000013.jpg", + "action_path": "path_to_validation/still.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "24 Val-06: key+camera excl rand Frame 4", + "image_path": "path_to_validation/000006.jpg", + "action_path": "path_to_validation/key_camera_excl_1_action_rand_1_f4.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "25 Val-07: key+camera excl rand Frame 4", + "image_path": "path_to_validation/000007.jpg", + "action_path": "path_to_validation/key_camera_excl_1_action_rand_2_f4.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "26 Train-00", + "image_path": "path_to_validation/000500.jpg", + "action_path": "path_to_validation/000500_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "27 Train-01", + "image_path": "path_to_validation/001000.jpg", + "action_path": "path_to_validation/001000_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "28 Doom-00: W", + "image_path": "path_to_validation/000000.jpg", + "action_path": "path_to_validation/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "29 Doom-01: key rand", + "image_path": "path_to_validation/000001.jpg", + "action_path": "path_to_validation/key_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "30 Doom-02: camera rand", + "image_path": "path_to_validation/000002.jpg", + "action_path": "path_to_validation/camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "31 Doom-03: key+camera excl rand", + "image_path": "path_to_validation/000003.jpg", + "action_path": "path_to_validation/key_camera_excl_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + } + ] +} diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py index 604fec1bb..aa5786f9d 100644 --- a/fastvideo/dataset/dataloader/record_schema.py +++ b/fastvideo/dataset/dataloader/record_schema.py @@ -191,3 +191,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 \ No newline at end of file diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index 048c7686c..33ea3428d 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/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' +]) \ No newline at end of file diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index ec937520a..324e04f35 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/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) diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index d2ea11925..65de38065 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -150,6 +150,13 @@ def _update_kv_cache_and_attend( else int(kv_cache["local_end_index"]) ) + if torch.is_grad_enabled(): + kv_cache["k"] = kv_cache["k"].detach().clone() + kv_cache["v"] = kv_cache["v"].detach().clone() + else: + kv_cache["k"] = kv_cache["k"].detach() + kv_cache["v"] = kv_cache["v"].detach() + # Check if we need to evict tokens if (current_end > original_global_end_index) and ( num_new_tokens + original_local_end_index > kv_cache_size diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 99bb8e134..6c1308c41 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -278,6 +278,13 @@ class CausalMatrixGameSelfAttention(nn.Module): else self.local_attn_size * frame_seqlen ) + if torch.is_grad_enabled(): + kv_cache["k"] = kv_cache["k"].detach().clone() + kv_cache["v"] = kv_cache["v"].detach().clone() + else: + kv_cache["k"] = kv_cache["k"].detach() + kv_cache["v"] = kv_cache["v"].detach() + kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.shape[1] global_end_index = ( @@ -905,13 +912,12 @@ class CausalMatrixGameWanModel(BaseDiT): encoder_hidden_states, torch.Tensor ): encoder_hidden_states = encoder_hidden_states[0] - if ( - isinstance(encoder_hidden_states_image, list) - and len(encoder_hidden_states_image) > 0 - ): - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = ( + encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 + else None + ) ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) @@ -1024,45 +1030,36 @@ class CausalMatrixGameWanModel(BaseDiT): else: encoder_hidden_states = encoder_hidden_states_image - if self.block_mask is None: - self.block_mask = self._prepare_blockwise_causal_attn_mask( + block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + if self.use_rope_keyboard: + block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + else: + block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( device=hidden_states.device, num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size, ) - if self.block_mask_keyboard is None: - if not self.use_rope_keyboard: - self.block_mask_keyboard = ( - self._prepare_blockwise_causal_attn_mask_keyboard( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - ) - else: - self.block_mask_keyboard = ( - self._prepare_blockwise_causal_attn_mask_action( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=1, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - ) - if self.block_mask_mouse is None: - self.block_mask_mouse = ( - self._prepare_blockwise_causal_attn_mask_action( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=1, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - ) + block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) if kv_cache is None: kv_cache = [None] * len(self.blocks) if kv_cache_mouse is None: @@ -1073,58 +1070,40 @@ class CausalMatrixGameWanModel(BaseDiT): crossattn_cache = [None] * len(self.blocks) for block_index, block in enumerate(self.blocks): + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "kv_cache_mouse": kv_cache_mouse[block_index] + if kv_cache_mouse + else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] + if kv_cache_keyboard + else None, + "crossattn_cache": crossattn_cache[block_index] + if crossattn_cache + else None, + "current_start": current_start, + "cache_start": cache_start, + } if torch.is_grad_enabled() and self.gradient_checkpointing: - kwargs.update( - { - "kv_cache": kv_cache[block_index], - "kv_cache_mouse": kv_cache_mouse[block_index] - if kv_cache_mouse - else None, - "kv_cache_keyboard": kv_cache_keyboard[block_index] - if kv_cache_keyboard - else None, - "crossattn_cache": crossattn_cache[block_index] - if crossattn_cache - else None, - "current_start": current_start, - "cache_start": cache_start, - "num_frame_per_block": effective_num_frame_per_block, - } - ) + causal_kwargs["num_frame_per_block"] = effective_num_frame_per_block hidden_states = self._gradient_checkpointing_func( - block, hidden_states, **kwargs + block, hidden_states, **causal_kwargs ) else: - kwargs.update( - { - "kv_cache": kv_cache[block_index], - "kv_cache_mouse": kv_cache_mouse[block_index] - if kv_cache_mouse - else None, - "kv_cache_keyboard": kv_cache_keyboard[block_index] - if kv_cache_keyboard - else None, - "crossattn_cache": crossattn_cache[block_index] - if crossattn_cache - else None, - "current_start": current_start, - "cache_start": cache_start, - } - ) hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, - block_mask=self.block_mask, + block_mask=block_mask, grid_sizes=grid_sizes, mouse_cond=mouse_cond, keyboard_cond=keyboard_cond, - block_mask_mouse=self.block_mask_mouse, - block_mask_keyboard=self.block_mask_keyboard, + block_mask_mouse=block_mask_mouse, + block_mask_keyboard=block_mask_keyboard, num_frame_per_block=effective_num_frame_per_block, use_rope_keyboard=self.use_rope_keyboard, - **kwargs, + **causal_kwargs, ) temb = temb.unflatten( @@ -1161,13 +1140,12 @@ class CausalMatrixGameWanModel(BaseDiT): encoder_hidden_states, torch.Tensor ): encoder_hidden_states = encoder_hidden_states[0] - if ( - isinstance(encoder_hidden_states_image, list) - and len(encoder_hidden_states_image) > 0 - ): - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = ( + encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 + else None + ) ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) @@ -1236,35 +1214,29 @@ class CausalMatrixGameWanModel(BaseDiT): local_attn_size=self.local_attn_size, ) if self.block_mask_keyboard is None: - if not self.use_rope_keyboard: - self.block_mask_keyboard = ( - self._prepare_blockwise_causal_attn_mask_keyboard( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - ) - else: - self.block_mask_keyboard = ( - self._prepare_blockwise_causal_attn_mask_action( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=1, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - ) - if self.block_mask_mouse is None: - self.block_mask_mouse = ( - self._prepare_blockwise_causal_attn_mask_action( + if self.use_rope_keyboard: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( device=hidden_states.device, num_frames=num_frames, frame_seqlen=1, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size, ) + else: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + if self.block_mask_mouse is None: + self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, ) hidden_states = self.patch_embedding(hidden_states) diff --git a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py index f5199db03..6031b5cc4 100644 --- a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py +++ b/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) diff --git a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline_ode_trajectory.py new file mode 100644 index 000000000..0896cde4a --- /dev/null +++ b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline_ode_trajectory.py @@ -0,0 +1,385 @@ +# 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'] + keyboard_cond = extra_features.get('keyboard_cond') + mouse_cond = extra_features.get('mouse_cond') + + 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 diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index d71354f5b..1f9c9c79e 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/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) diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index 861482fc2..7cfe5f4df 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -35,6 +35,7 @@ class DiffusionForcingSFTMethod(TrainingMethod): self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None)) self._timestep_index_range = (self._parse_timestep_index_range()) + self._training_weights = self._build_training_weights() # Initialize preprocessors on student. self.student.init_preprocessors(self.training_config) @@ -146,10 +147,24 @@ class DiffusionForcingSFTMethod(TrainingMethod): sigmas = schedule_sigmas[timestep_indices] sigmas = sigmas.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) pred_x0 = noisy_latents - pred * sigmas - loss = F.mse_loss(pred_x0.float(), clean_latents.float()) + per_frame_loss = F.mse_loss( + pred_x0.float(), + clean_latents.float(), + reduction="none", + ).mean(dim=(2, 3, 4)) else: target = noise - clean_latents - loss = F.mse_loss(pred.float(), target.float()) + per_frame_loss = F.mse_loss( + pred.float(), + target.float(), + reduction="none", + ).mean(dim=(2, 3, 4)) + + weight = self._get_training_weight( + timestep_indices, + clean_latents.device, + ).reshape(batch_size, num_latents) + loss = (per_frame_loss * weight).mean() attn_metadata = training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata @@ -317,3 +332,27 @@ class DiffusionForcingSFTMethod(TrainingMethod): ) expanded = chunk_indices.repeat_interleave(chunk_size, dim=1) return expanded[:, :num_latents] + + def _build_training_weights(self) -> torch.Tensor: + """Gaussian weighting over timestep indices. + + Emphasizes mid-noise timesteps, down-weights extremes + (near-clean and pure-noise). Matches Causal-Forcing's + bsmntw weighting scheme. + """ + scheduler = self.student.noise_scheduler + if scheduler is None: + raise ValueError("DFSFT requires student.noise_scheduler") + n = float(len(scheduler.timesteps)) + x = torch.arange(n, dtype=torch.float32) + y = torch.exp(-2 * ((x - n / 2) / n)**2) + y_shifted = y - y.min() + return y_shifted * (n / y_shifted.sum()) + + def _get_training_weight( + self, + timestep_indices: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + """Look up per-frame gaussian weights by timestep index.""" + return self._training_weights.to(device)[timestep_indices] diff --git a/fastvideo/train/models/matrixgame/matrixgame.py b/fastvideo/train/models/matrixgame/matrixgame.py new file mode 100644 index 000000000..1ae674790 --- /dev/null +++ b/fastvideo/train/models/matrixgame/matrixgame.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MatrixGame training model plugin.""" + +from __future__ import annotations + +import copy +from typing import Any, Literal + +import torch + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame +from fastvideo.distributed import ( + get_sp_group, + get_world_group, +) +from fastvideo.pipelines import TrainingBatch +from fastvideo.training.training_utils import normalize_dit_input + +from fastvideo.train.models.wan.wan import WanModel +from fastvideo.train.utils.dataloader import ( + build_parquet_matrixgame_train_dataloader, ) +from fastvideo.train.utils.moduleloader import ( + load_module_from_path, ) + + +class MatrixGameModel(WanModel): + """MatrixGame per-role model for finetuning in the new trainer.""" + + _transformer_cls_name: str = "MatrixGameWanModel" + + def init_preprocessors(self, training_config: Any) -> None: + self.vae = load_module_from_path( + model_path=str(training_config.model_path), + module_type="vae", + training_config=training_config, + ) + self.world_group = get_world_group() + self.sp_group = get_sp_group() + self._init_timestep_mechanics() + self.dataloader = build_parquet_matrixgame_train_dataloader( + training_config.data, + parquet_schema=pyarrow_schema_matrixgame, + ) + self.start_step = 0 + + def on_train_start(self) -> None: + # MatrixGame finetuning does not use negative text conditioning. + return + + def prepare_batch( + self, + raw_batch: dict[str, Any], + *, + generator: torch.Generator, + latents_source: Literal["data", "zeros"] = "data", + ) -> TrainingBatch: + assert self.training_config is not None + tc = self.training_config + dtype = self._get_training_dtype() + device = self.device + + training_batch = TrainingBatch() + infos = raw_batch.get("info_list") + batch_size = self._infer_batch_size(raw_batch) + + if latents_source == "zeros": + latents = self._make_zero_latents(batch_size=batch_size) + elif latents_source == "data": + latents = raw_batch["vae_latent"][:, :, :tc.data.num_latent_t] + latents = latents.to(device=device, dtype=dtype) + else: + raise ValueError(f"Unknown latents_source: {latents_source!r}") + + clip_feature = raw_batch["clip_feature"].to(device=device, dtype=dtype) + first_frame_latent = raw_batch["first_frame_latent"] + first_frame_latent = first_frame_latent[:, :, :tc.data.num_latent_t] + first_frame_latent = first_frame_latent.to(device=device, dtype=dtype) + + pil_image = raw_batch.get("pil_image") + if pil_image is not None: + pil_image = pil_image.to(device=device) + + keyboard_cond = self._get_optional_action( + raw_batch, + key="keyboard_cond", + expected_frames=self._expected_action_frames(tc.data.num_latent_t), + dtype=dtype, + ) + mouse_cond = self._get_optional_action( + raw_batch, + key="mouse_cond", + expected_frames=self._expected_action_frames(tc.data.num_latent_t), + dtype=dtype, + ) + + training_batch.latents = latents + training_batch.encoder_hidden_states = None + training_batch.encoder_attention_mask = None + training_batch.preprocessed_image = pil_image + training_batch.image_embeds = clip_feature + training_batch.image_latents = first_frame_latent + training_batch.keyboard_cond = keyboard_cond + training_batch.mouse_cond = mouse_cond + training_batch.infos = infos + + training_batch.latents = normalize_dit_input( + "wan", + training_batch.latents, + self.vae, + ) + training_batch = self._prepare_dit_inputs(training_batch, generator) + training_batch = self._build_attention_metadata(training_batch) + + training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata) + if training_batch.attn_metadata is not None: + training_batch.attn_metadata.VSA_sparsity = 0.0 # type: ignore[attr-defined] + + return training_batch + + def _prepare_dit_inputs( + self, + training_batch: TrainingBatch, + generator: torch.Generator, + ) -> TrainingBatch: + training_batch = super()._prepare_dit_inputs(training_batch, generator) + + image_latents = training_batch.image_latents + image_embeds = training_batch.image_embeds + if image_latents is None or image_embeds is None: + raise RuntimeError("MatrixGame requires image_latents and image_embeds") + + cond_latents = self._build_matrixgame_cond_concat(image_latents) + training_batch.image_latents = cond_latents + training_batch.noisy_model_input = torch.cat( + [training_batch.noisy_model_input, cond_latents], + dim=1, + ) + training_batch.conditional_dict = { + "encoder_hidden_states": None, + "encoder_attention_mask": None, + "encoder_hidden_states_image": image_embeds, + "keyboard_cond": training_batch.keyboard_cond, + "mouse_cond": training_batch.mouse_cond, + "image_latents": cond_latents, + } + training_batch.unconditional_dict = dict(training_batch.conditional_dict) + return training_batch + + def _get_uncond_text_dict( + self, + batch: TrainingBatch, + *, + cfg_uncond: dict[str, Any] | None, + ) -> dict[str, Any]: + del cfg_uncond + cond_dict = batch.conditional_dict + if cond_dict is None: + raise RuntimeError("Missing conditional_dict in TrainingBatch") + return batch.unconditional_dict or cond_dict + + def _build_distill_input_kwargs( + self, + noise_input: torch.Tensor, + timestep: torch.Tensor, + text_dict: dict[str, Any] | None, + ) -> dict[str, Any]: + if text_dict is None: + raise ValueError("text_dict cannot be None for MatrixGame") + hidden_states = noise_input.permute(0, 2, 1, 3, 4) + if hidden_states.shape[1] == 16: + cond_latents = text_dict.get("image_latents") + if cond_latents is None: + raise RuntimeError("MatrixGame requires image_latents in conditional_dict " + "when noise_input has 16 channels") + num_t = hidden_states.shape[2] + cond_latents = cond_latents[:, :, :num_t] + hidden_states = torch.cat([hidden_states, cond_latents], dim=1) + return { + "hidden_states": hidden_states, + "encoder_hidden_states": None, + "timestep": timestep.to(device=self.device, dtype=torch.bfloat16), + "encoder_hidden_states_image": text_dict["encoder_hidden_states_image"], + "keyboard_cond": text_dict["keyboard_cond"], + "mouse_cond": text_dict["mouse_cond"], + "return_dict": False, + } + + def _build_matrixgame_cond_concat( + self, + image_latents: torch.Tensor, + ) -> torch.Tensor: + if image_latents.ndim != 5: + raise ValueError("first_frame_latent must have shape [B, C, T, H, W], " + f"got {tuple(image_latents.shape)}") + if image_latents.shape[1] == 20: + return image_latents + if image_latents.shape[1] != 16: + raise ValueError("MatrixGame expects first_frame_latent with 16 or 20 channels, " + f"got {image_latents.shape[1]}") + + temporal_compression_ratio = self._temporal_compression_ratio() + batch_size, _, num_latent_t, latent_height, latent_width = image_latents.shape + num_frames = (num_latent_t - 1) * temporal_compression_ratio + 1 + + mask_lat_size = torch.ones( + batch_size, + 1, + num_frames, + latent_height, + latent_width, + device=image_latents.device, + dtype=image_latents.dtype, + ) + 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, + ).transpose(1, 2) + return torch.cat([mask_lat_size, image_latents], dim=1) + + def _get_optional_action( + self, + raw_batch: dict[str, Any], + *, + key: str, + expected_frames: int, + dtype: torch.dtype, + ) -> torch.Tensor | None: + value = raw_batch.get(key) + if value is None or value.numel() == 0: + return None + if value.shape[1] < expected_frames: + raise ValueError(f"{key} has {value.shape[1]} frames but requires at least {expected_frames}") + return value[:, :expected_frames].to(device=self.device, dtype=dtype) + + def _expected_action_frames(self, num_latent_t: int) -> int: + return (num_latent_t - 1) * self._temporal_compression_ratio() + 1 + + def _temporal_compression_ratio(self) -> int: + assert self.training_config is not None + return int(self.training_config.pipeline_config.vae_config.arch_config. + temporal_compression_ratio # type: ignore[union-attr] + ) + + def _infer_batch_size(self, raw_batch: dict[str, Any]) -> int: + if "vae_latent" in raw_batch: + return int(raw_batch["vae_latent"].shape[0]) + if "clip_feature" in raw_batch: + return int(raw_batch["clip_feature"].shape[0]) + raise ValueError("Unable to infer batch size from MatrixGame batch") + + def _make_zero_latents(self, *, batch_size: int) -> torch.Tensor: + assert self.training_config is not None + vae_config = self.training_config.pipeline_config.vae_config.arch_config # type: ignore[union-attr] + num_channels = vae_config.z_dim + spatial_compression_ratio = vae_config.spatial_compression_ratio + latent_height = self.training_config.data.num_height // spatial_compression_ratio + latent_width = self.training_config.data.num_width // spatial_compression_ratio + return torch.zeros( + batch_size, + num_channels, + self.training_config.data.num_latent_t, + latent_height, + latent_width, + device=self.device, + dtype=self._get_training_dtype(), + ) diff --git a/fastvideo/train/models/matrixgame/matrixgame_causal.py b/fastvideo/train/models/matrixgame/matrixgame_causal.py new file mode 100644 index 000000000..0e3c3d2cf --- /dev/null +++ b/fastvideo/train/models/matrixgame/matrixgame_causal.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MatrixGame causal training model plugin.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +import torch + +from fastvideo.forward_context import set_forward_context + +from fastvideo.train.models.matrixgame.matrixgame import MatrixGameModel +from fastvideo.train.models.wan.wan_causal import WanCausalModel + + +@dataclass(slots=True) +class _StreamingCaches: + kv_cache: list[dict[str, Any]] + kv_cache_mouse: list[dict[str, Any] | None] + kv_cache_keyboard: list[dict[str, Any] | None] + crossattn_cache: list[dict[str, Any]] | None + frame_seq_length: int + local_attn_size: int + sliding_window_num_frames: int + batch_size: int + dtype: torch.dtype + device: torch.device + + +class MatrixGameCausalModel(MatrixGameModel, WanCausalModel): + """MatrixGame per-role model with causal/streaming primitives.""" + + _transformer_cls_name: str = "CausalMatrixGameWanModel" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._streaming_caches: dict[tuple[int, str], _StreamingCaches] = {} + + def clear_caches(self, *, cache_tag: str = "pos") -> None: + self._streaming_caches.pop((id(self), str(cache_tag)), None) + + def predict_noise_streaming( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + batch: Any, + *, + conditional: bool, + cache_tag: str = "pos", + store_kv: bool = False, + cur_start_frame: int = 0, + cfg_uncond: dict[str, Any] | None = None, + attn_kind: Literal["dense", "vsa"] = "dense", + ) -> torch.Tensor | None: + del cfg_uncond + + if attn_kind == "dense": + attn_metadata = batch.attn_metadata + elif attn_kind == "vsa": + attn_metadata = batch.attn_metadata_vsa + else: + raise ValueError(f"Unknown attn_kind: {attn_kind!r}") + + cache_tag = str(cache_tag) + cur_start_frame = int(cur_start_frame) + if cur_start_frame < 0: + raise ValueError("cur_start_frame must be >= 0") + + batch_size = int(noisy_latents.shape[0]) + num_frames = int(noisy_latents.shape[1]) + timestep_full = self._ensure_per_frame_timestep( + timestep=timestep, + batch_size=batch_size, + num_frames=num_frames, + device=noisy_latents.device, + ) + + transformer = self._get_transformer(timestep_full) + caches = self._get_or_init_streaming_caches( + cache_tag=cache_tag, + transformer=transformer, + noisy_latents=noisy_latents, + ) + + frame_seq_length = int(caches.frame_seq_length) + kv_cache = caches.kv_cache + kv_cache_mouse = caches.kv_cache_mouse + kv_cache_keyboard = caches.kv_cache_keyboard + crossattn_cache = caches.crossattn_cache + + if self._should_snapshot_streaming_cache() and torch.is_grad_enabled(): + kv_cache = self._snapshot_kv_cache_indices(kv_cache) + kv_cache_mouse = self._snapshot_action_kv_cache_indices(kv_cache_mouse) + kv_cache_keyboard = self._snapshot_action_kv_cache_indices(kv_cache_keyboard) + + if conditional: + cond_dict = batch.conditional_dict + if cond_dict is None: + raise RuntimeError("Missing conditional_dict in TrainingBatch") + else: + cond_dict = batch.unconditional_dict or batch.conditional_dict + if cond_dict is None: + raise RuntimeError("Missing conditional_dict in TrainingBatch") + + device_type = self.device.type + dtype = noisy_latents.dtype + with ( + torch.autocast(device_type, dtype=dtype), + set_forward_context( + current_timestep=batch.timesteps, + attn_metadata=attn_metadata, + ), + ): + input_kwargs = self._build_streaming_input_kwargs( + noisy_latents=noisy_latents, + timestep=timestep_full, + cond_dict=cond_dict, + batch=batch, + cur_start_frame=cur_start_frame, + ) + input_kwargs["timestep"] = timestep_full.to( + device=self.device, + dtype=torch.long, + ) + input_kwargs.update({ + "kv_cache": kv_cache, + "kv_cache_mouse": kv_cache_mouse, + "kv_cache_keyboard": kv_cache_keyboard, + "crossattn_cache": crossattn_cache, + "current_start": cur_start_frame * frame_seq_length, + "start_frame": cur_start_frame, + "is_cache": bool(store_kv), + "num_frame_per_block": num_frames, + }) + + if store_kv: + with torch.no_grad(): + _ = transformer(**input_kwargs) + return None + + pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + return pred_noise + + def _get_or_init_streaming_caches( + self, + *, + cache_tag: str, + transformer: torch.nn.Module, + noisy_latents: torch.Tensor, + ) -> _StreamingCaches: + key = (id(self), cache_tag) + cached = self._streaming_caches.get(key) + + batch_size = int(noisy_latents.shape[0]) + dtype = noisy_latents.dtype + device = noisy_latents.device + + frame_seq_length = self._compute_frame_seq_length( + transformer, + noisy_latents, + ) + local_attn_size = self._get_local_attn_size(transformer) + sliding_window_num_frames = self._get_sliding_window_num_frames(transformer) + + meta = ( + frame_seq_length, + local_attn_size, + sliding_window_num_frames, + batch_size, + dtype, + device, + ) + + if cached is not None: + cached_meta = ( + cached.frame_seq_length, + cached.local_attn_size, + cached.sliding_window_num_frames, + cached.batch_size, + cached.dtype, + cached.device, + ) + if cached_meta == meta: + return cached + + kv_cache = self._initialize_kv_cache( + transformer=transformer, + batch_size=batch_size, + dtype=dtype, + device=device, + frame_seq_length=frame_seq_length, + local_attn_size=local_attn_size, + sliding_window_num_frames=sliding_window_num_frames, + checkpoint_safe=self._should_use_checkpoint_safe_kv_cache(), + ) + kv_cache_mouse = self._initialize_action_kv_cache( + transformer=transformer, + batch_size=batch_size, + dtype=dtype, + device=device, + frame_seq_length=frame_seq_length, + local_attn_size=local_attn_size, + channel="mouse", + ) + kv_cache_keyboard = self._initialize_action_kv_cache( + transformer=transformer, + batch_size=batch_size, + dtype=dtype, + device=device, + frame_seq_length=frame_seq_length, + local_attn_size=local_attn_size, + channel="keyboard", + ) + crossattn_cache = self._initialize_crossattn_cache( + transformer=transformer, + device=device, + ) + + caches = _StreamingCaches( + kv_cache=kv_cache, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + frame_seq_length=frame_seq_length, + local_attn_size=local_attn_size, + sliding_window_num_frames=sliding_window_num_frames, + batch_size=batch_size, + dtype=dtype, + device=device, + ) + self._streaming_caches[key] = caches + return caches + + def _build_streaming_input_kwargs( + self, + *, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + cond_dict: dict[str, Any], + batch: Any, + cur_start_frame: int, + ) -> dict[str, Any]: + if batch.image_latents is None: + raise RuntimeError("MatrixGame causal rollout requires image_latents") + if batch.image_embeds is None: + raise RuntimeError("MatrixGame causal rollout requires image_embeds") + + num_frames = int(noisy_latents.shape[1]) + frame_end = cur_start_frame + num_frames + cond_latents = self._slice_image_latents_for_chunk( + batch.image_latents, + start=cur_start_frame, + end=frame_end, + ) + hidden_states = torch.cat( + [noisy_latents.permute(0, 2, 1, 3, 4), cond_latents], + dim=1, + ) + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": None, + "timestep": timestep, + "encoder_hidden_states_image": cond_dict["encoder_hidden_states_image"], + "keyboard_cond": self._slice_action_prefix( + cond_dict.get("keyboard_cond"), + frame_end=frame_end, + ), + "mouse_cond": self._slice_action_prefix( + cond_dict.get("mouse_cond"), + frame_end=frame_end, + ), + "return_dict": False, + } + + def _slice_image_latents_for_chunk( + self, + image_latents: torch.Tensor, + *, + start: int, + end: int, + ) -> torch.Tensor: + num_frames = end - start + if image_latents.ndim != 5: + raise ValueError("image_latents must have shape [B, C, T, H, W], " + f"got {tuple(image_latents.shape)}") + if image_latents.shape[2] >= end: + return image_latents[:, :, start:end] + if image_latents.shape[2] > start: + cond = image_latents[:, :, start:] + pad_frames = num_frames - int(cond.shape[2]) + if pad_frames <= 0: + return cond + pad = torch.zeros( + cond.shape[0], + cond.shape[1], + pad_frames, + cond.shape[3], + cond.shape[4], + device=cond.device, + dtype=cond.dtype, + ) + return torch.cat([cond, pad], dim=2) + return torch.zeros( + image_latents.shape[0], + image_latents.shape[1], + num_frames, + image_latents.shape[3], + image_latents.shape[4], + device=image_latents.device, + dtype=image_latents.dtype, + ) + + def _slice_action_prefix( + self, + action: torch.Tensor | None, + *, + frame_end: int, + ) -> torch.Tensor | None: + if action is None: + return None + action_frame_end = ((frame_end - 1) * self._temporal_compression_ratio()) + 1 + if action.shape[1] < action_frame_end: + raise ValueError("Action tensor is shorter than required for causal rollout: " + f"got={action.shape[1]}, required>={action_frame_end}") + return action[:, :action_frame_end] + + def _initialize_action_kv_cache( + self, + *, + transformer: torch.nn.Module, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + frame_seq_length: int, + local_attn_size: int, + channel: Literal["mouse", "keyboard"], + ) -> list[dict[str, Any] | None]: + blocks = list(getattr(transformer, "blocks", [])) + num_blocks = len(blocks) + if num_blocks <= 0: + return [] + + action_config = getattr(transformer, "action_config", {}) or {} + action_blocks = {int(block_idx) for block_idx in action_config.get("blocks", [])} + if local_attn_size <= 0: + raise ValueError("MatrixGame causal streaming requires " + "transformer.local_attn_size > 0") + + action_heads_num = int(action_config.get("heads_num", 0) or 0) + if action_blocks and action_heads_num <= 0: + raise ValueError("MatrixGame causal action_config.heads_num must be > 0") + + if channel == "mouse": + enabled = bool(action_config.get("enable_mouse", False)) + hidden_dim = int(action_config.get("mouse_hidden_dim", 0) or 0) + batch_dim = batch_size * frame_seq_length + else: + enabled = bool(action_config.get("enable_keyboard", False)) + hidden_dim = int(action_config.get("keyboard_hidden_dim", 0) or 0) + batch_dim = batch_size + + caches: list[dict[str, Any] | None] = [] + for block_idx in range(num_blocks): + if block_idx not in action_blocks or not enabled: + caches.append(None) + continue + + if hidden_dim <= 0 or hidden_dim % action_heads_num != 0: + raise ValueError(f"Invalid {channel} action hidden size for causal " + f"cache initialization: hidden_dim={hidden_dim}, " + f"heads={action_heads_num}") + head_dim = hidden_dim // action_heads_num + + caches.append({ + "k": + torch.zeros( + [ + batch_dim, + local_attn_size, + action_heads_num, + head_dim, + ], + dtype=dtype, + device=device, + ), + "v": + torch.zeros( + [ + batch_dim, + local_attn_size, + action_heads_num, + head_dim, + ], + dtype=dtype, + device=device, + ), + "global_end_index": + torch.tensor( + [0], + dtype=torch.long, + device=device, + ), + "local_end_index": + torch.tensor( + [0], + dtype=torch.long, + device=device, + ), + }) + + return caches + + def _snapshot_action_kv_cache_indices( + self, + caches: list[dict[str, Any] | None], + ) -> list[dict[str, Any] | None]: + snapshot: list[dict[str, Any] | None] = [] + for cache in caches: + if cache is None: + snapshot.append(None) + continue + + global_end_index = cache.get("global_end_index") + local_end_index = cache.get("local_end_index") + if (not isinstance(global_end_index, torch.Tensor) or not isinstance(local_end_index, torch.Tensor)): + raise ValueError("Unexpected action kv_cache index tensors; expected " + "tensors at kv_cache_*[*].{global_end_index, " + "local_end_index}") + + copied = dict(cache) + copied["global_end_index"] = global_end_index.detach().clone() + copied["local_end_index"] = local_end_index.detach().clone() + snapshot.append(copied) + return snapshot diff --git a/fastvideo/train/utils/dataloader.py b/fastvideo/train/utils/dataloader.py index a1b22d3ba..5c2ce6794 100644 --- a/fastvideo/train/utils/dataloader.py +++ b/fastvideo/train/utils/dataloader.py @@ -31,3 +31,26 @@ def build_parquet_t2v_train_dataloader( seed=int(data_config.seed or 0), )) return dataloader + + +def build_parquet_matrixgame_train_dataloader( + data_config: DataConfig, + *, + parquet_schema: Any, +) -> Any: + """Build a parquet dataloader for MatrixGame datasets.""" + + from fastvideo.dataset import ( + build_parquet_map_style_dataloader, ) + + _dataset, dataloader = (build_parquet_map_style_dataloader( + data_config.data_path, + data_config.train_batch_size, + num_data_workers=(data_config.dataloader_num_workers), + parquet_schema=parquet_schema, + cfg_rate=float(data_config.training_cfg_rate or 0.0), + drop_last=True, + text_padding_length=512, + seed=int(data_config.seed or 0), + )) + return dataloader \ No newline at end of file diff --git a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py new file mode 100644 index 000000000..cdbabc731 --- /dev/null +++ b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py @@ -0,0 +1,451 @@ +# 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.api.sampling_param 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_self_forcing_flow_match import ( + SelfForcingFlowMatchScheduler, ) +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): + scheduler = SelfForcingFlowMatchScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True, + ) + scheduler.set_timesteps(num_inference_steps=1000, training=True) + self.modules["scheduler"] = scheduler + + def _log_training_info(self) -> None: + self.noise_scheduler = self.modules["scheduler"] + super()._log_training_info() + + 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) + + logger.info( + "MatrixGame AR diffusion pipeline initialized with " + "num_frame_per_block=%d, diffusion_forcing_shift=%.1f", + self.num_frame_per_block, + training_args.pipeline_config.flow_shift, + ) + + 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 = 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=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, + ) -> torch.Tensor: + """Sample one schedule index per chunk and broadcast to frames. + + Returns the shifted timestep, shape [B, num_frame]. + """ + device = get_local_torch_device() + num_schedule_steps = len(self.noise_scheduler.timesteps) + chunk_size = int(num_frame_per_block) + if chunk_size <= 0: + raise ValueError(f"num_frame_per_block must be > 0, got {chunk_size}") + num_chunks = (num_frame + chunk_size - 1) // chunk_size + chunk_indices = torch.randint( + 0, + num_schedule_steps, + (batch_size, num_chunks), + device=device, + dtype=torch.long, + ) + indices = chunk_indices.repeat_interleave(chunk_size, dim=1)[:, :num_frame] + schedule_t = self.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, + ) + + 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.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.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) + + per_frame_loss = F.mse_loss( + model_pred_btchw.float(), + training_target.float(), + reduction="none", + ).mean(dim=(2, 3, 4)) + timesteps = training_batch.timesteps + weight = self.noise_scheduler.training_weight(timesteps.flatten(0, 1)).to(per_frame_loss.dtype).reshape( + per_frame_loss.shape) + loss = (per_frame_loss * weight).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) diff --git a/fastvideo/training/matrixgame_ode_causal_pipeline.py b/fastvideo/training/matrixgame_ode_causal_pipeline.py new file mode 100644 index 000000000..b97725e62 --- /dev/null +++ b/fastvideo/training/matrixgame_ode_causal_pipeline.py @@ -0,0 +1,416 @@ +# 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.api.sampling_param 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 _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)) + + indexes = self._get_timestep( # [B, num_frames] + 0, len(self.dmd_denoising_steps), B, num_frames, 3, uniform_timestep=False) + 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] + + # 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) + + 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) + + 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) diff --git a/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py b/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py new file mode 100644 index 000000000..3960fe5d0 --- /dev/null +++ b/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py @@ -0,0 +1,873 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any, cast + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from einops import rearrange + +from fastvideo.api.sampling_param 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_self_forcing_flow_match import (SelfForcingFlowMatchScheduler) +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.pipelines import ComposedPipelineBase +from fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline import (MatrixGameCausalDMDPipeline) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.training.self_forcing_distillation_pipeline import (SelfForcingDistillationPipeline) +from fastvideo.training.training_utils import shift_timestep +from fastvideo.utils import is_vsa_available, shallow_asdict + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class MatrixGameSelfForcingDistillationPipeline(SelfForcingDistillationPipeline): + """ + A self-forcing distillation pipeline for MatrixGame that uses the self-forcing methodology + with DMD for video generation. + """ + _required_config_modules = [ + "scheduler", + "transformer", + "vae", + "image_encoder", + "image_processor", + ] + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_matrixgame + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> dict[str, Any]: + modules = ComposedPipelineBase.load_modules( + self, + fastvideo_args, + loaded_modules, + ) + training_args = cast(TrainingArgs, fastvideo_args) + old_override = training_args.override_transformer_cls_name + training_args.override_transformer_cls_name = "MatrixGameWanModel" + try: + if loaded_modules is not None and "real_score_transformer" in loaded_modules: + self.real_score_transformer = loaded_modules["real_score_transformer"] + elif training_args.real_score_model_path: + logger.info( + "Loading real score transformer from: %s", + training_args.real_score_model_path, + ) + self.real_score_transformer = self.load_module_from_path( + training_args.real_score_model_path, + "transformer", + training_args, + ) + else: + raise ValueError("real_score_model_path is required for DMD distillation pipeline") + modules["real_score_transformer"] = self.real_score_transformer + + if loaded_modules is not None and "fake_score_transformer" in loaded_modules: + self.fake_score_transformer = loaded_modules["fake_score_transformer"] + elif training_args.fake_score_model_path: + logger.info( + "Loading fake score transformer from: %s", + training_args.fake_score_model_path, + ) + self.fake_score_transformer = self.load_module_from_path( + training_args.fake_score_model_path, + "transformer", + training_args, + ) + else: + raise ValueError("fake_score_model_path is required for DMD distillation pipeline") + modules["fake_score_transformer"] = self.fake_score_transformer + finally: + training_args.override_transformer_cls_name = old_override + + self.real_score_transformer_2 = None + self.fake_score_transformer_2 = None + return modules + + def _build_matrixgame_cond_concat( + self, + image_latents: torch.Tensor, + ) -> torch.Tensor: + if image_latents.ndim != 5: + raise ValueError("first_frame_latent must have shape [B, C, T, H, W], got " + f"{tuple(image_latents.shape)}") + if image_latents.shape[1] != 16: + raise ValueError("MatrixGame cond conversion expects a 16-channel " + f"first_frame_latent, got {image_latents.shape[1]} channels.") + + temporal_compression_ratio = ( + self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio) + num_latent_t = image_latents.shape[2] + num_frames = ((num_latent_t - 1) * temporal_compression_ratio + 1) + batch_size, _, _, latent_height, latent_width = image_latents.shape + + mask_lat_size = torch.ones( + batch_size, + 1, + num_frames, + latent_height, + latent_width, + device=image_latents.device, + dtype=image_latents.dtype, + ) + 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).to( + image_latents.device, + dtype=image_latents.dtype, + ) + + return torch.cat([mask_lat_size, image_latents], dim=1) + + def _initialize_simulation_caches( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any] | None], list[dict[str, Any] | None]]: + """Initialize KV cache and cross-attention cache for multi-step simulation.""" + num_transformer_blocks = len(self.transformer.blocks) + latent_shape = self.video_latent_shape_sp + _, num_frames, _, height, width = latent_shape + + _, p_h, p_w = self.transformer.patch_size + post_patch_height = height // p_h + post_patch_width = width // p_w + + frame_seq_length = post_patch_height * post_patch_width + self.frame_seq_length = frame_seq_length + + # Get model configuration parameters - handle FSDP wrapping + num_attention_heads = getattr(self.transformer, 'num_attention_heads', None) + attention_head_dim = getattr(self.transformer, 'attention_head_dim', None) + + # 1 CLS token + 256 patch tokens = 257 + text_len = 257 + + action_config = getattr(self.transformer, 'action_config', {}) + action_blocks = action_config.get('blocks', []) if action_config else [] + local_attn_size = getattr(self.transformer, "local_attn_size", + action_config.get('local_attn_size', 6) if action_config else 6) + + if local_attn_size <= 0: + raise ValueError("MatrixGame self-forcing requires transformer.local_attn_size > 0") + kv_cache_size = local_attn_size * frame_seq_length + + kv_cache = [] + for _ in range(num_transformer_blocks): + kv_cache.append({ + "k": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device) + }) + + # Initialize cross-attention cache + crossattn_cache = [] + for _ in range(num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, + device=device), + "is_init": + False + }) + + # Initialize action module KV caches + action_heads_num = action_config.get('heads_num', 16) if action_config else 16 + mouse_hidden_dim = action_config.get('mouse_hidden_dim', 1024) if action_config else 1024 + keyboard_hidden_dim = action_config.get('keyboard_hidden_dim', 1024) if action_config else 1024 + + mouse_head_dim = mouse_hidden_dim // action_heads_num + keyboard_head_dim = keyboard_hidden_dim // action_heads_num + + kv_cache_mouse: list[dict[str, Any] | None] = [] + kv_cache_keyboard: list[dict[str, Any] | None] = [] + for block_idx in range(num_transformer_blocks): + if block_idx in action_blocks: + kv_cache_mouse.append({ + "k": + torch.zeros([batch_size * frame_seq_length, local_attn_size, action_heads_num, mouse_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size * frame_seq_length, local_attn_size, action_heads_num, mouse_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device) + }) + kv_cache_keyboard.append({ + "k": + torch.zeros([batch_size, local_attn_size, action_heads_num, keyboard_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, local_attn_size, action_heads_num, keyboard_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device) + }) + else: + kv_cache_mouse.append(None) + kv_cache_keyboard.append(None) + + return kv_cache, crossattn_cache, kv_cache_mouse, kv_cache_keyboard + + def _reset_simulation_caches(self, kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], + kv_cache_mouse: list[dict[str, Any] + | None] | None, kv_cache_keyboard: list[dict[str, Any] + | None] | None) -> None: + """Reset KV cache, cross-attention cache, and action caches to clean state.""" + if kv_cache is not None: + for cache_dict in kv_cache: + cache_dict["global_end_index"].fill_(0) + cache_dict["local_end_index"].fill_(0) + cache_dict["k"].zero_() + cache_dict["v"].zero_() + + if crossattn_cache is not None: + for cache_dict in crossattn_cache: + cache_dict["is_init"] = False + cache_dict["k"].zero_() + cache_dict["v"].zero_() + + for opt_caches in (kv_cache_mouse, kv_cache_keyboard): + if opt_caches is None: + continue + for opt_cache in opt_caches: + if opt_cache is None: + continue + opt_cache["global_end_index"].fill_(0) + opt_cache["local_end_index"].fill_(0) + opt_cache["k"].zero_() + opt_cache["v"].zero_() + + def _generator_multi_step_simulation_forward(self, + training_batch: TrainingBatch, + return_sim_steps: bool = False) -> torch.Tensor: + """Forward pass through student transformer matching inference procedure with KV cache management. + + This function is adapted from the reference self-forcing implementation's inference_with_trajectory + and includes gradient masking logic for dynamic frame generation. + """ + latents = training_batch.latents + dtype = latents.dtype + batch_size = latents.shape[0] + + num_training_frames = getattr(self.training_args, 'num_latent_t', 21) + min_num_frames = 20 if self.independent_first_frame else 21 + max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames + assert max_num_frames % self.num_frame_per_block == 0 + assert min_num_frames % self.num_frame_per_block == 0 + max_num_blocks = max_num_frames // self.num_frame_per_block + min_num_blocks = min_num_frames // self.num_frame_per_block + + # Sample number of blocks and sync across processes + num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1, ), device=self.device) + if dist.is_initialized(): + dist.broadcast(num_generated_blocks, src=0) + num_generated_blocks = num_generated_blocks.item() + num_generated_frames = num_generated_blocks * self.num_frame_per_block + if self.independent_first_frame: + num_generated_frames += 1 + min_num_frames += 1 + + noise_shape = [batch_size, num_generated_frames, *self.video_latent_shape[2:]] + noise = torch.randn(noise_shape, device=self.device, dtype=dtype) + if self.sp_world_size > 1: + noise = rearrange(noise, "b (n t) c h w -> b n t c h w", n=self.sp_world_size).contiguous() + noise = noise[:, self.rank_in_sp_group, :, :, :, :] + + batch_size, num_frames, num_channels, height, width = noise.shape + + if self.independent_first_frame: + assert (num_frames - 1) % self.num_frame_per_block == 0 + num_blocks = (num_frames - 1) // self.num_frame_per_block + else: + assert num_frames % self.num_frame_per_block == 0 + num_blocks = num_frames // self.num_frame_per_block + + num_output_frames = num_frames + output = torch.zeros([batch_size, num_output_frames, num_channels, height, width], + device=noise.device, + dtype=noise.dtype) + + # Step 1: Initialize KV cache to all zeros + (self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse, + self.kv_cache_keyboard) = self._initialize_simulation_caches(batch_size, dtype, self.device) + + # Step 2: Temporal denoising loop + current_start_frame = 0 + all_num_frames = [self.num_frame_per_block] * num_blocks + if self.independent_first_frame: + all_num_frames = [1] + all_num_frames + num_denoising_steps = len(self.denoising_step_list) + exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device) + start_gradient_frame_index = max(0, num_output_frames - 21) + + for block_index, current_num_frames in enumerate(all_num_frames): + noisy_input = noise[:, current_start_frame:current_start_frame + current_num_frames] + + # Step 3.1: Spatial denoising loop + for index, current_timestep in enumerate(self.denoising_step_list): + exit_flag = (index == exit_flags[0]) if self.same_step_across_blocks else (index + == exit_flags[block_index]) + + timestep = torch.ones([batch_size, current_num_frames], device=noise.device, + dtype=torch.int64) * current_timestep + + if self.boundary_timestep is not None and current_timestep < self.boundary_timestep and self.transformer_2 is not None: + current_model = self.transformer_2 + else: + current_model = self.transformer + + if not exit_flag: + with torch.no_grad(): + # Build input kwargs + training_batch_temp = self._build_distill_input_kwargs(noisy_input, + timestep, + training_batch.conditional_dict, + training_batch, + frame_start=current_start_frame, + frame_end=current_start_frame + + current_num_frames, + num_frame_per_block=current_num_frames) + + pred_flow = current_model(**training_batch_temp.input_kwargs, + kv_cache=self.kv_cache1, + kv_cache_mouse=self.kv_cache_mouse, + kv_cache_keyboard=self.kv_cache_keyboard, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame).permute(0, 2, 1, 3, 4) + + denoised_pred = pred_noise_to_pred_video(pred_noise=pred_flow.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, pred_flow.shape[:2]) + + next_timestep = self.denoising_step_list[index + 1] + noisy_input = self.noise_scheduler.add_noise( + denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), + next_timestep * torch.ones( + [batch_size * current_num_frames], device=noise.device, dtype=torch.long)).unflatten( + 0, denoised_pred.shape[:2]) + else: + # Final prediction with gradient control + if current_start_frame < start_gradient_frame_index: + with torch.no_grad(): + training_batch_temp = self._build_distill_input_kwargs( + noisy_input, + timestep, + training_batch.conditional_dict, + training_batch, + frame_start=current_start_frame, + frame_end=current_start_frame + current_num_frames, + num_frame_per_block=current_num_frames) + + pred_flow = current_model(**training_batch_temp.input_kwargs, + kv_cache=self.kv_cache1, + kv_cache_mouse=self.kv_cache_mouse, + kv_cache_keyboard=self.kv_cache_keyboard, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame).permute(0, 2, 1, 3, 4) + else: + training_batch_temp = self._build_distill_input_kwargs(noisy_input, + timestep, + training_batch.conditional_dict, + training_batch, + frame_start=current_start_frame, + frame_end=current_start_frame + + current_num_frames, + num_frame_per_block=current_num_frames) + + pred_flow = current_model(**training_batch_temp.input_kwargs, + kv_cache=self.kv_cache1, + kv_cache_mouse=self.kv_cache_mouse, + kv_cache_keyboard=self.kv_cache_keyboard, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame).permute(0, 2, 1, 3, 4) + + denoised_pred = pred_noise_to_pred_video(pred_noise=pred_flow.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, pred_flow.shape[:2]) + break + + # Step 3.2: record the model's output + output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred + + # Step 3.3: rerun with timestep zero to update the cache + context_timestep = torch.ones_like(timestep) * self.context_noise + denoised_pred = self.noise_scheduler.add_noise(denoised_pred.flatten(0, 1), + torch.randn_like(denoised_pred.flatten(0, 1)), + context_timestep).unflatten(0, denoised_pred.shape[:2]) + + with torch.no_grad(): + training_batch_temp = self._build_distill_input_kwargs(denoised_pred, + context_timestep, + training_batch.conditional_dict, + training_batch, + frame_start=current_start_frame, + frame_end=current_start_frame + + current_num_frames, + num_frame_per_block=current_num_frames) + + # context_timestep is 0 so we use transformer_2 + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + current_model(**training_batch_temp.input_kwargs, + kv_cache=self.kv_cache1, + kv_cache_mouse=self.kv_cache_mouse, + kv_cache_keyboard=self.kv_cache_keyboard, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame) + + # Step 3.4: update the start and end frame indices + current_start_frame += current_num_frames + + # Handle last 21 frames logic + pred_image_or_video = output + + # Slice last 21 frames if we generated more + gradient_mask = None + if pred_image_or_video.shape[1] > 21: + with torch.no_grad(): + # Re-encode to get image latent + latent_to_decode = pred_image_or_video[:, :-20, ...] + # Decode to video + latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] + + # Apply VAE scaling and shift factors + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent_to_decode = latent_to_decode / self.vae.scaling_factor.to( + latent_to_decode.device, latent_to_decode.dtype) + else: + latent_to_decode = latent_to_decode / self.vae.scaling_factor + + if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: + if isinstance(self.vae.shift_factor, torch.Tensor): + latent_to_decode += self.vae.shift_factor.to(latent_to_decode.device, latent_to_decode.dtype) + else: + latent_to_decode += self.vae.shift_factor + + # Decode to pixels + pixels = self.vae.decode(latent_to_decode) + frame = pixels[:, :, -1:, :, :].to(dtype) # Last frame [B, C, 1, H, W] + + # Encode frame back to get image latent + image_latent = self.vae.encode(frame).to(dtype) + image_latent = image_latent.permute(0, 2, 1, 3, 4) # [B, F, C, H, W] + + pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1) + else: + pred_image_or_video_last_21 = pred_image_or_video + + # Set up gradient mask if we generated more than minimum frames + if num_generated_frames != min_num_frames: + # Currently, we do not use gradient for the first chunk, since it contains image latents + gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool) + if self.independent_first_frame: + gradient_mask[:, :1] = False + else: + gradient_mask[:, :self.num_frame_per_block] = False + + # Apply gradient masking if needed + final_output = pred_image_or_video_last_21.to(dtype) + if gradient_mask is not None: + # Apply gradient masking: detach frames that shouldn't contribute gradients + final_output = torch.where( + gradient_mask, + pred_image_or_video_last_21, # Keep original values where gradient_mask is True + pred_image_or_video_last_21.detach() # Detach where gradient_mask is False + ) + + # Store visualization data + training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.as_tensor( + self.denoising_step_list[exit_flags[0]], + dtype=torch.float32, + device=self.device, + ).detach().clone() + + # Store gradient mask information for debugging + if gradient_mask is not None: + training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float() + training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(num_generated_frames, + dtype=torch.float32, + device=self.device) + training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(min_num_frames, + dtype=torch.float32, + device=self.device) + + # Clean up caches + assert self.kv_cache1 is not None + assert self.crossattn_cache is not None + self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse, self.kv_cache_keyboard) + + return final_output if gradient_mask is not None else pred_image_or_video + + 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) + + if not hasattr(self.validation_pipeline, "prompt_encoding_stage"): + # validation expects a prompt stage + def _prompt_encoding_stage(batch: ForwardBatch, _args: TrainingArgs) -> ForwardBatch: + if not batch.prompt_embeds: + batch.prompt_embeds = [None] + if batch.prompt_attention_mask is None or not batch.prompt_attention_mask: + batch.prompt_attention_mask = [None] + return batch + + self.validation_pipeline.prompt_encoding_stage = _prompt_encoding_stage # type: ignore[attr-defined] + + 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 + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + 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'] + + batch_size = clip_feature.shape[0] + vae_config = self.training_args.pipeline_config.vae_config.arch_config + num_channels = vae_config.z_dim + spatial_compression_ratio = vae_config.spatial_compression_ratio + + latent_height = self.training_args.num_height // spatial_compression_ratio + latent_width = self.training_args.num_width // spatial_compression_ratio + + latents = torch.randn(batch_size, num_channels, self.training_args.num_latent_t, latent_height, + latent_width).to(get_local_torch_device(), dtype=torch.bfloat16) + + 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.image_embeds = clip_feature.to(get_local_torch_device(), dtype=torch.bfloat16) + training_batch.image_latents = first_frame_latent.to(get_local_torch_device(), dtype=torch.bfloat16) + # Action conditioning + if keyboard_cond is not None and keyboard_cond.numel() > 0: + keyboard_cond_full = keyboard_cond.to(get_local_torch_device(), dtype=torch.bfloat16) + training_batch.keyboard_cond = keyboard_cond_full # For Teacher/Critic (dim=6) + else: + training_batch.keyboard_cond = None + if mouse_cond is not None and mouse_cond.numel() > 0: + training_batch.mouse_cond = mouse_cond.to(get_local_torch_device(), dtype=torch.bfloat16) + else: + training_batch.mouse_cond = None + training_batch.infos = infos + return training_batch + + 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) + + # cond_concat = [mask(4), image_latent(16)] with 20 channels. + expected_cond_channels = 20 + if image_latents.shape[1] == 16: + image_latents = self._build_matrixgame_cond_concat(image_latents) + elif image_latents.shape[1] != expected_cond_channels: + raise ValueError("Unexpected first_frame_latent channels, " + "Expected {expected_cond_channels} (cond_concat), got {image_latents.shape[1]}.") + + if self.sp_world_size > 1: + total_frames = image_latents.shape[2] + # Split cond latents to local SP shard only when tensor is still global. + if total_frames == self.training_args.num_latent_t: + if total_frames % self.sp_world_size != 0: + raise ValueError("image_latents temporal dim is not divisible by SP world size: " + f"frames={total_frames}, sp_world_size={self.sp_world_size}") + image_latents = rearrange(image_latents, "b c (n t) h w -> b c n t h w", + n=self.sp_world_size).contiguous() + image_latents = image_latents[:, :, self.rank_in_sp_group, :, :, :] + + training_batch.image_latents = image_latents + + return training_batch + + def _build_distill_input_kwargs(self, + noise_input: torch.Tensor, + timestep: torch.Tensor, + text_dict: dict[str, torch.Tensor] | None, + training_batch: TrainingBatch, + frame_start: int | None = None, + frame_end: int | None = None, + num_frame_per_block: int | None = None) -> TrainingBatch: + # Image Embeds for conditioning + image_embeds = training_batch.image_embeds + assert image_embeds is not None + assert torch.isnan(image_embeds).sum() == 0 + image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16) + image_embeds = [image_embeds] + + image_latents = training_batch.image_latents + if frame_start is not None and frame_end is not None: + image_latents = image_latents[:, :, frame_start:frame_end, :, :] + + vae_temporal_compression_ratio = 4 + if frame_end is not None: + action_frame_end = (frame_end - 1) * vae_temporal_compression_ratio + 1 + keyboard_cond_sliced = training_batch.keyboard_cond[:, : + action_frame_end, :] if training_batch.keyboard_cond is not None else None + mouse_cond_sliced = training_batch.mouse_cond[:, : + action_frame_end, :] if training_batch.mouse_cond is not None else None + else: + keyboard_cond_sliced = training_batch.keyboard_cond + mouse_cond_sliced = training_batch.mouse_cond + + noisy_model_input = torch.cat([noise_input, image_latents.permute(0, 2, 1, 3, 4)], dim=2) + + training_batch.input_kwargs = { + "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4), # bs, c, t, h, w + "encoder_hidden_states": None, + "timestep": timestep, + "encoder_hidden_states_image": image_embeds, + "keyboard_cond": keyboard_cond_sliced, + "mouse_cond": mouse_cond_sliced, + "num_frame_per_block": num_frame_per_block if num_frame_per_block is not None else self.num_frame_per_block, + } + training_batch.noise_latents = noise_input + + return training_batch + + def _dmd_forward(self, generator_pred_video: torch.Tensor, training_batch: TrainingBatch) -> torch.Tensor: + """Compute DMD (Diffusion Model Distillation) loss for MatrixGame.""" + original_latent = generator_pred_video + with torch.no_grad(): + timestep = torch.randint(0, self.num_train_timestep, [1], device=self.device, dtype=torch.long) + + timestep = shift_timestep(timestep, self.timestep_shift, self.num_train_timestep) + + timestep = timestep.clamp(self.min_timestep, self.max_timestep) + + noise = torch.randn(self.video_latent_shape, device=self.device, dtype=generator_pred_video.dtype) + + noisy_latent = self.noise_scheduler.add_noise(generator_pred_video.flatten(0, 1), noise.flatten( + 0, 1), timestep).detach().unflatten(0, (generator_pred_video.shape[0], generator_pred_video.shape[1])) + + # Non-causal models expect 1D timestep (batch_size,) + critic_timestep = timestep.expand(noisy_latent.shape[0]) + + self._build_distill_input_kwargs(noisy_latent, critic_timestep, None, training_batch) + + # fake_score_transformer forward + current_fake_score_transformer = self._get_fake_score_transformer(timestep) + fake_score_pred_noise = current_fake_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + + faker_score_pred_video = pred_noise_to_pred_video(pred_noise=fake_score_pred_noise.flatten(0, 1), + noise_input_latent=noisy_latent.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, fake_score_pred_noise.shape[:2]) + + # real_score_transformer forward + current_real_score_transformer = self._get_real_score_transformer(timestep) + real_score_pred_noise = current_real_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + + real_score_pred_video = pred_noise_to_pred_video(pred_noise=real_score_pred_noise.flatten(0, 1), + noise_input_latent=noisy_latent.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler).unflatten( + 0, real_score_pred_noise.shape[:2]) + + # No CFG for MatrixGame - use real_score_pred_video directly + grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(original_latent - + real_score_pred_video).mean() + grad = torch.nan_to_num(grad) + + dmd_loss = 0.5 * F.mse_loss(original_latent.float(), (original_latent.float() - grad.float()).detach()) + + training_batch.dmd_latent_vis_dict.update({ + "training_batch_dmd_fwd_clean_latent": training_batch.latents, + "generator_pred_video": original_latent.detach(), + "real_score_pred_video": real_score_pred_video.detach(), + "faker_score_pred_video": faker_score_pred_video.detach(), + "dmd_timestep": timestep.detach(), + }) + + return dmd_loss + + def faker_score_forward(self, training_batch: TrainingBatch) -> tuple[TrainingBatch, torch.Tensor]: + """Forward pass for critic training with MatrixGame action conditioning.""" + with torch.no_grad(), set_forward_context(current_timestep=training_batch.timesteps, + attn_metadata=training_batch.attn_metadata_vsa): + if self.training_args.simulate_generator_forward: + generator_pred_video = self._generator_multi_step_simulation_forward(training_batch) + else: + generator_pred_video = self._generator_forward(training_batch) + + fake_score_timestep = torch.randint(0, self.num_train_timestep, [1], device=self.device, dtype=torch.long) + + fake_score_timestep = shift_timestep(fake_score_timestep, self.timestep_shift, self.num_train_timestep) + + fake_score_timestep = fake_score_timestep.clamp(self.min_timestep, self.max_timestep) + + fake_score_noise = torch.randn(self.video_latent_shape, device=self.device, dtype=generator_pred_video.dtype) + + noisy_generator_pred_video = self.noise_scheduler.add_noise( + generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1), + fake_score_timestep).unflatten(0, (generator_pred_video.shape[0], generator_pred_video.shape[1])) + + # Non-causal critic expects 1D timestep (batch_size,), not 2D (batch_size, num_frames). + expanded_fake_score_timestep = fake_score_timestep.expand(noisy_generator_pred_video.shape[0]) + + self._build_distill_input_kwargs(noisy_generator_pred_video, expanded_fake_score_timestep, None, training_batch) + + with set_forward_context(current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata): + current_fake_score_transformer = self._get_fake_score_transformer(fake_score_timestep) + fake_score_pred_noise = current_fake_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + + target = fake_score_noise - generator_pred_video + flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2) + + training_batch.fake_score_latent_vis_dict = { + "training_batch_fakerscore_fwd_clean_latent": training_batch.latents, + "generator_pred_video": generator_pred_video, + "fake_score_timestep": fake_score_timestep, + } + + return training_batch, flow_matching_loss + + 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 self-forcing distillation pipeline...") + + pipeline = MatrixGameSelfForcingDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args) + + args = pipeline.training_args + pipeline.train() + logger.info("MatrixGame self-forcing distillation pipeline completed") + + +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() + main(args) diff --git a/fastvideo/training/trackers.py b/fastvideo/training/trackers.py index fbb29d99c..f1ccfaa42 100644 --- a/fastvideo/training/trackers.py +++ b/fastvideo/training/trackers.py @@ -188,6 +188,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, diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 8e4b3ae25..575d6dc5d 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -35,6 +35,7 @@ from fastvideo.distributed import (cleanup_dist_env_and_memory, get_local_torch_ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger +from fastvideo.models.vision_utils import load_video from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch, LoRAPipeline, TrainingBatch) from fastvideo.platforms import current_platform from fastvideo.training.activation_checkpoint import (apply_activation_checkpointing) @@ -80,6 +81,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") @@ -741,6 +743,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] = [] @@ -756,6 +759,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) @@ -789,6 +793,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 @@ -797,11 +802,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) @@ -830,10 +837,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 ref_filename, caption in zip(all_ref_videos, all_captions, strict=True): + if ref_filename is None: + continue + ref_frames = np.stack([np.asarray(frame) for frame in load_video(ref_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)