[feat] FastVideo World Model Training (#1179)
Co-authored-by: mignonjia <mhuo@ucsd.edu> Co-authored-by: alexzms <3036648523@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Kaiqin Kong
co-authored by
mignonjia
alexzms
mergify[bot]
mignonjia
alexzms
mergify[bot]
parent
636d3b743e
commit
d6dfe95466
@@ -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)
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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[@]}"
|
||||
@@ -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[@]}"
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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[@]}"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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'
|
||||
])
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
+385
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user