diff --git a/examples/inference/basic/basic_matrixgame.py b/examples/inference/basic/basic_matrixgame.py index b88fb5fcc..57be0e10e 100644 --- a/examples/inference/basic/basic_matrixgame.py +++ b/examples/inference/basic/basic_matrixgame.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -from fastvideo.configs.pipelines.wan import MatrixGameI2V480PConfig -from fastvideo.models.dits.matrix_game.utils import create_action_presets +from fastvideo.models.dits.matrixgame.utils import create_action_presets import torch diff --git a/examples/inference/basic/basic_matrixgame_streaming.py b/examples/inference/basic/basic_matrixgame_streaming.py index ea3f2abd4..40dac0b9a 100644 --- a/examples/inference/basic/basic_matrixgame_streaming.py +++ b/examples/inference/basic/basic_matrixgame_streaming.py @@ -1,5 +1,5 @@ from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator -from fastvideo.models.dits.matrix_game.utils import get_current_action_async, expand_action_to_frames +from fastvideo.models.dits.matrixgame.utils import get_current_action_async, expand_action_to_frames import torch import asyncio diff --git a/examples/inference/gradio/local/gradio_local_demo_matrixgame.py b/examples/inference/gradio/local/gradio_local_demo_matrixgame.py index ae8333a04..9c177f2ab 100644 --- a/examples/inference/gradio/local/gradio_local_demo_matrixgame.py +++ b/examples/inference/gradio/local/gradio_local_demo_matrixgame.py @@ -10,7 +10,7 @@ from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse, FileResponse from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator -from fastvideo.models.dits.matrix_game.utils import expand_action_to_frames +from fastvideo.models.dits.matrixgame.utils import expand_action_to_frames VARIANT_CONFIG = { diff --git a/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh b/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh new file mode 100644 index 000000000..ce00f73f3 --- /dev/null +++ b/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh @@ -0,0 +1,93 @@ +#!/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" + --max_train_steps 1500 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 4 + --num_latent_t 20 + --num_height 352 + --num_width 640 + --num_frames 77 + --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_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh b/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh new file mode 100644 index 000000000..7db267adb --- /dev/null +++ b/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Matrix-Game-2.0-Foundation-Diffusers" +DATA_MERGE_PATH="footsies-dataset/merge.txt" +OUTPUT_DIR="footsies-dataset/preprocessed/" + +# export CUDA_VISIBLE_DEVICES=0 +export MASTER_ADDR=localhost +export MASTER_PORT=29500 +export RANK=0 +export WORLD_SIZE=1 + +python fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 4 \ + --seed 42 \ + --max_height 352 \ + --max_width 640 \ + --num_frames 77 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --samples_per_file 4 \ + --train_fps 25 \ + --flush_frequency 4 \ + --preprocess_task matrixgame diff --git a/examples/training/finetune/MatrixGame2.0/validation.json b/examples/training/finetune/MatrixGame2.0/validation.json new file mode 100644 index 000000000..b3a1a5326 --- /dev/null +++ b/examples/training/finetune/MatrixGame2.0/validation.json @@ -0,0 +1,5020 @@ +{ + "data": [ + { + "caption": "all_left", + "image_path": "../../../../footsies-dataset/validate/0.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "all_right", + "image_path": "../../../../footsies-dataset/validate/0.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "all_idle", + "image_path": "../../../../footsies-dataset/validate/0.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "alternating_left_idle_right", + "image_path": "../../../../footsies-dataset/validate/0.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "all_left_with_left_position", + "image_path": "../../../../footsies-dataset/validate/1.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "all_right_with_right_position", + "image_path": "../../../../footsies-dataset/validate/2.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "alternating_left_right_with_1:2 ratio", + "image_path": "../../../../footsies-dataset/validate/1.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ] + ] + }, + { + "caption": "alternating_left_right_with_2:1 ratio", + "image_path": "../../../../footsies-dataset/validate/1.png", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77, + "keyboard_cond": [ + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 1, + 0, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ], + [ + 0, + 0, + 1, + 0, + 0, + 0 + ] + ] + } + ] +} \ No newline at end of file diff --git a/fastvideo/configs/models/dits/matrixgame.py b/fastvideo/configs/models/dits/matrixgame.py index f0f7e0e24..4b39f1601 100644 --- a/fastvideo/configs/models/dits/matrixgame.py +++ b/fastvideo/configs/models/dits/matrixgame.py @@ -8,8 +8,8 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig): # because MatrixGame checkpoints already have patch_embedding.proj format param_names_mapping: dict = field( default_factory=lambda: { - # Removed: r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1" - # because checkpoint already has correct format + r"^patch_embedding\.(?!proj\.)(.*)$": + r"patch_embedding.proj.\1", r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 55d917fb0..434839dcc 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -196,6 +196,12 @@ class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config): # ============================================= # ============= Matrix Game =================== # ============================================= +@dataclass +class MatrixGameBaseI2V480PConfig(WanI2V480PConfig): + dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig) + flow_shift: float | None = 5.0 + + @dataclass class MatrixGameI2V480PConfig(WanI2V480PConfig): dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig) diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index 02ad63a3f..048c7686c 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -116,3 +116,44 @@ pyarrow_schema_text_only = pa.schema([ # --- Metadata --- pa.field("caption", pa.string()), ]) + + +pyarrow_schema_matrixgame = pa.schema([ + pa.field("id", pa.string()), + # --- Image/Video VAE latents --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("vae_latent_bytes", pa.binary()), + # e.g., [C, T, H, W] or [C, H, W] + pa.field("vae_latent_shape", pa.list_(pa.int64())), + # e.g., 'float32' + pa.field("vae_latent_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()), + # --- 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 Validation + pa.field("pil_image_bytes", pa.binary()), + pa.field("pil_image_shape", pa.list_(pa.int64())), + pa.field("pil_image_dtype", pa.string()), + # --- Metadata --- + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), # 'image' or 'video' + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + # -- Video-specific (can be null/default for images) --- + # Number of frames processed (e.g., 1 for image, N for video) + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), +]) diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 48f8cf29d..d94727bd6 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -40,6 +40,7 @@ class PreprocessBatch: num_frames: int | None = None sample_frame_index: list[int] | None = None sample_num_frames: int | None = None + action_path: str | None = None # Processed data pixel_values: torch.Tensor | None = None @@ -470,8 +471,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Initialize tokenizer tokenizer_path = os.path.join(args.model_path, "tokenizer") - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, - cache_dir=args.cache_dir) + if os.path.exists(tokenizer_path): + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, + cache_dir=args.cache_dir) + else: + tokenizer = None # Initialize processing stages self._init_stages(args, transform, transform_topcrop, tokenizer) @@ -495,11 +499,14 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, self.video_transform_stage = VideoTransformStage(transform) self.image_transform_stage = ImageTransformStage( transform, transform_topcrop) - self.text_encoding_stage = TextEncodingStage( - tokenizer=tokenizer, - text_max_length=args.text_max_length, - cfg_rate=args.training_cfg_rate, - seed=self.seed) + if tokenizer is not None: + self.text_encoding_stage = TextEncodingStage( + tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=args.training_cfg_rate, + seed=self.seed) + else: + self.text_encoding_stage = None def _load_raw_data(self) -> list[dict]: """Load raw data from JSON files.""" @@ -521,6 +528,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Update paths with folder prefix for item in data_items: item["path"] = opj(folder, item["path"]) + if "action_path" in item and item["action_path"]: + item["action_path"] = opj(folder, item["action_path"]) return data_items @@ -542,7 +551,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, cap=item["cap"], resolution=item.get("resolution"), fps=item.get("fps"), - duration=item.get("duration")) + duration=item.get("duration"), + action_path=item.get("action_path")) # Apply filtering stages if not self._apply_filter_stages(batch, filter_counts): @@ -604,20 +614,27 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Apply transformation stages batch = self.video_transform_stage.process(batch) batch = self.image_transform_stage.process(batch) - batch = self.text_encoding_stage.process(batch) + if self.text_encoding_stage is not None: + batch = self.text_encoding_stage.process(batch) # Build result dictionary result = { "pixel_values": batch.pixel_values, - "text": batch.text, - "input_ids": batch.input_ids, - "cond_mask": batch.cond_mask, "path": batch.path, } + if batch.text is not None: + result["text"] = batch.text + result["input_ids"] = batch.input_ids + result["cond_mask"] = batch.cond_mask + # Add video-specific fields if batch.is_video: result.update({"fps": batch.fps, "duration": batch.duration}) + + # Add action_path + if batch.action_path: + result["action_path"] = batch.action_path return result diff --git a/fastvideo/models/dits/matrix_game/__init__.py b/fastvideo/models/dits/matrixgame/__init__.py similarity index 100% rename from fastvideo/models/dits/matrix_game/__init__.py rename to fastvideo/models/dits/matrixgame/__init__.py diff --git a/fastvideo/models/dits/matrix_game/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py similarity index 95% rename from fastvideo/models/dits/matrix_game/action_module.py rename to fastvideo/models/dits/matrixgame/action_module.py index ae4d01b4b..6c06babae 100644 --- a/fastvideo/models/dits/matrix_game/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -257,6 +257,17 @@ class ActionModule(nn.Module): ''' assert use_rope_keyboard + target_device = x.device + target_dtype = x.dtype + if mouse_condition is not None: + mouse_condition = mouse_condition.to(device=target_device, + dtype=target_dtype) + if keyboard_condition is not None: + keyboard_condition = keyboard_condition.to( + device=target_device, dtype=target_dtype) + else: + return x + B, N_frames, C = keyboard_condition.shape assert tt*th*tw == x.shape[1] assert ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 @@ -272,7 +283,9 @@ class ActionModule(nn.Module): # Defined freqs_cis early so it's available for both mouse and keyboard freqs_cis = (self._freqs_cos, self._freqs_sin) - assert (N_feats == tt and ((is_causal and kv_cache_mouse is None) or not is_causal)) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block and is_causal) + if is_causal: + assert (N_feats == tt and kv_cache_mouse is None) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block) + # For non-causal (training), we trust that the caller provides correctly shaped inputs if self.enable_mouse and mouse_condition is not None: hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th*tw) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 @@ -289,13 +302,15 @@ class ActionModule(nn.Module): mouse_condition = mouse_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)] else: - group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)] + local_num_frames = tt + group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)] group_mouse = torch.stack(group_mouse, dim = 1) + actual_num_frames = group_mouse.shape[1] # Use actual stacked frame count S = th * tw - group_mouse = group_mouse.unsqueeze(-1).expand(B, num_frame_per_block, pad_t, C, S) - group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, num_frame_per_block, pad_t * C) + group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C) group_mouse = torch.cat([hidden_states, group_mouse], dim = -1) group_mouse = self.mouse_mlp(group_mouse) @@ -314,7 +329,7 @@ class ActionModule(nn.Module): ## TODO: adding cache here if is_causal: if kv_cache_mouse is None: - assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 # == 880, f"{q.shape[0]},{k.shape[0]}" + assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 # == 880, f"{q.shape[0]},{k.shape[0]}" padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] padded_q = torch.cat( [q, @@ -395,7 +410,8 @@ class ActionModule(nn.Module): group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)] else: keyboard_condition = self.keyboard_embed(keyboard_condition) - group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)] + local_num_frames = tt + group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)] group_keyboard = torch.stack(group_keyboard, dim = 1) # B F RW C group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1)) # apply cross attn @@ -415,7 +431,7 @@ class ActionModule(nn.Module): q = self.key_attn_q_norm(q).to(v) k = self.key_attn_k_norm(k).to(v) S = th * tw - assert S == 880 + # assert S == 880 # position embed if use_rope_keyboard: B, TS, H, D = q.shape @@ -424,13 +440,13 @@ class ActionModule(nn.Module): q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) k1, k2, k3, k4 = k.shape - k = k.expand(S, k2, k3, k4) - v = v.expand(S, k2, k3, k4) + k = k.repeat_interleave(S, dim=0) + v = v.repeat_interleave(S, dim=0) if is_causal: if kv_cache_keyboard is None: - assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 + assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] padded_q = torch.cat( @@ -480,7 +496,7 @@ class ActionModule(nn.Module): else: local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() local_start_index = local_end_index - num_new_tokens - assert k.shape[0] == 880 # BS == 1 or the cache should not be saved/ load method should be modified + assert k.shape[0] == S # BS == 1 or the cache should not be saved/ load method should be modified kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1] kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1] diff --git a/fastvideo/models/dits/matrix_game/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py similarity index 99% rename from fastvideo/models/dits/matrix_game/causal_model.py rename to fastvideo/models/dits/matrixgame/causal_model.py index fc859e5a8..13806d2ac 100644 --- a/fastvideo/models/dits/matrix_game/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -1096,4 +1096,4 @@ class CausalMatrixGameWanModel(BaseDiT): if block.action_model.proj_keyboard.bias is not None: nn.init.zeros_(block.action_model.proj_keyboard.bias) except AttributeError: - pass + pass \ No newline at end of file diff --git a/fastvideo/models/dits/matrix_game/model.py b/fastvideo/models/dits/matrixgame/model.py similarity index 97% rename from fastvideo/models/dits/matrix_game/model.py rename to fastvideo/models/dits/matrixgame/model.py index 30891164d..201d03015 100644 --- a/fastvideo/models/dits/matrix_game/model.py +++ b/fastvideo/models/dits/matrixgame/model.py @@ -61,14 +61,17 @@ class MatrixGameTimeImageEmbedding(nn.Module): timestep_proj = self.time_modulation(temb) # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states - # and return None for the text embedding part if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( encoder_hidden_states_image) - return temb, timestep_proj, None, encoder_hidden_states_image + encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image class MatrixGameCrossAttention(WanSelfAttention): @@ -313,6 +316,7 @@ class MatrixGameWanModel(BaseDiT): inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads + self.num_channels_latents = config.out_channels self.patch_size = config.patch_size # 1. Patch & position embedding @@ -395,7 +399,8 @@ class MatrixGameWanModel(BaseDiT): self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) + rope_theta=10000, + do_sp_sharding=True) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = (freqs_cos.float(), diff --git a/fastvideo/models/dits/matrix_game/utils.py b/fastvideo/models/dits/matrixgame/utils.py similarity index 97% rename from fastvideo/models/dits/matrix_game/utils.py rename to fastvideo/models/dits/matrixgame/utils.py index 7f365eed7..4dd937699 100644 --- a/fastvideo/models/dits/matrix_game/utils.py +++ b/fastvideo/models/dits/matrixgame/utils.py @@ -157,8 +157,8 @@ def load_initial_image(image_path: str = None) -> Image.Image: return Image.new("RGB", (640, 352), (128, 128, 128)) def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None): - if keyboard_dim not in (2, 4, 7): - raise ValueError(f"keyboard_dim must be 2, 4, or 7, got {keyboard_dim}") + if keyboard_dim not in (2, 4, 6, 7): + raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}") if num_frames % 4 != 1: raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.") @@ -181,6 +181,11 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No actions_double_action = [] actions_single_camera = ["camera_l", "camera_r"] keyboard_idx = {"forward": 0, "back": 1} + elif keyboard_dim == 6: + actions_single_action = ["forward", "back", "left", "right"] + actions_double_action = ["forward_left", "forward_right"] + actions_single_camera = ["camera_l", "camera_r"] + keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5} else: # keyboard_dim == 7 # Temple Run model: still, w, s, left, right, a, d (no mouse) actions_single_action = ["forward", "back", "left", "right"] diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 67a2b6ede..ef98576bd 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -76,7 +76,7 @@ class WanTimeTextImageEmbedding(nn.Module): dim, dim, bias=True, - act_type="gelu_pytorch_tanh") + act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None self.image_embedder = None if image_embed_dim is not None: @@ -92,7 +92,10 @@ class WanTimeTextImageEmbedding(nn.Module): temb = self.time_embedder(timestep, timestep_seq_len) timestep_proj = self.time_modulation(temb) - encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if self.text_embedder is not None: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + else: + encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), device=temb.device, dtype=temb.dtype) if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( @@ -179,7 +182,10 @@ class WanT2VCrossAttention(WanSelfAttention): v = self.to_v(context)[0].view(b, -1, n, d) # compute attention - x = self.attn(q, k, v) + if k.size(1) > 0: + x = self.attn(q, k, v) + else: + x = torch.zeros_like(q) # output x = x.flatten(2) @@ -227,7 +233,10 @@ class WanI2VCrossAttention(WanSelfAttention): v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) img_x = self.attn(q, k_img, v_img) # compute attention - x = self.attn(q, k, v) + if k.size(1) > 0: + x = self.attn(q, k, v) + else: + x = torch.zeros_like(q) # output x = x.flatten(2) @@ -633,7 +642,7 @@ class WanTransformer3DModel(CachableDiT): enable_teacache = forward_batch is not None and forward_batch.enable_teacache orig_dtype = hidden_states.dtype - if not isinstance(encoder_hidden_states, torch.Tensor): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: @@ -705,8 +714,11 @@ class WanTransformer3DModel(CachableDiT): timestep_proj = timestep_proj.unflatten(1, (6, -1)) if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + if encoder_hidden_states is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + else: + encoder_hidden_states = encoder_hidden_states_image if current_platform.is_mps() or current_platform.is_npu(): encoder_hidden_states = encoder_hidden_states.to(orig_dtype) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 2ca789307..954b96049 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -42,8 +42,8 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), - "MatrixGameWanModel": ("dits", "matrix_game", "MatrixGameWanModel"), - "CausalMatrixGameWanModel": ("dits", "matrix_game", "CausalMatrixGameWanModel"), + "MatrixGameWanModel": ("dits", "matrixgame", "MatrixGameWanModel"), + "CausalMatrixGameWanModel": ("dits", "matrixgame", "CausalMatrixGameWanModel"), } _TEXT_ENCODER_MODELS = { diff --git a/fastvideo/pipelines/preprocess/matrixgame/__init__.py b/fastvideo/pipelines/preprocess/matrixgame/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py new file mode 100644 index 000000000..5079ec264 --- /dev/null +++ b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines.preprocess.preprocess_pipeline_base import ( + BasePreprocessPipeline) +from fastvideo.pipelines.stages import ImageEncodingStage + + +class PreprocessPipeline_MatrixGame(BasePreprocessPipeline): + """I2V preprocessing pipeline implementation.""" + + _required_config_modules = ["vae", "image_encoder", "image_processor"] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + self.add_stage(stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + def get_pyarrow_schema(self): + """Return the PyArrow schema for I2V pipeline.""" + return pyarrow_schema_matrixgame + + 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 + """Get VAE 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) + + latent_condition = encoder_outputs.mean + if (hasattr(self.get_module("vae"), "shift_factor") + and self.get_module("vae").shift_factor is not None): + if isinstance(self.get_module("vae").shift_factor, torch.Tensor): + latent_condition -= self.get_module("vae").shift_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition -= self.get_module("vae").shift_factor + + if isinstance(self.get_module("vae").scaling_factor, torch.Tensor): + latent_condition = latent_condition * self.get_module( + "vae").scaling_factor.to(latent_condition.device, + latent_condition.dtype) + else: + latent_condition = latent_condition * self.get_module( + "vae").scaling_factor + + # mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, + # latent_width) + # mask_lat_size[:, :, list(range(1, num_frames))] = 0 + # first_frame_mask = mask_lat_size[:, :, 0:1] + # first_frame_mask = torch.repeat_interleave( + # first_frame_mask, + # dim=2, + # repeats=self.get_module("vae").temporal_compression_ratio) + # mask_lat_size = torch.concat( + # [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + # mask_lat_size = mask_lat_size.view( + # batch_size, -1, + # self.get_module("vae").temporal_compression_ratio, latent_height, + # latent_width) + # mask_lat_size = mask_lat_size.transpose(1, 2) + # mask_lat_size = mask_lat_size.to(latent_condition.device) + + # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1) + + features["first_frame_latent"] = latent_condition + + if "action_path" in valid_data and valid_data["action_path"]: + keyboard_cond_list = [] + mouse_cond_list = [] + num_bits = 6 + 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_raw = action_dict["keyboard"] + # Convert 1D bit-flag values to 2D multi-hot encoding + if isinstance(keyboard_raw, np.ndarray): + if keyboard_raw.ndim == 1: + # [T] -> [T, num_bits] + T = len(keyboard_raw) + multi_hot = np.zeros((T, num_bits), + dtype=np.float32) + action_values = keyboard_raw.astype(int) + for bit_idx in range(num_bits): + target_idx = ( + 2 - + (bit_idx % 3)) + 3 * (bit_idx // 3) + if target_idx < num_bits: + multi_hot[:, target_idx] = ( + (action_values >> bit_idx) + & 1).astype(np.float32) + keyboard_cond_list.append(multi_hot) + else: + # If already 2D, pad to num_bits if necessary + k_data = keyboard_raw.astype(np.float32) + if k_data.ndim == 2 and k_data.shape[ + -1] < num_bits: + padding = np.zeros( + (k_data.shape[0], + num_bits - k_data.shape[-1]), + dtype=np.float32) + k_data = np.concatenate( + [k_data, padding], axis=-1) + keyboard_cond_list.append(k_data) + else: + keyboard_cond_list.append(keyboard_raw) + if "mouse" in action_dict: + mouse_cond_list.append(action_dict["mouse"]) + else: + if isinstance(action_data, + np.ndarray) and action_data.ndim == 1: + T = len(action_data) + multi_hot = np.zeros((T, num_bits), + dtype=np.float32) + action_values = action_data.astype(int) + for bit_idx in range(num_bits): + target_idx = ( + 2 - (bit_idx % 3)) + 3 * (bit_idx // 3) + if target_idx < num_bits: + multi_hot[:, target_idx] = ( + (action_values >> bit_idx) & 1).astype( + np.float32) + keyboard_cond_list.append(multi_hot) + else: + # If already 2D, pad to num_bits if necessary + k_data = action_data.astype(np.float32) + if k_data.ndim == 2 and k_data.shape[-1] < num_bits: + padding = np.zeros( + (k_data.shape[0], + num_bits - k_data.shape[-1]), + dtype=np.float32) + k_data = np.concatenate([k_data, padding], + axis=-1) + keyboard_cond_list.append(k_data) + if keyboard_cond_list: + features["keyboard_cond"] = keyboard_cond_list + if mouse_cond_list: + features["mouse_cond"] = mouse_cond_list + + return features + + def create_record( + self, + video_name: str, + vae_latent: np.ndarray, + text_embedding: np.ndarray, + valid_data: dict[str, Any], + idx: int, + extra_features: dict[str, Any] | None = None) -> dict[str, Any]: + """Create a record for the Parquet dataset with CLIP features.""" + record = super().create_record(video_name=video_name, + vae_latent=vae_latent, + text_embedding=text_embedding, + valid_data=valid_data, + idx=idx, + extra_features=extra_features) + + if extra_features and "clip_feature" in extra_features: + clip_feature = extra_features["clip_feature"] + record.update({ + "clip_feature_bytes": clip_feature.tobytes(), + "clip_feature_shape": list(clip_feature.shape), + "clip_feature_dtype": str(clip_feature.dtype), + }) + else: + record.update({ + "clip_feature_bytes": b"", + "clip_feature_shape": [], + "clip_feature_dtype": "", + }) + + if extra_features and "first_frame_latent" in extra_features: + first_frame_latent = extra_features["first_frame_latent"] + 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), + }) + else: + record.update({ + "first_frame_latent_bytes": b"", + "first_frame_latent_shape": [], + "first_frame_latent_dtype": "", + }) + + if extra_features and "pil_image" in extra_features: + pil_image = extra_features["pil_image"] + 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": "", + }) + + if extra_features and "keyboard_cond" in extra_features: + keyboard_cond = extra_features["keyboard_cond"] + 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 extra_features and "mouse_cond" in extra_features: + mouse_cond = extra_features["mouse_cond"] + 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": "", + }) + + return record + + +EntryClass = PreprocessPipeline_MatrixGame diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 78e169315..083277bc2 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -320,12 +320,18 @@ class BasePreprocessPipeline(ComposedPipelineBase): "pixel_values": torch.stack( [data["pixel_values"][i] for i in valid_indices]), - "text": [data["text"][i] for i in valid_indices], + "text": [data["text"][i] for i in valid_indices] + if "text" in data else ["" for _ in valid_indices], "path": [data["path"][i] for i in valid_indices], "fps": [data["fps"][i] for i in valid_indices], "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 + ] + # VAE with torch.autocast("cuda", dtype=torch.float32): latents = self.get_module("vae").encode( @@ -343,28 +349,35 @@ class BasePreprocessPipeline(ComposedPipelineBase): prompt_embeds=[], prompt_attention_mask=[], ) - assert hasattr(self, "prompt_encoding_stage") - result_batch = self.prompt_encoding_stage(batch, fastvideo_args) - prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[ - 0], result_batch.prompt_attention_mask[0] - assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0] - # Get sequence lengths from attention masks (number of 1s) - seq_lens = prompt_attention_mask.sum(dim=1) + if hasattr(self, "prompt_encoding_stage"): + result_batch = self.prompt_encoding_stage( + batch, fastvideo_args) + prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[ + 0], result_batch.prompt_attention_mask[0] + assert prompt_embeds.shape[ + 0] == prompt_attention_mask.shape[0] - non_padded_embeds = [] - non_padded_masks = [] + # Get sequence lengths from attention masks (number of 1s) + seq_lens = prompt_attention_mask.sum(dim=1) - # Process each item in the batch - for i in range(prompt_embeds.size(0)): - seq_len = seq_lens[i].item() - # Slice the embeddings and masks to keep only non-padding parts - non_padded_embeds.append(prompt_embeds[i, :seq_len]) - non_padded_masks.append(prompt_attention_mask[i, :seq_len]) + non_padded_embeds = [] + non_padded_masks = [] - # Update the tensors with non-padded versions - prompt_embeds = non_padded_embeds - prompt_attention_mask = non_padded_masks + # Process each item in the batch + for i in range(prompt_embeds.size(0)): + seq_len = seq_lens[i].item() + # Slice the embeddings and masks to keep only non-padding parts + non_padded_embeds.append(prompt_embeds[i, :seq_len]) + non_padded_masks.append( + prompt_attention_mask[i, :seq_len]) + + # Update the tensors with non-padded versions + prompt_embeds = non_padded_embeds + prompt_attention_mask = non_padded_masks + else: + bs = len(valid_indices) + prompt_embeds = [torch.zeros(0) for _ in range(bs)] # Prepare batch data for Parquet dataset batch_data = [] diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 058cb2e5d..18455d70f 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -16,6 +16,8 @@ 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.utils import maybe_download_model logger = init_logger(__name__) @@ -60,9 +62,12 @@ def main(args) -> None: assert args.flow_shift is not None, "flow_shift is required for ode_trajectory" fastvideo_args.pipeline_config.flow_shift = args.flow_shift PreprocessPipeline = PreprocessPipeline_ODE_Trajectory + elif args.preprocess_task == "matrixgame": + PreprocessPipeline = PreprocessPipeline_MatrixGame else: - raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. " - f"Valid options: t2v, i2v, ode_trajectory, text_only") + raise ValueError( + f"Invalid preprocess task: {args.preprocess_task}. " + f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame") logger.info("Preprocess task: %s using %s", args.preprocess_task, PreprocessPipeline.__name__) @@ -106,11 +111,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"], - help="Type of preprocessing task to run") + 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("--train_fps", type=int, default=30) parser.add_argument("--use_image_num", type=int, default=0) parser.add_argument("--text_max_length", type=int, default=256) diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 089fd7192..2768668fc 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -210,9 +210,9 @@ class InputValidationStage(PipelineStage): f"keyboard_cond must have 3 dimensions (B, T, K), but got {batch.keyboard_cond.dim()}" ) keyboard_dim = batch.keyboard_cond.shape[-1] - if keyboard_dim not in {2, 4, 6, 7}: + if keyboard_dim not in {2, 3, 4, 6, 7}: raise ValueError( - f"keyboard_cond last dimension must be 2, 4, 6, or 7, but got {keyboard_dim}" + f"keyboard_cond last dimension must be 2, 3, 4, 6, or 7, but got {keyboard_dim}" ) logger.info( "Action control: keyboard_cond validated - shape %s (dim=%d)", diff --git a/fastvideo/tests/ssim/test_matrixgame_similarity.py b/fastvideo/tests/ssim/test_matrixgame_similarity.py index e8c394a8e..c720a922b 100644 --- a/fastvideo/tests/ssim/test_matrixgame_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame_similarity.py @@ -5,7 +5,7 @@ import torch import pytest from fastvideo import VideoGenerator -from fastvideo.models.dits.matrix_game.utils import create_action_presets +from fastvideo.models.dits.matrixgame.utils import create_action_presets from fastvideo.logger import init_logger from fastvideo.tests.utils import ( compute_video_ssim_torchvision, diff --git a/fastvideo/training/matrixgame_training_pipeline.py b/fastvideo/training/matrixgame_training_pipeline.py new file mode 100644 index 000000000..78abe2d8d --- /dev/null +++ b/fastvideo/training/matrixgame_training_pipeline.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any + +import torch + +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) +from fastvideo.pipelines.basic.matrixgame.matrixgame_i2v_pipeline import ( + MatrixGamePipeline) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.utils import is_vsa_available, shallow_asdict + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class MatrixGameTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Matrix-Game-2.0. + """ + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_matrixgame + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + args_copy.dit_cpu_offload = True + # args_copy.pipeline_config.vae_config.load_encoder = False + # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + self.validation_pipeline = MatrixGamePipeline.from_pretrained( + training_args.model_path, + args=None, + inference_mode=True, + loaded_modules={ + "transformer": self.get_module("transformer"), + }, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + dit_cpu_offload=True) + + 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) + # 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) + + latents = batch['vae_latent'] + latents = latents[:, :, :self.training_args.num_latent_t] + # encoder_hidden_states = batch['text_embedding'] + # encoder_attention_mask = batch['text_attention_mask'] + 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 + # MatrixGame doesn't use text encoder + 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 + + # Action conditioning + 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 + + 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) + + temporal_compression_ratio = self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (self.training_args.num_latent_t - + 1) * temporal_compression_ratio + 1 + batch_size, num_channels, _, latent_height, latent_width = image_latents.shape + mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, + latent_width) + mask_lat_size[:, :, 1:] = 0 + + first_frame_mask = mask_lat_size[:, :, :1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, dim=2, repeats=temporal_compression_ratio) + mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], + dim=2) + mask_lat_size = mask_lat_size.view(batch_size, -1, + temporal_compression_ratio, + latent_height, latent_width) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to( + image_latents.device).to(dtype=torch.bfloat16) + + training_batch.noisy_model_input = torch.cat( + [training_batch.noisy_model_input, mask_lat_size, image_latents], + dim=1) + + return training_batch + + def _build_input_kwargs(self, + training_batch: TrainingBatch) -> TrainingBatch: + + # Image Embeds for conditioning + image_embeds = training_batch.image_embeds + assert torch.isnan(image_embeds).sum() == 0 + image_embeds = image_embeds.to(get_local_torch_device(), + dtype=torch.bfloat16) + encoder_hidden_states_image = image_embeds + + # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs + training_batch.input_kwargs = { + "hidden_states": + training_batch.noisy_model_input, + "encoder_hidden_states": + training_batch.encoder_hidden_states, # None for MatrixGame + "timestep": + training_batch.timesteps.to(get_local_torch_device(), + dtype=torch.bfloat16), + # "encoder_attention_mask": + # training_batch.encoder_attention_mask, + "encoder_hidden_states_image": + encoder_hidden_states_image, + # Action conditioning + "mouse_cond": + training_batch.mouse_cond, + "keyboard_cond": + training_batch.keyboard_cond, + "return_dict": + False, + } + 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 training pipeline...") + + pipeline = MatrixGameTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 8b0c8d262..58723241e 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -654,7 +654,11 @@ class TrainingPipeline(LoRAPipeline, ABC): seq_len = (training_batch.raw_latent_shape[2] // patch_t) * ( training_batch.raw_latent_shape[3] // patch_h) * (training_batch.raw_latent_shape[4] // patch_w) - context_len = int(training_batch.encoder_hidden_states.shape[1]) + if training_batch.encoder_hidden_states is not None: + context_len = int( + training_batch.encoder_hidden_states.shape[1]) + else: + context_len = 0 metrics["dit_seq_len"] = int(seq_len) metrics["context_len"] = context_len