diff --git a/examples/train/configs/dfsft_wangame_causal_v3.yaml b/examples/train/configs/dfsft_wangame_causal_v3.yaml new file mode 100644 index 000000000..f09cb6117 --- /dev/null +++ b/examples/train/configs/dfsft_wangame_causal_v3.yaml @@ -0,0 +1,95 @@ +# V3 config: WanGame causal Diffusion-Forcing SFT (DFSFT). +# +# Uses _target_-based instantiation — each model role is an independent +# class instance; the method class is resolved directly from the YAML. + +models: + student: + _target_: fastvideo.train.models.wangame.WanGameCausalModel + init_from: /mnt/weka/home/hao.zhang/kaiqin/wg_models/WanGame-2.1-0306-61500steps + trainable: true + transformer_override_safetensor: /mnt/weka/home/hao.zhang/mhuo/FastVideo-hyw/outputs/wangame_dfsft_causal_v3/checkpoint-best-step-36500/transformer/model.safetensors + +method: + _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod + attn_kind: dense + # use_ema: true + chunk_size: 3 + min_timestep_ratio: 0.02 + max_timestep_ratio: 0.98 + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + data: + data_path: >- + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_2130/preprocessed:1, + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_1600/preprocessed:0, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/0_static_plus_w_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/1_wasd_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/wasdonly_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/camera/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/camera4hold_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/preprocessed:3 + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 18 + num_height: 352 + num_width: 640 + num_frames: 69 + apply_bot_died_filter: true + + optimizer: + learning_rate: 1e-4 + betas: [0.9, 0.95] + weight_decay: 1e-5 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 2 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wangame_dfsft_causal_v3 + # resume_from_checkpoint: /mnt/weka/home/hao.zhang/mhuo/FastVideo-hyw/outputs/wangame_dfsft_causal_v3/checkpoint-best-step-36500 + training_state_checkpointing_steps: 5000 + weight_only_checkpointing_steps: 5000 + checkpoints_total_limit: 10 + best_checkpoint_start_step: 2000 + best_checkpoint_top_k: 5 + + tracker: + project_name: distillation_wangame_r + run_name: wangame_dfsft_causal_v3 + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + # ema: + # _target_: fastvideo.train.callbacks.ema.EMACallback + # beta: 0.9999 + validation: + _target_: fastvideo.train.callbacks.validation.ValidationCallback + pipeline_target: fastvideo.pipelines.basic.wan.wangame_causal_dmd_pipeline.WangameCausalOdeDMDPipeline + dataset_file: examples/training/finetune/WanGame2.1_1.3b_i2v/validation_8.json + every_steps: 500 + sampling_steps: [40] + scheduler_target: fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler + guidance_scale: 1.0 + num_frames: 69 + evaluate_ptlflow: false + +pipeline: + flow_shift: 3 diff --git a/examples/train/configs/self_forcing_wangame.yaml b/examples/train/configs/self_forcing_wangame.yaml new file mode 100644 index 000000000..2e7e8ff38 --- /dev/null +++ b/examples/train/configs/self_forcing_wangame.yaml @@ -0,0 +1,100 @@ +models: + student: + _target_: fastvideo.train.models.wangame.WanGameCausalModel + init_from: /mnt/weka/home/hao.zhang/kaiqin/wg_models/SFWanGame-2.1-0308-10000steps + trainable: true + teacher: + _target_: fastvideo.train.models.wangame.WanGameModel + init_from: /mnt/weka/home/hao.zhang/kaiqin/wg_models/WanGame-2.1-0223-9000steps + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.wangame.WanGameModel + init_from: /mnt/weka/home/hao.zhang/kaiqin/wg_models/WanGame-2.1-0223-9000steps + 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: 1.0 + 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: >- + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_2130/preprocessed:1, + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_1600/preprocessed:0, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/0_static_plus_w_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/1_wasd_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/wasdonly_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/camera/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/camera4hold_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/preprocessed:3 + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 18 + num_height: 352 + num_width: 640 + num_frames: 69 + apply_bot_died_filter: true + + optimizer: + learning_rate: 1e-5 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 6 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wangame_1.3b_self_forcing + training_state_checkpointing_steps: 500 + checkpoints_total_limit: 5 + + tracker: + project_name: wangame_sf + run_name: wangame_1.3b_self_forcing + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.wan.wangame_causal_dmd_pipeline.WangameCausalSdeDMDPipeline + dataset_file: examples/training/finetune/WanGame2.1_1.3b_i2v/validation_4.json + every_steps: 5 + sampling_steps: [4] + sampling_timesteps: [1000, 750, 500, 250] + num_frames: 69 + guidance_scale: 1.0 + evaluate_ptlflow: false + +pipeline: + flow_shift: 3 diff --git a/examples/train/configs/tfsft_wangame_causal.yaml b/examples/train/configs/tfsft_wangame_causal.yaml new file mode 100644 index 000000000..79f2c5819 --- /dev/null +++ b/examples/train/configs/tfsft_wangame_causal.yaml @@ -0,0 +1,94 @@ +# V3 config: WanGame causal Diffusion-Forcing SFT (DFSFT). +# +# Uses _target_-based instantiation — each model role is an independent +# class instance; the method class is resolved directly from the YAML. + +models: + student: + _target_: fastvideo.train.models.wangame.WanGameCausalModel + init_from: /mnt/weka/home/hao.zhang/kaiqin/wg_models/WanGame-2.1-0306-61500steps + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod + attn_kind: dense + # use_ema: true + chunk_size: 3 + min_timestep_ratio: 0.02 + max_timestep_ratio: 0.98 + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: >- + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_2130/preprocessed:1, + /mnt/weka/home/hao.zhang/mhuo/traindata_0204_1600/preprocessed:0, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/0_static_plus_w_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0205_1330/data/1_wasd_only/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/wasdonly_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0206_1200/data/camera/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/camera4hold_alpha1/preprocessed:3, + /mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/preprocessed:3 + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 18 + num_height: 352 + num_width: 640 + num_frames: 69 + apply_bot_died_filter: true + + optimizer: + learning_rate: 1e-5 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 10000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wangame_tfsft_causal + # resume_from_checkpoint: /mnt/weka/home/hao.zhang/mhuo/FastVideo-hyw/outputs/wangame_dfsft_causal_v3/checkpoint-best-step-36500 + training_state_checkpointing_steps: 5000 + weight_only_checkpointing_steps: 5000 + checkpoints_total_limit: 10 + best_checkpoint_start_step: 2000 + best_checkpoint_top_k: 5 + + tracker: + project_name: distillation_wangame_r + run_name: wangame_tfsft_causal + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + # ema: + # _target_: fastvideo.train.callbacks.ema.EMACallback + # beta: 0.9999 + validation: + _target_: fastvideo.train.callbacks.validation.ValidationCallback + pipeline_target: fastvideo.pipelines.basic.wan.wangame_causal_dmd_pipeline.WangameCausalOdeDMDPipeline + dataset_file: examples/training/finetune/WanGame2.1_1.3b_i2v/validation_4.json + every_steps: 500 + sampling_steps: [40] + scheduler_target: fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler + guidance_scale: 1.0 + num_frames: 69 + evaluate_ptlflow: false + +pipeline: + flow_shift: 3 diff --git a/examples/train/run.sh b/examples/train/run.sh index 263a6713a..df116b380 100755 --- a/examples/train/run.sh +++ b/examples/train/run.sh @@ -22,16 +22,17 @@ shift # ── GPU / node settings ────────────────────────────────────────── NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}" -NUM_GPUS="${NUM_GPUS:-1}" +NUM_GPUS="${NUM_GPUS:-4}" NNODES="${NNODES:-1}" NODE_RANK="${NODE_RANK:-0}" MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" MASTER_PORT="${MASTER_PORT:-29501}" export TOKENIZERS_PARALLELISM=false # ── W&B ────────────────────────────────────────────────────────── -export WANDB_API_KEY="${WANDB_API_KEY:-}" +export WANDB_API_KEY="${WANDB_API_KEY:-7ff8b6e8356924f7a6dd51a0342dd1a422ea9352}" export WANDB_MODE="${WANDB_MODE:-online}" + # ── Log file ───────────────────────────────────────────────────── CONFIG_NAME="$(basename "${CONFIG}" .yaml)" TIMESTAMP="$(date +%Y%m%d_%H%M%S)" @@ -39,6 +40,12 @@ LOG_DIR="${LOG_DIR:-examples/train}" mkdir -p "${LOG_DIR}" LOG_FILE="${LOG_DIR}/${CONFIG_NAME}_${TIMESTAMP}.log" +set +u +source ~/conda/miniconda/bin/activate +conda activate mhuo-fv +set -u +export PYTHONPATH="/mnt/weka/home/hao.zhang/mhuo/FastVideo-refactor:${PYTHONPATH:-}" + echo "=== Train Training ===" echo "Config: ${CONFIG}" echo "Num GPUs: ${NUM_GPUS}" @@ -49,7 +56,7 @@ echo "Extra args: $*" echo "Log file: ${LOG_FILE}" echo "==============================" -python -m torch.distributed.run \ +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.run \ --nnodes "${NNODES}" \ --node_rank "${NODE_RANK}" \ --nproc_per_node "${NUM_GPUS}" \ diff --git a/examples/train/run.slurm b/examples/train/run.slurm new file mode 100644 index 000000000..cfafcefb0 --- /dev/null +++ b/examples/train/run.slurm @@ -0,0 +1,80 @@ +#!/bin/bash +#SBATCH --job-name=wg-sf +#SBATCH --partition=main +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=examples/train/slurm_%j.out +#SBATCH --error=examples/train/slurm_%j.err +#SBATCH --exclusive + +set -euo pipefail + +CONFIG="${1:?Usage: sbatch examples/train/run.slurm [extra flags...]}" +shift +EXTRA_ARGS=("$@") +set -- + +cd /mnt/weka/home/hao.zhang/mhuo/FastVideo-refactor + +get_num_gpus() { + if [[ -n "${NUM_GPUS:-}" ]]; then + echo "${NUM_GPUS}" + return + fi + if [[ -n "${SLURM_GPUS_ON_NODE:-}" ]]; then + echo "${SLURM_GPUS_ON_NODE%%(*}" + return + fi + if [[ -n "${SLURM_GPUS_PER_NODE:-}" ]]; then + echo "${SLURM_GPUS_PER_NODE%%(*}" + return + fi + if command -v nvidia-smi >/dev/null 2>&1; then + nvidia-smi -L 2>/dev/null | wc -l | tr -d " " + else + echo 8 + fi +} + +export NNODES="${NNODES:-${SLURM_JOB_NUM_NODES:-1}}" +export NUM_GPUS="$(get_num_gpus)" +export MASTER_PORT="${MASTER_PORT:-29501}" +if [[ -z "${MASTER_ADDR:-}" ]]; then + nodes=( $(scontrol show hostnames "${SLURM_JOB_NODELIST}") ) + export MASTER_ADDR="${nodes[0]}" +fi + +export NCCL_P2P_DISABLE="${NCCL_P2P_DISABLE:-1}" +export TORCH_NCCL_ENABLE_MONITORING="${TORCH_NCCL_ENABLE_MONITORING:-0}" +export NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}" +export TOKENIZERS_PARALLELISM=false +export WANDB_API_KEY="${WANDB_API_KEY:-7ff8b6e8356924f7a6dd51a0342dd1a422ea9352}" +export WANDB_MODE="${WANDB_MODE:-online}" + +set +u +source ~/conda/miniconda/bin/activate +conda activate mhuo-fv +set -u +export PYTHONPATH="/mnt/weka/home/hao.zhang/mhuo/FastVideo-refactor:${PYTHONPATH:-}" + +echo "=== Distillation Training (Slurm) ===" +echo "Config: ${CONFIG}" +echo "Num GPUs: ${NUM_GPUS}" +echo "Num Nodes: ${NNODES}" +echo "Master: ${MASTER_ADDR}:${MASTER_PORT}" +echo "Extra args: ${EXTRA_ARGS[*]:-}" +echo "=====================================" + +srun torchrun \ + --nnodes "${NNODES}" \ + --nproc_per_node "${NUM_GPUS}" \ + --rdzv_backend c10d \ + --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ + --node_rank "${SLURM_PROCID}" \ + fastvideo/train/entrypoint/train.py \ + --config "${CONFIG}" \ + "${EXTRA_ARGS[@]}" diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_4.json b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_4.json new file mode 100644 index 000000000..eb3210229 --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_4.json @@ -0,0 +1,44 @@ +{ + "data": [ + { + "caption": "00 Val-00: W", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01 Val-01: S", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/S.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02 Val-02: A", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000004.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/A.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03 Val-03: D", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/D.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_8.json b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_8.json new file mode 100644 index 000000000..0bde93715 --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_8.json @@ -0,0 +1,84 @@ +{ + "data": [ + { + "caption": "00 Val-00: W", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01 Val-01: S", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/S.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02 Val-02: A", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000004.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/A.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03 Val-03: D", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/D.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "04 Val-04: u", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000000.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/u.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "05 Val-05: d", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000001.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/d.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "06 Val-06: l", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000006.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/l.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "07 Val-07: r", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000007.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/r.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_random.json b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_random.json new file mode 100644 index 000000000..41e51fc79 --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation_random.json @@ -0,0 +1,324 @@ +{ + "data": [ + { + "caption": "00 Val-00: W", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01 Val-01: S", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/S.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02 Val-02: A", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000004.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/A.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03 Val-03: D", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/D.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "04 Val-04: u", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000000.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/u.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "05 Val-05: d", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000001.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/d.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "06 Val-06: l", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000006.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/l.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "07 Val-07: r", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000007.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/r.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "08 Val-00: key rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "09 Val-01: key rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "10 Val-02: camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000004.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "11 Val-03: camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/camera_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "12 Val-00: key+camera excl rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "13 Val-01: key+camera excl rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "14 Val-02: key+camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000004.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "15 Val-03: key+camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_1_action_rand_2.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "16 Val-04: (simultaneous) key rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000000.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "17 Val-05: (simultaneous) camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000001.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/camera_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "18 Val-06: (simultaneous) key+camera excl rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000006.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "19 Val-07: (simultaneous) key+camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000007.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_2_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "20 Val-08: W+A", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/humanplay/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/WA.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "21 Val-09: S+u", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/humanplay/000013.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/S_u.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "22 Val-08: Still", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/humanplay/000005.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/still.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "23 Val-09: Still", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/humanplay/000013.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/still.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "24 Val-06: key+camera excl rand Frame 4", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000006.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_1_action_rand_1_f4.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "25 Val-07: key+camera excl rand Frame 4", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/mc_wasd_10/validate/000007.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_1_action_rand_2_f4.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "26 Train-00", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/first_frame/000500.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/videos/000500_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "27 Train-01", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/first_frame/001000.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/traindata_0208_2000/data/wasd4holdrandview_simple_1key1mouse1/videos/001000_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "28 Doom-00: W", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/doom/000000.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/W.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "29 Doom-01: key rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/doom/000001.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "30 Doom-02: camera rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/doom/000002.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/camera_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "31 Doom-03: key+camera excl rand", + "image_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/doom/000003.jpg", + "action_path": "/mnt/weka/home/hao.zhang/mhuo/FastVideo/examples/training/finetune/WanGame2.1_1.3b_i2v/actions/key_camera_excl_1_action_rand_1.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 7ea6d1284..abf0162f0 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -7,9 +7,11 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", - "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig" + "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", + "WanGameVideoConfig" ] diff --git a/fastvideo/configs/models/dits/wangamevideo.py b/fastvideo/configs/models/dits/wangamevideo.py new file mode 100644 index 000000000..62fd5dfac --- /dev/null +++ b/fastvideo/configs/models/dits/wangamevideo.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.wanvideo import ( + WanVideoArchConfig, + WanVideoConfig, +) + + +@dataclass +class WanGameVideoArchConfig(WanVideoArchConfig): + """Wangame keeps WanVideo architecture defaults and checkpoint mappings.""" + + +@dataclass +class WanGameVideoConfig(WanVideoConfig): + arch_config: WanGameVideoArchConfig = field( + default_factory=WanGameVideoArchConfig + ) + prefix: str = "WanGame" diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 83df65f8d..889be6f5d 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -69,6 +69,8 @@ class PipelineConfig: # DMD parameters dmd_denoising_steps: list[int] | None = field(default=None) + ode_solver: str = "unipc" + # Wan2.2 TI2V parameters ti2v_task: bool = False boundary_ratio: float | None = None @@ -175,6 +177,14 @@ class PipelineConfig: help= "Comma-separated list of denoising steps (e.g., '1000,757,522')", ) + parser.add_argument( + f"--{prefix_with_dot}ode-solver", + type=str, + choices=["unipc", "euler"], + dest=f"{prefix_with_dot.replace('-', '_')}ode_solver", + default=PipelineConfig.ode_solver, + help="ODE solver selection for ode sampling.", + ) # Add VAE configuration arguments from fastvideo.configs.models.vaes.base import VAEConfig diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index 048c7686c..101b046d0 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -157,3 +157,5 @@ pyarrow_schema_matrixgame = pa.schema([ pa.field("duration_sec", pa.float64()), pa.field("fps", pa.float64()), ]) + +pyarrow_schema_wangame = pyarrow_schema_matrixgame diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index c14eee93e..5724c850d 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import hashlib import os import pickle import random @@ -36,7 +37,9 @@ class DP_SP_BatchSampler(Sampler[list[int]]): global_rank: int, drop_last: bool = True, drop_first_row: bool = False, + reshuffle_each_epoch: bool = True, seed: int = 0, + candidate_indices: list[int] | None = None, ): self.batch_size = batch_size self.dataset_size = dataset_size @@ -45,38 +48,43 @@ class DP_SP_BatchSampler(Sampler[list[int]]): self.num_sp_groups = num_sp_groups self.global_rank = global_rank self.sp_world_size = sp_world_size + self.drop_first_row = drop_first_row + self.reshuffle_each_epoch = reshuffle_each_epoch + self.candidate_indices = ( + torch.as_tensor(candidate_indices, dtype=torch.long) + if candidate_indices is not None else None + ) - # ── epoch-level RNG ──────────────────────────────────────────────── - rng = torch.Generator().manual_seed(self.seed) - # Create a random permutation of all indices - global_indices = torch.randperm(self.dataset_size, generator=rng) + self._build_indices(0) - if drop_first_row: - # drop 0 in global_indices + def _build_indices(self, epoch: int) -> None: + rng = torch.Generator().manual_seed(self.seed + epoch) + if self.candidate_indices is None: + global_indices = torch.randperm(self.dataset_size, generator=rng) + else: + perm = torch.randperm(len(self.candidate_indices), generator=rng) + global_indices = self.candidate_indices[perm] + + dataset_size = len(global_indices) + if self.drop_first_row: global_indices = global_indices[global_indices != 0] - self.dataset_size = self.dataset_size - 1 + dataset_size = len(global_indices) if self.drop_last: - # For drop_last=True, we: - # 1. Ensure total samples is divisible by (batch_size * num_sp_groups) - # 2. This guarantees each SP group gets same number of complete batches - # 3. Prevents uneven batch sizes across SP groups at end of epoch - num_batches = self.dataset_size // self.batch_size + num_batches = dataset_size // self.batch_size num_global_batches = num_batches // self.num_sp_groups global_indices = global_indices[:num_global_batches * self.num_sp_groups * self.batch_size] else: - if self.dataset_size % (self.num_sp_groups * self.batch_size) != 0: - # add more indices to make it divisible by (batch_size * num_sp_groups) + if dataset_size % (self.num_sp_groups * self.batch_size) != 0: padding_size = self.num_sp_groups * self.batch_size - ( - self.dataset_size % (self.num_sp_groups * self.batch_size)) + dataset_size % (self.num_sp_groups * self.batch_size)) logger.info("Padding the dataset from %d to %d", - self.dataset_size, self.dataset_size + padding_size) + dataset_size, dataset_size + padding_size) global_indices = torch.cat( [global_indices, global_indices[:padding_size]]) - # shard the indices to each sp group ith_sp_group = self.global_rank // self.sp_world_size sp_group_local_indices = global_indices[ith_sp_group::self. num_sp_groups] @@ -84,6 +92,22 @@ class DP_SP_BatchSampler(Sampler[list[int]]): logger.info("Dataset size for each sp group: %d", len(sp_group_local_indices)) + def set_candidate_indices( + self, + candidate_indices: list[int] | None, + epoch: int = 0, + ) -> None: + self.candidate_indices = ( + torch.as_tensor(candidate_indices, dtype=torch.long) + if candidate_indices is not None else None + ) + self._build_indices(epoch) + + def set_epoch(self, epoch: int) -> None: + if not self.reshuffle_each_epoch: + return + self._build_indices(epoch) + def __iter__(self): indices = self.sp_group_local_indices for i in range(0, len(indices), self.batch_size): @@ -94,19 +118,89 @@ class DP_SP_BatchSampler(Sampler[list[int]]): return len(self.sp_group_local_indices) // self.batch_size +def _parse_data_path_specs(path: str) -> list[tuple[str, int]]: + """ + Parse data_path into a list of (directory, repeat_count). + Syntax: comma-separated entries; each entry is "path" (default 1) or "path:N" (N = repeat count). + N=0 means skip that path (convenience to disable without removing). If no ":" present, default is 1. + Example: "/dir1:2,/dir2,/dir3:0" -> dir1 2x, dir2 1x, dir3 skipped. + """ + specs: list[tuple[str, int]] = [] + for part in path.split(","): + part = part.strip() + if not part: + continue + if ":" in part: + p, _, count_str = part.rpartition(":") + p = p.strip() + try: + count = int(count_str.strip()) + except ValueError: + raise ValueError( + f"data_path repeat count must be an integer, got {count_str!r}" + ) from None + if count < 0: + raise ValueError( + f"data_path repeat count must be >= 0, got {count}" + ) + specs.append((p, count)) + else: + specs.append((part, 1)) + return specs + + +def _scan_parquet_files_for_path(p: str) -> tuple[list[str], list[int]]: + """Return (file_paths, row_lengths) for a single directory.""" + file_names: list[str] = [] + for root, _, files in os.walk(p): + for file in sorted(files): + if file.endswith(".parquet"): + file_names.append(os.path.join(root, file)) + lengths = [] + for file_path in tqdm.tqdm( + file_names, desc="Reading parquet files to get lengths"): + lengths.append(pq.ParquetFile(file_path).metadata.num_rows) + logger.info("Found %d parquet files with %d total rows", len(file_names), + sum(lengths)) + return file_names, lengths + + def get_parquet_files_and_length(path: str): - dataset_root = os.path.realpath(os.path.expanduser(path)) - # Check if cached info exists - cache_dir = os.path.join(dataset_root, "map_style_cache") - cache_file = os.path.join(cache_dir, "file_info.pkl") + """ + Collect parquet file paths and row lengths from one or more directories. + path: single directory, or comma-separated "path" or "path:N" (N = repeat count). + E.g. "/dir1:2,/dir2:1" -> dir1's files appear 2x (oversampled), dir2 once. + """ + path_specs = _parse_data_path_specs(path) + if not path_specs: + raise ValueError( + "data_path must be a non-empty path or comma-separated path specs" + ) + + first_path = next((p for p, c in path_specs if c > 0), path_specs[0][0]) + is_single_no_repeat = len(path_specs) == 1 and path_specs[0][1] == 1 + effective_path = path.strip() + + if is_single_no_repeat: + cache_dir = os.path.join(first_path, "map_style_cache") + cache_suffix = "file_info.pkl" + else: + neutral_root = os.environ.get( + "FASTVIDEO_MAP_STYLE_CACHE_DIR", + os.path.join(os.path.expanduser("~"), ".cache", "fastvideo", + "map_style_cache"), + ) + cache_dir = neutral_root + cache_suffix = ("file_info_" + + hashlib.md5(effective_path.encode()).hexdigest()[:16] + + ".pkl") + cache_file = os.path.join(cache_dir, cache_suffix) - # Only rank 0 checks for cache and scans files if needed if get_world_rank() == 0: cache_loaded = False file_names_sorted = None lengths_sorted = None - # First try to load existing cache if os.path.exists(cache_file): logger.info("Loading cached file info from %s", cache_file) try: @@ -117,24 +211,11 @@ def get_parquet_files_and_length(path: str): os.path.join(os.getcwd(), p) if not os.path.isabs(p) else p) for p in file_names_sorted) - files_outside_dataset_root = [ - file_path for file_path in file_names_sorted - if os.path.commonpath([dataset_root, file_path - ]) != dataset_root - ] missing_files = [ file_path for file_path in file_names_sorted if not os.path.exists(file_path) ] - if files_outside_dataset_root: - logger.warning( - "Cached parquet file list points outside dataset root " - "(%s). Cache will be rebuilt. First out-of-root file: %s", - dataset_root, - files_outside_dataset_root[0], - ) - cache_loaded = False - elif missing_files: + if missing_files: logger.warning( "Cached parquet file list contains %d missing files. " "Cache will be rebuilt. First missing file: %s", @@ -150,43 +231,43 @@ def get_parquet_files_and_length(path: str): logger.info("Falling back to scanning files") cache_loaded = False - # If cache not loaded (either doesn't exist or failed to load), scan files if not cache_loaded: - logger.info("Scanning parquet files to get lengths") - lengths = [] - file_names = [] - for root, _, files in os.walk(dataset_root): - for file in sorted(files): - if file.endswith('.parquet'): - file_path = os.path.realpath(os.path.join(root, file)) - file_names.append(file_path) - if len(file_names) == 0: + logger.info( + "Scanning parquet files (path specs: %s)", + [(p, c) for p, c in path_specs], + ) + combined: list[tuple[str, int, int]] = [] + sort_index = 0 + for p, count in path_specs: + if count == 0: + continue + fnames, lens = _scan_parquet_files_for_path(p) + if not fnames: + logger.warning("No parquet files found under path spec %s", p) + continue + for _ in range(count): + for f, ln in zip(fnames, lens, strict=True): + combined.append((f, ln, sort_index)) + sort_index += 1 + + if len(combined) == 0: raise FileNotFoundError( "No parquet files found under dataset path: " f"{path}. " "Please verify this path points to preprocessed parquet " "data.") - for file_path in tqdm.tqdm( - file_names, desc="Reading parquet files to get lengths"): - num_rows = pq.ParquetFile(file_path).metadata.num_rows - lengths.append(num_rows) - # sort according to file name to ensure all rank has the same order - file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names, - lengths, - strict=True), - key=lambda x: x[0]), - strict=True) - # Save the cache + + file_names_sorted, lengths_sorted, _ = zip( + *sorted(combined, key=lambda x: (x[0], x[2])), strict=True) + os.makedirs(cache_dir, exist_ok=True) with open(cache_file, "wb") as f: pickle.dump((file_names_sorted, lengths_sorted), f) logger.info("Saved file info to %s", cache_file) - # Wait for rank 0 to finish creating/loading cache world_group = get_world_group() world_group.barrier() - # Now all ranks load the cache (it should exist and be valid now) logger.info("Loading cached file info from %s after barrier", cache_file) with open(cache_file, "rb") as f: file_names_sorted, lengths_sorted = pickle.load(f) @@ -366,6 +447,84 @@ def passthrough(batch): return batch +def build_bot_died_excluded_indices( + data_path: str, + parquet_files: list[str], + lengths: list[int], +) -> set[int]: + """Build global row indices to exclude based on per-dir bot_died.json.""" + import json + + path_specs = [ + (os.path.realpath(os.path.expanduser(p)), count) + for p, count in _parse_data_path_specs(data_path) + ] + + bot_died_per_dir: dict[str, set[int]] = {} + for p, count in path_specs: + if count == 0: + continue + candidates = [ + os.path.join(os.path.dirname(p), "filter", "blue_water_random_half.json"), + os.path.join( + os.path.dirname(os.path.abspath(os.path.expanduser(p))), + "filter", + "blue_water_random_half.json", + ), + ] + filter_path = next((fp for fp in candidates if os.path.exists(fp)), None) + if filter_path is None: + continue + with open(filter_path, "r", encoding="utf-8") as f: + bot_died_per_dir[p] = set(json.load(f)) + logger.info( + "Loaded bot_died filter from %s: %d entries to exclude", + filter_path, + len(bot_died_per_dir[p]), + ) + + if not bot_died_per_dir: + return set() + + excluded: set[int] = set() + matched_file_count: dict[str, int] = {k: 0 for k in bot_died_per_dir} + global_offset = 0 + for file_path, length in zip(parquet_files, lengths, strict=True): + file_path_real = os.path.realpath(file_path) + matching_dir = None + for dir_path in bot_died_per_dir: + if (file_path.startswith(dir_path) + or file_path_real.startswith(dir_path)): + matching_dir = dir_path + break + + if matching_dir is not None: + matched_file_count[matching_dir] += 1 + bot_died_set = bot_died_per_dir[matching_dir] + try: + table = pq.read_table(file_path, columns=["file_name"]) + file_names = table.column("file_name").to_pylist() + for local_idx, fn in enumerate(file_names): + if int(str(fn).strip()) in bot_died_set: + excluded.add(global_offset + local_idx) + except Exception as e: + logger.warning( + "Failed to read file_name from %s for bot_died filter: %s", + file_path, + e, + ) + global_offset += length + + for dir_path, count in matched_file_count.items(): + logger.info( + "bot_died matching: dir=%s matched_parquet_files=%d", + dir_path, + count, + ) + + return excluded + + def build_parquet_map_style_dataloader( path, batch_size, diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 5ab467d75..cf97e8bc0 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -4,6 +4,7 @@ import os import pathlib import datasets +import numpy as np from torch.utils.data import IterableDataset from fastvideo.distributed import (get_sp_world_size, get_world_rank, @@ -16,8 +17,9 @@ logger = init_logger(__name__) class ValidationDataset(IterableDataset): - def __init__(self, filename: str): + def __init__(self, filename: str, num_samples: int | None = None): super().__init__() + self.num_samples = num_samples self.filename = pathlib.Path(filename) # get directory of filename @@ -58,6 +60,12 @@ class ValidationDataset(IterableDataset): # Convert to list to get total samples self.all_samples = list(data) + + # Limit number of samples if specified + if self.num_samples is not None and self.num_samples < len(self.all_samples): + self.all_samples = self.all_samples[:self.num_samples] + logger.info("Limiting validation samples to %s", self.num_samples) + self.original_total_samples = len(self.all_samples) # Extend samples to be a multiple of DP degree (num_sp_groups) @@ -160,5 +168,25 @@ class ValidationDataset(IterableDataset): else: sample["control_video"] = load_video(control_video_path) + if sample.get("action_path", None) is not None: + action_path = sample["action_path"] + action_path = os.path.join(self.dir, action_path) + sample["action_path"] = action_path + if not pathlib.Path(action_path).is_file(): + logger.warning("Action file %s does not exist.", action_path) + else: + try: + action_data = np.load(action_path, allow_pickle=True) + num_frames = sample["num_frames"] + if action_data.dtype == object: action_data = action_data.item() + if isinstance(action_data, dict): + sample["keyboard_cond"] = action_data["keyboard"][:num_frames] + sample["mouse_cond"] = action_data["mouse"][:num_frames] + else: + sample["keyboard_cond"] = action_data[:num_frames] + except Exception as e: + logger.error("Error loading action file %s: %s", + action_path, e) + sample = {k: v for k, v in sample.items() if v is not None} yield sample diff --git a/fastvideo/models/dits/hyworld/pose.py b/fastvideo/models/dits/hyworld/pose.py index b1b3f5df3..ceadfc016 100644 --- a/fastvideo/models/dits/hyworld/pose.py +++ b/fastvideo/models/dits/hyworld/pose.py @@ -10,6 +10,7 @@ Adapted from HY-WorldPlay: https://github.com/Tencent-Hunyuan/HY-WorldPlay """ import json +import logging import numpy as np import torch from scipy.spatial.transform import Rotation as R @@ -18,6 +19,9 @@ from typing import Union, Optional from .trajectory import generate_camera_trajectory_local +logger = logging.getLogger(__name__) + + # Mapping from one-hot action encoding to single label mapping = { (0, 0, 0, 0): 0, @@ -305,6 +309,152 @@ def camera_center_normalization(w2c: np.ndarray) -> np.ndarray: return np.linalg.inv(c2w_aligned) +def reformat_keyboard_and_mouse_tensors( + keyboard_tensor: torch.Tensor, + mouse_tensor: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Collapse frame-level keyboard/mouse controls to latent-level controls. + + The first frame is the anchor frame. The remaining frames are grouped into + chunks of 4 and each chunk is expected to be constant. + """ + num_frames = keyboard_tensor.shape[0] + assert (num_frames - 1) % 4 == 0, "num_frames must be a multiple of 4" + assert mouse_tensor.shape[0] == num_frames, ( + "mouse_tensor must have the same number of frames as keyboard_tensor" + ) + keyboard_tensor = keyboard_tensor[1:, :] + mouse_tensor = mouse_tensor[1:, :] + + groups = keyboard_tensor.view(-1, 4, keyboard_tensor.shape[1]) + if not (groups == groups[:, 0:1]).all(dim=1).all(): + logger.warning("keyboard_tensor has different values per 4-frame group") + + groups = mouse_tensor.view(-1, 4, mouse_tensor.shape[1]) + if not (groups == groups[:, 0:1]).all(dim=1).all(): + logger.warning("mouse_tensor has different values per 4-frame group") + + return keyboard_tensor[::4], mouse_tensor[::4] + + +def process_custom_actions( + keyboard_tensor: torch.Tensor, + mouse_tensor: torch.Tensor, + forward_speed: float = DEFAULT_FORWARD_SPEED, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Convert custom keyboard/mouse controls into viewmats, intrinsics, and labels. + """ + if keyboard_tensor.ndim == 3: + keyboard_tensor = keyboard_tensor.squeeze(0) + if mouse_tensor.ndim == 3: + mouse_tensor = mouse_tensor.squeeze(0) + + keyboard_tensor, mouse_tensor = reformat_keyboard_and_mouse_tensors( + keyboard_tensor, mouse_tensor + ) + + motions: list[dict[str, float]] = [] + for t in range(keyboard_tensor.shape[0]): + frame_motion: dict[str, float] = {} + + fwd = 0.0 + if keyboard_tensor[t, 0] > 0.5: + fwd += forward_speed + if keyboard_tensor[t, 1] > 0.5: + fwd -= forward_speed + if fwd != 0.0: + frame_motion["forward"] = fwd + + rgt = 0.0 + if keyboard_tensor[t, 2] > 0.5: + rgt -= forward_speed + if keyboard_tensor[t, 3] > 0.5: + rgt += forward_speed + if rgt != 0.0: + frame_motion["right"] = rgt + + pitch = mouse_tensor[t, 0].item() + yaw = mouse_tensor[t, 1].item() + if abs(pitch) > 1e-4: + frame_motion["pitch"] = pitch + if abs(yaw) > 1e-4: + frame_motion["yaw"] = yaw + + motions.append(frame_motion) + + poses = generate_camera_trajectory_local(motions) + + w2c_list = [] + intrinsic_list = [] + K = np.array(DEFAULT_INTRINSIC) + K[0, 0] /= K[0, 2] * 2 + K[1, 1] /= K[1, 2] * 2 + K[0, 2] = 0.5 + K[1, 2] = 0.5 + + for pose in poses: + c2w = np.array(pose) + w2c = np.linalg.inv(c2w) + w2c_list.append(w2c) + intrinsic_list.append(K) + + viewmats = torch.as_tensor(np.array(w2c_list)) + intrinsics = torch.as_tensor(np.array(intrinsic_list)) + + c2ws = np.linalg.inv(np.array(w2c_list)) + c_inv = np.linalg.inv(c2ws[:-1]) + relative_c2w = np.zeros_like(c2ws) + relative_c2w[0, ...] = c2ws[0, ...] + relative_c2w[1:, ...] = c_inv @ c2ws[1:, ...] + + trans_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) + rotate_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) + move_norm_valid = 0.0001 + + for i in range(1, relative_c2w.shape[0]): + move_dirs = relative_c2w[i, :3, 3] + move_norms = np.linalg.norm(move_dirs) + + if move_norms > move_norm_valid: + move_norm_dirs = move_dirs / move_norms + angles_rad = np.arccos(move_norm_dirs.clip(-1.0, 1.0)) + trans_angles_deg = angles_rad * (180.0 / np.pi) + else: + trans_angles_deg = np.zeros(3) + + r_rel = relative_c2w[i, :3, :3] + rot_angles_deg = R.from_matrix(r_rel).as_euler("xyz", degrees=True) + + if move_norms > move_norm_valid: + if trans_angles_deg[2] < 60: + trans_one_hot[i, 0] = 1 + elif trans_angles_deg[2] > 120: + trans_one_hot[i, 1] = 1 + + if trans_angles_deg[0] < 60: + trans_one_hot[i, 2] = 1 + elif trans_angles_deg[0] > 120: + trans_one_hot[i, 3] = 1 + + if rot_angles_deg[1] > 5e-2: + rotate_one_hot[i, 0] = 1 + elif rot_angles_deg[1] < -5e-2: + rotate_one_hot[i, 1] = 1 + + if rot_angles_deg[0] > 5e-2: + rotate_one_hot[i, 2] = 1 + elif rot_angles_deg[0] < -5e-2: + rotate_one_hot[i, 3] = 1 + + trans_label = one_hot_to_one_dimension(torch.tensor(trans_one_hot)) + rotate_label = one_hot_to_one_dimension(torch.tensor(rotate_one_hot)) + action_labels = trans_label * 9 + rotate_label + + return viewmats, intrinsics, action_labels + + def parse_pose_string_to_actions(pose_string: str, fps: int = 24) -> list[dict]: """ diff --git a/fastvideo/models/dits/matrixgame/utils.py b/fastvideo/models/dits/matrixgame/utils.py index 4dd937699..f9cc96c08 100644 --- a/fastvideo/models/dits/matrixgame/utils.py +++ b/fastvideo/models/dits/matrixgame/utils.py @@ -299,6 +299,209 @@ def parse_config(config, mode="universal"): ) return key_data, mouse_data + +def _get_cv2(): + import cv2 + + return cv2 + + +def draw_rounded_rectangle( + image, + top_left, + bottom_right, + color, + radius=10, + alpha=0.5, +): + cv2 = _get_cv2() + overlay = image.copy() + x1, y1 = top_left + x2, y2 = bottom_right + + cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1) + cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1) + cv2.ellipse( + overlay, + (x1 + radius, y1 + radius), + (radius, radius), + 180, + 0, + 90, + color, + -1, + ) + cv2.ellipse( + overlay, + (x2 - radius, y1 + radius), + (radius, radius), + 270, + 0, + 90, + color, + -1, + ) + cv2.ellipse( + overlay, + (x1 + radius, y2 - radius), + (radius, radius), + 90, + 0, + 90, + color, + -1, + ) + cv2.ellipse( + overlay, + (x2 - radius, y2 - radius), + (radius, radius), + 0, + 0, + 90, + color, + -1, + ) + cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image) + + +def draw_keys_on_frame( + frame, + keys, + key_size=(30, 30), + spacing=5, + top_margin=15, + mode="universal", +): + """Draw keyboard action badges on the top-left of the frame.""" + cv2 = _get_cv2() + del spacing # Preserved for compatibility with the original helper. + + left_margin = 15 + gap = 3 + + key_positions = { + "W": (left_margin + key_size[0] + gap, top_margin), + "A": (left_margin, top_margin + key_size[1] + gap), + "S": (left_margin + key_size[0] + gap, top_margin + key_size[1] + gap), + "D": ( + left_margin + (key_size[0] + gap) * 2, + top_margin + key_size[1] + gap, + ), + } + key_icon = { + "W": "W", + "A": "A", + "S": "S", + "D": "D", + "left": "L", + "right": "R", + } + if mode == "templerun": + key_positions.update( + { + "left": ( + left_margin + (key_size[0] + gap) * 3 + 10, + top_margin + key_size[1] + gap, + ), + "right": ( + left_margin + (key_size[0] + gap) * 4 + 15, + top_margin + key_size[1] + gap, + ), + } + ) + + for key, (x, y) in key_positions.items(): + is_pressed = keys.get(key, False) + top_left = (x, y) + bottom_right = (x + key_size[0], y + key_size[1]) + + color = (0, 255, 0) if is_pressed else (200, 200, 200) + alpha = 0.8 if is_pressed else 0.5 + draw_rounded_rectangle( + frame, + top_left, + bottom_right, + color, + radius=5, + alpha=alpha, + ) + + text = key_icon[key] + text_size = cv2.getTextSize( + text, + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + 1, + )[0] + text_x = x + (key_size[0] - text_size[0]) // 2 + text_y = y + (key_size[1] + text_size[1]) // 2 + cv2.putText( + frame, + text, + (text_x, text_y), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (0, 0, 0), + 1, + ) + + +def draw_mouse_on_frame(frame, pitch, yaw, top_margin=15): + """Draw a mouse-look crosshair with a direction arrow.""" + cv2 = _get_cv2() + h, w, _ = frame.shape + + right_margin = 15 + crosshair_radius = 25 + crosshair_x = w - right_margin - crosshair_radius + crosshair_y = top_margin + crosshair_radius + + dx = int(yaw * crosshair_radius * 8) + dy = int(-pitch * crosshair_radius * 8) + + max_arrow = crosshair_radius - 5 + dx = max(-max_arrow, min(max_arrow, dx)) + dy = max(-max_arrow, min(max_arrow, dy)) + + cv2.circle( + frame, + (crosshair_x, crosshair_y), + crosshair_radius, + (50, 50, 50), + -1, + ) + cv2.circle( + frame, + (crosshair_x, crosshair_y), + crosshair_radius, + (200, 200, 200), + 1, + ) + cv2.line( + frame, + (crosshair_x - crosshair_radius + 5, crosshair_y), + (crosshair_x + crosshair_radius - 5, crosshair_y), + (100, 100, 100), + 1, + ) + cv2.line( + frame, + (crosshair_x, crosshair_y - crosshair_radius + 5), + (crosshair_x, crosshair_y + crosshair_radius - 5), + (100, 100, 100), + 1, + ) + + if abs(dx) > 1 or abs(dy) > 1: + cv2.arrowedLine( + frame, + (crosshair_x, crosshair_y), + (crosshair_x + dx, crosshair_y + dy), + (0, 255, 0), + 2, + tipLength=0.3, + ) + # NOTE: drawing functions are commented out to avoid cv2/libGL dependency. # # def draw_rounded_rectangle(image, top_left, bottom_right, color, radius=10, alpha=0.5): diff --git a/fastvideo/models/dits/wangame/__init__.py b/fastvideo/models/dits/wangame/__init__.py new file mode 100644 index 000000000..7d391cbed --- /dev/null +++ b/fastvideo/models/dits/wangame/__init__.py @@ -0,0 +1,16 @@ +from .model import WanGameActionTransformer3DModel +from .causal_model import CausalWanGameActionTransformer3DModel +from .hyworld_action_module import WanGameActionTimeImageEmbedding, WanGameActionSelfAttention + +__all__ = [ + "WanGameActionTransformer3DModel", + "CausalWanGameActionTransformer3DModel", + "WanGameActionTimeImageEmbedding", + "WanGameActionSelfAttention", +] + +# Entry point for model registry +EntryClass = [ + WanGameActionTransformer3DModel, + CausalWanGameActionTransformer3DModel, +] diff --git a/fastvideo/models/dits/wangame/causal_model.py b/fastvideo/models/dits/wangame/causal_model.py new file mode 100644 index 000000000..e2c7f4a30 --- /dev/null +++ b/fastvideo/models/dits/wangame/causal_model.py @@ -0,0 +1,863 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +import torch.nn as nn + +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from torch.nn.attention.flex_attention import BlockMask +# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention +# see https://github.com/pytorch/pytorch/issues/133254 +# change to default for other models +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") +import torch.distributed as dist + +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import PatchEmbed +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanI2VCrossAttention +from fastvideo.platforms import AttentionBackendEnum, current_platform + +from .hyworld_action_module import ( + WanGameActionTimeImageEmbedding, + WanGameActionSelfAttention, +) +from fastvideo.models.dits.hyworld.camera_rope import prope_qkv + +logger = init_logger(__name__) + + +_DEFAULT_WANGAME_CONFIG = WanGameVideoConfig() + + +class CausalWanGameCrossAttention(WanI2VCrossAttention): + """Cross-attention for WanGame causal model""" + + def forward(self, x, context, context_lens=None, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache: Optional cache dict for inference + """ + context_img = context + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + crossattn_cache["k"] = k_img + crossattn_cache["v"] = v_img + else: + k_img = crossattn_cache["k"] + v_img = crossattn_cache["v"] + else: + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + + img_x = self.attn(q, k_img, v_img) + + # output + x = img_x.flatten(2) + x, _ = self.to_out(x) + return x + + +class CausalWanGameActionSelfAttention(WanGameActionSelfAttention): + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm=True, + eps=1e-6) -> None: + super().__init__( + dim=dim, + num_heads=num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps, + ) + self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 + + # Local attention for KV-cache inference + self.local_attn = LocalAttention( + num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + @staticmethod + def _masked_flex_attn( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_mask: BlockMask, + ) -> torch.Tensor: + padded_length = math.ceil(query.shape[1] / 128) * 128 - query.shape[1] + if padded_length > 0: + query = torch.cat( + [ + query, + torch.zeros( + [query.shape[0], padded_length, query.shape[2], query.shape[3]], + device=query.device, + dtype=value.dtype, + ), + ], + dim=1, + ) + key = torch.cat( + [ + key, + torch.zeros( + [key.shape[0], padded_length, key.shape[2], key.shape[3]], + device=key.device, + dtype=value.dtype, + ), + ], + dim=1, + ) + value = torch.cat( + [ + value, + torch.zeros( + [value.shape[0], padded_length, value.shape[2], value.shape[3]], + device=value.device, + dtype=value.dtype, + ), + ], + dim=1, + ) + + out = flex_attention( + query=query.transpose(2, 1), + key=key.transpose(2, 1), + value=value.transpose(2, 1), + block_mask=block_mask, + ).transpose(2, 1) + + if padded_length > 0: + out = out[:, :-padded_length] + return out + + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask | None = None, + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False): + """ + Forward pass with causal attention. + """ + if cache_start is None: + cache_start = current_start + + if kv_cache is None: + if block_mask is None: + raise ValueError( + "block_mask must be provided for causal training attention") + if viewmats is None or Ks is None: + raise ValueError( + "viewmats and Ks must be provided for WanGame causal attention") + + cos, sin = freqs_cis + query_rope = _apply_rotary_emb( + q, cos, sin, is_neox_style=False).type_as(v) + key_rope = _apply_rotary_emb( + k, cos, sin, is_neox_style=False).type_as(v) + rope_output = self._masked_flex_attn( + query_rope, key_rope, v, block_mask) + + # PRoPE path with the same causal mask. + query_prope, key_prope, value_prope, apply_fn_o = prope_qkv( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + viewmats=viewmats, + Ks=Ks, + patches_x=40, + patches_y=22, + ) + query_prope = query_prope.transpose(1, 2) + key_prope = key_prope.transpose(1, 2) + value_prope = value_prope.transpose(1, 2) + prope_output = self._masked_flex_attn( + query_prope, key_prope, value_prope, block_mask) + prope_output = apply_fn_o( + prope_output.transpose(1, 2)).transpose(1, 2) + + return rope_output, prope_output + else: + # Inference mode with KV cache + if viewmats is None or Ks is None: + raise ValueError( + "viewmats and Ks must be provided for WanGame causal attention") + + cos, sin = freqs_cis + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) + roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) + query_prope, key_prope, value_prope, apply_fn_o = prope_qkv( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + viewmats=viewmats, + Ks=Ks, + patches_x=40, + patches_y=22, + ) + query_prope = query_prope.transpose(1, 2).type_as(v) + key_prope = key_prope.transpose(1, 2).type_as(v) + value_prope = value_prope.transpose(1, 2).type_as(v) + + frame_seqlen = q.shape[1] + current_end = current_start + roped_query.shape[1] + sink_tokens = self.sink_size * frame_seqlen + # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = roped_query.shape[1] + + # rope+prope + cache_head_dim = kv_cache["k"].shape[-1] + local_end_index = kv_cache["local_end_index"].item() + + # read cache but never mutate it. + if not is_cache: + if cache_head_dim not in (self.head_dim, self.head_dim * 2): + raise ValueError( + f"Unexpected kv_cache head dim: {cache_head_dim}, " + f"expected {self.head_dim} or {self.head_dim * 2}") + + cache_k_rope = kv_cache["k"][..., :self.head_dim] + cache_v_rope = kv_cache["v"][..., :self.head_dim] + rope_k = torch.cat( + [cache_k_rope[:, :local_end_index], roped_key], dim=1) + rope_v = torch.cat( + [cache_v_rope[:, :local_end_index], v], dim=1) + rope_k = rope_k[:, -self.max_attention_size:] + rope_v = rope_v[:, -self.max_attention_size:] + rope_x = self.local_attn(roped_query, rope_k, rope_v) + + if cache_head_dim == self.head_dim * 2: + cache_k_prope = kv_cache["k"][..., self.head_dim:] + cache_v_prope = kv_cache["v"][..., self.head_dim:] + prope_k = torch.cat( + [cache_k_prope[:, :local_end_index], key_prope], dim=1) + prope_v = torch.cat( + [cache_v_prope[:, :local_end_index], value_prope], dim=1) + prope_k = prope_k[:, -self.max_attention_size:] + prope_v = prope_v[:, -self.max_attention_size:] + prope_x = self.local_attn(query_prope, prope_k, prope_v) + else: + prope_x = self.local_attn( + query_prope, key_prope, value_prope) + + prope_x = apply_fn_o(prope_x.transpose(1, 2)).transpose(1, 2) + return rope_x, prope_x + + # update cache. + if cache_head_dim == self.head_dim: + kv_cache["k"] = torch.cat( + [kv_cache["k"], torch.zeros_like(kv_cache["k"])], dim=-1) + kv_cache["v"] = torch.cat( + [kv_cache["v"], torch.zeros_like(kv_cache["v"])], dim=-1) + elif cache_head_dim != self.head_dim * 2: + raise ValueError( + f"Unexpected kv_cache head dim: {cache_head_dim}, " + f"expected {self.head_dim} or {self.head_dim * 2}") + + cache_k_rope = kv_cache["k"][..., :self.head_dim] + cache_k_prope = kv_cache["k"][..., self.head_dim:] + cache_v_rope = kv_cache["v"][..., :self.head_dim] + cache_v_prope = kv_cache["v"][..., self.head_dim:] + + if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): + # Calculate the number of new tokens added in this step + # Shift existing cache content left to discard oldest tokens + # Clone the source slice to avoid overlapping memory error + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens + cache_k_rope[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache_k_rope[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + cache_v_rope[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache_v_rope[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + cache_k_prope[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache_k_prope[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + cache_v_prope[:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache_v_prope[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache["local_end_index"].item() + current_end - \ + kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + cache_k_rope[:, local_start_index:local_end_index] = roped_key + cache_v_rope[:, local_start_index:local_end_index] = v + cache_k_prope[:, local_start_index:local_end_index] = key_prope + cache_v_prope[:, local_start_index:local_end_index] = value_prope + else: + # Assign new keys/values directly up to current_end + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache["k"] = kv_cache["k"].detach() + kv_cache["v"] = kv_cache["v"].detach() + cache_k_rope = kv_cache["k"][..., :self.head_dim] + cache_k_prope = kv_cache["k"][..., self.head_dim:] + cache_v_rope = kv_cache["v"][..., :self.head_dim] + cache_v_prope = kv_cache["v"][..., self.head_dim:] + # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) + cache_k_rope[:, local_start_index:local_end_index] = roped_key + cache_v_rope[:, local_start_index:local_end_index] = v + cache_k_prope[:, local_start_index:local_end_index] = key_prope + cache_v_prope[:, local_start_index:local_end_index] = value_prope + + rope_x = self.local_attn( + roped_query, + cache_k_rope[:, max(0, local_end_index - self.max_attention_size):local_end_index], + cache_v_rope[:, max(0, local_end_index - self.max_attention_size):local_end_index] + ) + prope_x = self.local_attn( + query_prope, + cache_k_prope[:, max(0, local_end_index - self.max_attention_size):local_end_index], + cache_v_prope[:, max(0, local_end_index - self.max_attention_size):local_end_index] + ) + prope_x = apply_fn_o(prope_x.transpose(1, 2)).transpose(1, 2) + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) + + return rope_x, prope_x + + +class CausalWanGameActionTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ReplicatedLinear(dim, dim, bias=True) + + self.attn1 = CausalWanGameActionSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) + + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + dim_head = dim // num_heads + + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + raise ValueError(f"QK Norm type {qk_norm} not supported") + + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + compute_dtype=torch.float32) + + # 2. Cross-attention (I2V only) + self.attn2 = CausalWanGameCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + # norm3 for FFN input + self.norm3 = LayerNormScaleShift(dim, norm_type="layer", eps=eps, + elementwise_affine=False) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # PRoPE output projection (initialized via add_discrete_action_parameters on the model) + self.to_out_prope = ReplicatedLinear(dim, dim, bias=True) + nn.init.zeros_(self.to_out_prope.weight) + if self.to_out_prope.bias is not None: + nn.init.zeros_(self.to_out_prope.bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask | None = None, + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + + num_frames = temb.shape[1] + frame_seqlen = hidden_states.shape[1] // num_frames + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + + # Cast temb to float32 for scale/shift computation + e = self.scale_shift_table + temb.float() + assert e.shape == (bs, num_frames, 6, self.hidden_dim) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).to(orig_dtype).flatten(1, 2) + + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q.forward_native(query) + if self.norm_k is not None: + key = self.norm_k.forward_native(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + # Self-attention with camera PRoPE + attn_output_rope, attn_output_prope = self.attn1( + query, key, value, freqs_cis, + block_mask, kv_cache, current_start, cache_start, + viewmats, Ks, is_cache=is_cache + ) + # Combine rope and prope outputs + attn_output_rope = attn_output_rope.flatten(2) + attn_output_rope, _ = self.to_out(attn_output_rope) + attn_output_prope = attn_output_prope.flatten(2) + attn_output_prope, _ = self.to_out_prope(attn_output_prope) + attn_output = attn_output_rope.squeeze(1) + attn_output_prope.squeeze(1) + + # Self-attention residual + norm in float32 + null_shift = null_scale = torch.zeros(1, device=hidden_states.device, dtype=torch.float32) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states.float(), attn_output.float(), gate_msa, null_shift, null_scale) + hidden_states = hidden_states.type_as(attn_output) + norm_hidden_states = norm_hidden_states.type_as(attn_output) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states.to(orig_dtype), + context=encoder_hidden_states, + context_lens=None, + crossattn_cache=crossattn_cache) + # Cross-attention residual in bfloat16 + hidden_states = hidden_states + attn_output + + # norm3 for FFN input in float32 + norm_hidden_states = self.norm3( + hidden_states.float(), c_shift_msa, c_scale_msa + ).type_as(hidden_states) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states.to(orig_dtype)) + hidden_states = self.mlp_residual(hidden_states.float(), ff_output.float(), c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +class CausalWanGameActionTransformer3DModel(BaseDiT): + supports_action_input = True + + _fsdp_shard_conditions = _DEFAULT_WANGAME_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_WANGAME_CONFIG._compile_conditions + _supported_attention_backends = ( + _DEFAULT_WANGAME_CONFIG._supported_attention_backends + ) + param_names_mapping = _DEFAULT_WANGAME_CONFIG.param_names_mapping + reverse_param_names_mapping = ( + _DEFAULT_WANGAME_CONFIG.reverse_param_names_mapping + ) + lora_param_names_mapping = _DEFAULT_WANGAME_CONFIG.lora_param_names_mapping + + def __init__(self, config: WanGameVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + 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.attention_head_dim = config.attention_head_dim + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.local_attn_size = config.local_attn_size + self.inner_dim = inner_dim + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanGameActionTimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + CausalWanGameActionTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}") + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.block_mask = None + self.num_frame_per_block = config.arch_config.num_frames_per_block + assert self.num_frame_per_block <= 3 + + self.__post_init__() + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=0, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for tmp in frame_indices: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + + def _forward_inference( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + action: torch.Tensor | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + kv_cache: list[dict] | None = None, + crossattn_cache: list[dict] | None = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + is_cache: bool = False, + **kwargs + ) -> torch.Tensor: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + """ + orig_dtype = hidden_states.dtype + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, action, encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image) + + # condition_embedder returns: + # - temb: [B*T, dim] where T = post_patch_num_frames + # - timestep_proj: [B*T, 6*dim] + # Reshape to [B, T, 6, dim] for transformer blocks + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)) # [B*T, 6, dim] + timestep_proj = timestep_proj.view(batch_size, post_patch_num_frames, 6, self.hidden_size) # [B, T, 6, dim] + + encoder_hidden_states = encoder_hidden_states_image + + # Transformer blocks + for block_idx, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + self.block_mask, + kv_cache[block_idx] if kv_cache else None, + crossattn_cache[block_idx] if crossattn_cache else None, + current_start, cache_start, + viewmats, Ks, is_cache) + else: + hidden_states = block( + hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + block_mask=self.block_mask, + kv_cache=kv_cache[block_idx] if kv_cache else None, + crossattn_cache=crossattn_cache[block_idx] if crossattn_cache else None, + current_start=current_start, cache_start=cache_start, + viewmats=viewmats, Ks=Ks, is_cache=is_cache) + + # If cache-only mode, return early + if is_cache: + return kv_cache + + # Output norm, projection & unpatchify + # temb is [B*T, dim], reshape to [B, T, 1, dim] + temb = temb.view(batch_size, post_patch_num_frames, -1).unsqueeze(2) # [B, T, 1, dim] + + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def _forward_train( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + action: torch.Tensor | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + start_frame: int = 0, + **kwargs + ) -> torch.Tensor: + + orig_dtype = hidden_states.dtype + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + # Construct blockwise causal attn mask + if self.block_mask is None: + self.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 + ) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, action, encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image) + + # condition_embedder returns: + # - temb: [B*T, dim] where T = post_patch_num_frames + # - timestep_proj: [B*T, 6*dim] + # Reshape to [B, T, 6, dim] for transformer blocks + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)) # [B*T, 6, dim] + timestep_proj = timestep_proj.view(batch_size, post_patch_num_frames, 6, self.hidden_size) # [B, T, 6, dim] + + encoder_hidden_states = encoder_hidden_states_image + + # Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + self.block_mask, + None, None, # kv_cache, crossattn_cache + 0, None, # current_start, cache_start + viewmats, Ks, False) # viewmats, Ks, is_cache + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + block_mask=self.block_mask, + viewmats=viewmats, Ks=Ks) + + # Output norm, projection & unpatchify + # temb is [B*T, dim], reshape to [B, T, 1, dim] + temb = temb.view(batch_size, post_patch_num_frames, -1).unsqueeze(2) # [B, T, 1, dim] + + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) diff --git a/fastvideo/models/dits/wangame/hyworld_action_module.py b/fastvideo/models/dits/wangame/hyworld_action_module.py new file mode 100644 index 000000000..22206d501 --- /dev/null +++ b/fastvideo/models/dits/wangame/hyworld_action_module.py @@ -0,0 +1,284 @@ +import math + +import torch +import torch.nn as nn + +from fastvideo.layers.visual_embedding import TimestepEmbedder, ModulateProjection, timestep_embedding +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.attention import DistributedAttention +from fastvideo.forward_context import set_forward_context +from fastvideo.models.dits.wanvideo import WanImageEmbedding + +from fastvideo.models.dits.hyworld.camera_rope import prope_qkv +from fastvideo.layers.rotary_embedding import _apply_rotary_emb +from fastvideo.layers.mlp import MLP + +class WanGameActionTimeImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + image_embed_dim: int | None = None, + ): + super().__init__() + + self.time_freq_dim = time_freq_dim + self.time_embedder = TimestepEmbedder( + dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, + factor=6, + act_layer="silu") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + self.action_embedder = MLP( + time_freq_dim, + dim, + dim, + bias=True, + act_type="silu" + ) + # Initialize fc_in with kaiming_uniform (same as nn.Linear default) + nn.init.kaiming_uniform_(self.action_embedder.fc_in.weight, a=math.sqrt(5)) + # Initialize fc_out with zeros for residual-like behavior + nn.init.zeros_(self.action_embedder.fc_out.weight) + if self.action_embedder.fc_out.bias is not None: + nn.init.zeros_(self.action_embedder.fc_out.bias) + + def forward( + self, + timestep: torch.Tensor, + action: torch.Tensor, + encoder_hidden_states: torch.Tensor, # Kept for interface compatibility + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + ): + """ + Args: + timestep: [B] diffusion timesteps (one per batch sample) + action: [B, T] action labels (one per frame per batch sample) + + Returns: + temb: [B*T, dim] combined timestep + action embedding + timestep_proj: [B*T, 6*dim] modulation projection + """ + # timestep may be [B] (one per sample) or [B*T] (one per frame, from causal training) + temb = self.time_embedder(timestep, timestep_seq_len) + + # Handle action embedding for batch > 1 + # action shape: [B, T] where B=batch_size, T=num_frames + batch_size = action.shape[0] + num_frames = action.shape[1] + + # Compute action embeddings: [B, T] -> [B*T] -> [B*T, dim] + action_flat = action.flatten() # [B*T] + action_emb = timestep_embedding(action_flat, self.time_freq_dim) + action_embedder_dtype = next(iter(self.action_embedder.parameters())).dtype + if ( + action_emb.dtype != action_embedder_dtype + and action_embedder_dtype != torch.int8 + ): + action_emb = action_emb.to(action_embedder_dtype) + action_emb = self.action_embedder(action_emb).type_as(temb) # [B*T, dim] + + # temb is [B*T, dim] when timestep was already per-frame (causal training), + # or [B, dim] when timestep is per-sample (inference). + # Only expand if temb is per-sample [B, dim]. + if temb.shape[0] == batch_size and num_frames > 1: + # Expand temb: [B, dim] -> [B, T, dim] -> [B*T, dim] + temb_expanded = temb.unsqueeze(1).expand(-1, num_frames, -1) # [B, T, dim] + temb_expanded = temb_expanded.reshape(batch_size * num_frames, -1) # [B*T, dim] + else: + # temb is already [B*T, dim] (per-frame timesteps) + temb_expanded = temb + + # Add action embedding to expanded temb + temb = temb_expanded + action_emb # [B*T, dim] + + timestep_proj = self.time_modulation(temb) # [B*T, 6*dim] + + # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states + + 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) + + encoder_hidden_states = torch.zeros((batch_size, 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + +class WanGameActionSelfAttention(nn.Module): + """ + Self-attention module with support for: + - Standard RoPE-based attention + - Camera PRoPE-based attention (when viewmats and Ks are provided) + - KV caching for autoregressive generation + """ + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm=True, + eps=1e-6) -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 + + # Scaled dot product attention (using DistributedAttention for SP support) + self.attn = DistributedAttention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False, + attention_mask: torch.Tensor | None = None): + """ + Forward pass with camera PRoPE attention combining standard RoPE and projective positional encoding. + + Args: + q, k, v: Query, key, value tensors [B, L, num_heads, head_dim] + freqs_cis: RoPE frequency cos/sin tensors + kv_cache: KV cache dict (may have None values for training) + current_start: Current position for KV cache + cache_start: Cache start position + viewmats: Camera view matrices for PRoPE [B, cameras, 4, 4] + Ks: Camera intrinsics for PRoPE [B, cameras, 3, 3] + is_cache: Whether to store to KV cache (for inference) + attention_mask: Attention mask [B, L] (1 = attend, 0 = mask) + """ + if cache_start is None: + cache_start = current_start + + # Apply RoPE manually + cos, sin = freqs_cis + query_rope = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) + key_rope = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) + value_rope = v + + # # DEBUG: Check camera matrices + # if self.training and torch.distributed.get_rank() == 0: + # vm_info = f"viewmats={viewmats.shape if viewmats is not None else None}" + # ks_info = f"Ks={Ks.shape if Ks is not None else None}" + # vm_nonzero = (viewmats != 0).sum().item() if viewmats is not None else 0 + # ks_nonzero = (Ks != 0).sum().item() if Ks is not None else 0 + # print(f"[DEBUG] PRoPE input: {vm_info} nonzero={vm_nonzero}, {ks_info} nonzero={ks_nonzero}", flush=True) + + # Get PRoPE transformed q, k, v + query_prope, key_prope, value_prope, apply_fn_o = prope_qkv( + q.transpose(1, 2), # [B, num_heads, L, head_dim] + k.transpose(1, 2), + v.transpose(1, 2), + viewmats=viewmats, + Ks=Ks, + patches_x=40, # hardcoded for now + patches_y=22, + ) + # PRoPE returns [B, num_heads, L, head_dim], convert to [B, L, num_heads, head_dim] + query_prope = query_prope.transpose(1, 2) + key_prope = key_prope.transpose(1, 2) + value_prope = value_prope.transpose(1, 2) + + # # DEBUG: Check prope_qkv output + # if self.training and torch.distributed.get_rank() == 0: + # q_nz = (query_prope != 0).sum().item() + # k_nz = (key_prope != 0).sum().item() + # v_nz = (value_prope != 0).sum().item() + # print(f"[DEBUG] prope_qkv output: q_nonzero={q_nz}, k_nonzero={k_nz}, v_nonzero={v_nz}", flush=True) + + # KV cache handling + if kv_cache is not None: + cache_key = kv_cache.get("k", None) + cache_value = kv_cache.get("v", None) + + if cache_value is not None and not is_cache: + cache_key_rope, cache_key_prope = cache_key.chunk(2, dim=-1) + cache_value_rope, cache_value_prope = cache_value.chunk(2, dim=-1) + + key_rope = torch.cat([cache_key_rope, key_rope], dim=1) + value_rope = torch.cat([cache_value_rope, value_rope], dim=1) + key_prope = torch.cat([cache_key_prope, key_prope], dim=1) + value_prope = torch.cat([cache_value_prope, value_prope], dim=1) + + if is_cache: + # Store to cache (update input dict directly) + kv_cache["k"] = torch.cat([key_rope, key_prope], dim=-1) + kv_cache["v"] = torch.cat([value_rope, value_prope], dim=-1) + + # Concatenate rope and prope paths (matching original) + query_all = torch.cat([query_rope, query_prope], dim=0) + key_all = torch.cat([key_rope, key_prope], dim=0) + value_all = torch.cat([value_rope, value_prope], dim=0) + + # Check if Q and KV have different sequence lengths (KV cache mode) + # In this case, use LocalAttention (supports different Q/KV lengths) + if query_all.shape[1] != key_all.shape[1]: + raise ValueError("Q and KV have different sequence lengths") + else: + # Same sequence length: use DistributedAttention (supports SP) + # Create default attention mask if not provided + # NOTE: query_all has shape [2*B, L, ...] (rope+prope concatenated), so mask needs 2*B + if attention_mask is None: + batch_size, seq_len = q.shape[0], q.shape[1] + attention_mask = torch.ones(batch_size * 2, seq_len, device=q.device, dtype=q.dtype) + + if q.dtype == torch.float32: + from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder + attn_metadata_builder = SDPAMetadataBuilder + else: + from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder + attn_metadata_builder = FlashAttnMetadataBuilder + attn_metadata = attn_metadata_builder().build( + current_timestep=0, + attn_mask=attention_mask, + ) + with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): + hidden_states_all, _ = self.attn( + query_all, + key_all, + value_all, + ) + + hidden_states_rope, hidden_states_prope = hidden_states_all.chunk(2, dim=0) + + # # DEBUG: Check attention output and apply_fn_o + # if self.training and torch.distributed.get_rank() == 0: + # attn_all_nz = (hidden_states_all != 0).sum().item() + # rope_nz = (hidden_states_rope != 0).sum().item() + # prope_before = (hidden_states_prope != 0).sum().item() + # print(f"[DEBUG] attn output: all_nonzero={attn_all_nz}, rope_nonzero={rope_nz}, prope_before_apply={prope_before}", flush=True) + + hidden_states_prope = apply_fn_o(hidden_states_prope.transpose(1, 2)).transpose(1, 2) + + # # DEBUG: Check after apply_fn_o + # if self.training and torch.distributed.get_rank() == 0: + # prope_after = (hidden_states_prope != 0).sum().item() + # print(f"[DEBUG] prope_after_apply_fn_o={prope_after}", flush=True) + + return hidden_states_rope, hidden_states_prope diff --git a/fastvideo/models/dits/wangame/model.py b/fastvideo/models/dits/wangame/model.py new file mode 100644 index 000000000..a2ed9e357 --- /dev/null +++ b/fastvideo/models/dits/wangame/model.py @@ -0,0 +1,433 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import PatchEmbed +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanI2VCrossAttention +from fastvideo.platforms import AttentionBackendEnum, current_platform + +from .hyworld_action_module import ( + WanGameActionSelfAttention, + WanGameActionTimeImageEmbedding, +) + +logger = init_logger(__name__) + + +_DEFAULT_WANGAME_CONFIG = WanGameVideoConfig() + + +class WanGameCrossAttention(WanI2VCrossAttention): + def forward(self, x, context, context_lens=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + context_img = context + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + img_x = self.attn(q, k_img, v_img) + + # output + x = img_x.flatten(2) + x, _ = self.to_out(x) + return x + +class WanGameActionTransformerBlock(nn.Module): + """ + Transformer block for WAN Action model with support for: + - Self-attention with RoPE and camera PRoPE + - Cross-attention with text/image context + - Feed-forward network with AdaLN modulation + """ + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ReplicatedLinear(dim, dim, bias=True) + + self.attn1 = WanGameActionSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) + + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + dim_head = dim // num_heads + + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + raise ValueError(f"QK Norm type {qk_norm} not supported") + + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + compute_dtype=torch.float32) + + # 2. Cross-attention (I2V only for now) + self.attn2 = WanGameCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + # norm3 for FFN input + self.norm3 = LayerNormScaleShift(dim, norm_type="layer", eps=eps, + elementwise_affine=False) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # PRoPE output projection (initialized via add_discrete_action_parameters on the model) + self.to_out_prope = ReplicatedLinear(dim, dim, bias=True) + nn.init.zeros_(self.to_out_prope.weight) + if self.to_out_prope.bias is not None: + nn.init.zeros_(self.to_out_prope.bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + + num_frames = temb.shape[1] + frame_seqlen = hidden_states.shape[1] // num_frames + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + + # Cast temb to float32 for scale/shift computation + e = self.scale_shift_table + temb.float() + assert e.shape == (bs, num_frames, 6, self.hidden_dim) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).to(orig_dtype).flatten(1, 2) + + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q.forward_native(query) + if self.norm_k is not None: + key = self.norm_k.forward_native(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + # Self-attention with camera PRoPE + attn_output_rope, attn_output_prope = self.attn1( + query, key, value, freqs_cis, + kv_cache, current_start, cache_start, viewmats, Ks, + is_cache=is_cache + ) + # Combine rope and prope outputs + attn_output_rope = attn_output_rope.flatten(2) + attn_output_rope, _ = self.to_out(attn_output_rope) + attn_output_prope = attn_output_prope.flatten(2) + + # # DEBUG: Check if prope input is zero + # if self.training and torch.distributed.get_rank() == 0: + # prope_nonzero = (attn_output_prope != 0).sum().item() + # prope_total = attn_output_prope.numel() + # if prope_nonzero == 0: + # print(f"[DEBUG] to_out_prope INPUT is ALL ZEROS! shape={attn_output_prope.shape}", flush=True) + + attn_output_prope, _ = self.to_out_prope(attn_output_prope) + attn_output = attn_output_rope.squeeze(1) + attn_output_prope.squeeze(1) + + # Self-attention residual + norm in float32 + null_shift = null_scale = torch.zeros(1, device=hidden_states.device, dtype=torch.float32) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states.float(), attn_output.float(), gate_msa, null_shift, null_scale) + hidden_states = hidden_states.type_as(attn_output) + norm_hidden_states = norm_hidden_states.type_as(attn_output) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states.to(orig_dtype), + context=encoder_hidden_states, + context_lens=None) + # Cross-attention residual in bfloat16 + hidden_states = hidden_states + attn_output + + # norm3 for FFN input in float32 + norm_hidden_states = self.norm3( + hidden_states.float(), c_shift_msa, c_scale_msa + ).type_as(hidden_states) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states.to(orig_dtype)) + hidden_states = self.mlp_residual(hidden_states.float(), ff_output.float(), c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) # Cast back to original dtype + + return hidden_states + +class WanGameActionTransformer3DModel(BaseDiT): + """ + WAN Action Transformer 3D Model for video generation with action conditioning. + + Extends the base WAN video model with: + - Action embedding support for controllable generation + - camera PRoPE attention for 3D-aware generation + - KV caching for autoregressive inference + """ + supports_action_input = True + + _fsdp_shard_conditions = _DEFAULT_WANGAME_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_WANGAME_CONFIG._compile_conditions + _supported_attention_backends = ( + _DEFAULT_WANGAME_CONFIG._supported_attention_backends + ) + param_names_mapping = _DEFAULT_WANGAME_CONFIG.param_names_mapping + reverse_param_names_mapping = ( + _DEFAULT_WANGAME_CONFIG.reverse_param_names_mapping + ) + lora_param_names_mapping = _DEFAULT_WANGAME_CONFIG.lora_param_names_mapping + + def __init__(self, config: WanGameVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + 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.attention_head_dim = config.attention_head_dim + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.local_attn_size = config.local_attn_size + self.inner_dim = inner_dim + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings (with action support) + self.condition_embedder = WanGameActionTimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + WanGameActionTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}") + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.num_frame_per_block = config.arch_config.num_frames_per_block + assert self.num_frame_per_block <= 3 + + self.__post_init__() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor], + guidance=None, + action: torch.Tensor | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + kv_cache: list[dict] | None = None, + crossattn_cache: list[dict] | None = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + is_cache: bool = False, + **kwargs + ) -> torch.Tensor: + """ + Forward pass for both training and inference with KV caching. + + Args: + hidden_states: Video latents [B, C, T, H, W] + encoder_hidden_states: Text embeddings [B, L, D] + timestep: Timestep tensor + encoder_hidden_states_image: Optional image embeddings + action: Action tensor [B, T] for per-frame conditioning + viewmats: Camera view matrices for PRoPE [B, T, 4, 4] + Ks: Camera intrinsics for PRoPE [B, T, 3, 3] + kv_cache: KV cache for autoregressive inference (list of dicts per layer) + crossattn_cache: Cross-attention cache for inference + current_start: Current position for KV cache + cache_start: Cache start position + start_frame: RoPE offset for new frames in autoregressive mode + is_cache: If True, populate KV cache and return early (cache-only mode) + """ + orig_dtype = hidden_states.dtype + # if not isinstance(encoder_hidden_states, torch.Tensor): + # encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + # else: + # encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, action, encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image) + + # condition_embedder returns: + # - temb: [B*T, dim] where T = post_patch_num_frames + # - timestep_proj: [B*T, 6*dim] + # Reshape to [B, T, 6, dim] for transformer blocks + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)) # [B*T, 6, dim] + timestep_proj = timestep_proj.view(batch_size, post_patch_num_frames, 6, self.hidden_size) # [B, T, 6, dim] + + encoder_hidden_states = encoder_hidden_states_image + + # Transformer blocks + for block_idx, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + kv_cache[block_idx] if kv_cache else None, + crossattn_cache[block_idx] if crossattn_cache else None, + current_start, cache_start, + viewmats, Ks, is_cache) + else: + hidden_states = block( + hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + kv_cache[block_idx] if kv_cache else None, + crossattn_cache[block_idx] if crossattn_cache else None, + current_start, cache_start, + viewmats, Ks, is_cache) + + # If cache-only mode, return early + if is_cache: + return kv_cache + + # Output norm, projection & unpatchify + # temb is [B*T, dim], reshape to [B, T, 1, dim] + temb = temb.view(batch_size, post_patch_num_frames, -1).unsqueeze(2) # [B, T, 1, dim] + + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 6ee1b28c3..dd502f54f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -844,6 +844,15 @@ class TransformerLoader(ComponentLoader): cls_name.startswith("Cosmos25") or cls_name == "Cosmos25Transformer3DModel" or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25" + ) and not ( + cls_name.startswith("WanGame") + or cls_name == "WanGameActionTransformer3DModel" + or cls_name.startswith("CausalWan") + or getattr(fastvideo_args.pipeline_config, "prefix", "") == "WanGame" + or cls_name.startswith("WanLingBot") + or cls_name == "WanLingBotTransformer3DModel" + or getattr(fastvideo_args.pipeline_config, "prefix", "") == "WanLingBot" + or cls_name.startswith("CausalWanGameActionTransformer3DModel") ) model = maybe_load_fsdp_model( model_cls=model_cls, diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 9ba60320a..7ca9e4a2b 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -290,8 +290,31 @@ def load_model_from_full_model_state_dict( """ meta_sd = model.state_dict() sharded_sd = {} - custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( - full_sd_iterator, param_names_mapping) # type: ignore + if param_names_mapping is None: + custom_param_sd = dict(full_sd_iterator) + reverse_param_names_mapping = { + name: (name, None, None) + for name in custom_param_sd + } + else: + def _mapping_with_passthrough( + source_param_name: str, + ) -> tuple[str, Any, Any]: + target_param_name, merge_index, num_params_to_merge = ( + param_names_mapping(source_param_name) + ) + # Custom override safetensors may already use FastVideo-internal + # names. In that case, keep the source name instead of remapping it + # a second time through the HF -> custom regex rules. + if (target_param_name not in meta_sd + and source_param_name in meta_sd): + return source_param_name, None, None + return target_param_name, merge_index, num_params_to_merge + + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( + full_sd_iterator, + _mapping_with_passthrough, + ) for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index a22a582e0..f2de89eff 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -46,6 +46,12 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "CausalWanGameTransformer3DModel": + ("dits", "wangame", "CausalWanGameActionTransformer3DModel"), + "CausalWanGameActionTransformer3DModel": + ("dits", "wangame", "CausalWanGameActionTransformer3DModel"), + "WanGameActionTransformer3DModel": + ("dits", "wangame", "WanGameActionTransformer3DModel"), "MatrixGameWanModel": ("dits", "matrixgame", "MatrixGameWanModel"), "CausalMatrixGameWanModel": ("dits", "matrixgame", "CausalMatrixGameWanModel"), } diff --git a/fastvideo/pipelines/basic/wan/__init__.py b/fastvideo/pipelines/basic/wan/__init__.py index e69de29bb..98aba65ef 100644 --- a/fastvideo/pipelines/basic/wan/__init__.py +++ b/fastvideo/pipelines/basic/wan/__init__.py @@ -0,0 +1,3 @@ +from fastvideo.pipelines.basic.wan.wangame_causal_dmd_pipeline import ( + WangameCausalOdeDMDPipeline as WangameCausalOdeDMDPipeline, + WangameCausalSdeDMDPipeline as WangameCausalSdeDMDPipeline, ) diff --git a/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py new file mode 100644 index 000000000..c5bd1c50b --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wangame causal DMD pipeline implementations.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline + +from fastvideo.pipelines.stages import ( + ConditioningStage, DecodingStage, MatrixGameCausalDenoisingStage, + MatrixGameCausalOdeDenoisingStage, MatrixGameImageEncodingStage, + InputValidationStage, LatentPreparationStage, TextEncodingStage, + TimestepPreparationStage) +from fastvideo.pipelines.stages.image_encoding import ( + MatrixGameImageVAEEncodingStage) + +logger = init_logger(__name__) + + +class _WangameCausalDMDPipelineBase(LoRAPipeline, ComposedPipelineBase): + requires_timestep_preparation: bool + requires_dmd_denoising_steps: bool + denoising_stage_cls: type[MatrixGameCausalDenoisingStage] + + _required_config_modules = [ + "vae", "transformer", "scheduler", "image_encoder", "image_processor" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + del fastvideo_args + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + if (self.get_module("text_encoder", None) is not None + and self.get_module("tokenizer", None) is not None): + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + if (self.get_module("image_encoder", None) is not None + and self.get_module("image_processor", None) is not None): + 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="conditioning_stage", + stage=ConditioningStage()) + + if self.requires_timestep_preparation: + 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="image_latent_preparation_stage", + stage=MatrixGameImageVAEEncodingStage(vae=self.get_module("vae"))) + + denoising_stage = self.denoising_stage_cls( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self, + vae=self.get_module("vae"), + ) + + self.add_stage(stage_name="denoising_stage", stage=denoising_stage) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + logger.info("%s initialized with action support", + self.__class__.__name__) + +class WangameCausalOdeDMDPipeline(_WangameCausalDMDPipelineBase): + requires_timestep_preparation = True + requires_dmd_denoising_steps = False + denoising_stage_cls = MatrixGameCausalOdeDenoisingStage + + +class WangameCausalSdeDMDPipeline(_WangameCausalDMDPipelineBase): + requires_timestep_preparation = False + requires_dmd_denoising_steps = True + denoising_stage_cls = MatrixGameCausalDenoisingStage + +EntryClass = [WangameCausalOdeDMDPipeline, WangameCausalSdeDMDPipeline] diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index ee9086e82..ae1d445eb 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -160,6 +160,10 @@ class ForwardBatch: # Timesteps timesteps: torch.Tensor | None = None + # Optional explicit denoising-loop timesteps (sampler-specific). + # When set, some samplers iterate this list instead of `timesteps` + # produced by the timestep preparation stage. + sampling_timesteps: torch.Tensor | None = None timestep: torch.Tensor | float | int | None = None step_index: int | None = None boundary_ratio: float | None = None diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 9bd5280b8..330561fbd 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -34,7 +34,7 @@ from fastvideo.pipelines.stages.ltx2_latent_preparation import ( LTX2LatentPreparationStage) from fastvideo.pipelines.stages.ltx2_text_encoding import LTX2TextEncodingStage from fastvideo.pipelines.stages.matrixgame_denoising import ( - MatrixGameCausalDenoisingStage) + MatrixGameCausalDenoisingStage, MatrixGameCausalOdeDenoisingStage) from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage from fastvideo.pipelines.stages.text_encoding import (Cosmos25TextEncodingStage, @@ -66,6 +66,7 @@ __all__ = [ "CausalDMDDenosingStage", "CausalDenoisingStage", "MatrixGameCausalDenoisingStage", + "MatrixGameCausalOdeDenoisingStage", "HYWorldDenoisingStage", "GameCraftDenoisingStage", "CosmosDenoisingStage", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 1b11478ce..2acf1bdfe 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -1237,6 +1237,37 @@ class DmdDenoisingStage(DenoisingStage): }, ) + if batch.mouse_cond is not None and batch.keyboard_cond is not None: + from fastvideo.models.dits.hyworld.pose import process_custom_actions + + viewmats, intrinsics, action_labels = process_custom_actions( + batch.keyboard_cond, batch.mouse_cond) + camera_action_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "viewmats": + viewmats.unsqueeze(0).to( + get_local_torch_device(), dtype=target_dtype), + "Ks": + intrinsics.unsqueeze(0).to( + get_local_torch_device(), dtype=target_dtype), + "action": + action_labels.unsqueeze(0).to( + get_local_torch_device(), dtype=target_dtype), + }, + ) + else: + camera_action_kwargs = {} + + action_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "mouse_cond": batch.mouse_cond, + "keyboard_cond": batch.keyboard_cond, + "c2ws_plucker_emb": batch.c2ws_plucker_emb, + }, + ) + # Get latents and embeddings assert batch.latents is not None, "latents must be provided" latents = batch.latents @@ -1245,14 +1276,29 @@ class DmdDenoisingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds assert not torch.isnan( prompt_embeds[0]).any(), "prompt_embeds contains nan" - timesteps = torch.tensor( - fastvideo_args.pipeline_config.dmd_denoising_steps, - dtype=torch.long, - device=get_local_torch_device()) + loop_timesteps = batch.sampling_timesteps + if loop_timesteps is None: + legacy = getattr(fastvideo_args.pipeline_config, + "dmd_denoising_steps", None) + if legacy is not None: + loop_timesteps = torch.tensor(legacy, dtype=torch.long) + else: + loop_timesteps = batch.timesteps + + if loop_timesteps is None: + raise ValueError( + "SDE sampling requires `batch.sampling_timesteps` " + "(preferred) or `pipeline_config.dmd_denoising_steps`.") + if not isinstance(loop_timesteps, torch.Tensor): + loop_timesteps = torch.tensor(loop_timesteps, dtype=torch.long) + if loop_timesteps.ndim != 1: + raise ValueError("Expected 1D `sampling_timesteps`, got shape " + f"{tuple(loop_timesteps.shape)}") + loop_timesteps = loop_timesteps.to(get_local_torch_device()) # Run denoising loop - with self.progress_bar(total=len(timesteps)) as progress_bar: - for i, t in enumerate(timesteps): + with self.progress_bar(total=len(loop_timesteps)) as progress_bar: + for i, t in enumerate(loop_timesteps): # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: continue @@ -1326,6 +1372,8 @@ class DmdDenoisingStage(DenoisingStage): guidance=guidance_expand, **image_kwargs, **pos_cond_kwargs, + **action_kwargs, + **camera_action_kwargs, ).permute(0, 2, 1, 3, 4) pred_video = pred_noise_to_pred_video( @@ -1335,8 +1383,8 @@ class DmdDenoisingStage(DenoisingStage): scheduler=self.scheduler).unflatten( 0, pred_noise.shape[:2]) - if i < len(timesteps) - 1: - next_timestep = timesteps[i + 1] * torch.ones( + if i < len(loop_timesteps) - 1: + next_timestep = loop_timesteps[i + 1] * torch.ones( [1], dtype=torch.long, device=pred_video.device) noise = torch.randn(video_raw_latent_shape, dtype=pred_video.dtype, @@ -1349,7 +1397,7 @@ class DmdDenoisingStage(DenoisingStage): latents = pred_video # Update progress bar - if i == len(timesteps) - 1 or ( + if i == len(loop_timesteps) - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None): diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index cd9599bbb..4b058089c 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -1,6 +1,8 @@ from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass +import hashlib +import os from typing import Any import torch # type: ignore @@ -26,6 +28,22 @@ except ImportError: logger = init_logger(__name__) +def _should_debug_timesteps(prompt: object) -> bool: + if not os.environ.get("FASTVIDEO_DEBUG_TIMESTEPS"): + return False + if isinstance(prompt, str): + return prompt.startswith("00 Val-00:") + if isinstance(prompt, list) and prompt: + first = prompt[0] + return isinstance(first, str) and first.startswith("00 Val-00:") + return False + + +def _tensor_md5(tensor: torch.Tensor) -> str: + array = tensor.detach().cpu().to(torch.int64).contiguous().numpy() + return hashlib.md5(array.tobytes()).hexdigest() + + @dataclass class BlockProcessingContext: """Dataclass contains for block processing.""" @@ -54,6 +72,9 @@ class BlockProcessingContext: image_kwargs: dict[str, Any] pos_cond_kwargs: dict[str, Any] + viewmats_full: torch.Tensor | None = None + intrinsics_full: torch.Tensor | None = None + action_full: torch.Tensor | None = None def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]: if self.boundary_timestep is not None: @@ -97,10 +118,12 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): -1) except Exception: self.local_attn_size = -1 + try: + self.local_attn_size = getattr(self.transformer.model, + "local_attn_size", -1) + except Exception: + self.local_attn_size = -1 - assert self.local_attn_size != -1, ( - f"local_attn_size must be set for Matrix-Game causal inference, " - f"got {self.local_attn_size}. Check MatrixGameWanVideoArchConfig.") assert self.num_frame_per_block > 0, ( f"num_frame_per_block must be positive, got {self.num_frame_per_block}" ) @@ -126,7 +149,10 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): ) and not fastvideo_args.disable_autocast latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] - patch_size = self.transformer.patch_size + if hasattr(self.transformer, "patch_size"): + patch_size = self.transformer.patch_size + else: + patch_size = self.transformer.config.arch_config.patch_size patch_ratio = patch_size[-1] * patch_size[-2] self.frame_seq_length = latent_seq_length // patch_ratio @@ -166,6 +192,31 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 + viewmats_full = None + intrinsics_full = None + action_full = None + if batch.mouse_cond is not None and batch.keyboard_cond is not None: + from fastvideo.models.dits.hyworld.pose import process_custom_actions + + viewmats_list = [] + intrinsics_list = [] + action_list = [] + for bi in range(b): + vm, ks, action = process_custom_actions(batch.keyboard_cond[bi], + batch.mouse_cond[bi]) + viewmats_list.append(vm) + intrinsics_list.append(ks) + action_list.append(action) + viewmats_full = torch.stack(viewmats_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + intrinsics_full = torch.stack(intrinsics_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + action_full = torch.stack(action_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) @@ -225,6 +276,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): "context_noise", 0), image_kwargs=image_kwargs, pos_cond_kwargs=pos_cond_kwargs, + viewmats_full=viewmats_full, + intrinsics_full=intrinsics_full, + action_full=action_full, ) context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", @@ -240,6 +294,8 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): action_kwargs = self._prepare_action_kwargs( batch, start_index, current_num_frames) + camera_action_kwargs = self._prepare_camera_action_kwargs( + ctx, start_index, current_num_frames) current_latents = self._process_single_block( current_latents=current_latents, @@ -249,6 +305,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): timesteps=timesteps, ctx=ctx, action_kwargs=action_kwargs, + camera_action_kwargs=camera_action_kwargs, progress_bar=progress_bar, ) @@ -263,6 +320,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): current_num_frames=current_num_frames, ctx=ctx, action_kwargs=action_kwargs, + camera_action_kwargs=camera_action_kwargs, context_noise=context_noise, ) @@ -324,9 +382,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), }) return kv_cache @@ -362,9 +420,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), }) kv_cache_mouse.append({ "k": @@ -382,9 +440,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=dtype, device=device), "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), + torch.zeros((), dtype=torch.long, device=device), }) return kv_cache_mouse, kv_cache_keyboard @@ -418,6 +476,18 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): }) return crossattn_cache + def _prepare_camera_action_kwargs( + self, ctx: BlockProcessingContext, start_index: int, + current_num_frames: int) -> dict[str, Any]: + if ctx.action_full is None or ctx.viewmats_full is None or ctx.intrinsics_full is None: + return {} + end_index = start_index + current_num_frames + return { + "viewmats": ctx.viewmats_full[:, start_index:end_index], + "Ks": ctx.intrinsics_full[:, start_index:end_index], + "action": ctx.action_full[:, start_index:end_index], + } + def _process_single_block( self, current_latents: torch.Tensor, @@ -427,6 +497,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): timesteps: torch.Tensor, ctx: BlockProcessingContext, action_kwargs: dict[str, Any], + camera_action_kwargs: dict[str, Any], noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor] | None = None, progress_bar: Any | None = None, @@ -445,7 +516,16 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): independent_first_frame = getattr(self.transformer, 'independent_first_frame', False) - if batch.image_latent is not None and independent_first_frame and start_index == 0: + if batch.image_latent is not None and not independent_first_frame: + image_latent_chunk = batch.image_latent[:, :, start_index: + start_index + + current_num_frames, :, :] + latent_model_input = torch.cat([ + latent_model_input, + image_latent_chunk.to(ctx.target_dtype) + ], + dim=1) + elif batch.image_latent is not None and independent_first_frame and start_index == 0: latent_model_input = torch.cat([ latent_model_input, batch.image_latent.to(ctx.target_dtype) @@ -495,6 +575,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): "crossattn_cache": ctx.crossattn_cache, "current_start": start_index * self.frame_seq_length, "start_frame": start_index, + "is_cache": False, } if self.use_action_module and current_model == self.transformer: @@ -510,6 +591,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): latent_model_input, prompt_embeds, t_expanded_noise, + **camera_action_kwargs, **ctx.image_kwargs, **ctx.pos_cond_kwargs, **model_kwargs, @@ -582,6 +664,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): current_num_frames: int, ctx: BlockProcessingContext, action_kwargs: dict[str, Any], + camera_action_kwargs: dict[str, Any], context_noise: float, ) -> None: prompt_embeds = batch.prompt_embeds @@ -592,6 +675,17 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): device=latents_device, dtype=torch.long) * int(context_noise) context_bcthw = current_latents.to(ctx.target_dtype) + context_input = context_bcthw + independent_first_frame = getattr(self.transformer, + "independent_first_frame", False) + if batch.image_latent is not None and not independent_first_frame: + image_context_chunk = batch.image_latent[:, :, + start_index:start_index + + current_num_frames, :, :] + context_input = torch.cat( + [context_input, + image_context_chunk.to(ctx.target_dtype)], + dim=1) with torch.autocast(device_type="cuda", dtype=ctx.target_dtype, @@ -605,6 +699,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): "crossattn_cache": ctx.crossattn_cache, "current_start": start_index * self.frame_seq_length, "start_frame": start_index, + "is_cache": True, } if self.use_action_module: @@ -617,26 +712,409 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): context_model_kwargs.update(action_kwargs) if ctx.boundary_timestep is not None and self.transformer_2 is not None: - self.transformer_2( - context_bcthw, + cache_update_ret_2 = self.transformer_2( + context_input, prompt_embeds, t_context, kv_cache=ctx.kv_cache2, crossattn_cache=ctx.crossattn_cache, current_start=start_index * self.frame_seq_length, start_frame=start_index, + is_cache=True, + **camera_action_kwargs, **ctx.image_kwargs, **ctx.pos_cond_kwargs, ) + if isinstance(cache_update_ret_2, + list) and len(cache_update_ret_2) > 0: + ctx.kv_cache2 = cache_update_ret_2 - self.transformer( - context_bcthw, + cache_update_ret = self.transformer( + context_input, prompt_embeds, t_context, + **camera_action_kwargs, **ctx.image_kwargs, **ctx.pos_cond_kwargs, **context_model_kwargs, ) + if isinstance(cache_update_ret, list) and len(cache_update_ret) > 0: + ctx.kv_cache1 = cache_update_ret + + +class MatrixGameCausalOdeDenoisingStage(MatrixGameCausalDenoisingStage): + """Causal ODE denoising for WanGame/MatrixGame. + + This is the deterministic counterpart of `MatrixGameCausalDenoisingStage`. + It performs block-by-block causal rollout, but uses the scheduler's ODE-style + `step()` update (no re-noising between steps). + """ + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + timesteps = batch.timesteps + if timesteps is None: + raise ValueError( + "MatrixGameCausalOdeDenoisingStage requires batch.timesteps. " + "Make sure TimestepPreparationStage runs before this stage.") + + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + if hasattr(self.transformer, "patch_size"): + patch_size = self.transformer.patch_size + else: + patch_size = self.transformer.config.arch_config.patch_size + patch_ratio = patch_size[-1] * patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + + timesteps = timesteps.to(get_local_torch_device()) + if _should_debug_timesteps(batch.prompt): + logger.info( + "DEBUG_TIMESTEPS stage=ode_denoising rank=%s prompt=%r " + "len=%s md5=%s head=%s tail=%s", + os.environ.get("RANK", "?"), + batch.prompt, + int(timesteps.numel()), + _tensor_md5(timesteps), + timesteps[:10].detach().cpu().tolist(), + timesteps[-10:].detach().cpu().tolist(), + ) + + boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, + "boundary_ratio", None) + if boundary_ratio is not None: + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + else: + boundary_timestep = None + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + # directly set the kwarg. + image_kwargs = {"encoder_hidden_states_image": image_embeds} + pos_cond_kwargs: dict[str, Any] = {} + + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + b, c, t, h, w = latents.shape + + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + viewmats_full = None + intrinsics_full = None + action_full = None + if batch.mouse_cond is not None and batch.keyboard_cond is not None: + from fastvideo.models.dits.hyworld.pose import process_custom_actions + + viewmats_list = [] + intrinsics_list = [] + action_list = [] + for bi in range(b): + vm, ks, action = process_custom_actions(batch.keyboard_cond[bi], + batch.mouse_cond[bi]) + viewmats_list.append(vm) + intrinsics_list.append(ks) + action_list.append(action) + viewmats_full = torch.stack(viewmats_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + intrinsics_full = torch.stack(intrinsics_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + action_full = torch.stack(action_list, + dim=0).to(device=latents.device, + dtype=target_dtype) + + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + kv_cache_mouse = None + kv_cache_keyboard = None + if self.use_action_module: + kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache( + batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=257, # 1 CLS + 256 patch tokens + dtype=target_dtype, + device=latents.device) + + if t % self.num_frame_per_block != 0: + raise ValueError( + "num_frames must be divisible by num_frame_per_block for causal denoising" + ) + num_blocks = t // self.num_frame_per_block + block_sizes = [self.num_frame_per_block] * num_blocks + start_index = 0 + + if boundary_timestep is not None: + block_sizes[0] = 1 + + ctx = BlockProcessingContext( + batch=batch, + block_idx=0, + start_index=0, + kv_cache1=kv_cache1, + kv_cache2=kv_cache2, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + timesteps=timesteps, + block_sizes=block_sizes, + noise_pool=None, + fastvideo_args=fastvideo_args, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + boundary_timestep=boundary_timestep, + high_noise_timesteps=None, + context_noise=getattr(fastvideo_args.pipeline_config, + "context_noise", 0), + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + viewmats_full=viewmats_full, + intrinsics_full=intrinsics_full, + action_full=action_full, + ) + + context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", + 0) + + with self.progress_bar(total=len(block_sizes) * + len(timesteps)) as progress_bar: + for block_idx, current_num_frames in enumerate(block_sizes): + ctx.block_idx = block_idx + ctx.start_index = start_index + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + + # The scheduler maintains an internal `step_index` (and potentially + # additional multistep state, e.g. UniPC). Since causal streaming runs + # a full denoising trajectory *per block*, reset that state before + # each block rollout. + self._reset_scheduler_state_for_new_rollout() + + action_kwargs = self._prepare_action_kwargs( + batch, start_index, current_num_frames) + camera_action_kwargs = self._prepare_camera_action_kwargs( + ctx, start_index, current_num_frames) + + current_latents = self._process_single_block_ode( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timesteps=timesteps, + ctx=ctx, + action_kwargs=action_kwargs, + camera_action_kwargs=camera_action_kwargs, + progress_bar=progress_bar, + ) + + latents[:, :, start_index:start_index + + current_num_frames, :, :] = current_latents + + # Update KV caches with clean context + self._update_context_cache( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=action_kwargs, + camera_action_kwargs=camera_action_kwargs, + context_noise=context_noise, + ) + + start_index += current_num_frames + + if boundary_timestep is not None: + num_frames_to_remove = self.num_frame_per_block - 1 + if num_frames_to_remove > 0: + latents = latents[:, :, :-num_frames_to_remove, :, :] + + batch.latents = latents + return batch + + def _reset_scheduler_state_for_new_rollout(self) -> None: + scheduler = self.scheduler + + # Common diffusers-like state. + if hasattr(scheduler, "_step_index"): + scheduler._step_index = None # type: ignore[attr-defined] + if hasattr(scheduler, "_begin_index"): + scheduler._begin_index = None # type: ignore[attr-defined] + + # UniPC multistep state (FlowUniPCMultistepScheduler) needs additional reset + # between independent trajectories. + if hasattr(scheduler, "model_outputs") and hasattr(scheduler, "config"): + try: + solver_order = int( + getattr(scheduler.config, "solver_order", 0) or 0) + except Exception: + solver_order = 0 + if solver_order > 0: + scheduler.model_outputs = [ + None + ] * solver_order # type: ignore[attr-defined] + if hasattr(scheduler, "timestep_list") and hasattr(scheduler, "config"): + try: + solver_order = int( + getattr(scheduler.config, "solver_order", 0) or 0) + except Exception: + solver_order = 0 + if solver_order > 0: + scheduler.timestep_list = [ + None + ] * solver_order # type: ignore[attr-defined] + if hasattr(scheduler, "lower_order_nums"): + scheduler.lower_order_nums = 0 # type: ignore[attr-defined] + if hasattr(scheduler, "last_sample"): + scheduler.last_sample = None # type: ignore[attr-defined] + + def _process_single_block_ode( + self, + *, + current_latents: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + timesteps: torch.Tensor, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + camera_action_kwargs: dict[str, Any], + progress_bar: Any | None = None, + ) -> torch.Tensor: + prompt_embeds = batch.prompt_embeds + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": batch.generator, + "eta": batch.eta, + }, + ) + + for i, t_cur in enumerate(timesteps): + if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep: + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + else: + current_model = self.transformer + + latent_model_input = current_latents.to(ctx.target_dtype) + + independent_first_frame = getattr(self.transformer, + "independent_first_frame", False) + if batch.image_latent is not None and not independent_first_frame: + image_latent_chunk = batch.image_latent[:, :, start_index: + start_index + + current_num_frames, :, :] + latent_model_input = torch.cat([ + latent_model_input, + image_latent_chunk.to(ctx.target_dtype) + ], + dim=1) + elif (batch.image_latent is not None and independent_first_frame + and start_index == 0): + latent_model_input = torch.cat([ + latent_model_input, + batch.image_latent.to(ctx.target_dtype) + ], + dim=2) + + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t_cur) + + # Build attention metadata if VSA is available + if vsa_available and self.attn_backend == VideoSparseAttentionBackend: + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + h, w = current_latents.shape[-2:] + attn_metadata = self.attn_metadata_builder.build( + current_timestep=i, + raw_latent_shape=(current_num_frames, h, w), + patch_size=ctx.fastvideo_args.pipeline_config. + dit_config.patch_size, + VSA_sparsity=ctx.fastvideo_args.VSA_sparsity, + device=get_local_torch_device(), + ) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + t_expanded = t_cur * torch.ones( + (latent_model_input.shape[0], current_num_frames), + device=latent_model_input.device, + dtype=t_cur.dtype) + + model_kwargs = { + "kv_cache": ctx.get_kv_cache(t_cur), + "crossattn_cache": ctx.crossattn_cache, + "current_start": start_index * self.frame_seq_length, + "start_frame": start_index, + "is_cache": False, + } + + if self.use_action_module and current_model == self.transformer: + model_kwargs.update({ + "kv_cache_mouse": + ctx.kv_cache_mouse, + "kv_cache_keyboard": + ctx.kv_cache_keyboard, + }) + model_kwargs.update(action_kwargs) + + noise_pred = current_model( + latent_model_input, + prompt_embeds, + t_expanded, + **camera_action_kwargs, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + **model_kwargs, + ) + + current_latents = self.scheduler.step( + noise_pred, + t_cur, + current_latents, + **extra_step_kwargs, + return_dict=False, + )[0] + + if progress_bar is not None: + progress_bar.update() + + return current_latents def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: @@ -645,7 +1123,10 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): ) and not fastvideo_args.disable_autocast latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] - patch_size = self.transformer.patch_size + if hasattr(self.transformer, "patch_size"): + patch_size = self.transformer.patch_size + else: + patch_size = self.transformer.config.arch_config.patch_size patch_ratio = patch_size[-1] * patch_size[-2] self.frame_seq_length = latent_seq_length // patch_ratio @@ -821,6 +1302,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): timesteps=ctx.timesteps, ctx=ctx, action_kwargs=action_kwargs, + camera_action_kwargs={}, noise_generator=streaming_noise_generator, ) @@ -835,6 +1317,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): current_num_frames=current_num_frames, ctx=ctx, action_kwargs=action_kwargs, + camera_action_kwargs={}, context_noise=ctx.context_noise, ) diff --git a/fastvideo/train/callbacks/callback.py b/fastvideo/train/callbacks/callback.py index 024a92860..a59491945 100644 --- a/fastvideo/train/callbacks/callback.py +++ b/fastvideo/train/callbacks/callback.py @@ -173,3 +173,6 @@ class CallbackDict: fn(*args, **kwargs) return _dispatch + + def get_callback(self, name: str) -> Callback | None: + return self._callbacks.get(name) diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index 6bcf2da4f..9da6e0145 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -8,14 +8,17 @@ section. The pipeline class is resolved from from __future__ import annotations -import contextlib +import gc import os +import sys from dataclasses import dataclass +from pathlib import Path from typing import Any, TYPE_CHECKING import imageio import numpy as np import torch +import torch.distributed as dist import torchvision from einops import rearrange from torch.utils.data import DataLoader @@ -38,6 +41,8 @@ from fastvideo.utils import shallow_asdict if TYPE_CHECKING: from fastvideo.train.methods.base import TrainingMethod + from fastvideo.train.utils.training_config import ( + TrainingConfig, ) logger = init_logger(__name__) @@ -46,6 +51,7 @@ logger = init_logger(__name__) class _ValidationStepResult: videos: list[list[np.ndarray]] captions: list[str] + action_paths: list[str | None] class ValidationCallback(Callback): @@ -57,6 +63,18 @@ class ValidationCallback(Callback): contract. """ + _FLOW_EVAL_SCALAR_KEYS = ( + "mf_epe_mean", + "mf_angle_err_mean", + "mf_cosine_mean", + "mf_mag_ratio_mean", + "pixel_epe_mean_mean", + "px_angle_rmse_mean", + "fl_all_mean", + "foe_dist_mean", + "flow_kl_2d_mean", + ) + def __init__( self, *, @@ -64,20 +82,54 @@ class ValidationCallback(Callback): dataset_file: str, every_steps: int = 100, sampling_steps: list[int] | None = None, + scheduler_target: str | None = None, guidance_scale: float | None = None, num_frames: int | None = None, output_dir: str | None = None, sampling_timesteps: list[int] | None = None, + evaluate_ptlflow: bool = False, + ptlflow_dir: str = "/mnt/weka/home/hao.zhang/mhuo/FastVideo/benchmarks/ptlflow", + ptlflow_ckpt: str | None = None, + ptlflow_calibration_path: str | None = None, + ptlflow_use_depth: bool = True, **pipeline_kwargs: Any, ) -> None: self.pipeline_target = str(pipeline_target) self.dataset_file = str(dataset_file) self.every_steps = int(every_steps) self.sampling_steps = ([int(s) for s in sampling_steps] if sampling_steps else [40]) + + pipeline_cls = resolve_target(self.pipeline_target) + self._requires_dmd_denoising_steps = bool( + getattr( + pipeline_cls, + "requires_dmd_denoising_steps", + False, + ) + ) + self.scheduler_target = ( + str(scheduler_target) + if scheduler_target is not None + else None + ) self.guidance_scale = (float(guidance_scale) if guidance_scale is not None else None) self.num_frames = (int(num_frames) if num_frames is not None else None) self.output_dir = (str(output_dir) if output_dir is not None else None) self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None) + + self.evaluate_ptlflow = bool(evaluate_ptlflow) + self.ptlflow_dir = str(ptlflow_dir) + self.ptlflow_ckpt = ( + str(ptlflow_ckpt) + if ptlflow_ckpt is not None + else str(Path(self.ptlflow_dir) / "dpflow-things-2012b5d6.ckpt") + ) + self.ptlflow_calibration_path = ( + str(ptlflow_calibration_path) + if ptlflow_calibration_path is not None + else str(Path(self.ptlflow_dir) / "calibration.json") + ) + self.ptlflow_use_depth = bool(ptlflow_use_depth) self.pipeline_kwargs = dict(pipeline_kwargs) # Set after on_train_start. @@ -87,6 +139,12 @@ class ValidationCallback(Callback): self.tracker: Any = DummyTracker() self.validation_random_generator: (torch.Generator | None) = None self.seed: int = 0 + self._flow_eval_init_done: bool = False + self._flow_eval_ready: bool = False + self._flow_eval_fn: Any | None = None + self._last_scalar_metrics: dict[str, float] = {} + self._last_mf_angle_err_mean: float = float("inf") + self._last_eval_step: int = -1 # ---------------------------------------------------------- # Callback hooks @@ -138,80 +196,72 @@ class ValidationCallback(Callback): method: TrainingMethod, step: int, ) -> None: - - transformer = method.student.transformer - # Look for an EMA callback to temporarily swap - # EMA weights during validation. - ema_cb = self._find_ema_callback() - if ema_cb is not None: - ctx = ema_cb.ema_context(transformer) - else: - ctx = contextlib.nullcontext(transformer) - with ctx as t: - self._run_validation_inner( - method, - step, - t, - ) - - def _find_ema_callback(self) -> Any | None: - """Find the EMA callback in the callback dict.""" - from fastvideo.train.callbacks.ema import ( - EMACallback, ) - - cb_dict = getattr(self, "_callback_dict", None) - if cb_dict is not None: - for cb in cb_dict._callbacks.values(): - if isinstance(cb, EMACallback): - return cb - return None - - def _run_validation_inner( - self, - method: TrainingMethod, - step: int, - transformer: torch.nn.Module, - ) -> None: tc = self.training_config - was_training = bool(getattr(transformer, "training", False)) + self._last_scalar_metrics = {} + self._last_mf_angle_err_mean = float("inf") + self._last_eval_step = int(step) + # Use EMA transformer for validation when available. + transformer = method.transformer_inference + was_training = bool( + getattr(transformer, "training", False) + ) - output_dir = (self.output_dir or tc.checkpoint.output_dir) + output_dir = ( + self.output_dir + or tc.checkpoint.output_dir + or os.getcwd() + ) + # For streaming SDE pipelines we may need to + # temporarily set dmd_denoising_steps on + # pipeline_config. + old_dmd_denoising_steps = getattr( + tc.pipeline_config, + "dmd_denoising_steps", + None, + ) try: transformer.eval() num_sp_groups = (self.world_group.world_size // self.sp_group.world_size) for num_inference_steps in self.sampling_steps: + self._maybe_set_dmd_denoising_steps( + tc, + num_inference_steps, + ) + result = self._run_validation_for_steps( num_inference_steps, transformer=transformer, ) + flow_metric_keys = list( + self._FLOW_EVAL_SCALAR_KEYS + ) + local_video_filenames: list[str] = [] + local_video_captions: list[str] = [] + local_metric_sums = { + k: 0.0 for k in flow_metric_keys + } + local_metric_counts = { + k: 0.0 for k in flow_metric_keys + } + local_eval_errors: list[str] = [] + flow_eval_enabled = bool( + self.evaluate_ptlflow + ) - if self.rank_in_sp_group != 0: - continue - - if self.global_rank == 0: - all_videos = list(result.videos) - all_captions = list(result.captions) - for sp_idx in range(1, num_sp_groups): - src = (sp_idx * self.sp_world_size) - recv_v = (self.world_group.recv_object(src=src)) - recv_c = (self.world_group.recv_object(src=src)) - all_videos.extend(recv_v) - all_captions.extend(recv_c) - + if self.rank_in_sp_group == 0: os.makedirs( - output_dir, - exist_ok=True, + output_dir, exist_ok=True, ) - video_filenames: list[str] = [] sp = self._get_sampling_param() - for i, video in enumerate(all_videos): + for i, video in enumerate(result.videos): fname = os.path.join( output_dir, f"validation_step_{step}" f"_inference_steps_" f"{num_inference_steps}" + f"_rank_{self.global_rank}" f"_video_{i}.mp4", ) imageio.mimsave( @@ -219,13 +269,131 @@ class ValidationCallback(Callback): video, fps=sp.fps, ) - video_filenames.append(fname) + local_video_filenames.append( + fname + ) + local_video_captions = list( + result.captions + ) + + if flow_eval_enabled: + self._init_flow_eval_module() + if not self._flow_eval_ready: + logger.warning( + "PTLFlow evaluation is " + "enabled but evaluator " + "initialization failed. " + "Skipping flow metrics " + "for this validation run." + ) + flow_eval_enabled = False + + if flow_eval_enabled: + for fname, cap, action_path in zip( + local_video_filenames, + result.captions, + result.action_paths, + strict=True, + ): + try: + sample_metrics = ( + self._evaluate_validation_video( + video_path=fname, + caption=cap, + action_path=action_path, + global_step=step, + num_inference_steps=num_inference_steps, + ) + ) + if sample_metrics: + for key in ( + flow_metric_keys + ): + val = ( + sample_metrics.get( + key + ) + ) + if not isinstance( + val, + ( + float, + int, + np.floating, + np.integer, + ), + ): + continue + val_float = float( + val + ) + if not np.isfinite( + val_float + ): + continue + local_metric_sums[ + key + ] += val_float + local_metric_counts[ + key + ] += 1.0 + except Exception as e: + err = ( + "Validation flow " + "evaluation failed " + f"on rank " + f"{self.global_rank} " + f"for {fname}: {e}" + ) + logger.exception(err) + local_eval_errors.append( + err + ) + finally: + # PTLFlow allocates large CUDA + # buffers. Free caches aggressively + # so training backward can allocate + # cuBLAS handles. + self._release_validation_cuda_memory() + + if ( + self.rank_in_sp_group == 0 + and self.global_rank == 0 + ): + all_video_filenames = list( + local_video_filenames + ) + all_captions = list( + local_video_captions + ) + for sp_idx in range( + 1, num_sp_groups + ): + src = ( + sp_idx * self.sp_world_size + ) + recv_video_filenames = ( + self.world_group.recv_object( + src=src + ) + ) + recv_captions = ( + self.world_group.recv_object( + src=src + ) + ) + all_video_filenames.extend( + recv_video_filenames + ) + all_captions.extend( + recv_captions + ) video_logs = [] for fname, cap in zip( - video_filenames, - all_captions, - strict=True, + all_video_filenames, + all_captions, + strict=True, ): art = self.tracker.video( fname, @@ -243,19 +411,221 @@ class ValidationCallback(Callback): logs, step, ) - else: + elif self.rank_in_sp_group == 0: self.world_group.send_object( - result.videos, - dst=0, + local_video_filenames, dst=0, ) self.world_group.send_object( - result.captions, - dst=0, + local_video_captions, dst=0, ) + + use_cuda_tensor = ( + dist.is_available() + and dist.is_initialized() + and dist.get_backend() == "nccl" + and torch.cuda.is_available() + ) + tensor_device = ( + torch.device("cuda") + if use_cuda_tensor + else torch.device("cpu") + ) + + if ( + dist.is_available() + and dist.is_initialized() + ): + eval_failed = torch.tensor( + [ + 1 + if local_eval_errors + else 0 + ], + device=tensor_device, + dtype=torch.int32, + ) + dist.all_reduce( + eval_failed, + op=dist.ReduceOp.MAX, + ) + if int(eval_failed.item()) > 0: + if local_eval_errors: + raise RuntimeError( + "Validation flow " + "evaluation failed:\n" + + "\n".join( + local_eval_errors + ) + ) + raise RuntimeError( + "Validation flow " + "evaluation failed on at " + "least one rank. Check " + "per-rank logs." + ) + elif local_eval_errors: + raise RuntimeError( + "Validation flow evaluation " + "failed:\n" + + "\n".join(local_eval_errors) + ) + + if not self.evaluate_ptlflow: + continue + + metric_sum_tensor = torch.zeros( + len(flow_metric_keys), + device=tensor_device, + dtype=torch.float64, + ) + metric_count_tensor = torch.zeros( + len(flow_metric_keys), + device=tensor_device, + dtype=torch.float64, + ) + if ( + self.rank_in_sp_group == 0 + and flow_eval_enabled + ): + metric_sum_tensor = torch.tensor( + [ + local_metric_sums[k] + for k in flow_metric_keys + ], + device=tensor_device, + dtype=torch.float64, + ) + metric_count_tensor = torch.tensor( + [ + local_metric_counts[k] + for k in flow_metric_keys + ], + device=tensor_device, + dtype=torch.float64, + ) + + if ( + dist.is_available() + and dist.is_initialized() + ): + dist.all_reduce( + metric_sum_tensor, + op=dist.ReduceOp.SUM, + ) + dist.all_reduce( + metric_count_tensor, + op=dist.ReduceOp.SUM, + ) + + if self.global_rank == 0: + metric_logs: dict[str, float] = {} + for i, metric_key in enumerate( + flow_metric_keys + ): + count = float( + metric_count_tensor[ + i + ].item() + ) + if count <= 0: + continue + value = float( + metric_sum_tensor[ + i + ].item() + / count + ) + if not np.isfinite(value): + continue + metric_logs[ + f"metrics/{metric_key}" + ] = value + self._last_scalar_metrics[ + metric_key + ] = value + + if metric_logs: + self.tracker.log( + metric_logs, step + ) + + mf_val = metric_logs.get( + "metrics/mf_angle_err_mean" + ) + if mf_val is not None: + self._last_mf_angle_err_mean = ( + float(mf_val) + ) finally: + self._release_validation_cuda_memory() + if hasattr(tc.pipeline_config, "dmd_denoising_steps"): + tc.pipeline_config.dmd_denoising_steps = ( + old_dmd_denoising_steps + ) if was_training: transformer.train() + if ( + dist.is_available() + and dist.is_initialized() + ): + backend = dist.get_backend() + use_cuda = ( + backend == "nccl" + and torch.cuda.is_available() + ) + tensor_device = ( + torch.device("cuda") + if use_cuda + else torch.device("cpu") + ) + mf_tensor = torch.tensor( + [self._last_mf_angle_err_mean], + dtype=torch.float32, + device=tensor_device, + ) + dist.broadcast(mf_tensor, src=0) + self._last_mf_angle_err_mean = float( + mf_tensor.item() + ) + if np.isfinite( + self._last_mf_angle_err_mean + ): + self._last_scalar_metrics[ + "mf_angle_err_mean" + ] = self._last_mf_angle_err_mean + + def _maybe_set_dmd_denoising_steps( + self, + tc: TrainingConfig, + num_inference_steps: int, + ) -> None: + """Set dmd_denoising_steps on pipeline_config for + pipelines that require explicit DMD timesteps.""" + if not self._requires_dmd_denoising_steps: + return + if self.sampling_timesteps is not None: + tc.pipeline_config.dmd_denoising_steps = ( # type: ignore[union-attr] + list(self.sampling_timesteps) + ) + else: + timesteps = np.linspace( + 1000, 0, int(num_inference_steps), + ) + tc.pipeline_config.dmd_denoising_steps = [ # type: ignore[union-attr] + int(max(0, min(1000, round(t)))) + for t in timesteps + ] + + # Also set any pipeline-specific kwargs from + # YAML (e.g. dmd_denoising_steps override). + pk = self.pipeline_kwargs + if "dmd_denoising_steps" in pk: + tc.pipeline_config.dmd_denoising_steps = [ # type: ignore[union-attr] + int(s) + for s in pk["dmd_denoising_steps"] + ] + # ---------------------------------------------------------- # Pipeline management # ---------------------------------------------------------- @@ -270,7 +640,11 @@ class ValidationCallback(Callback): *, transformer: torch.nn.Module, ) -> Any: - key = (id(transformer), ) + key = ( + id(transformer), + self.pipeline_target, + self.scheduler_target, + ) if (self._pipeline is not None and self._pipeline_key == key): return self._pipeline @@ -296,6 +670,13 @@ class ValidationCallback(Callback): if flow_shift is not None: kwargs["flow_shift"] = float(flow_shift) + # Build and inject a scheduler if target is set. + scheduler = self._build_scheduler(flow_shift) + if scheduler is not None: + kwargs["loaded_modules"]["scheduler"] = ( + scheduler + ) + self._pipeline = PipelineCls.from_pretrained( tc.model_path, **kwargs, @@ -303,6 +684,20 @@ class ValidationCallback(Callback): self._pipeline_key = key return self._pipeline + def _build_scheduler( + self, flow_shift: float | None, + ) -> Any | None: + """Build scheduler from ``scheduler_target``.""" + if self.scheduler_target is None: + return None + if flow_shift is None: + return None + + SchedulerCls = resolve_target( + self.scheduler_target + ) + return SchedulerCls(shift=float(flow_shift)) + # ---------------------------------------------------------- # Batch preparation # ---------------------------------------------------------- @@ -363,6 +758,7 @@ class ValidationCallback(Callback): eta=0.0, VSA_sparsity=tc.vsa_sparsity, timesteps=sampling_timesteps_tensor, + sampling_timesteps=sampling_timesteps_tensor, ) batch._inference_args = inference_args # type: ignore[attr-defined] @@ -370,8 +766,167 @@ class ValidationCallback(Callback): if ("image" in validation_batch and validation_batch["image"] is not None): batch.pil_image = validation_batch["image"] + self._maybe_set_action_conds( + batch, validation_batch, sampling_param, + ) return batch + def _maybe_set_action_conds( + self, + batch: ForwardBatch, + validation_batch: dict[str, Any], + sampling_param: SamplingParam, + ) -> None: + """Set keyboard_cond / mouse_cond on the batch if + present in the dataset.""" + target_len = int(sampling_param.num_frames) + + if ( + "keyboard_cond" in validation_batch + and validation_batch["keyboard_cond"] + is not None + ): + kb = torch.as_tensor( + validation_batch["keyboard_cond"] + ).to(dtype=torch.bfloat16) + if kb.ndim == 3 and kb.shape[0] == 1: + kb = kb.squeeze(0) + if kb.ndim != 2: + raise ValueError( + "validation keyboard_cond must have" + " shape (T, K), got " + f"{tuple(kb.shape)}" + ) + if kb.shape[0] > target_len: + kb = kb[:target_len] + elif kb.shape[0] < target_len: + pad = torch.zeros( + ( + target_len - kb.shape[0], + kb.shape[1], + ), + dtype=kb.dtype, + device=kb.device, + ) + kb = torch.cat([kb, pad], dim=0) + batch.keyboard_cond = kb.unsqueeze(0) + + if ( + "mouse_cond" in validation_batch + and validation_batch["mouse_cond"] + is not None + ): + mc = torch.as_tensor( + validation_batch["mouse_cond"] + ).to(dtype=torch.bfloat16) + if mc.ndim == 3 and mc.shape[0] == 1: + mc = mc.squeeze(0) + if mc.ndim != 2: + raise ValueError( + "validation mouse_cond must have " + "shape (T, 2), got " + f"{tuple(mc.shape)}" + ) + if mc.shape[0] > target_len: + mc = mc[:target_len] + elif mc.shape[0] < target_len: + pad = torch.zeros( + ( + target_len - mc.shape[0], + mc.shape[1], + ), + dtype=mc.dtype, + device=mc.device, + ) + mc = torch.cat([mc, pad], dim=0) + batch.mouse_cond = mc.unsqueeze(0) + + # ---------------------------------------------------------- + # Post-processing + # ---------------------------------------------------------- + + def _post_process_validation_frames( + self, + frames: list[np.ndarray], + batch: ForwardBatch, + ) -> list[np.ndarray]: + """Overlay action indicators if conditions present.""" + keyboard_cond = getattr(batch, "keyboard_cond", None) + mouse_cond = getattr(batch, "mouse_cond", None) + if keyboard_cond is None and mouse_cond is None: + return frames + + try: + from fastvideo.models.dits.matrixgame.utils import ( + draw_keys_on_frame, + draw_mouse_on_frame, + ) + except Exception as e: + logger.warning( + "Action overlay unavailable: %s", e, + ) + return frames + + if ( + keyboard_cond is not None + and torch.is_tensor(keyboard_cond) + ): + keyboard_np = ( + keyboard_cond.squeeze(0) + .detach() + .cpu() + .float() + .numpy() + ) + else: + keyboard_np = None + + if ( + mouse_cond is not None + and torch.is_tensor(mouse_cond) + ): + mouse_np = ( + mouse_cond.squeeze(0) + .detach() + .cpu() + .float() + .numpy() + ) + else: + mouse_np = None + + key_names = ["W", "S", "A", "D", "left", "right"] + processed: list[np.ndarray] = [] + for fi, frame in enumerate(frames): + frame = np.ascontiguousarray(frame.copy()) + if ( + keyboard_np is not None + and fi < len(keyboard_np) + ): + keys = { + key_names[i]: bool( + keyboard_np[fi, i] + ) + for i in range( + min( + len(key_names), + int(keyboard_np.shape[1]), + ) + ) + } + draw_keys_on_frame( + frame, keys, mode="universal", + ) + if ( + mouse_np is not None + and fi < len(mouse_np) + ): + pitch = float(mouse_np[fi, 0]) + yaw = float(mouse_np[fi, 1]) + draw_mouse_on_frame(frame, pitch, yaw) + processed.append(frame) + return processed + # ---------------------------------------------------------- # Validation loop # ---------------------------------------------------------- @@ -398,13 +953,9 @@ class ValidationCallback(Callback): model_path=tc.model_path, ) - # Propagate sampling_timesteps to pipeline_config so - # causal/DMD denoising stages can read them. - if (self.sampling_timesteps is not None and inference_args.pipeline_config.dmd_denoising_steps is None): - inference_args.pipeline_config.dmd_denoising_steps = ([int(s) for s in self.sampling_timesteps]) - videos: list[list[np.ndarray]] = [] captions: list[str] = [] + action_paths: list[str | None] = [] for validation_batch in dataloader: batch = self._prepare_validation_batch( @@ -412,9 +963,13 @@ class ValidationCallback(Callback): validation_batch, num_inference_steps, ) + action_path = validation_batch.get( + "action_path" + ) + if not isinstance(action_path, str): + action_path = None assert (batch.prompt is not None and isinstance(batch.prompt, str)) - captions.append(batch.prompt) with torch.no_grad(): output_batch = pipeline.forward( @@ -425,6 +980,8 @@ class ValidationCallback(Callback): samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue + captions.append(batch.prompt) + action_paths.append(action_path) video = rearrange( samples, @@ -438,13 +995,305 @@ class ValidationCallback(Callback): ) x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1)) frames.append((x * 255).numpy().astype(np.uint8)) + frames = ( + self._post_process_validation_frames( + frames, batch, + ) + ) videos.append(frames) return _ValidationStepResult( videos=videos, captions=captions, + action_paths=action_paths, ) + def _init_flow_eval_module(self) -> None: + if self._flow_eval_init_done: + return + self._flow_eval_init_done = True + self._flow_eval_ready = False + + ptlflow_dir = Path(self.ptlflow_dir) + try: + ptlflow_dir_str = str(ptlflow_dir.resolve()) + if ptlflow_dir_str not in sys.path: + sys.path.insert(0, ptlflow_dir_str) + + from eval_flow_divergence import ( + evaluate_pair_synthetic, ) + + self._flow_eval_fn = ( + evaluate_pair_synthetic + ) + self._flow_eval_ready = True + logger.info( + "Initialized flow divergence " + "evaluator: %s", + ptlflow_dir_str, + ) + except Exception as e: + logger.warning( + "Failed to initialize flow " + "divergence evaluator: %s", + e, + ) + + def _count_video_frames( + self, + video_path: str, + ) -> int: + reader: Any | None = None + try: + reader = imageio.get_reader(video_path) + try: + return int(reader.count_frames()) + except Exception: + n_frames = 0 + for _ in reader: + n_frames += 1 + return int(n_frames) + except Exception as e: + logger.warning( + "Failed to count frames for " + "validation video %s: %s", + video_path, + e, + ) + return 0 + finally: + if reader is not None: + try: + reader.close() + except Exception: + pass + + def _prepare_aligned_action_path( + self, + *, + action_path: str, + video_path: str, + eval_output_dir: str, + ) -> str: + try: + actions_obj = np.load( + action_path, + allow_pickle=True, + ).item() + except Exception as e: + logger.warning( + "Failed to load action file " + "%s: %s", + action_path, + e, + ) + return action_path + + if not isinstance(actions_obj, dict): + logger.warning( + "Action file %s is not a dict; " + "skipping action alignment.", + action_path, + ) + return action_path + + keyboard = actions_obj.get("keyboard") + mouse = actions_obj.get("mouse") + if keyboard is None or mouse is None: + logger.warning( + "Action file %s is missing " + "'keyboard' or 'mouse'; " + "skipping action alignment.", + action_path, + ) + return action_path + + keyboard_arr = np.asarray(keyboard) + mouse_arr = np.asarray(mouse) + if ( + keyboard_arr.ndim < 1 + or mouse_arr.ndim < 1 + ): + logger.warning( + "Action file %s has invalid " + "shapes keyboard=%s mouse=%s; " + "skipping action alignment.", + action_path, + keyboard_arr.shape, + mouse_arr.shape, + ) + return action_path + + n_frames = self._count_video_frames(video_path) + if n_frames <= 0: + return action_path + + n_keyboard = int(keyboard_arr.shape[0]) + n_mouse = int(mouse_arr.shape[0]) + n_use = min(n_frames, n_keyboard, n_mouse) + if n_use <= 0: + return action_path + + if ( + n_use == n_keyboard + and n_use == n_mouse + ): + return action_path + + aligned_actions = dict(actions_obj) + aligned_actions["keyboard"] = keyboard_arr[:n_use] + aligned_actions["mouse"] = mouse_arr[:n_use] + + os.makedirs(eval_output_dir, exist_ok=True) + aligned_path = os.path.join( + eval_output_dir, + "actions_aligned.npy", + ) + np.save( + aligned_path, + aligned_actions, + allow_pickle=True, + ) + logger.info( + "Aligned action file for PTLFlow: " + "video_frames=%d keyboard=%d " + "mouse=%d used=%d path=%s", + n_frames, + n_keyboard, + n_mouse, + n_use, + aligned_path, + ) + return aligned_path + + def _release_validation_cuda_memory(self) -> None: + if not torch.cuda.is_available(): + return + try: + gc.collect() + except Exception: + pass + try: + torch.cuda.empty_cache() + except Exception: + pass + try: + torch.cuda.ipc_collect() + except Exception: + pass + + def _evaluate_validation_video( + self, + *, + video_path: str, + caption: str, + action_path: str | None, + global_step: int, + num_inference_steps: int, + ) -> dict[str, float]: + del caption + self._init_flow_eval_module() + if not self._flow_eval_ready: + raise RuntimeError( + "ptlflow evaluator is not " + "initialized; cannot compute " + "flow metrics." + ) + if ( + not isinstance(action_path, str) + or not os.path.isfile(action_path) + ): + raise FileNotFoundError( + "Validation sample is missing a " + f"valid action_path: {action_path}" + ) + + output_root = ( + self.output_dir + or self.training_config.checkpoint.output_dir + or os.getcwd() + ) + eval_output_dir = os.path.join( + output_root, + "flow_eval", + f"step_{global_step}", + f"inference_steps_{num_inference_steps}", + Path(video_path).stem, + ) + + aligned_action_path = ( + self._prepare_aligned_action_path( + action_path=action_path, + video_path=video_path, + eval_output_dir=eval_output_dir, + ) + ) + with torch.inference_mode(): + summary = self._flow_eval_fn( + gen_video=video_path, + action_path=aligned_action_path, + calibration_path=self.ptlflow_calibration_path, + output_dir=eval_output_dir, + model_name="dpflow", + ckpt=self.ptlflow_ckpt, + no_viz=True, + use_depth=self.ptlflow_use_depth, + ) + if not isinstance(summary, dict): + raise RuntimeError( + "ptlflow returned invalid summary " + f"type: {type(summary)}" + ) + + metrics: dict[str, float] = {} + missing_or_invalid: list[str] = [] + for key in self._FLOW_EVAL_SCALAR_KEYS: + value = summary.get(key) + if not isinstance( + value, + ( + float, + int, + np.floating, + np.integer, + ), + ): + missing_or_invalid.append(key) + continue + value_float = float(value) + if not np.isfinite(value_float): + missing_or_invalid.append(key) + continue + metrics[key] = value_float + + if missing_or_invalid: + raise RuntimeError( + "ptlflow summary missing/invalid " + "metrics: " + + ", ".join(missing_or_invalid) + ) + return metrics + + def get_latest_metric( + self, + metric_key: str, + *, + step: int | None = None, + ) -> float | None: + if ( + step is not None + and int(step) != self._last_eval_step + ): + return None + value = self._last_scalar_metrics.get( + str(metric_key) + ) + if value is None: + return None + if not np.isfinite(value): + return None + return float(value) + # ---------------------------------------------------------- # State management # ---------------------------------------------------------- @@ -453,6 +1302,13 @@ class ValidationCallback(Callback): state: dict[str, Any] = {} if self.validation_random_generator is not None: state["validation_rng"] = (self.validation_random_generator.get_state()) + + state["last_eval_step"] = int( + self._last_eval_step + ) + state["last_mf_angle_err_mean"] = float( + self._last_mf_angle_err_mean + ) return state def load_state_dict( @@ -462,3 +1318,20 @@ class ValidationCallback(Callback): rng_state = state_dict.get("validation_rng") if (rng_state is not None and self.validation_random_generator is not None): self.validation_random_generator.set_state(rng_state) + + last_step = state_dict.get("last_eval_step") + if isinstance(last_step, int): + self._last_eval_step = int(last_step) + last_mf = state_dict.get( + "last_mf_angle_err_mean" + ) + if isinstance( + last_mf, (float, int, np.floating, np.integer) + ): + self._last_mf_angle_err_mean = float( + last_mf + ) + if np.isfinite(self._last_mf_angle_err_mean): + self._last_scalar_metrics[ + "mf_angle_err_mean" + ] = self._last_mf_angle_err_mean diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py index bd20f0bd2..ecdfae2d3 100644 --- a/fastvideo/train/entrypoint/train.py +++ b/fastvideo/train/entrypoint/train.py @@ -33,6 +33,10 @@ def run_training_from_config( config_path: str, *, dry_run: bool = False, + resume_from_checkpoint: str | None = None, + override_output_dir: str | None = None, + best_checkpoint_start_step: int | None = None, + best_checkpoint_top_k: int | None = None, overrides: list[str] | None = None, ) -> None: """YAML-only training entrypoint (schema v2).""" @@ -54,6 +58,15 @@ def run_training_from_config( cfg = load_run_config(config_path, overrides=overrides) tc = cfg.training + if resume_from_checkpoint is not None: + tc.checkpoint.resume_from_checkpoint = str(resume_from_checkpoint) + if override_output_dir is not None: + tc.checkpoint.output_dir = str(override_output_dir) + if best_checkpoint_start_step is not None: + tc.checkpoint.best_checkpoint_start_step = int(best_checkpoint_start_step) + if best_checkpoint_top_k is not None: + tc.checkpoint.best_checkpoint_top_k = max(1, int(best_checkpoint_top_k)) + # Auto-set attention backend for VSA when sparsity is configured. if tc.vsa_sparsity > 0.0: os.environ.setdefault( @@ -97,6 +110,7 @@ def run_training_from_config( output_dir=tc.checkpoint.output_dir, config=ckpt_config, callbacks=trainer.callbacks, + tracker=trainer.tracker, raw_config=cfg.raw, ) @@ -115,6 +129,18 @@ def main( ) -> None: config_path = str(args.config) dry_run = bool(args.dry_run) + resume_from_checkpoint = getattr( + args, "resume_from_checkpoint", None + ) + override_output_dir = getattr( + args, "override_output_dir", None + ) + best_checkpoint_start_step = getattr( + args, "best_checkpoint_start_step", None + ) + best_checkpoint_top_k = getattr( + args, "best_checkpoint_top_k", None + ) logger.info( "Starting training from config=%s", config_path, @@ -122,6 +148,10 @@ def main( run_training_from_config( config_path, dry_run=dry_run, + resume_from_checkpoint=resume_from_checkpoint, + override_output_dir=override_output_dir, + best_checkpoint_start_step=best_checkpoint_start_step, + best_checkpoint_top_k=best_checkpoint_top_k, overrides=overrides, ) logger.info("Training completed") @@ -142,5 +172,43 @@ if __name__ == "__main__": help=("Parse config and build runtime, " "but do not start training."), ) + parser.add_argument( + "--resume-from-checkpoint", + type=str, + default=None, + help=( + "Path to a checkpoint directory " + "(checkpoint-), its 'dcp/' subdir, " + "or an output_dir containing checkpoints " + "(auto-picks latest)." + ), + ) + parser.add_argument( + "--override-output-dir", + type=str, + default=None, + help=( + "Override training.output_dir from YAML " + "(useful for repeated runs)." + ), + ) + parser.add_argument( + "--best-checkpoint-start-step", + type=int, + default=None, + help=( + "Override training.checkpoint.best_checkpoint_start_step " + "(0 disables best-checkpoint saving)." + ), + ) + parser.add_argument( + "--best-checkpoint-top-k", + type=int, + default=None, + help=( + "Override training.checkpoint.best_checkpoint_top_k " + "(minimum 1)." + ), + ) args, unknown = parser.parse_known_args(argv[1:]) main(args, overrides=unknown if unknown else None) diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index 61fd6ef2e..a99e85fe3 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -8,6 +8,7 @@ __all__ = [ "FineTuneMethod", "SelfForcingMethod", "DiffusionForcingSFTMethod", + "TeacherForcingSFTMethod", ] @@ -24,4 +25,7 @@ def __getattr__(name: str) -> object: if name == "DiffusionForcingSFTMethod": from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod return DiffusionForcingSFTMethod + if name == "TeacherForcingSFTMethod": + from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod + return TeacherForcingSFTMethod raise AttributeError(name) diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index 8561f167f..9b35cd696 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -73,6 +73,10 @@ class TrainingMethod(torch.nn.Module, ABC): def set_tracker(self, tracker: Any) -> None: self.tracker = tracker + @property + def transformer_inference(self) -> torch.nn.Module: + return self.student.transformer + @abstractmethod def single_train_step( self, diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py index 6bfbb44ad..6cf14344f 100644 --- a/fastvideo/train/methods/distribution_matching/dmd2.py +++ b/fastvideo/train/methods/distribution_matching/dmd2.py @@ -126,6 +126,12 @@ class DMD2Method(TrainingMethod): } outputs: dict[str, Any] = dict(critic_outputs) + if training_batch.dmd_latent_vis_dict: + outputs["dmd_latent_vis_dict"] = training_batch.dmd_latent_vis_dict + if training_batch.fake_score_latent_vis_dict: + outputs["fake_score_latent_vis_dict"] = ( + training_batch.fake_score_latent_vis_dict + ) outputs["_fv_backward"] = { "update_student": update_student, "student_ctx": student_ctx, @@ -623,6 +629,12 @@ class DMD2Method(TrainingMethod): grad = (faker_x0 - real_cfg_x0) / denom grad = torch.nan_to_num(grad) + batch.dmd_latent_vis_dict.update({ + "generator_pred_video": generator_pred_x0.detach(), + "real_score_pred_video": real_cfg_x0.detach(), + "faker_score_pred_video": faker_x0.detach(), + "dmd_timestep": timestep.float().detach(), + }) loss = 0.5 * F.mse_loss( generator_pred_x0.float(), (generator_pred_x0.float() - grad.float()).detach(), diff --git a/fastvideo/train/methods/distribution_matching/self_forcing.py b/fastvideo/train/methods/distribution_matching/self_forcing.py index 4939a85bd..637803bb5 100644 --- a/fastvideo/train/methods/distribution_matching/self_forcing.py +++ b/fastvideo/train/methods/distribution_matching/self_forcing.py @@ -572,5 +572,11 @@ class SelfForcingMethod(DMD2Method): grad = (faker_x0 - real_cfg_x0) / denom grad = torch.nan_to_num(grad) + batch.dmd_latent_vis_dict.update({ + "generator_pred_video": generator_pred_x0.detach(), + "real_score_pred_video": real_cfg_x0.detach(), + "faker_score_pred_video": faker_x0.detach(), + "dmd_timestep": timestep.float().detach(), + }) loss = 0.5 * torch.mean((generator_pred_x0.float() - (generator_pred_x0.float() - grad.float()).detach())**2) return loss diff --git a/fastvideo/train/methods/fine_tuning/__init__.py b/fastvideo/train/methods/fine_tuning/__init__.py index 6f862df4e..96da179f7 100644 --- a/fastvideo/train/methods/fine_tuning/__init__.py +++ b/fastvideo/train/methods/fine_tuning/__init__.py @@ -7,10 +7,12 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod + from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod __all__ = [ "DiffusionForcingSFTMethod", "FineTuneMethod", + "TeacherForcingSFTMethod", ] @@ -25,4 +27,9 @@ def __getattr__(name: str) -> object: from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod return FineTuneMethod + if name == "TeacherForcingSFTMethod": + from fastvideo.train.methods.fine_tuning.tfsft import ( + TeacherForcingSFTMethod, ) + + return TeacherForcingSFTMethod raise AttributeError(name) diff --git a/fastvideo/train/methods/fine_tuning/tfsft.py b/fastvideo/train/methods/fine_tuning/tfsft.py new file mode 100644 index 000000000..32b1786a3 --- /dev/null +++ b/fastvideo/train/methods/fine_tuning/tfsft.py @@ -0,0 +1,500 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Teacher-forcing SFT method (TFSFT; algorithm layer).""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Literal + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import create_block_mask + +from fastvideo.train.methods.base import TrainingMethod, LogScalar +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.optimizer import ( + build_optimizer_and_scheduler, ) + + +class TeacherForcingSFTMethod(TrainingMethod): + """Teacher-forcing SFT (TFSFT) for causal students. + + Training uses a concatenated ``[clean | noisy]`` latent sequence plus a + custom block mask so noisy chunks can attend to a clean prefix while the + denoising loss is applied only on the noisy half. + """ + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, ModelBase], + ) -> None: + super().__init__(cfg=cfg, role_models=role_models) + + if "student" not in role_models: + raise ValueError("TFSFT requires role 'student'") + if not self.student._trainable: + raise ValueError("TFSFT requires student to be trainable") + self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) + + self._chunk_size = self._parse_chunk_size( + self.method_config.get("chunk_size", None)) + self._timestep_index_range = self._parse_timestep_index_range() + + # Initialize preprocessors on student. + self.student.init_preprocessors(self.training_config) + + self._block_mask_cache: dict[tuple[str, int, int, int], Any] = {} + self._init_optimizers_and_schedulers() + + @property + def _optimizer_dict(self) -> dict[str, Any]: + return {"student": self._student_optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self._student_lr_scheduler} + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + del iteration + training_batch = self.student.prepare_batch( + batch, + generator=self.cuda_generator, + latents_source="data", + ) + + if training_batch.latents is None: + raise RuntimeError("prepare_batch() must set TrainingBatch.latents") + + clean_latents = training_batch.latents + if not torch.is_tensor(clean_latents): + raise TypeError("TrainingBatch.latents must be a torch.Tensor") + if clean_latents.ndim != 5: + raise ValueError("TrainingBatch.latents must be " + "[B, T, C, H, W], got " + f"shape={tuple(clean_latents.shape)}") + + batch_size, num_latents = ( + int(clean_latents.shape[0]), + int(clean_latents.shape[1]), + ) + + expected_chunk = getattr( + self.student.transformer, + "num_frame_per_block", + None, + ) + if (expected_chunk is not None + and int(expected_chunk) != int(self._chunk_size)): + raise ValueError("TFSFT chunk_size must match " + "transformer.num_frame_per_block for " + f"causal training (got {self._chunk_size}, " + f"expected {expected_chunk}).") + + timestep_indices = self._sample_t_inhom_indices( + batch_size=batch_size, + num_latents=num_latents, + device=clean_latents.device, + ) + sp_size = int(self.training_config.distributed.sp_size) + sp_group = getattr(self.student, "sp_group", None) + if (sp_size > 1 and sp_group is not None + and hasattr(sp_group, "broadcast")): + sp_group.broadcast(timestep_indices, src=0) + + scheduler = self.student.noise_scheduler + if scheduler is None: + raise ValueError("TFSFT requires student.noise_scheduler") + + schedule_timesteps = scheduler.timesteps.to( + device=clean_latents.device, + dtype=torch.float32, + ) + schedule_sigmas = scheduler.sigmas.to( + device=clean_latents.device, + dtype=clean_latents.dtype, + ) + t_inhom = schedule_timesteps[timestep_indices] + + noise = getattr(training_batch, "noise", None) + if noise is None: + noise = torch.randn_like(clean_latents) + else: + if not torch.is_tensor(noise): + raise TypeError("TrainingBatch.noise must be a " + "torch.Tensor when set") + noise = noise.permute(0, 2, 1, 3, 4).to(dtype=clean_latents.dtype) + + noisy_latents = self.student.add_noise( + clean_latents, + noise, + t_inhom.flatten(), + ) + + concat_latents = torch.cat([clean_latents, noisy_latents], dim=1) + concat_timesteps = torch.cat([torch.zeros_like(t_inhom), t_inhom], dim=1) + + saved_batch_state = self._expand_batch_for_concat( + training_batch, + num_concat_latents=int(concat_latents.shape[1]), + ) + training_batch.timesteps = concat_timesteps + self._refresh_attention_metadata(training_batch) + + transformer = self.student.transformer + custom_block_mask = self._get_teacher_forcing_block_mask( + transformer=transformer, + latents=concat_latents, + ) + prev_block_mask = getattr(transformer, "block_mask", None) + + try: + transformer.block_mask = custom_block_mask + pred = self.student.predict_noise( + concat_latents, + concat_timesteps, + training_batch, + conditional=True, + attn_kind=self._attn_kind, + ) + finally: + transformer.block_mask = prev_block_mask + self._restore_batch_after_concat(training_batch, saved_batch_state) + + pred_noisy = pred[:, num_latents:] + + if bool(self.training_config.model.precondition_outputs): + sigmas = schedule_sigmas[timestep_indices] + sigmas = sigmas.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + pred_x0 = noisy_latents - pred_noisy * sigmas + total_loss = F.mse_loss(pred_x0.float(), clean_latents.float()) + else: + target = noise - clean_latents + total_loss = F.mse_loss(pred_noisy.float(), target.float()) + + if self._attn_kind == "vsa": + attn_metadata = training_batch.attn_metadata_vsa + else: + attn_metadata = training_batch.attn_metadata + + loss_map = {"total_loss": total_loss, "tfsft_loss": total_loss} + outputs: dict[str, Any] = { + "_fv_backward": ( + concat_timesteps, + attn_metadata, + ) + } + metrics: dict[str, LogScalar] = {} + return loss_map, outputs, metrics + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int = 1, + ) -> None: + grad_accum_rounds = max(1, int(grad_accum_rounds)) + ctx = outputs.get("_fv_backward") + if ctx is None: + super().backward( + loss_map, + outputs, + grad_accum_rounds=grad_accum_rounds, + ) + return + self.student.backward( + loss_map["total_loss"], + ctx, + grad_accum_rounds=grad_accum_rounds, + ) + + def get_optimizers( + self, + iteration: int, + ) -> list[torch.optim.Optimizer]: + del iteration + return [self._student_optimizer] + + def get_lr_schedulers( + self, + iteration: int, + ) -> list[Any]: + del iteration + return [self._student_lr_scheduler] + + def _refresh_attention_metadata(self, training_batch: Any) -> None: + build_fn = getattr(self.student, "_build_attention_metadata", None) + if not callable(build_fn): + return + + build_fn(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] + + def _expand_batch_for_concat( + self, + training_batch: Any, + *, + num_concat_latents: int, + ) -> dict[str, Any]: + saved: dict[str, Any] = {} + + raw_latent_shape = getattr(training_batch, "raw_latent_shape", None) + if raw_latent_shape is not None and len(raw_latent_shape) == 5: + saved["raw_latent_shape"] = raw_latent_shape + batch_size, channels, _num_frames, height, width = raw_latent_shape + training_batch.raw_latent_shape = ( + batch_size, + channels, + num_concat_latents, + height, + width, + ) + + temporal_dims = { + "image_latents": 2, + "mask_lat_size": 2, + "viewmats": 1, + "Ks": 1, + "action": 1, + "mouse_cond": 1, + "keyboard_cond": 1, + } + for name, dim in temporal_dims.items(): + value = getattr(training_batch, name, None) + if isinstance(value, torch.Tensor): + saved[name] = value + setattr(training_batch, name, torch.cat([value, value], dim=dim)) + + return saved + + def _restore_batch_after_concat( + self, + training_batch: Any, + saved: dict[str, Any], + ) -> None: + for name, value in saved.items(): + setattr(training_batch, name, value) + + def _get_teacher_forcing_block_mask( + self, + *, + transformer: Any, + latents: torch.Tensor, + ) -> Any: + num_frames = int(latents.shape[1]) + if num_frames % 2 != 0: + raise ValueError("TFSFT concat-and-mask requires an even " + f"number of frames, got {num_frames}") + + patch_size = getattr(transformer, "patch_size", (1, 1, 1)) + if len(patch_size) != 3: + raise ValueError("Unexpected transformer.patch_size: " + f"{patch_size!r}") + + patch_t, patch_h, patch_w = [int(v) for v in patch_size] + post_patch_num_frames = num_frames // patch_t + if post_patch_num_frames * patch_t != num_frames: + raise ValueError("TFSFT requires num_frames divisible by " + f"patch temporal size {patch_t}, got {num_frames}") + + frame_seqlen = (int(latents.shape[-2]) // patch_h) * ( + int(latents.shape[-1]) // patch_w) + clean_frames = post_patch_num_frames // 2 + if clean_frames <= 0: + raise ValueError("TFSFT requires at least one clean frame") + + device_key = str(latents.device) + cache_key = ( + device_key, + post_patch_num_frames, + frame_seqlen, + self._chunk_size, + ) + cached = self._block_mask_cache.get(cache_key) + if cached is not None: + return cached + + total_length = post_patch_num_frames * frame_seqlen + padded_length = math.ceil(total_length / 128) * 128 - total_length + total_tokens = total_length + padded_length + + prefix_starts = torch.zeros(total_tokens, device=latents.device, dtype=torch.long) + prefix_ends = torch.zeros_like(prefix_starts) + noisy_starts = torch.zeros_like(prefix_starts) + noisy_ends = torch.zeros_like(prefix_starts) + + for frame_idx in range(post_patch_num_frames): + token_start = frame_idx * frame_seqlen + token_end = min(total_length, token_start + frame_seqlen) + if frame_idx < clean_frames: + chunk_end = min(clean_frames, ((frame_idx // self._chunk_size) + 1) * self._chunk_size) + prefix_start = 0 + prefix_end = chunk_end * frame_seqlen + noisy_start = 0 + noisy_end = 0 + else: + noisy_frame = frame_idx - clean_frames + chunk_start = (noisy_frame // self._chunk_size) * self._chunk_size + chunk_end = min(clean_frames, chunk_start + self._chunk_size) + prefix_start = 0 + prefix_end = chunk_start * frame_seqlen + noisy_start = (clean_frames + chunk_start) * frame_seqlen + noisy_end = (clean_frames + chunk_end) * frame_seqlen + + prefix_starts[token_start:token_end] = prefix_start + prefix_ends[token_start:token_end] = prefix_end + noisy_starts[token_start:token_end] = noisy_start + noisy_ends[token_start:token_end] = noisy_end + + def attention_mask(b, h, q_idx, kv_idx): + valid = (q_idx < total_length) & (kv_idx < total_length) + prefix_ok = ((kv_idx >= prefix_starts[q_idx]) + & (kv_idx < prefix_ends[q_idx])) + noisy_ok = ((kv_idx >= noisy_starts[q_idx]) + & (kv_idx < noisy_ends[q_idx])) + return valid & (prefix_ok | noisy_ok) + + block_mask = create_block_mask( + attention_mask, + B=None, + H=None, + Q_LEN=total_tokens, + KV_LEN=total_tokens, + _compile=False, + device=latents.device, + ) + self._block_mask_cache[cache_key] = block_mask + return block_mask + + def _parse_chunk_size(self, raw: Any) -> int: + if raw in (None, ""): + return 3 + if isinstance(raw, bool): + raise ValueError("method_config.chunk_size must be an int, " + "got bool") + if isinstance(raw, float) and not raw.is_integer(): + raise ValueError("method_config.chunk_size must be an int, " + "got float") + if isinstance(raw, str) and not raw.strip(): + raise ValueError("method_config.chunk_size must be an int, " + "got empty string") + try: + value = int(raw) + except (TypeError, ValueError) as e: + raise ValueError("method_config.chunk_size must be an int, " + f"got {type(raw).__name__}") from e + if value <= 0: + raise ValueError("method_config.chunk_size must be > 0") + return value + + def _parse_ratio( + self, + raw: Any, + *, + where: str, + default: float, + ) -> float: + if raw in (None, ""): + return float(default) + if isinstance(raw, bool): + raise ValueError(f"{where} must be a number/string, got bool") + if isinstance(raw, int | float): + return float(raw) + if isinstance(raw, str) and raw.strip(): + return float(raw) + raise ValueError(f"{where} must be a number/string, " + f"got {type(raw).__name__}") + + def _parse_timestep_index_range(self) -> tuple[int, int]: + scheduler = self.student.noise_scheduler + if scheduler is None: + raise ValueError("TFSFT requires student.noise_scheduler") + num_steps = int( + getattr(scheduler, "config", scheduler).num_train_timesteps) + + min_ratio = self._parse_ratio( + self.method_config.get("min_timestep_ratio", None), + where="method.min_timestep_ratio", + default=0.0, + ) + max_ratio = self._parse_ratio( + self.method_config.get("max_timestep_ratio", None), + where="method.max_timestep_ratio", + default=1.0, + ) + + if not (0.0 <= min_ratio <= 1.0 and 0.0 <= max_ratio <= 1.0): + raise ValueError("TFSFT timestep ratios must be in [0,1], " + f"got min={min_ratio}, max={max_ratio}") + if max_ratio < min_ratio: + raise ValueError("method_config.max_timestep_ratio must be " + ">= min_timestep_ratio") + + min_index = int(min_ratio * num_steps) + max_index = int(max_ratio * num_steps) + min_index = max(0, min(min_index, num_steps - 1)) + max_index = max(0, min(max_index, num_steps - 1)) + + if max_index <= min_index: + max_index = min(num_steps - 1, min_index + 1) + + return min_index, max_index + 1 + + def _init_optimizers_and_schedulers(self) -> None: + tc = self.training_config + student_lr = float(tc.optimizer.learning_rate) + if student_lr <= 0.0: + raise ValueError("training.learning_rate must be > 0 " + "for tfsft") + + student_betas = tc.optimizer.betas + student_sched = str(tc.optimizer.lr_scheduler) + student_params = [ + p for p in self.student.transformer.parameters() if p.requires_grad + ] + ( + self._student_optimizer, + self._student_lr_scheduler, + ) = build_optimizer_and_scheduler( + params=student_params, + optimizer_config=tc.optimizer, + loop_config=tc.loop, + learning_rate=student_lr, + betas=student_betas, + scheduler_name=student_sched, + ) + + def _sample_t_inhom_indices( + self, + *, + batch_size: int, + num_latents: int, + device: torch.device, + ) -> torch.Tensor: + chunk_size = self._chunk_size + num_chunks = ((num_latents + chunk_size - 1) // chunk_size) + low, high = self._timestep_index_range + chunk_indices = torch.randint( + low=low, + high=high, + size=(batch_size, num_chunks), + device=device, + dtype=torch.long, + generator=self.cuda_generator, + ) + expanded = chunk_indices.repeat_interleave(chunk_size, dim=1) + return expanded[:, :num_latents] diff --git a/fastvideo/train/models/wangame/__init__.py b/fastvideo/train/models/wangame/__init__.py new file mode 100644 index 000000000..101e3e7a8 --- /dev/null +++ b/fastvideo/train/models/wangame/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +"""WanGame model plugin package.""" + +from fastvideo.train.models.wangame.wangame import ( + WanGameModel as WanGameModel, ) +from fastvideo.train.models.wangame.wangame_causal import ( + WanGameCausalModel as WanGameCausalModel, ) diff --git a/fastvideo/train/models/wangame/wangame.py b/fastvideo/train/models/wangame/wangame.py new file mode 100644 index 000000000..bbb0763e6 --- /dev/null +++ b/fastvideo/train/models/wangame/wangame.py @@ -0,0 +1,862 @@ +# SPDX-License-Identifier: Apache-2.0 +"""WanGame bidirectional model plugin (per-role instance).""" + +from __future__ import annotations + +import copy +from typing import Any, Literal, TYPE_CHECKING + +import torch + +import fastvideo.envs as envs +from fastvideo.distributed import ( + get_local_torch_device, + get_sp_group, + get_world_group, +) +from fastvideo.forward_context import set_forward_context +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.pipelines import TrainingBatch +from fastvideo.training.activation_checkpoint import ( + apply_activation_checkpointing, ) +from fastvideo.training.training_utils import ( + compute_density_for_timestep_sampling, + get_sigmas, + normalize_dit_input, + shift_timestep, +) +from fastvideo.utils import ( + is_vmoba_available, + is_vsa_available, + set_random_seed, +) + +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.module_state import ( + apply_trainable, ) +from fastvideo.train.utils.moduleloader import ( + load_module_from_path, ) + +if TYPE_CHECKING: + from fastvideo.train.utils.training_config import ( + TrainingConfig, ) + +try: + from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionMetadataBuilder, ) + from fastvideo.attention.backends.vmoba import ( + VideoMobaAttentionMetadataBuilder, ) +except Exception: + VideoSparseAttentionMetadataBuilder = None # type: ignore[assignment] + VideoMobaAttentionMetadataBuilder = None # type: ignore[assignment] + + +class WanGameModel(ModelBase): + """WanGame per-role model: owns transformer + noise_scheduler.""" + + _transformer_cls_name: str = ("WanGameActionTransformer3DModel") + + def __init__( + self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + disable_custom_init_weights: bool = False, + flow_shift: float = 3.0, + enable_gradient_checkpointing_type: str | None = None, + transformer_override_safetensor: str | None = None, + ) -> None: + self._init_from = str(init_from) + self._trainable = bool(trainable) + + self.transformer = self._load_transformer( + init_from=self._init_from, + trainable=self._trainable, + disable_custom_init_weights=(disable_custom_init_weights), + enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type), + training_config=training_config, + transformer_override_safetensor=(transformer_override_safetensor), + ) + + self.noise_scheduler = (FlowMatchEulerDiscreteScheduler(shift=float(flow_shift))) + + # Filled by init_preprocessors (student only). + self.vae: Any = None + self.training_config: TrainingConfig = training_config + self.dataloader: Any = None + self.validator: Any = None + self.start_step: int = 0 + + self.world_group: Any = None + self.sp_group: Any = None + self.device: Any = get_local_torch_device() + + self.noise_random_generator: (torch.Generator | None) = None + self.noise_gen_cuda: torch.Generator | None = None + + self.timestep_shift: float = float(flow_shift) + self.num_train_timestep: int = int(self.noise_scheduler.num_train_timesteps) + self.min_timestep: int = 0 + self.max_timestep: int = self.num_train_timestep + + def _load_transformer( + self, + *, + init_from: str, + trainable: bool, + disable_custom_init_weights: bool, + enable_gradient_checkpointing_type: str | None, + training_config: TrainingConfig, + transformer_override_safetensor: str | None = None, + ) -> torch.nn.Module: + transformer = load_module_from_path( + model_path=init_from, + module_type="transformer", + training_config=training_config, + disable_custom_init_weights=(disable_custom_init_weights), + override_transformer_cls_name=(self._transformer_cls_name), + transformer_override_safetensor=(transformer_override_safetensor), + ) + transformer = apply_trainable(transformer, trainable=trainable) + ckpt_type = (enable_gradient_checkpointing_type or getattr( + getattr(training_config, "model", None), + "enable_gradient_checkpointing_type", + None, + )) + if trainable and ckpt_type: + transformer = apply_activation_checkpointing( + transformer, + checkpointing_type=ckpt_type, + ) + return transformer + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def init_preprocessors(self, training_config: TrainingConfig) -> None: + """Load VAE, build dataloader, seed RNGs.""" + 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() + + from fastvideo.dataset.dataloader.schema import ( + pyarrow_schema_wangame, ) + from fastvideo.train.utils.dataloader import ( + build_parquet_wangame_train_dataloader, ) + + self.dataloader = (build_parquet_wangame_train_dataloader( + training_config.data, + parquet_schema=pyarrow_schema_wangame, + )) + self.start_step = 0 + + # ------------------------------------------------------------------ + # ModelBase overrides: timestep helpers + # ------------------------------------------------------------------ + + @property + def num_train_timesteps(self) -> int: + return int(self.num_train_timestep) + + def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + timestep = shift_timestep( + timestep, + self.timestep_shift, + self.num_train_timestep, + ) + return timestep.clamp(self.min_timestep, self.max_timestep) + + # ------------------------------------------------------------------ + # ModelBase overrides: lifecycle hooks + # ------------------------------------------------------------------ + + def on_train_start(self) -> None: + assert self.training_config is not None + tc = self.training_config + seed = tc.data.seed + if seed is None: + raise ValueError("training.data.seed must be set " + "for training") + + global_rank = int(getattr(self.world_group, "rank", 0)) + sp_world_size = int(tc.distributed.sp_size or 1) + if sp_world_size > 1: + sp_group_seed = int(seed) + (global_rank // sp_world_size) + set_random_seed(sp_group_seed) + else: + set_random_seed(int(seed) + global_rank) + + self.noise_random_generator = torch.Generator(device="cpu").manual_seed(int(seed)) + self.noise_gen_cuda = torch.Generator(device=self.device).manual_seed(int(seed)) + + def get_rng_generators(self, ) -> dict[str, torch.Generator]: + generators: dict[str, torch.Generator] = {} + if self.noise_random_generator is not None: + generators["noise_cpu"] = (self.noise_random_generator) + if self.noise_gen_cuda is not None: + generators["noise_cuda"] = self.noise_gen_cuda + return generators + + # ------------------------------------------------------------------ + # ModelBase overrides: runtime primitives + # ------------------------------------------------------------------ + + def _trim_temporal_prefix( + self, + value: torch.Tensor, + *, + target_length: int, + where: str, + dim: int, + ) -> torch.Tensor: + current_length = int(value.shape[dim]) + if current_length < target_length: + raise ValueError( + f"{where} temporal dim mismatch: got {current_length}, " + f"expected at least {target_length}" + ) + if current_length == target_length: + return value + + index = [slice(None)] * value.ndim + index[dim] = slice(0, target_length) + return value[tuple(index)] + + def prepare_batch( + self, + raw_batch: dict[str, Any], + *, + generator: torch.Generator | None = None, + current_vsa_sparsity: float = 0.0, + latents_source: Literal["data", "zeros"] = "data", + ) -> TrainingBatch: + del generator + assert self.training_config is not None + tc = self.training_config + dtype = self._get_training_dtype() + device = self.device + + training_batch = TrainingBatch(current_vsa_sparsity=current_vsa_sparsity) + infos = raw_batch.get("info_list") + + if latents_source == "zeros": + clip_feature = raw_batch["clip_feature"] + batch_size = int(clip_feature.shape[0]) + vae_config = ( + tc.pipeline_config.vae_config.arch_config # type: ignore[union-attr] + ) + num_channels = int(vae_config.z_dim) + spatial_compression_ratio = int(vae_config.spatial_compression_ratio) + latent_height = (int(tc.data.num_height) // spatial_compression_ratio) + latent_width = (int(tc.data.num_width) // spatial_compression_ratio) + latents = torch.zeros( + batch_size, + num_channels, + int(tc.data.num_latent_t), + latent_height, + latent_width, + device=device, + dtype=dtype, + ) + elif latents_source == "data": + if "vae_latent" not in raw_batch: + raise ValueError("vae_latent not found in batch " + "and latents_source='data'") + latents = raw_batch["vae_latent"] + latents = self._trim_temporal_prefix( + latents, + target_length=int(tc.data.num_latent_t), + where="vae_latent", + dim=2, + ) + latents = latents.to(device, dtype=dtype) + else: + raise ValueError(f"Unknown latents_source: " + f"{latents_source!r}") + + if "clip_feature" not in raw_batch: + raise ValueError("clip_feature must be present for WanGame") + image_embeds = raw_batch["clip_feature"].to(device, dtype=dtype) + + if "first_frame_latent" not in raw_batch: + raise ValueError("first_frame_latent must be present " + "for WanGame") + image_latents = raw_batch["first_frame_latent"] + image_latents = self._trim_temporal_prefix( + image_latents, + target_length=int(tc.data.num_latent_t), + where="first_frame_latent", + dim=2, + ) + image_latents = image_latents.to(device, dtype=dtype) + + pil_image = raw_batch.get("pil_image") + if isinstance(pil_image, torch.Tensor): + training_batch.preprocessed_image = (pil_image.to(device=device)) + else: + training_batch.preprocessed_image = pil_image + + keyboard_cond = raw_batch.get("keyboard_cond") + if (isinstance(keyboard_cond, torch.Tensor) and keyboard_cond.numel() > 0): + training_batch.keyboard_cond = keyboard_cond.to(device, dtype=dtype) + else: + training_batch.keyboard_cond = None + + mouse_cond = raw_batch.get("mouse_cond") + if (isinstance(mouse_cond, torch.Tensor) and mouse_cond.numel() > 0): + training_batch.mouse_cond = mouse_cond.to(device, dtype=dtype) + else: + training_batch.mouse_cond = None + + temporal_compression_ratio = ( + tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio # type: ignore[union-attr] + ) + expected_num_frames = ((tc.data.num_latent_t - 1) * temporal_compression_ratio + 1) + if training_batch.keyboard_cond is not None: + training_batch.keyboard_cond = self._trim_temporal_prefix( + training_batch.keyboard_cond, + target_length=int(expected_num_frames), + where="keyboard_cond", + dim=1, + ) + if training_batch.mouse_cond is not None: + training_batch.mouse_cond = self._trim_temporal_prefix( + training_batch.mouse_cond, + target_length=int(expected_num_frames), + where="mouse_cond", + dim=1, + ) + + training_batch.latents = latents + training_batch.encoder_hidden_states = None + training_batch.encoder_attention_mask = None + training_batch.image_embeds = image_embeds + training_batch.image_latents = image_latents + training_batch.infos = infos + + training_batch.latents = normalize_dit_input("wan", training_batch.latents, self.vae) + training_batch = self._prepare_dit_inputs(training_batch) + 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] + + training_batch.mask_lat_size = (self._build_i2v_mask_latents(image_latents)) + viewmats, intrinsics, action_labels = (self._process_actions(training_batch)) + training_batch.viewmats = viewmats + training_batch.Ks = intrinsics + training_batch.action = action_labels + + return training_batch + + def add_noise( + self, + clean_latents: torch.Tensor, + noise: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + b, t = clean_latents.shape[:2] + noisy = self.noise_scheduler.add_noise( + clean_latents.flatten(0, 1), + noise.flatten(0, 1), + timestep, + ).unflatten(0, (b, t)) + return noisy + + def predict_x0( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + conditional: bool, + cfg_uncond: dict[str, Any] | None = None, + attn_kind: Literal["dense", "vsa"] = "dense", + ) -> torch.Tensor: + device_type = self.device.type + dtype = noisy_latents.dtype + + 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}") + + with torch.autocast(device_type, dtype=dtype), set_forward_context( + current_timestep=batch.timesteps, + attn_metadata=attn_metadata, + ): + cond_inputs = (self._select_cfg_condition_inputs( + batch, + conditional=conditional, + cfg_uncond=cfg_uncond, + )) + input_kwargs = (self._build_distill_input_kwargs( + noisy_latents, + timestep, + image_embeds=cond_inputs["image_embeds"], + image_latents=cond_inputs["image_latents"], + mask_lat_size=cond_inputs["mask_lat_size"], + viewmats=cond_inputs["viewmats"], + Ks=cond_inputs["Ks"], + action=cond_inputs["action"], + mouse_cond=cond_inputs["mouse_cond"], + keyboard_cond=cond_inputs["keyboard_cond"], + )) + transformer = self._get_transformer(timestep) + pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + pred_x0 = pred_noise_to_pred_video( + pred_noise=pred_noise.flatten(0, 1), + noise_input_latent=noisy_latents.flatten(0, 1), + timestep=timestep, + scheduler=self.noise_scheduler, + ).unflatten(0, pred_noise.shape[:2]) + return pred_x0 + + def predict_noise( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + conditional: bool, + cfg_uncond: dict[str, Any] | None = None, + attn_kind: Literal["dense", "vsa"] = "dense", + ) -> torch.Tensor: + device_type = self.device.type + dtype = noisy_latents.dtype + + 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}") + + with torch.autocast(device_type, dtype=dtype), set_forward_context( + current_timestep=batch.timesteps, + attn_metadata=attn_metadata, + ): + cond_inputs = (self._select_cfg_condition_inputs( + batch, + conditional=conditional, + cfg_uncond=cfg_uncond, + )) + input_kwargs = (self._build_distill_input_kwargs( + noisy_latents, + timestep, + image_embeds=cond_inputs["image_embeds"], + image_latents=cond_inputs["image_latents"], + mask_lat_size=cond_inputs["mask_lat_size"], + viewmats=cond_inputs["viewmats"], + Ks=cond_inputs["Ks"], + action=cond_inputs["action"], + mouse_cond=cond_inputs["mouse_cond"], + keyboard_cond=cond_inputs["keyboard_cond"], + )) + transformer = self._get_transformer(timestep) + pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + return pred_noise + + def backward( + self, + loss: torch.Tensor, + ctx: Any, + *, + grad_accum_rounds: int, + ) -> None: + timesteps, attn_metadata = ctx + with set_forward_context( + current_timestep=timesteps, + attn_metadata=attn_metadata, + ): + (loss / max(1, int(grad_accum_rounds))).backward() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _get_training_dtype(self) -> torch.dtype: + return torch.bfloat16 + + def _init_timestep_mechanics(self) -> None: + assert self.training_config is not None + tc = self.training_config + self.timestep_shift = float(tc.pipeline_config.flow_shift # type: ignore[union-attr] + ) + self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps) + self.min_timestep = 0 + self.max_timestep = self.num_train_timestep + + def _sample_timesteps(self, batch_size: int, device: torch.device) -> torch.Tensor: + if self.noise_random_generator is None: + raise RuntimeError("on_train_start() must be called " + "before prepare_batch()") + assert self.training_config is not None + tc = self.training_config + + u = compute_density_for_timestep_sampling( + weighting_scheme=tc.model.weighting_scheme, + batch_size=batch_size, + generator=self.noise_random_generator, + logit_mean=tc.model.logit_mean, + logit_std=tc.model.logit_std, + mode_scale=tc.model.mode_scale, + ) + indices = (u * self.noise_scheduler.config.num_train_timesteps).long() + return self.noise_scheduler.timesteps[indices].to(device=device) + + def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_config is not None + tc = self.training_config + latents_shape = training_batch.raw_latent_shape + patch_size = ( + tc.pipeline_config.dit_config.patch_size # type: ignore[union-attr] + ) + current_vsa_sparsity = (training_batch.current_vsa_sparsity) + assert latents_shape is not None + assert training_batch.timesteps is not None + + if (envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN"): + if (not is_vsa_available() or VideoSparseAttentionMetadataBuilder is None): + raise ImportError("FASTVIDEO_ATTENTION_BACKEND is " + "VIDEO_SPARSE_ATTN, but " + "fastvideo_kernel is not correctly " + "installed or detected.") + training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder().build( # type: ignore[misc] + raw_latent_shape=latents_shape[2:5], + current_timestep=(training_batch.timesteps), + patch_size=patch_size, + VSA_sparsity=current_vsa_sparsity, + device=self.device, + ) + elif (envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN"): + if (not is_vmoba_available() or VideoMobaAttentionMetadataBuilder is None): + raise ImportError("FASTVIDEO_ATTENTION_BACKEND is " + "VMOBA_ATTN, but fastvideo_kernel " + "(or flash_attn>=2.7.4) is not " + "correctly installed.") + moba_params = tc.model.moba_config.copy() + moba_params.update({ + "current_timestep": (training_batch.timesteps), + "raw_latent_shape": (training_batch.raw_latent_shape[2:5]), + "patch_size": patch_size, + "device": self.device, + }) + training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder().build(** + moba_params) # type: ignore[misc] + else: + training_batch.attn_metadata = None + + return training_batch + + def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_config is not None + tc = self.training_config + latents = training_batch.latents + assert isinstance(latents, torch.Tensor) + batch_size = latents.shape[0] + + if self.noise_gen_cuda is None: + raise RuntimeError("on_train_start() must be called " + "before prepare_batch()") + + noise = torch.randn( + latents.shape, + generator=self.noise_gen_cuda, + device=latents.device, + dtype=latents.dtype, + ) + timesteps = self._sample_timesteps(batch_size, latents.device) + if int(tc.distributed.sp_size or 1) > 1: + self.sp_group.broadcast(timesteps, src=0) + + sigmas = get_sigmas( + self.noise_scheduler, + latents.device, + timesteps, + n_dim=latents.ndim, + dtype=latents.dtype, + ) + noisy_model_input = ((1.0 - sigmas) * latents + sigmas * noise) + + training_batch.noisy_model_input = (noisy_model_input) + training_batch.timesteps = timesteps + training_batch.sigmas = sigmas + training_batch.noise = noise + training_batch.raw_latent_shape = latents.shape + + training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4)) + return training_batch + + def _build_i2v_mask_latents(self, image_latents: torch.Tensor) -> torch.Tensor: + assert self.training_config is not None + tc = self.training_config + temporal_compression_ratio = ( + tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio # type: ignore[union-attr] + ) + num_frames = ((tc.data.num_latent_t - 1) * temporal_compression_ratio + 1) + + ( + batch_size, + _num_channels, + _t, + 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) + return mask_lat_size.to( + device=image_latents.device, + dtype=image_latents.dtype, + ) + + def _process_actions(self, training_batch: TrainingBatch) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + keyboard_cond = getattr(training_batch, "keyboard_cond", None) + mouse_cond = getattr(training_batch, "mouse_cond", None) + if keyboard_cond is None or mouse_cond is None: + raise ValueError("WanGame batch must provide " + "keyboard_cond and mouse_cond") + + from fastvideo.models.dits.hyworld.pose import ( + process_custom_actions, ) + + batch_size = int(training_batch.noisy_model_input.shape[0] # type: ignore[union-attr] + ) + viewmats_list: list[torch.Tensor] = [] + intrinsics_list: list[torch.Tensor] = [] + action_labels_list: list[torch.Tensor] = [] + for b in range(batch_size): + v, i, a = process_custom_actions(keyboard_cond[b], mouse_cond[b]) + viewmats_list.append(v) + intrinsics_list.append(i) + action_labels_list.append(a) + + viewmats = torch.stack(viewmats_list, dim=0).to(device=self.device, dtype=torch.bfloat16) + intrinsics = torch.stack(intrinsics_list, dim=0).to(device=self.device, dtype=torch.bfloat16) + action_labels = torch.stack(action_labels_list, dim=0).to(device=self.device, dtype=torch.bfloat16) + + num_latent_t = int(training_batch.noisy_model_input.shape[2] # type: ignore[union-attr] + ) + if int(action_labels.shape[1]) != num_latent_t: + raise ValueError("Action conditioning temporal dim " + "mismatch: " + f"action={tuple(action_labels.shape)} " + f"vs latent_t={num_latent_t}") + if int(viewmats.shape[1]) != num_latent_t: + raise ValueError("Viewmats temporal dim mismatch: " + f"viewmats={tuple(viewmats.shape)} " + f"vs latent_t={num_latent_t}") + + return viewmats, intrinsics, action_labels + + def _build_distill_input_kwargs( + self, + noisy_video_latents: torch.Tensor, + timestep: torch.Tensor, + *, + image_embeds: torch.Tensor, + image_latents: torch.Tensor, + mask_lat_size: torch.Tensor, + viewmats: torch.Tensor | None, + Ks: torch.Tensor | None, + action: torch.Tensor | None, + mouse_cond: torch.Tensor | None, + keyboard_cond: torch.Tensor | None, + ) -> dict[str, Any]: + hidden_states = torch.cat( + [ + noisy_video_latents.permute(0, 2, 1, 3, 4), + mask_lat_size, + image_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": image_embeds, + "viewmats": viewmats, + "Ks": Ks, + "action": action, + "mouse_cond": mouse_cond, + "keyboard_cond": keyboard_cond, + "return_dict": False, + } + + def _select_cfg_condition_inputs( + self, + batch: TrainingBatch, + *, + conditional: bool, + cfg_uncond: dict[str, Any] | None, + ) -> dict[str, Any]: + image_embeds = batch.image_embeds + image_latents = batch.image_latents + mask_lat_size = batch.mask_lat_size + if image_embeds is None: + raise RuntimeError("WanGameModel requires " + "TrainingBatch.image_embeds") + if image_latents is None: + raise RuntimeError("WanGameModel requires " + "TrainingBatch.image_latents") + if mask_lat_size is None: + raise RuntimeError("WanGameModel requires " + "TrainingBatch.mask_lat_size") + + viewmats = getattr(batch, "viewmats", None) + Ks = getattr(batch, "Ks", None) + action = getattr(batch, "action", None) + mouse_cond = getattr(batch, "mouse_cond", None) + keyboard_cond = getattr(batch, "keyboard_cond", None) + + if conditional or cfg_uncond is None: + return { + "image_embeds": image_embeds, + "image_latents": image_latents, + "mask_lat_size": mask_lat_size, + "viewmats": viewmats, + "Ks": Ks, + "action": action, + "mouse_cond": mouse_cond, + "keyboard_cond": keyboard_cond, + } + + on_missing_raw = cfg_uncond.get("on_missing", "error") + if not isinstance(on_missing_raw, str): + raise ValueError("method_config.cfg_uncond.on_missing " + "must be a string, got " + f"{type(on_missing_raw).__name__}") + on_missing = on_missing_raw.strip().lower() + if on_missing not in {"error", "ignore"}: + raise ValueError("method_config.cfg_uncond.on_missing " + "must be one of {error, ignore}, got " + f"{on_missing_raw!r}") + + supported_channels = {"image", "action"} + for channel, policy_raw in cfg_uncond.items(): + if channel in {"on_missing"}: + continue + if channel in supported_channels: + continue + if policy_raw is None: + continue + if not isinstance(policy_raw, str): + raise ValueError("method_config.cfg_uncond values " + "must be strings, got " + f"{channel}=" + f"{type(policy_raw).__name__}") + policy = policy_raw.strip().lower() + if policy == "keep": + continue + if on_missing == "ignore": + continue + raise ValueError("WanGameModel does not support " + "cfg_uncond channel " + f"{channel!r} (policy={policy!r}). " + "Set cfg_uncond.on_missing=ignore or " + "remove the channel.") + + def _get_policy(channel: str) -> str: + raw = cfg_uncond.get(channel, "keep") + if raw is None: + return "keep" + if not isinstance(raw, str): + raise ValueError("method_config.cfg_uncond values " + "must be strings, got " + f"{channel}={type(raw).__name__}") + policy = raw.strip().lower() + if policy not in {"keep", "zero", "drop"}: + raise ValueError("method_config.cfg_uncond values " + "must be one of " + "{keep, zero, drop}, got " + f"{channel}={raw!r}") + return policy + + image_policy = _get_policy("image") + if image_policy == "zero": + image_embeds = torch.zeros_like(image_embeds) + image_latents = torch.zeros_like(image_latents) + mask_lat_size = torch.zeros_like(mask_lat_size) + elif image_policy == "drop": + raise ValueError("cfg_uncond.image=drop is not supported " + "for WanGame I2V; use " + "cfg_uncond.image=zero or keep.") + + action_policy = _get_policy("action") + if action_policy == "zero": + if (viewmats is None or Ks is None or action is None): + if on_missing == "ignore": + pass + else: + raise ValueError("cfg_uncond.action=zero requires " + "action conditioning tensors, " + "but TrainingBatch is missing " + "{viewmats, Ks, action}.") + else: + viewmats = torch.zeros_like(viewmats) + Ks = torch.zeros_like(Ks) + action = torch.zeros_like(action) + if mouse_cond is not None: + mouse_cond = torch.zeros_like(mouse_cond) + if keyboard_cond is not None: + keyboard_cond = torch.zeros_like(keyboard_cond) + elif action_policy == "drop": + viewmats = None + Ks = None + action = None + mouse_cond = None + keyboard_cond = None + + return { + "image_embeds": image_embeds, + "image_latents": image_latents, + "mask_lat_size": mask_lat_size, + "viewmats": viewmats, + "Ks": Ks, + "action": action, + "mouse_cond": mouse_cond, + "keyboard_cond": keyboard_cond, + } + + def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module: + return self.transformer diff --git a/fastvideo/train/models/wangame/wangame_causal.py b/fastvideo/train/models/wangame/wangame_causal.py new file mode 100644 index 000000000..6e472f828 --- /dev/null +++ b/fastvideo/train/models/wangame/wangame_causal.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +"""WanGame causal model plugin (per-role instance, streaming/cache).""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal, TYPE_CHECKING + +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.models.utils import pred_noise_to_pred_video + +from fastvideo.train.models.base import CausalModelBase +from fastvideo.train.models.wangame.wangame import WanGameModel + +if TYPE_CHECKING: + from fastvideo.train.utils.training_config import ( + TrainingConfig, ) + + +@dataclass(slots=True) +class _StreamingCaches: + kv_cache: list[dict[str, Any]] + 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 WanGameCausalModel(WanGameModel, CausalModelBase): + """WanGame per-role model with causal/streaming primitives.""" + + _transformer_cls_name: str = ("CausalWanGameActionTransformer3DModel") + + def __init__( + self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + disable_custom_init_weights: bool = False, + flow_shift: float = 3.0, + enable_gradient_checkpointing_type: str | None = None, + transformer_override_safetensor: str | None = None, + ) -> None: + super().__init__( + init_from=init_from, + trainable=trainable, + disable_custom_init_weights=disable_custom_init_weights, + flow_shift=flow_shift, + enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type), + training_config=training_config, + transformer_override_safetensor=(transformer_override_safetensor), + ) + self._streaming_caches: dict[tuple[int, str], _StreamingCaches] = {} + + # --- CausalModelBase override: clear_caches --- + def clear_caches(self, *, cache_tag: str = "pos") -> None: + self._streaming_caches.pop((id(self), str(cache_tag)), None) + + # --- CausalModelBase override: predict_noise_streaming --- + 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: + 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 + 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) + + model_kwargs: dict[str, Any] = { + "kv_cache": kv_cache, + "crossattn_cache": crossattn_cache, + "current_start": cur_start_frame * frame_seq_length, + "start_frame": cur_start_frame, + "is_cache": bool(store_kv), + } + + 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, + ): + cond_inputs = self._select_cfg_condition_inputs( + batch, + conditional=conditional, + cfg_uncond=cfg_uncond, + ) + cond_inputs = self._slice_cond_inputs_for_streaming( + cond_inputs=cond_inputs, + cur_start_frame=cur_start_frame, + num_frames=num_frames, + ) + input_kwargs = self._build_distill_input_kwargs( + noisy_latents, + timestep_full, + image_embeds=cond_inputs["image_embeds"], + image_latents=cond_inputs["image_latents"], + mask_lat_size=cond_inputs["mask_lat_size"], + viewmats=cond_inputs["viewmats"], + Ks=cond_inputs["Ks"], + action=cond_inputs["action"], + mouse_cond=cond_inputs["mouse_cond"], + keyboard_cond=cond_inputs["keyboard_cond"], + ) + + input_kwargs["timestep"] = timestep_full.to(device=self.device, dtype=torch.long) + input_kwargs.update(model_kwargs) + + 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 + + # --- CausalModelBase override: predict_x0_streaming --- + def predict_x0_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: + pred_noise = self.predict_noise_streaming( + noisy_latents, + timestep, + batch, + conditional=conditional, + cache_tag=cache_tag, + store_kv=store_kv, + cur_start_frame=cur_start_frame, + cfg_uncond=cfg_uncond, + attn_kind=attn_kind, + ) + if pred_noise is None: + return None + + pred_x0 = pred_noise_to_pred_video( + pred_noise=pred_noise.flatten(0, 1), + noise_input_latent=noisy_latents.flatten(0, 1), + timestep=self.shift_and_clamp_timestep( + self._ensure_per_frame_timestep( + timestep=timestep, + batch_size=int(noisy_latents.shape[0]), + num_frames=int(noisy_latents.shape[1]), + device=noisy_latents.device, + ).flatten()), + scheduler=self.noise_scheduler, + ).unflatten(0, pred_noise.shape[:2]) + return pred_x0 + + # --- internal helpers --- + + def _ensure_per_frame_timestep( + self, + *, + timestep: torch.Tensor, + batch_size: int, + num_frames: int, + device: torch.device, + ) -> torch.Tensor: + if timestep.ndim == 0: + return (timestep.view(1, 1).expand(batch_size, num_frames).to(device=device)) + if timestep.ndim == 1: + if int(timestep.shape[0]) == batch_size: + return (timestep.view(batch_size, 1).expand(batch_size, num_frames).to(device=device)) + raise ValueError("streaming timestep must be scalar, [B], or " + f"[B, T]; got shape={tuple(timestep.shape)}") + if timestep.ndim == 2: + return timestep.to(device=device) + raise ValueError("streaming timestep must be scalar, [B], or [B, T]; " + f"got ndim={int(timestep.ndim)}") + + def _slice_cond_inputs_for_streaming( + self, + *, + cond_inputs: dict[str, Any], + cur_start_frame: int, + num_frames: int, + ) -> dict[str, Any]: + start = int(cur_start_frame) + num_frames = int(num_frames) + if num_frames <= 0: + raise ValueError("num_frames must be positive for streaming") + if start < 0: + raise ValueError("cur_start_frame must be >= 0 for streaming") + end = start + num_frames + + sliced: dict[str, Any] = dict(cond_inputs) + + image_latents = cond_inputs.get("image_latents") + if isinstance(image_latents, torch.Tensor): + sliced["image_latents"] = image_latents[:, :, start:end] + + mask_lat_size = cond_inputs.get("mask_lat_size") + if isinstance(mask_lat_size, torch.Tensor): + sliced["mask_lat_size"] = mask_lat_size[:, :, start:end] + + viewmats = cond_inputs.get("viewmats") + if isinstance(viewmats, torch.Tensor): + sliced["viewmats"] = viewmats[:, start:end] + + Ks = cond_inputs.get("Ks") + if isinstance(Ks, torch.Tensor): + sliced["Ks"] = Ks[:, start:end] + + action = cond_inputs.get("action") + if isinstance(action, torch.Tensor): + sliced["action"] = action[:, start:end] + + temporal_compression_ratio = int( + self.training_config.pipeline_config.vae_config.arch_config.temporal_compression_ratio) + raw_end_frame_idx = (1 + temporal_compression_ratio * max(0, end - 1)) + + mouse_cond = cond_inputs.get("mouse_cond") + if isinstance(mouse_cond, torch.Tensor): + sliced["mouse_cond"] = mouse_cond[:, :raw_end_frame_idx] + + keyboard_cond = cond_inputs.get("keyboard_cond") + if isinstance(keyboard_cond, torch.Tensor): + sliced["keyboard_cond"] = keyboard_cond[:, :raw_end_frame_idx] + + return sliced + + 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()), + ) + crossattn_cache = self._initialize_crossattn_cache(transformer=transformer, device=device) + + caches = _StreamingCaches( + kv_cache=kv_cache, + 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 _compute_frame_seq_length( + self, + transformer: torch.nn.Module, + noisy_latents: torch.Tensor, + ) -> int: + latent_seq_length = int(noisy_latents.shape[-1]) * int(noisy_latents.shape[-2]) + patch_size = getattr(transformer, "patch_size", None) + if patch_size is None: + patch_size = getattr( + getattr(transformer, "config", None), + "arch_config", + None, + ) + patch_size = getattr(patch_size, "patch_size", None) + if patch_size is None: + raise ValueError("Unable to determine transformer.patch_size " + "for causal streaming") + patch_ratio = int(patch_size[-1]) * int(patch_size[-2]) + if patch_ratio <= 0: + raise ValueError("Invalid patch_size for causal streaming") + return latent_seq_length // patch_ratio + + def _get_sliding_window_num_frames(self, transformer: torch.nn.Module) -> int: + cfg = getattr(transformer, "config", None) + arch_cfg = getattr(cfg, "arch_config", None) + value = (getattr(arch_cfg, "sliding_window_num_frames", None) if arch_cfg is not None else None) + if value is None: + return 15 + return int(value) + + def _get_local_attn_size(self, transformer: torch.nn.Module) -> int: + try: + value = getattr(transformer, "local_attn_size", -1) + except Exception: + value = -1 + if value is None: + return -1 + return int(value) + + def _initialize_kv_cache( + self, + *, + transformer: torch.nn.Module, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + frame_seq_length: int, + local_attn_size: int, + sliding_window_num_frames: int, + checkpoint_safe: bool, + ) -> list[dict[str, Any]]: + num_blocks = len(getattr(transformer, "blocks", [])) + if num_blocks <= 0: + raise ValueError("Unexpected transformer.blocks for causal " + "streaming") + + try: + num_attention_heads = int(transformer.num_attention_heads # type: ignore[attr-defined] + ) + except AttributeError as e: + raise ValueError("Transformer is missing num_attention_heads") from e + + try: + attention_head_dim = int(transformer.attention_head_dim # type: ignore[attr-defined] + ) + except AttributeError: + try: + hidden_size = int(transformer.hidden_size # type: ignore[attr-defined] + ) + except AttributeError as e: + raise ValueError("Transformer is missing attention_head_dim " + "and hidden_size") from e + attention_head_dim = hidden_size // max(1, num_attention_heads) + + if local_attn_size != -1: + kv_cache_size = (int(local_attn_size) * int(frame_seq_length)) + else: + kv_cache_size = int(frame_seq_length) * int(sliding_window_num_frames) + + if checkpoint_safe: + tc = getattr(self, "training_config", None) + total_frames = int(tc.data.num_frames if tc is not None else 0) + if total_frames <= 0: + raise ValueError("training.num_frames must be set to enable " + "checkpoint-safe streaming KV cache; " + f"got {total_frames}") + kv_cache_size = max( + kv_cache_size, + int(frame_seq_length) * total_frames, + ) + + kv_cache: list[dict[str, Any]] = [] + for _ in range(num_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.zeros((), dtype=torch.long, device=device), + "local_end_index": + torch.zeros((), dtype=torch.long, device=device), + }) + + return kv_cache + + def _should_use_checkpoint_safe_kv_cache(self) -> bool: + tc = getattr(self, "training_config", None) + if tc is not None: + checkpointing_type = tc.model.enable_gradient_checkpointing_type + else: + checkpointing_type = None + return bool(checkpointing_type) and bool(self._trainable) + + def _should_snapshot_streaming_cache(self) -> bool: + return self._should_use_checkpoint_safe_kv_cache() + + def _snapshot_kv_cache_indices(self, kv_cache: list[dict[str, Any]]) -> list[dict[str, Any]]: + snapshot: list[dict[str, Any]] = [] + for block_cache in kv_cache: + global_end_index = block_cache.get("global_end_index") + local_end_index = block_cache.get("local_end_index") + if not isinstance(global_end_index, torch.Tensor) or not isinstance(local_end_index, torch.Tensor): + raise ValueError("Unexpected kv_cache index tensors; expected " + "tensors at kv_cache[*].{global_end_index, " + "local_end_index}") + + copied = dict(block_cache) + copied["global_end_index"] = (global_end_index.detach().clone()) + copied["local_end_index"] = (local_end_index.detach().clone()) + snapshot.append(copied) + return snapshot + + def _initialize_crossattn_cache( + self, + *, + transformer: torch.nn.Module, + device: torch.device, + ) -> list[dict[str, Any]] | None: + num_blocks = len(getattr(transformer, "blocks", [])) + if num_blocks <= 0: + return None + return [{ + "is_init": False, + "k": torch.empty(0, device=device), + "v": torch.empty(0, device=device), + } for _ in range(num_blocks)] diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index d814a6cfe..a3cae9ed3 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -2,15 +2,18 @@ from __future__ import annotations +import os import time from collections.abc import Iterator from dataclasses import dataclass from typing import Any, TYPE_CHECKING +import numpy as np import torch from tqdm.auto import tqdm from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.logger import init_logger from fastvideo.train.callbacks.callback import CallbackDict from fastvideo.train.methods.base import TrainingMethod from fastvideo.train.utils.tracking import build_tracker @@ -19,6 +22,8 @@ if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( TrainingConfig, ) +logger = init_logger(__name__) + def _coerce_log_scalar(value: Any, *, where: str) -> float: if isinstance(value, torch.Tensor): @@ -32,6 +37,78 @@ def _coerce_log_scalar(value: Any, *, where: str) -> float: f"{where}, got {type(value).__name__}") +def _maybe_log_resume_fingerprint( + method: TrainingMethod, + *, + global_rank: int, + step: int, +) -> None: + if os.getenv("FASTVIDEO_DEBUG_RESUME_HASH", "").lower() not in {"1", "true", "yes"}: + return + + transformer = method.transformer_inference + fingerprints: list[str] = [] + for idx, (name, param) in enumerate(transformer.named_parameters()): + if idx >= 3: + break + data = param.detach().reshape(-1) + if data.numel() == 0: + fingerprints.append(f"{name}:empty") + continue + sample = data[:16].float() + fingerprints.append( + f"{name}:shape={tuple(param.shape)} " + f"sample_sum={sample.sum().item():.8f} " + f"sample_mean={sample.mean().item():.8f} " + f"first={sample[0].item():.8f}" + ) + + print( + "DEBUG_RESUME_FINGERPRINT " + f"rank={global_rank} step={step} " + " | ".join(fingerprints), + flush=True, + ) + + +def _decode_latent_video( + vae: Any, + latents: torch.Tensor, +) -> np.ndarray: + with torch.no_grad(): + latents = latents.detach() + latents = latents.permute(0, 2, 1, 3, 4) + + scaling_factor = getattr(vae, "scaling_factor", None) + if scaling_factor is not None: + if isinstance(scaling_factor, torch.Tensor): + latents = latents / scaling_factor.to( + latents.device, + latents.dtype, + ) + else: + latents = latents / float(scaling_factor) + + shift_factor = getattr(vae, "shift_factor", None) + if shift_factor is not None: + if isinstance(shift_factor, torch.Tensor): + latents = latents + shift_factor.to( + latents.device, + latents.dtype, + ) + else: + latents = latents + float(shift_factor) + + with torch.autocast( + device_type="cuda", + dtype=torch.bfloat16, + enabled=latents.is_cuda, + ): + video = vae.decode(latents) + video = (video / 2 + 0.5).clamp(0, 1) + video = video.detach().cpu().float().permute(0, 2, 1, 3, 4) + return (video * 255).numpy().astype(np.uint8) + + @dataclass(slots=True) class TrainLoopState: step: int @@ -63,6 +140,65 @@ class Trainer: training_config, ) + def _should_log_train_artifacts( + self, + step: int, + ) -> bool: + validation_cb = self.callbacks.get_callback("validation") + every_steps = getattr(validation_cb, "every_steps", None) + if every_steps is None: + return False + every_steps = int(every_steps) + if every_steps <= 0: + return False + return step % every_steps == 0 + + def _log_train_artifacts( + self, + method: TrainingMethod, + outputs: dict[str, Any], + *, + step: int, + ) -> None: + if self.global_rank != 0: + return + if not self._should_log_train_artifacts(step): + return + + vae = getattr(method.student, "vae", None) + if vae is None: + return + + artifacts: dict[str, Any] = {} + dmd_latent_dict = outputs.get("dmd_latent_vis_dict") + if isinstance(dmd_latent_dict, dict) and dmd_latent_dict: + latents = dmd_latent_dict.get("real_score_pred_video") + if isinstance(latents, torch.Tensor): + video = _decode_latent_video(vae, latents) + artifact = self.tracker.video( + video, + fps=24, + format="mp4", + ) + if artifact is not None: + artifacts["dmd_real_score_pred_video"] = artifact + + for scalar_key in ("generator_timestep", "dmd_timestep"): + value = dmd_latent_dict.get(scalar_key) + if isinstance(value, torch.Tensor) and value.numel() == 1: + artifacts[scalar_key] = float(value.detach().item()) + + fake_score_latent_dict = outputs.get("fake_score_latent_vis_dict") + if isinstance(fake_score_latent_dict, dict) and fake_score_latent_dict: + value = fake_score_latent_dict.get("fake_score_timestep") + if isinstance(value, torch.Tensor) and value.numel() == 1: + artifacts["fake_score_timestep"] = float( + value.detach().item() + ) + + if artifacts: + self.tracker.log_artifacts(artifacts, step) + def _iter_dataloader(self, dataloader: Any) -> Iterator[dict[str, Any]]: data_iter = iter(dataloader) while True: @@ -89,18 +225,21 @@ class Trainer: method.set_tracker(self.tracker) method.on_train_start() + + resume_from_checkpoint = (tc.checkpoint.resume_from_checkpoint or "") + if checkpoint_manager is not None: + resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint))) + if resumed_step is not None: + start_step = int(resumed_step) + _maybe_log_resume_fingerprint( + method, + global_rank=self.global_rank, + step=start_step, + ) self.callbacks.on_train_start( method, iteration=start_step, ) - - resume_from_checkpoint = (tc.checkpoint.resume_from_checkpoint or "") - if checkpoint_manager is not None: - if resume_from_checkpoint: - method.seed_optimizer_state_for_resume() - resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint))) - if resumed_step is not None: - start_step = int(resumed_step) self.callbacks.on_validation_begin( method, iteration=start_step, @@ -108,14 +247,9 @@ class Trainer: method.optimizers_zero_grad(start_step) data_stream = self._iter_dataloader(dataloader) - - # Restore the RNG snapshot LAST — after dcp.load, - # after iter(dataloader), after everything that may - # have advanced the RNG as a side-effect. - if (checkpoint_manager is not None and resume_from_checkpoint): - checkpoint_manager.load_rng_snapshot(resume_from_checkpoint, ) progress = tqdm( range(start_step + 1, max_steps + 1), + total=max_steps, initial=start_step, desc="Steps", disable=self.local_rank > 0, @@ -125,12 +259,14 @@ class Trainer: loss_sums: dict[str, float] = {} metric_sums: dict[str, float] = {} + last_outputs: dict[str, Any] = {} for accum_iter in range(grad_accum): batch = next(data_stream) loss_map, outputs, step_metrics = method.single_train_step( batch, step, ) + last_outputs = outputs method.backward( loss_map, @@ -166,6 +302,11 @@ class Trainer: metrics["vsa_sparsity"] = float(tc.vsa_sparsity) if self.global_rank == 0 and metrics: self.tracker.log(metrics, step) + self._log_train_artifacts( + method, + last_outputs, + step=step, + ) self.callbacks.on_training_step_end( method, @@ -184,6 +325,33 @@ class Trainer: method, iteration=step, ) + if checkpoint_manager is not None: + validation_cb = self.callbacks.get_callback("validation") + latest_mf_metric: float | None = None + get_latest_metric = getattr( + validation_cb, + "get_latest_metric", + None, + ) + if callable(get_latest_metric): + latest_mf_metric = get_latest_metric( + "mf_angle_err_mean", + step=step, + ) + + checkpoint_manager.maybe_save_best( + step=step, + metric_value=latest_mf_metric, + metric_name="mf_angle_err_mean", + start_step=int( + tc.checkpoint.best_checkpoint_start_step + or 0 + ), + top_k=int( + tc.checkpoint.best_checkpoint_top_k + or 1 + ), + ) self.callbacks.on_train_end( method, diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 94453a196..c2028d825 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import math import os import random import re @@ -20,6 +21,9 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) _CHECKPOINT_DIR_RE = re.compile(r"^checkpoint-(\d+)$") +_BEST_CHECKPOINT_DIR_RE = re.compile( + r"^checkpoint-best-step-(\d+)$" +) def _is_stateful(obj: Any) -> bool: @@ -37,11 +41,38 @@ def _barrier() -> None: dist.barrier() +def _broadcast_int(value: int, *, src: int = 0) -> int: + if not (dist.is_available() and dist.is_initialized()): + return int(value) + backend = dist.get_backend() + use_cuda = ( + backend == "nccl" + and torch.cuda.is_available() + ) + device = ( + torch.device("cuda") + if use_cuda + else torch.device("cpu") + ) + tensor = torch.tensor( + [int(value)], + dtype=torch.int32, + device=device, + ) + dist.broadcast(tensor, src=src) + return int(tensor.item()) + + def _parse_step_from_dir(checkpoint_dir: Path) -> int: - match = _CHECKPOINT_DIR_RE.match(checkpoint_dir.name) + match = ( + _CHECKPOINT_DIR_RE.match(checkpoint_dir.name) + or _BEST_CHECKPOINT_DIR_RE.match(checkpoint_dir.name) + ) if not match: - raise ValueError(f"Invalid checkpoint directory name {checkpoint_dir.name!r}; " - "expected 'checkpoint-'") + raise ValueError( + f"Invalid checkpoint directory name {checkpoint_dir.name!r}; " + "expected 'checkpoint-' or 'checkpoint-best-step-'" + ) return int(match.group(1)) @@ -53,7 +84,10 @@ def _find_latest_checkpoint(output_dir: Path) -> Path | None: for child in output_dir.iterdir(): if not child.is_dir(): continue - if not _CHECKPOINT_DIR_RE.match(child.name): + if not ( + _CHECKPOINT_DIR_RE.match(child.name) + or _BEST_CHECKPOINT_DIR_RE.match(child.name) + ): continue if not (child / "dcp").is_dir(): continue @@ -100,7 +134,10 @@ def _resolve_resume_checkpoint(resume_from_checkpoint: str, *, output_dir: str) if path.is_dir() and path.name == "dcp": path = path.parent - if path.is_dir() and _CHECKPOINT_DIR_RE.match(path.name): + if path.is_dir() and ( + _CHECKPOINT_DIR_RE.match(path.name) + or _BEST_CHECKPOINT_DIR_RE.match(path.name) + ): if not (path / "dcp").is_dir(): raise FileNotFoundError(f"Missing dcp dir under checkpoint: {path / 'dcp'}") return path @@ -113,8 +150,9 @@ def _resolve_resume_checkpoint(resume_from_checkpoint: str, *, output_dir: str) # Give a clearer error message. out = Path(os.path.expanduser(str(output_dir))).resolve() raise ValueError("Could not resolve resume checkpoint. Expected a checkpoint directory " - f"named 'checkpoint-' (with 'dcp/' inside), or an output_dir " - f"containing such checkpoints. Got: {path} (output_dir={out}).") + "named 'checkpoint-' or 'checkpoint-best-step-' " + f"(with 'dcp/' inside), or an output_dir containing such " + f"checkpoints. Got: {path} (output_dir={out}).") class _RoleModuleContainer(torch.nn.Module): @@ -141,8 +179,7 @@ class _CallbackStateWrapper: return self._callbacks.state_dict() def load_state_dict( - self, - state_dict: dict[str, Any], + self, state_dict: dict[str, Any], ) -> None: self._callbacks.load_state_dict(state_dict) @@ -168,6 +205,7 @@ class CheckpointManager: output_dir: str, config: CheckpointConfig, callbacks: Any | None = None, + tracker: Any | None = None, raw_config: dict[str, Any] | None = None, ) -> None: self.method = method @@ -175,8 +213,10 @@ class CheckpointManager: self.output_dir = str(output_dir) self.config = config self._callbacks = callbacks + self._tracker = tracker self._raw_config = raw_config self._last_saved_step: int | None = None + self._last_best_saved_step: int | None = None def _build_states(self) -> dict[str, Any]: states: dict[str, Any] = self.method.checkpoint_state() @@ -187,7 +227,9 @@ class CheckpointManager: # Callback state (e.g. EMA shadow weights, validation RNG). if self._callbacks is not None and _is_stateful(self._callbacks): - states["callbacks"] = _CallbackStateWrapper(self._callbacks, ) + states["callbacks"] = _CallbackStateWrapper( + self._callbacks, + ) return states @@ -197,6 +239,15 @@ class CheckpointManager: def _dcp_dir(self, step: int) -> Path: return self._checkpoint_dir(step) / "dcp" + def _best_checkpoint_dir(self, step: int) -> Path: + return ( + Path(self.output_dir) + / f"checkpoint-best-step-{step}" + ) + + def _checkpoint_best_alias_path(self) -> Path: + return Path(self.output_dir) / "checkpoint-best" + def maybe_save(self, step: int) -> None: save_steps = int(self.config.save_steps or 0) if save_steps <= 0: @@ -215,120 +266,298 @@ class CheckpointManager: def save(self, step: int) -> None: checkpoint_dir = self._checkpoint_dir(step) - dcp_dir = self._dcp_dir(step) + self._save_checkpoint_dir( + checkpoint_dir, + step=step, + log_prefix="Saving checkpoint", + ) + self._last_saved_step = step + + self._cleanup_old_checkpoints() + + def _save_checkpoint_dir( + self, + checkpoint_dir: Path, + *, + step: int, + log_prefix: str, + metadata_extra: dict[str, Any] | None = None, + ) -> None: + dcp_dir = checkpoint_dir / "dcp" os.makedirs(dcp_dir, exist_ok=True) states = self._build_states() if _rank() == 0: - logger.info( - "Saving checkpoint to %s", + logger.info("%s to %s", log_prefix, checkpoint_dir) + self._write_metadata( checkpoint_dir, + step, + extra=metadata_extra, ) - self._write_metadata(checkpoint_dir, step) dcp.save(states, checkpoint_id=str(dcp_dir)) _barrier() - - # Save RNG state AFTER dcp.save so it captures the - # exact state the continuous run continues with. - # dcp.save triggers FSDP all-gather ops that can - # advance the RNG between when DCP captures it and - # when the save completes. self._save_rng_snapshot(checkpoint_dir) _barrier() - self._last_saved_step = step - - self._cleanup_old_checkpoints() - def _write_metadata( self, checkpoint_dir: Path, step: int, + *, + extra: dict[str, Any] | None = None, ) -> None: metadata: dict[str, Any] = {"step": step} if self._raw_config is not None: metadata["config"] = self._raw_config + if extra: + metadata.update(extra) meta_path = checkpoint_dir / "metadata.json" with open(meta_path, "w", encoding="utf-8") as f: json.dump(metadata, f, indent=2) @staticmethod - def load_metadata(checkpoint_dir: str | Path, ) -> dict[str, Any]: + def load_metadata( + checkpoint_dir: str | Path, + ) -> dict[str, Any]: """Read ``metadata.json`` from a checkpoint dir.""" meta_path = Path(checkpoint_dir) / "metadata.json" if not meta_path.is_file(): - raise FileNotFoundError(f"No metadata.json in {checkpoint_dir}") + raise FileNotFoundError( + f"No metadata.json in {checkpoint_dir}" + ) with open(meta_path, encoding="utf-8") as f: return json.load(f) # type: ignore[no-any-return] - def _save_rng_snapshot(self, checkpoint_dir: Path) -> None: - """Save per-rank RNG state to a separate file. - - Called AFTER ``dcp.save`` so the snapshot reflects - the exact state the continuous run continues with. - Each rank saves its own file because CUDA RNG and - custom generators differ across ranks. - """ - rank = _rank() - rng: dict[str, Any] = { - "torch_rng": torch.get_rng_state(), - "python_rng": random.getstate(), - "numpy_rng": np.random.get_state(), - } - rng["cuda_rng"] = torch.cuda.get_rng_state() - rng["gen_cuda"] = self.method.cuda_generator.get_state() - torch.save( - rng, - checkpoint_dir / f"rng_state_rank{rank}.pt", - ) - - def load_rng_snapshot( + def _list_best_checkpoint_entries( self, - checkpoint_path: str, - ) -> None: - """Restore per-rank RNG state from the snapshot file. + metric_name: str, + ) -> list[dict[str, Any]]: + output_dir = Path(self.output_dir) + if not output_dir.is_dir(): + return [] - Must be called AFTER ``dcp.load`` **and** after - ``iter(dataloader)`` so no later operation can - clobber the restored state. - """ - resolved = _resolve_resume_checkpoint( - checkpoint_path, - output_dir=self.output_dir, - ) - if resolved is None: - return - rank = _rank() - rng_path = resolved / f"rng_state_rank{rank}.pt" - if not rng_path.is_file(): - # Fall back to legacy single-file snapshot. - rng_path = resolved / "rng_state.pt" - if not rng_path.is_file(): - logger.warning( - "No rng_state in %s; skipping " - "RNG snapshot restore.", - resolved, + entries: list[dict[str, Any]] = [] + for child in output_dir.iterdir(): + if ( + not child.is_dir() + or not _BEST_CHECKPOINT_DIR_RE.match( + child.name + ) + ): + continue + metric_path = child / "best_metric.json" + if not metric_path.is_file(): + logger.warning( + "Skipping %s: best_metric.json " + "is missing", + child, + ) + continue + try: + with open( + metric_path, encoding="utf-8" + ) as f: + meta = json.load(f) + metric_raw = meta.get(metric_name) + if metric_raw is None: + metric_raw = meta.get( + "mf_angle_err_mean" + ) + if metric_raw is None: + metric_raw = meta.get( + "metric_value" + ) + + step_raw = meta.get("step") + if step_raw is None: + match = ( + _BEST_CHECKPOINT_DIR_RE.match( + child.name + ) + ) + if match is None: + continue + step_raw = int(match.group(1)) + + metric_val = float(metric_raw) + step_val = int(step_raw) + if not math.isfinite(metric_val): + raise ValueError( + "metric is non-finite" + ) + except Exception as e: + logger.warning( + "Skipping %s: invalid " + "best metric metadata (%s)", + child, + e, + ) + continue + entries.append( + { + "path": child, + "step": step_val, + "metric": metric_val, + } ) + + entries.sort( + key=lambda x: ( + float(x["metric"]), + int(x["step"]), + ) + ) + return entries + + def _update_best_checkpoint_alias( + self, + best_checkpoint_path: Path, + ) -> None: + alias_path = self._checkpoint_best_alias_path() + try: + if alias_path.is_symlink() or alias_path.is_file(): + alias_path.unlink() + elif alias_path.is_dir(): + shutil.rmtree(alias_path) + os.symlink( + os.path.basename(str(best_checkpoint_path)), + str(alias_path), + ) + except OSError as e: + logger.warning( + "Failed to update checkpoint-best " + "alias: %s", + e, + ) + + def maybe_save_best( + self, + *, + step: int, + metric_value: float | None, + metric_name: str = "mf_angle_err_mean", + start_step: int = 0, + top_k: int = 1, + ) -> None: + start_step = int(start_step or 0) + if start_step <= 0: + return + if int(step) < start_step: + return + if metric_value is None: + return + metric = float(metric_value) + if not math.isfinite(metric): + return + if self._last_best_saved_step == int(step): return - rng = torch.load( - rng_path, - map_location="cpu", - weights_only=False, - ) - if "torch_rng" in rng: - torch.set_rng_state(rng["torch_rng"]) - if "python_rng" in rng: - random.setstate(rng["python_rng"]) - if "numpy_rng" in rng: - np.random.set_state(rng["numpy_rng"]) + top_k = max(1, int(top_k or 1)) + metric_name = str(metric_name) - torch.cuda.set_rng_state(rng["cuda_rng"]) - self.method.cuda_generator.set_state(rng["gen_cuda"]) - logger.info( - "Restored RNG snapshot from %s", - rng_path, + should_save = 0 + if _rank() == 0: + entries = self._list_best_checkpoint_entries( + metric_name + ) + if len(entries) < top_k: + should_save = 1 + else: + worst = entries[-1] + if metric < float( + worst["metric"] + ): + should_save = 1 + should_save = _broadcast_int( + should_save, src=0 ) + if should_save == 0: + return + + checkpoint_dir = self._best_checkpoint_dir( + int(step) + ) + if _rank() == 0 and checkpoint_dir.exists(): + shutil.rmtree(checkpoint_dir, ignore_errors=True) + _barrier() + + logger.info( + "%s=%.6f at step %s entered top-%s " + "best checkpoints; saving.", + metric_name, + metric, + step, + top_k, + ) + self._save_checkpoint_dir( + checkpoint_dir, + step=int(step), + log_prefix="Saving best checkpoint", + metadata_extra={ + "kind": "best", + "metric_name": metric_name, + "metric_value": metric, + }, + ) + self._last_best_saved_step = int(step) + + if _rank() == 0: + metric_meta = { + "step": int(step), + metric_name: metric, + "metric_name": metric_name, + "metric_value": metric, + } + if metric_name == "mf_angle_err_mean": + metric_meta[ + "mf_angle_err_mean" + ] = metric + metric_path = ( + checkpoint_dir / "best_metric.json" + ) + with open( + metric_path, "w", encoding="utf-8" + ) as f: + json.dump(metric_meta, f, indent=2) + + best_entries = self._list_best_checkpoint_entries( + metric_name + ) + kept_entries = best_entries[:top_k] + for stale_entry in best_entries[top_k:]: + stale_path = Path( + stale_entry["path"] + ) + logger.info( + "Removing non-top-k best " + "checkpoint: %s", + stale_path, + ) + shutil.rmtree( + stale_path, ignore_errors=True + ) + + if kept_entries: + top1 = kept_entries[0] + self._update_best_checkpoint_alias( + Path(top1["path"]) + ) + if self._tracker is not None: + self._tracker.log( + { + f"best/{metric_name}": float( + top1["metric"] + ), + "best/step": int( + top1["step"] + ), + "best/topk_count": len( + kept_entries + ), + }, + int(step), + ) + _barrier() def maybe_resume(self, *, resume_from_checkpoint: str | None) -> int | None: if not resume_from_checkpoint: @@ -344,11 +573,111 @@ class CheckpointManager: states = self._build_states() logger.info("Loading Phase 2 checkpoint from %s", resolved) - dcp.load(states, checkpoint_id=str(resolved / "dcp")) + try: + dcp.load(states, checkpoint_id=str(resolved / "dcp")) + except BaseException as exc: + if not isinstance(exc, dcp.CheckpointException): + raise + msg = str(exc) + fallback_prefixes = ( + "optimizers.", + "schedulers.", + "dataloader", + "callbacks.", + "random_state", + ) + can_fallback = ( + "Missing key in checkpoint state_dict:" in msg + and any( + f"Missing key in checkpoint state_dict: {prefix}" + in msg + for prefix in fallback_prefixes + ) + ) + if not can_fallback: + raise + + model_only_states = { + key: value + for key, value in states.items() + if key.startswith("roles.") + } + if not model_only_states: + raise + + logger.warning( + "Resume checkpoint is missing non-model state " + "(optimizer/scheduler/dataloader/callback/RNG). " + "Falling back to model-only restore from %s " + "with %d role states. Optimizer/scheduler/etc. " + "will be reinitialized.", + resolved, + len(model_only_states), + ) + dcp.load( + model_only_states, + checkpoint_id=str(resolved / "dcp"), + ) _barrier() logger.info("Checkpoint loaded; resuming from step=%s", step) return step + def _save_rng_snapshot(self, checkpoint_dir: Path) -> None: + """Save per-rank RNG state after DCP save completes.""" + rank = _rank() + rng: dict[str, Any] = { + "torch_rng": torch.get_rng_state(), + "python_rng": random.getstate(), + "numpy_rng": np.random.get_state(), + } + if torch.cuda.is_available(): + rng["cuda_rng"] = torch.cuda.get_rng_state() + cuda_generator = getattr(self.method, "cuda_generator", None) + if cuda_generator is not None: + rng["gen_cuda"] = cuda_generator.get_state() + torch.save( + rng, + checkpoint_dir / f"rng_state_rank{rank}.pt", + ) + + def load_rng_snapshot( + self, + checkpoint_path: str, + ) -> None: + resolved = _resolve_resume_checkpoint( + checkpoint_path, + output_dir=self.output_dir, + ) + if resolved is None: + return + rank = _rank() + rng_path = resolved / f"rng_state_rank{rank}.pt" + if not rng_path.is_file(): + rng_path = resolved / "rng_state.pt" + if not rng_path.is_file(): + logger.warning( + "No rng_state in %s; skipping RNG snapshot restore.", + resolved, + ) + return + + rng = torch.load( + rng_path, + map_location="cpu", + weights_only=False, + ) + if "torch_rng" in rng: + torch.set_rng_state(rng["torch_rng"]) + if "python_rng" in rng: + random.setstate(rng["python_rng"]) + if "numpy_rng" in rng: + np.random.set_state(rng["numpy_rng"]) + if torch.cuda.is_available() and "cuda_rng" in rng: + torch.cuda.set_rng_state(rng["cuda_rng"]) + cuda_generator = getattr(self.method, "cuda_generator", None) + if cuda_generator is not None and "gen_cuda" in rng: + cuda_generator.set_state(rng["gen_cuda"]) + def _cleanup_old_checkpoints(self) -> None: keep_last = int(self.config.keep_last or 0) if keep_last <= 0: diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 86e9ebaf3..08bdd5018 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -321,6 +321,7 @@ def _build_training_config( data_path=str(da.get("data_path", "") or ""), train_batch_size=int(da.get("train_batch_size", 1) or 1), dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0), + apply_bot_died_filter=bool(da.get("apply_bot_died_filter", False)), training_cfg_rate=float(da.get("training_cfg_rate", 0.0) or 0.0), seed=int(da.get("seed", 0) or 0), num_height=int(da.get("num_height", 0) or 0), @@ -347,6 +348,8 @@ def _build_training_config( resume_from_checkpoint=str(ck.get("resume_from_checkpoint", "") or ""), training_state_checkpointing_steps=int(ck.get("training_state_checkpointing_steps", 0) or 0), checkpoints_total_limit=int(ck.get("checkpoints_total_limit", 0) or 0), + best_checkpoint_start_step=int(ck.get("best_checkpoint_start_step", 0) or 0), + best_checkpoint_top_k=max(1, int(ck.get("best_checkpoint_top_k", 1) or 1)), ), tracker=TrackerConfig( trackers=list(tr.get("trackers", []) or []), diff --git a/fastvideo/train/utils/dataloader.py b/fastvideo/train/utils/dataloader.py index a1b22d3ba..7e696950f 100644 --- a/fastvideo/train/utils/dataloader.py +++ b/fastvideo/train/utils/dataloader.py @@ -4,10 +4,68 @@ from __future__ import annotations from typing import Any, TYPE_CHECKING +from fastvideo.logger import init_logger + if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( DataConfig, ) +logger = init_logger(__name__) + + +def _maybe_apply_bot_died_filter( + *, + data_config: "DataConfig", + dataset: Any, +) -> None: + if not bool(getattr(data_config, "apply_bot_died_filter", False)): + return + + from fastvideo.dataset.parquet_dataset_map_style import ( + build_bot_died_excluded_indices, + ) + from fastvideo.distributed import ( + get_world_group, + get_world_rank, + ) + + world_group = get_world_group() + world_rank = int(get_world_rank()) + total = int(sum(dataset.lengths)) + bot_died_excluded: set[int] | None = None + if world_rank == 0: + bot_died_excluded = build_bot_died_excluded_indices( + data_path=str(data_config.data_path), + parquet_files=list(dataset.parquet_files), + lengths=list(dataset.lengths), + ) + + bot_died_excluded = world_group.broadcast_object(bot_died_excluded, src=0) + + excluded_count = int(len(bot_died_excluded or set())) + remaining_count = int(total - excluded_count) + + if bot_died_excluded: + valid_indices = [i for i in range(total) if i not in bot_died_excluded] + dataset.sampler.set_candidate_indices(valid_indices, epoch=0) + + local_samples = int(len(getattr(dataset.sampler, "sp_group_local_indices", []))) + local_batches = int(len(dataset.sampler)) + + if world_rank == 0: + logger.info( + "BOT_DIED_FILTER_SUMMARY enabled=true total=%d excluded=%d remaining=%d", + total, + excluded_count, + remaining_count, + ) + logger.info( + "BOT_DIED_FILTER_LOCAL rank=%d samples=%d batches=%d", + world_rank, + local_samples, + local_batches, + ) + def build_parquet_t2v_train_dataloader( data_config: DataConfig, @@ -30,4 +88,35 @@ def build_parquet_t2v_train_dataloader( text_padding_length=int(text_len), seed=int(data_config.seed or 0), )) + _maybe_apply_bot_died_filter( + data_config=data_config, + dataset=_dataset, + ) + return dataloader + + +def build_parquet_wangame_train_dataloader( + data_config: DataConfig, + *, + parquet_schema: Any, +) -> Any: + """Build a parquet dataloader for WanGame 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), + )) + _maybe_apply_bot_died_filter( + data_config=data_config, + dataset=_dataset, + ) return dataloader diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 8dd0f440b..d202545c9 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -25,6 +25,7 @@ class DataConfig: data_path: str = "" train_batch_size: int = 1 dataloader_num_workers: int = 0 + apply_bot_died_filter: bool = False training_cfg_rate: float = 0.0 seed: int = 0 num_height: int = 0 @@ -57,6 +58,8 @@ class CheckpointConfig: resume_from_checkpoint: str = "" training_state_checkpointing_steps: int = 0 checkpoints_total_limit: int = 0 + best_checkpoint_start_step: int = 0 + best_checkpoint_top_k: int = 1 @dataclass(slots=True) diff --git a/fastvideo/training/checkpointing_utils.py b/fastvideo/training/checkpointing_utils.py index bc6aeed55..6a249b729 100644 --- a/fastvideo/training/checkpointing_utils.py +++ b/fastvideo/training/checkpointing_utils.py @@ -21,10 +21,20 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful): state_dict = get_model_state_dict( self.model) # type: ignore[no-any-return] # filter out non-trainable parameters - param_requires_grad = set([ - k for k, v in dict(self.model.named_parameters()).items() - if v.requires_grad - ]) + param_requires_grad: set[str] = set() + for name, param in self.model.named_parameters(): + if not bool(param.requires_grad): + continue + param_requires_grad.add(name) + + # Activation checkpointing wraps modules with an internal attribute + # `_checkpoint_wrapped_module`, which changes the parameter names + # returned by `named_parameters()` but not the keys returned by + # `get_model_state_dict()`. + if "._checkpoint_wrapped_module." in name: + param_requires_grad.add( + name.replace("._checkpoint_wrapped_module.", ".") + ) state_dict = { k: v for k, v in state_dict.items() if k in param_requires_grad