From 3c71218987d322478ee3f2973a563a47bd6c9ae3 Mon Sep 17 00:00:00 2001 From: Wei Zhou Date: Wed, 1 Apr 2026 02:22:08 +0000 Subject: [PATCH] ckpt --- .../distill/SFWan2.1-T2V/distill_dmd.slurm | 172 +++++++++ .../validation_64.json | 0 examples/inference/basic/basic.py | 2 +- .../causal_ode_init/wan_df_init.slurm | 137 ++++++++ .../causal_ode_init/wan_tf_init.slurm | 138 ++++++++ .../causal_ode_init/wan_tf_ode_init.slurm | 138 ++++++++ .../finetune/wan_t2v_1.3B/mixkit/README.md | 1 + .../mixkit/preprocess_wan_data_t2v.sh | 25 ++ .../mixkit/preprocess_wan_data_t2v.slurm | 66 ++++ fastvideo/configs/pipelines/wan.py | 4 +- fastvideo/configs/sample/wan.py | 3 +- fastvideo/dataset/__init__.py | 3 +- fastvideo/dataset/dataloader/record_schema.py | 41 +-- fastvideo/dataset/dataloader/schema.py | 18 +- fastvideo/dataset/preprocessing_datasets.py | 46 ++- fastvideo/dataset/utils.py | 22 +- fastvideo/dataset/validation_dataset.py | 36 +- fastvideo/fastvideo_args.py | 18 + fastvideo/models/dits/causal_wanvideo.py | 160 ++++++++- fastvideo/models/dits/wanvideo.py | 178 +++++++++- fastvideo/models/loader/component_loader.py | 12 +- fastvideo/models/loader/fsdp_load.py | 19 +- .../scheduling_self_forcing_flow_match.py | 6 +- fastvideo/models/vaes/wanvae.py | 1 + fastvideo/pipelines/pipeline_batch_info.py | 1 + .../preprocess/preprocess_pipeline_base.py | 6 +- .../preprocess_pipeline_ode_trajectory.py | 20 +- .../preprocess/preprocess_pipeline_tf_ode.py | 329 ++++++++++++++++++ .../pipelines/preprocess/v1_preprocess.py | 11 +- .../pipelines/stages/causal_denoising.py | 40 ++- fastvideo/pipelines/stages/decoding.py | 31 +- fastvideo/pipelines/stages/denoising.py | 27 +- .../pipelines/stages/hy15_causal_denoising.py | 4 +- .../pipelines/stages/matrixgame_denoising.py | 4 +- fastvideo/training/ar_diffusion_pipeline.py | 321 +++++++++++++++++ fastvideo/training/distillation_pipeline.py | 325 ++++++++++++++++- fastvideo/training/ode_causal_pipeline.py | 219 +++++++++++- .../self_forcing_distillation_pipeline.py | 58 ++- fastvideo/training/training_pipeline.py | 22 +- fastvideo/training/training_utils.py | 8 + fastvideo/training/wan_training_pipeline.py | 5 +- wan_ode_preprocess/launch_slurm.sh | 21 ++ wan_ode_preprocess/syn.slurm | 99 ++++++ wan_preprocess/launch_slurm.sh | 21 ++ wan_preprocess/syn.slurm | 95 +++++ wan_text_preprocess/launch_slurm.sh | 21 ++ wan_text_preprocess/syn.slurm | 95 +++++ 47 files changed, 2866 insertions(+), 163 deletions(-) create mode 100644 examples/distill/SFWan2.1-T2V/distill_dmd.slurm rename examples/distill/{Wan2.2-TI2V-5B-Diffusers/Data-free => SFWan2.1-T2V}/validation_64.json (100%) create mode 100644 examples/training/consistency_finetune/causal_ode_init/wan_df_init.slurm create mode 100644 examples/training/consistency_finetune/causal_ode_init/wan_tf_init.slurm create mode 100644 examples/training/consistency_finetune/causal_ode_init/wan_tf_ode_init.slurm create mode 100644 examples/training/finetune/wan_t2v_1.3B/mixkit/README.md create mode 100644 examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.sh create mode 100644 examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.slurm create mode 100644 fastvideo/pipelines/preprocess/preprocess_pipeline_tf_ode.py create mode 100644 fastvideo/training/ar_diffusion_pipeline.py create mode 100644 wan_ode_preprocess/launch_slurm.sh create mode 100644 wan_ode_preprocess/syn.slurm create mode 100644 wan_preprocess/launch_slurm.sh create mode 100644 wan_preprocess/syn.slurm create mode 100644 wan_text_preprocess/launch_slurm.sh create mode 100644 wan_text_preprocess/syn.slurm diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd.slurm b/examples/distill/SFWan2.1-T2V/distill_dmd.slurm new file mode 100644 index 000000000..d7c918688 --- /dev/null +++ b/examples/distill/SFWan2.1-T2V/distill_dmd.slurm @@ -0,0 +1,172 @@ +#!/bin/bash +#SBATCH --job-name=distill_dmd_t2v_1.3b +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=distill_dmd_t2v_1.3b_output/distill_dmd_t2v_1.3b_%j.out +#SBATCH --error=distill_dmd_t2v_1.3b_output/distill_dmd_t2v_1.3b_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +export HF_HOME=/mnt/data/huggingface +export HF_DATASETS_CACHE="/tmp/hf_datasets_${SLURM_JOB_ID}_${SLURM_PROCID}" +mkdir -p "${HF_DATASETS_CACHE}" +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" + +source ~/miniconda3/bin/activate +conda activate fastvideo + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=64 + +# Model paths for Self-Forcing DMD distillation: +GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" # Teacher model +FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model + +DATA_DIR="/mnt/data/vidprom_filtered_extended_umt5_text_embed" +# DATA_DIR="/mnt/data/mixkit_wan_1.3b_processed_t2v_distributed" +DATA_DIR_2="/mnt/data/mixkit_wan_1.3b_processed_t2v_distributed" +VALIDATION_DATASET_FILE="/mnt/home/zhouw.jerry2017/FastVideo/examples/distill/SFWan2.1-T2V/validation_64.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +training_args=( + --tracker_project_name SFwan_t2v_distill_self_forcing_dmd + --wandb_run_name causal_forcing_vidprom_extended + --output_dir /mnt/data/wan_cf_df_distill_1.3b + --max_train_steps 4000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --enable_gradient_checkpointing_type "full" + --log_visualization + --simulate_generator_forward + --num_frames 81 + --num_frame_per_block 3 # Frame generation block size for self-forcing + --enable_gradient_masking + --gradient_mask_last_n_frames 21 +) + +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +model_args=( + --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --pretrained_model_name_or_path $GENERATOR_MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_MODEL_PATH +) + +dataset_args=( + --data_path "$DATA_DIR" + # --data_path_2 "$DATA_DIR_2" + --dataloader_num_workers 4 +) + +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" +) + +optimizer_args=( + --optimizer-type "adamw" + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 500 + --weight_only_checkpointing_steps 500 + --weight_decay 0.01 + --betas '0.0,0.999' + --max_grad_norm 1.0 +) + +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --flow_shift 5 + --seed 1000 + --use_ema True + --ema_decay 0.99 + --ema_start_step 200 +# --init_weights_from_safetensors "/mnt/data/wan_tf_ode_init_3333/checkpoint-1600/transformer/diffusion_pytorch_model.safetensors" +# --init_weights_from_safetensors "/mnt/data/wan_tf_init_3333/checkpoint-800/transformer/diffusion_pytorch_model.safetensors" + # --init_weights_from_safetensors "/mnt/data/Causal-Forcing-wan1.3b/causal_ode.safetensors" + # --init_weights_from_safetensors "/mnt/data/wan_ar_diffusion_3333_mixkit_shift5/checkpoint-7200/transformer/diffusion_pytorch_model.safetensors" + --init_weights_from_safetensors "/mnt/data/wan_ar_diffusion_3333_mixkit_shift5/checkpoint-4800/transformer/diffusion_pytorch_model.safetensors" + # --init_weights_from_safetensors "/mnt/data/wan_tf_ode_init_3333_mixkit_shift5/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors" + # --init_weights_from_safetensors "/mnt/data/wan_ar_df_diffusion_3333_mixkit_shift5/checkpoint-4800/transformer/diffusion_pytorch_model.safetensors" +) + +dmd_args=( + --dmd_denoising_steps '1000,750,500,250' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --dfake_gen_update_ratio 5 + --real_score_guidance_scale 3.0 + --fake_score_learning_rate 2e-6 + --fake_score_betas '0.0,0.999' + --warp_denoising_step + # --use-decoupled-dmd True + # --use-gan-loss True +) + +self_forcing_args=( + --independent_first_frame False # Whether to treat first frame independently + --same_step_across_blocks True # Whether to use same denoising step across all blocks + --last_step_only False # Whether to only use the last denoising step + --context_noise 0 # Amount of noise to add during context caching (0 = no noise) +) + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/wan_self_forcing_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ + "${self_forcing_args[@]}" diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json b/examples/distill/SFWan2.1-T2V/validation_64.json similarity index 100% rename from examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json rename to examples/distill/SFWan2.1-T2V/validation_64.json diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 3f00832b9..485740a37 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -15,7 +15,7 @@ def main(): use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=False, vae_cpu_offload=False, - text_encoder_cpu_offload=True, + text_encoder_cpu_offload=False, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, ) diff --git a/examples/training/consistency_finetune/causal_ode_init/wan_df_init.slurm b/examples/training/consistency_finetune/causal_ode_init/wan_df_init.slurm new file mode 100644 index 000000000..2c6b53aeb --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/wan_df_init.slurm @@ -0,0 +1,137 @@ +#!/bin/bash +#SBATCH --job-name=wan_tf_3333 +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=wan_tf_3333_output/wan_tf_3333_%j.out +#SBATCH --error=wan_tf_3333_output/wan_tf_3333_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" + +source ~/miniconda3/bin/activate +conda activate fastvideo + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=64 + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/data/mixkit_wan_1.3b_processed_t2v_distributed" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_tf_init" + --output_dir "/mnt/data/wan_ar_df_diffusion_3333_mixkit_shift5" + --wandb_run_name "wan_df_init_3333" + --warp_denoising_step + --dmd_denoising_steps '1000,750,500,250,0' + --log_visualization + --visualization_steps 100 + --max_train_steps 30000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --num_frames 81 + --override-transformer-cls-name "CausalWanTransformer3DModel" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --text-encoder-cpu-offload +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Optimizer arguments +optimizer_args=( + --optimizer-type "adamw" + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 800 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +validation_args=( + # --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 10 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" # not used for dmd inference + # --vae_cpu_offload True +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 6 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/ar_diffusion_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/examples/training/consistency_finetune/causal_ode_init/wan_tf_init.slurm b/examples/training/consistency_finetune/causal_ode_init/wan_tf_init.slurm new file mode 100644 index 000000000..1b774fb68 --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/wan_tf_init.slurm @@ -0,0 +1,138 @@ +#!/bin/bash +#SBATCH --job-name=wan_tf_3333 +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=wan_tf_3333_output/wan_tf_3333_%j.out +#SBATCH --error=wan_tf_3333_output/wan_tf_3333_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" + +source ~/miniconda3/bin/activate +conda activate fastvideo + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=64 + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/data/mixkit_wan_1.3b_processed_t2v_distributed" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_tf_init" + --output_dir "/mnt/data/wan_ar_diffusion_3333_mixkit_shift5" + --wandb_run_name "wan_tf_init_3333" + --warp_denoising_step + --dmd_denoising_steps '1000,750,500,250,0' + --log_visualization + --visualization_steps 100 + --max_train_steps 30000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --num_frames 81 + --override-transformer-cls-name "CausalWanTransformer3DModel" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --text-encoder-cpu-offload +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Optimizer arguments +optimizer_args=( + --optimizer-type "adamw" + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 800 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +validation_args=( + # --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 10 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" # not used for dmd inference + # --vae_cpu_offload True +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 6 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" + --use-tf True +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/ar_diffusion_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/examples/training/consistency_finetune/causal_ode_init/wan_tf_ode_init.slurm b/examples/training/consistency_finetune/causal_ode_init/wan_tf_ode_init.slurm new file mode 100644 index 000000000..3cb0244eb --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/wan_tf_ode_init.slurm @@ -0,0 +1,138 @@ +#!/bin/bash +#SBATCH --job-name=wan_tf_ode_init_3333 +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=wan_tf_ode_init_3333_output/wan_tf_ode_init_3333_%j.out +#SBATCH --error=wan_tf_ode_init_3333_output/wan_tf_ode_init_3333_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" + +source ~/miniconda3/bin/activate +conda activate fastvideo + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=64 + +MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/data/fv-tf-ode-preprocessing-6k-mixkit-wan-1.3b/" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_tf_ode_init" + --output_dir "/mnt/data/wan_tf_ode_init_3333_mixkit_shift5" + --wandb_run_name "wan_tf_ode_init_3333_mixkit_shift5" + --warp_denoising_step + --dmd_denoising_steps '1000,750,500,250' + --log_visualization + --visualization_steps 100 + --max_train_steps 30000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --num_frames 81 + --enable_gradient_checkpointing_type "full" + --init_weights_from_safetensors "/mnt/data/wan_ar_diffusion_3333_mixkit_shift5/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --text-encoder-cpu-offload +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Optimizer arguments +optimizer_args=( + --optimizer-type "adamw" + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 500 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +validation_args=( + # --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 10 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" # not used for dmd inference + # --vae_cpu_offload True +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 6 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + # --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md b/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md new file mode 100644 index 000000000..b1bd07aab --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md @@ -0,0 +1 @@ +The raw mixkit dataset is downloaded and preprocessed following instructions from https://github.com/tianweiy/CausVid/tree/master \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.sh new file mode 100644 index 000000000..03c1ca087 --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="/mnt/data/mixkit/merge.txt" +OUTPUT_DIR="/mnt/data/mixkit_wan_1.3b_processed_t2v/" + +torchrun --nproc_per_node=$GPU_NUM \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --drop_short_ratio 0.0 \ + --preprocess_task "t2v" \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.slurm b/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.slurm new file mode 100644 index 000000000..d1cac7caf --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/preprocess_wan_data_t2v.slurm @@ -0,0 +1,66 @@ +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=preprocess_mixkit_wan_1.3b_t2v/preprocess_wan_data_t2v.log +#SBATCH --error=preprocess_mixkit_wan_1.3b_t2v/preprocess_wan_data_t2v.error +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +# export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" + +source ~/miniconda3/bin/activate +conda activate fastvideo + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="/mnt/data/mixkit/merge.txt" +OUTPUT_DIR="/mnt/data/mixkit_wan_1.3b_processed_t2v/" + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $GPU_NUM \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --dataloader_num_workers 1 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --drop_short_ratio 0.0 \ + --preprocess_task "t2v" \ No newline at end of file diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 434839dcc..5d0f8e83c 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -41,7 +41,7 @@ class WanT2V480PConfig(PipelineConfig): vae_sp: bool = False # Denoising stage - flow_shift: float | None = 3.0 + flow_shift: float | None = 5.0 # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( @@ -176,6 +176,8 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig): flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 750, 500, 250]) + # dmd_denoising_steps: list[int] | None = field( + # default_factory=lambda: [1000, 680, 360, 180]) warp_denoising_step: bool = True diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 85f76a7ac..f37f10354 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -15,7 +15,8 @@ class WanT2V_1_3B_SamplingParam(SamplingParam): # Denoising stage guidance_scale: float = 3.0 - negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + # negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" num_inference_steps: int = 50 teacache_params: WanTeaCacheParams = field( diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index 9c1e98473..b5c7fb15d 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -10,7 +10,7 @@ from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, from fastvideo.dataset.validation_dataset import ValidationDataset -def getdataset(args) -> VideoCaptionMergedDataset: +def getdataset(args, start_idx: int = 0) -> VideoCaptionMergedDataset: if args.do_temporal_sample: temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x else: @@ -36,6 +36,7 @@ def getdataset(args) -> VideoCaptionMergedDataset: transform=transform, temporal_sample=temporal_sample, transform_topcrop=transform_topcrop, + start_idx=start_idx, seed=args.seed) diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py index fa8b8676a..b0bd1ea97 100644 --- a/fastvideo/dataset/dataloader/record_schema.py +++ b/fastvideo/dataset/dataloader/record_schema.py @@ -127,10 +127,7 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: def ode_text_only_record_creator( video_name: str, text_embedding: np.ndarray, caption: str, trajectory_latents: np.ndarray, - trajectory_timesteps: np.ndarray, - text_mask: np.ndarray | None = None, - text_embedding_2: np.ndarray | None = None, - text_mask_2: np.ndarray | None = None) -> dict[str, Any]: + trajectory_timesteps: np.ndarray) -> dict[str, Any]: """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only. Args: @@ -168,24 +165,24 @@ def ode_text_only_record_creator( "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype), }) - if text_embedding_2 is not None: - record.update({ - "text_embedding_2_bytes": text_embedding_2.tobytes(), - "text_embedding_2_shape": list(text_embedding_2.shape), - "text_embedding_2_dtype": str(text_embedding_2.dtype), - }) - if text_mask is not None: - record.update({ - "text_mask_bytes": text_mask.tobytes(), - "text_mask_shape": list(text_mask.shape), - "text_mask_dtype": str(text_mask.dtype), - }) - if text_mask_2 is not None: - record.update({ - "text_mask_2_bytes": text_mask_2.tobytes(), - "text_mask_2_shape": list(text_mask_2.shape), - "text_mask_2_dtype": str(text_mask_2.dtype), - }) + # if text_embedding_2 is not None: + # record.update({ + # "text_embedding_2_bytes": text_embedding_2.tobytes(), + # "text_embedding_2_shape": list(text_embedding_2.shape), + # "text_embedding_2_dtype": str(text_embedding_2.dtype), + # }) + # if text_mask is not None: + # record.update({ + # "text_mask_bytes": text_mask.tobytes(), + # "text_mask_shape": list(text_mask.shape), + # "text_mask_dtype": str(text_mask.dtype), + # }) + # if text_mask_2 is not None: + # record.update({ + # "text_mask_2_bytes": text_mask_2.tobytes(), + # "text_mask_2_shape": list(text_mask_2.shape), + # "text_mask_2_dtype": str(text_mask_2.dtype), + # }) return record diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index c0583e241..166a010dc 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -90,15 +90,15 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([ pa.field("text_embedding_shape", pa.list_(pa.int64())), # e.g., 'bfloat16' or 'float32' pa.field("text_embedding_dtype", pa.string()), - pa.field("text_embedding_2_bytes", pa.binary()), - pa.field("text_embedding_2_shape", pa.list_(pa.int64())), - pa.field("text_embedding_2_dtype", pa.string()), - pa.field("text_mask_bytes", pa.binary()), - pa.field("text_mask_shape", pa.list_(pa.int64())), - pa.field("text_mask_dtype", pa.string()), - pa.field("text_mask_2_bytes", pa.binary()), - pa.field("text_mask_2_shape", pa.list_(pa.int64())), - pa.field("text_mask_2_dtype", pa.string()), + # pa.field("text_embedding_2_bytes", pa.binary()), + # pa.field("text_embedding_2_shape", pa.list_(pa.int64())), + # pa.field("text_embedding_2_dtype", pa.string()), + # pa.field("text_mask_bytes", pa.binary()), + # pa.field("text_mask_shape", pa.list_(pa.int64())), + # pa.field("text_mask_dtype", pa.string()), + # pa.field("text_mask_2_bytes", pa.binary()), + # pa.field("text_mask_2_shape", pa.list_(pa.int64())), + # pa.field("text_mask_2_dtype", pa.string()), # --- ODE Trajectory --- pa.field("trajectory_latents_bytes", pa.binary()), pa.field("trajectory_latents_shape", pa.list_(pa.int64())), diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 563c3cc74..2e4b720af 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -470,15 +470,14 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, self.seed = seed # Initialize tokenizer - tokenizer_path = os.path.join(args.model_path, "tokenizer") - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, - cache_dir=args.cache_dir) + tokenizer_path = os.path.join(maybe_download_model(args.model_path), "tokenizer") + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) # Initialize processing stages self._init_stages(args, transform, transform_topcrop, tokenizer) # Process metadata - self.processed_batches = self._process_metadata() + self.processed_batches = self._process_metadata(start_idx=start_idx) def _init_stages(self, args, transform, transform_topcrop, tokenizer) -> None: @@ -496,13 +495,16 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, self.video_transform_stage = VideoTransformStage(transform) self.image_transform_stage = ImageTransformStage( transform, transform_topcrop) - self.text_encoding_stage = TextEncodingStage( - tokenizer=tokenizer, - text_max_length=args.text_max_length, - cfg_rate=args.training_cfg_rate, - seed=self.seed) + if tokenizer is not None: + self.text_encoding_stage = TextEncodingStage( + tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=args.training_cfg_rate, + seed=self.seed) + else: + self.text_encoding_stage = None - def _load_raw_data(self) -> list[dict]: + def _load_raw_data(self, start_idx: int = 0) -> list[dict]: """Load raw data from JSON files.""" # Read folder-annotation pairs with open(self.data_merge_path) as f: @@ -523,11 +525,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, for item in data_items: item["path"] = opj(folder, item["path"]) - return data_items + return data_items[start_idx:] - def _process_metadata(self) -> list[PreprocessBatch]: + def _process_metadata(self, start_idx: int = 0) -> list[PreprocessBatch]: """Process the raw metadata through all filtering stages.""" - raw_data = self._load_raw_data() + raw_data = self._load_raw_data(start_idx=start_idx) processed_batches = [] # Initialize counters @@ -605,17 +607,27 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Apply transformation stages batch = self.video_transform_stage.process(batch) batch = self.image_transform_stage.process(batch) - batch = self.text_encoding_stage.process(batch) + if self.text_encoding_stage is not None: + batch = self.text_encoding_stage.process(batch) + else: + raise ValueError("Text encoding stage is not initialized") # Build result dictionary result = { "pixel_values": batch.pixel_values, - "text": batch.text, - "input_ids": batch.input_ids, - "cond_mask": batch.cond_mask, + # "text": batch.text, + # "input_ids": batch.input_ids, + # "cond_mask": batch.cond_mask, "path": batch.path, } + if batch.text is not None: + result["text"] = batch.text + result["input_ids"] = batch.input_ids + result["cond_mask"] = batch.cond_mask + else: + raise ValueError("Text encoding stage is not initialized") + # Add video-specific fields if batch.is_video: result.update({"fps": batch.fps, "duration": batch.duration}) diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 142529cac..0fc07927e 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -154,14 +154,16 @@ def collate_rows_from_parquet_schema(rows, ) if rng else random.random()) < cfg_rate: data = np.zeros((512, 4096), dtype=np.float32) else: - if row[f"{tensor_name}_dtype"] == "float32": - data = np.frombuffer( - bytes_data, dtype=np.float32).reshape(shape).copy() - elif row[f"{tensor_name}_dtype"] == "int64": - data = np.frombuffer( - bytes_data, dtype=np.int64).reshape(shape).copy() - else: - raise ValueError(f"Unsupported dtype: {row[f"{tensor_name}_dtype"]}") + data = np.frombuffer( + bytes_data, dtype=np.float32).reshape(shape).copy() + # if row[f"{tensor_name}_dtype"] == "float32": + # data = np.frombuffer( + # bytes_data, dtype=np.float32).reshape(shape).copy() + # elif row[f"{tensor_name}_dtype"] == "int64": + # data = np.frombuffer( + # bytes_data, dtype=np.int64).reshape(shape).copy() + # else: + # raise ValueError(f"Unsupported dtype: {row[f"{tensor_name}_dtype"]}") tensor = torch.from_numpy(data) # if len(data.shape) == 3: # B, L, D = tensor.shape @@ -174,13 +176,15 @@ def collate_rows_from_parquet_schema(rows, tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) # Stack tensors with special handling for text embeddings - if tensor_name == 'null': + if tensor_name == 'text_embedding': # Handle text embeddings with padding padded_tensors = [] attention_masks = [] for tensor in tensor_list: if tensor.numel() > 0: + if tensor.ndim == 3: + tensor = tensor.squeeze(0) padded_tensor, mask = pad(tensor, text_padding_length) padded_tensors.append(padded_tensor) attention_masks.append(mask) diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 5ab467d75..746f0ffc9 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -2,6 +2,7 @@ # adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py import os import pathlib +import json import datasets from torch.utils.data import IterableDataset @@ -32,10 +33,7 @@ class ValidationDataset(IterableDataset): data_files=self.filename.as_posix(), split="train") elif self.filename.suffix == ".json": - data = datasets.load_dataset("json", - data_files=self.filename.as_posix(), - split="train", - field="data") + data = self._load_json_without_hf_locking(self.filename) elif self.filename.suffix == ".parquet": data = datasets.load_dataset("parquet", data_files=self.filename.as_posix(), @@ -162,3 +160,33 @@ class ValidationDataset(IterableDataset): sample = {k: v for k, v in sample.items() if v is not None} yield sample + + @staticmethod + def _load_json_without_hf_locking(filename: pathlib.Path) -> list[dict]: + """Load JSON validation files without using datasets.load_dataset lock files. + + This avoids distributed lock contention when many ranks initialize + validation simultaneously on shared filesystems. + """ + with filename.open("r", encoding="utf-8") as f: + payload = json.load(f) + + if isinstance(payload, dict): + if "data" not in payload: + raise ValueError( + f"Validation JSON {filename.as_posix()} must contain a 'data' field when top-level is an object." + ) + records = payload["data"] + elif isinstance(payload, list): + records = payload + else: + raise ValueError( + f"Validation JSON {filename.as_posix()} must be either a list of samples or an object containing a 'data' list." + ) + + if not isinstance(records, list): + raise ValueError( + f"Validation JSON records in {filename.as_posix()} must be a list." + ) + + return records diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 104b364b0..06ab45420 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -937,6 +937,9 @@ class TrainingArgs(FastVideoArgs): # simulate generator forward to match inference simulate_generator_forward: bool = False warp_denoising_step: bool = False + use_tf: bool = False + use_decoupled_dmd: bool = False # use decoupled DMD + use_gan_loss: bool = False # Self-forcing specific arguments num_frame_per_block: int = 3 @@ -1345,6 +1348,21 @@ class TrainingArgs(FastVideoArgs): help= "Whether to warp denoising step according to the scheduler time shift" ) + parser.add_argument( + "--use-tf", + action=StoreBoolean, + help="Whether to use Teacher-forcing for finetuning" + ) + parser.add_argument( + "--use-decoupled-dmd", + action=StoreBoolean, + help="Whether to use decoupled DMD" + ) + parser.add_argument( + "--use-gan-loss", + action=StoreBoolean, + help="Whether to use GAN loss" + ) # Self-forcing specific arguments parser.add_argument( diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 69b3e5430..42eaf4629 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -33,7 +33,7 @@ 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 WanT2VCrossAttention, WanTimeTextImageEmbedding -from fastvideo.platforms import AttentionBackendEnum +from fastvideo.platforms import AttentionBackendEnum, current_platform logger = init_logger(__name__) class CausalWanSelfAttention(nn.Module): @@ -88,10 +88,25 @@ class CausalWanSelfAttention(nn.Module): cache_start = current_start 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) if kv_cache is None: + is_tf = (q.shape[1] == cos.shape[0] * 2) + if is_tf: + q_chunk = torch.chunk(q, 2, dim=1) + k_chunk = torch.chunk(k, 2, dim=1) + roped_query = [] + roped_key = [] + for ii in range(2): + rq = _apply_rotary_emb(q_chunk[ii], cos, sin, is_neox_style=False).type_as(v) + rk = _apply_rotary_emb(k_chunk[ii], cos, sin, is_neox_style=False).type_as(v) + roped_query.append(rq) + roped_key.append(rk) + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + else: + 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) + # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( @@ -120,6 +135,8 @@ class CausalWanSelfAttention(nn.Module): block_mask=block_mask )[:, :, :-padded_length].transpose(2, 1) else: + 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) frame_seqlen = q.shape[1] current_end = current_start + roped_query.shape[1] sink_tokens = self.sink_size * frame_seqlen @@ -286,6 +303,8 @@ class CausalWanTransformerBlock(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, @@ -294,10 +313,13 @@ class CausalWanTransformerBlock(nn.Module): crossattn_cache=crossattn_cache) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) return hidden_states @@ -376,6 +398,94 @@ class CausalWanTransformer3DModel(BaseDiT): self.__post_init__() + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=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 + """ + # debug + DEBUG = False + if DEBUG: + num_frames = 9 + frame_seqlen = 256 + + total_length = num_frames * frame_seqlen * 2 + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + # for clean context frames, we can construct their flex attention mask based on a [start, end] interval + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end] + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_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 + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange( + start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, dtype=torch.long + ) + + # attention for clean context frames + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange( + num_frames * frame_seqlen, total_length, + step=attention_block_size, + device=device, dtype=torch.long + ) + noisy_image_end_list = noisy_image_start_list + attention_block_size + + # attention for noisy frames + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + # attend to noisy tokens within the same block + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + # attend to context tokens in previous blocks + # noise_context_starts[start:end] = 0 + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + # first design the mask for clean frames + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + # then design the mask for noisy frames + # noisy frames will attend to all clean preceeding clean frames + itself + C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (C1 | C2) + + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + 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 DEBUG: + print(block_mask) + import imageio + import numpy as np + from torch.nn.attention.flex_attention import create_mask + + mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + padded_length, KV_LEN=total_length + padded_length, device=device) + import cv2 + mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + @staticmethod def _prepare_blockwise_causal_attn_mask( device: torch.device | str, num_frames: int = 21, @@ -551,6 +661,7 @@ class CausalWanTransformer3DModel(BaseDiT): encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, start_frame: int = 0, + clean_hidden_states: torch.Tensor | None = None, **kwargs) -> torch.Tensor: orig_dtype = hidden_states.dtype @@ -562,6 +673,10 @@ class CausalWanTransformer3DModel(BaseDiT): else: encoder_hidden_states_image = None + if timestep.ndim == 1: + # [B] -> [B, F] + timestep = timestep.unsqueeze(1).repeat(1, hidden_states.shape[2]) + 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 @@ -588,13 +703,21 @@ class CausalWanTransformer3DModel(BaseDiT): # 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 - ) + if clean_hidden_states is not None: + self.block_mask = self._prepare_teacher_forcing_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, + ) + else: + 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) grid_sizes = torch.stack( @@ -607,6 +730,17 @@ class CausalWanTransformer3DModel(BaseDiT): timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) + if clean_hidden_states is not None: + clean_hidden_states = self.patch_embedding(clean_hidden_states) + clean_hidden_states = clean_hidden_states.flatten(2).transpose(1, 2) + hidden_states = torch.cat([clean_hidden_states, hidden_states], dim=1) + + aug_timestep = torch.zeros_like(timestep) + temb_clean = self.condition_embedder.time_embedder(aug_timestep.flatten(), None) + timestep_proj_clean = self.condition_embedder.time_modulation(temb_clean) + timestep_proj_clean = timestep_proj_clean.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=aug_timestep.shape) + timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1) + if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1) @@ -630,6 +764,9 @@ class CausalWanTransformer3DModel(BaseDiT): timestep_proj, freqs_cis, block_mask=self.block_mask) + if clean_hidden_states is not None: + hidden_states = hidden_states[:, hidden_states.shape[1] // 2:] + # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, @@ -644,12 +781,13 @@ class CausalWanTransformer3DModel(BaseDiT): def forward( self, *args, + clean_hidden_states: torch.Tensor | None = None, **kwargs ): if kwargs.get('kv_cache', None) is not None: return self._forward_inference(*args, **kwargs) else: - return self._forward_train(*args, **kwargs) + return self._forward_train(*args, clean_hidden_states=clean_hidden_states, **kwargs) def unpatchify(self, x, grid_sizes): diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index ef98576bd..ea696bc97 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -8,6 +8,8 @@ import numpy as np import torch import torch.nn as nn +from einops import repeat + import fastvideo.envs as envs from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) @@ -152,6 +154,32 @@ class WanSelfAttention(nn.Module): """ pass +class WanGanCrossAttention(WanSelfAttention): + + 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 (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + qq = self.norm_q(self.to_q(context)[0]).view(b, 1, -1, d) + + kk = self.norm_k(self.to_k(x)[0]).view(b, -1, n, d) + vv = self.to_v(x)[0].view(b, -1, n, d) + + # compute attention + x = self.attn(qq, kk, vv) + + # output + x = x.flatten(2) + x, _ = self.to_out(x) + return x + class WanT2VCrossAttention(WanSelfAttention): @@ -245,6 +273,98 @@ class WanI2VCrossAttention(WanSelfAttention): x, _ = self.to_out(x) return x +class GanAttentionBlock(nn.Module): + + def __init__(self, + dim=1536, + ffn_dim=8192, + num_heads=12, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + # self.norm1 = WanLayerNorm(dim, eps) + # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + # eps) + self.norm3 = FP32LayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + self.cross_attn = WanGanCrossAttention(dim, num_heads, + (-1, -1), + qk_norm, + eps) + + # modulation + # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states, + encoder_hidden_states, + # seq_lens, + # grid_sizes, + # freqs, + # context, + # context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + # e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # # self-attention + # y = self.self_attn( + # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + # freqs) + # # with amp.autocast(dtype=torch.float32): + # x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context): + token = context + self.cross_attn(self.norm3(x), context) + y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + # x = x + y * e[5] + return y + + x = cross_attn_ffn(hidden_states, encoder_hidden_states) + return x + +class RegisterTokens(nn.Module): + def __init__(self, num_registers: int, dim: int): + super().__init__() + self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02) + self.rms_norm = RMSNorm(dim, eps=1e-6) + + def forward(self): + return self.rms_norm(self.register_tokens) + + def reset_parameters(self): + nn.init.normal_(self.register_tokens, std=0.02) class WanTransformerBlock(nn.Module): @@ -630,6 +750,26 @@ class WanTransformer3DModel(CachableDiT): self.cnt = 0 self.__post_init__() + def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None: + self._cls_pred_branch = nn.Sequential( + # Input: [B, 384, 21, 60, 104] + nn.LayerNorm(atten_dim * 3 + time_embed_dim), + nn.Linear(atten_dim * 3 + time_embed_dim, 1536), + nn.SiLU(), + nn.Linear(atten_dim, num_class) + ) + self._cls_pred_branch.requires_grad_(True) + num_registers = 3 + self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim) + self._register_tokens.requires_grad_(True) + + gan_ca_blocks = [] + for _ in range(num_registers): + block = GanAttentionBlock() + gan_ca_blocks.append(block) + self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks) + self._gan_ca_blocks.requires_grad_(True) + def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], @@ -637,6 +777,8 @@ class WanTransformer3DModel(CachableDiT): encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, + classify_mode=False, + concat_time_embeddings=False, **kwargs) -> torch.Tensor: forward_batch = get_forward_context().forward_batch enable_teacache = forward_batch is not None and forward_batch.enable_teacache @@ -727,6 +869,16 @@ class WanTransformer3DModel(CachableDiT): assert encoder_hidden_states.dtype == orig_dtype + final_x = None + if classify_mode: + assert self._register_tokens is not None + assert self._gan_ca_blocks is not None + assert self._cls_pred_branch is not None + + final_x = [] + registers = repeat(self._register_tokens(), "n d -> b n d", b=hidden_states.shape[0]) + + gan_idx = 0 # 4. Transformer blocks # if caching is enabled, we might be able to skip the forward pass should_skip_forward = self.should_skip_forward_for_cached_states( @@ -740,19 +892,33 @@ class WanTransformer3DModel(CachableDiT): if enable_teacache: original_hidden_states = hidden_states.clone() - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: + for ii, 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, attention_mask) - else: - for block in self.blocks: + else: hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, attention_mask) + timestep_proj, freqs_cis, attention_mask) + + if classify_mode and ii in [13, 21, 29]: + gan_token = registers[:, gan_idx: gan_idx + 1] + final_x.append(self._gan_ca_blocks[gan_idx](hidden_states, gan_token)) + gan_idx += 1 # if teacache is enabled, we need to cache the original hidden states if enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) + + if classify_mode: + final_x = torch.cat(final_x, dim=1) + if concat_time_embeddings: + cls_input = torch.cat([final_x.float(), (10 * temb[:, None, :]).float()], dim=1).view(final_x.shape[0], -1) + final_x = self._cls_pred_branch(cls_input) + else: + cls_input = final_x.float().view(final_x.shape[0], -1) + final_x = self._cls_pred_branch(cls_input) + # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) @@ -778,6 +944,8 @@ class WanTransformer3DModel(CachableDiT): 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) + if classify_mode: + return output, torch.nan_to_num(final_x) return output def maybe_cache_states(self, hidden_states: torch.Tensor, diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index b8b1acb4a..78813ef20 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -809,6 +809,11 @@ class TransformerLoader(ComponentLoader): or cls_name == "Cosmos25Transformer3DModel" or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25" ) + + add_cls_branch = False + if hasattr(fastvideo_args, "_adding_cls_branch") and fastvideo_args._adding_cls_branch: + strict_load = False + add_cls_branch = True model = maybe_load_fsdp_model( model_cls=model_cls, init_params={"config": dit_config, "hf_config": hf_config}, @@ -828,6 +833,7 @@ class TransformerLoader(ComponentLoader): training_mode=fastvideo_args.training_mode, enable_torch_compile=fastvideo_args.enable_torch_compile, torch_compile_kwargs=fastvideo_args.torch_compile_kwargs, + add_cls_branch=add_cls_branch, ) total_params = sum(p.numel() for p in model.parameters()) @@ -869,9 +875,11 @@ class SchedulerLoader(ComponentLoader): scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) + if fastvideo_args.pipeline_config.flow_shift is not None and "shift" in config: + config["shift"] = fastvideo_args.pipeline_config.flow_shift scheduler = scheduler_cls(**config) - if fastvideo_args.pipeline_config.flow_shift is not None: - scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift) + # if fastvideo_args.pipeline_config.flow_shift is not None: + # scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift) if fastvideo_args.pipeline_config.timesteps_scale is not None: scheduler.set_timesteps_scale( fastvideo_args.pipeline_config.timesteps_scale diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 9ba60320a..cc94782af 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -75,6 +75,7 @@ def maybe_load_fsdp_model( pin_cpu_memory: bool = True, enable_torch_compile: bool = False, torch_compile_kwargs: dict[str, Any] | None = None, + add_cls_branch: bool = False, ) -> torch.nn.Module: """ Load the model with FSDP if is training, else load the model without FSDP. @@ -96,6 +97,8 @@ def maybe_load_fsdp_model( logger.info("Loading model with default_dtype: %s", default_dtype) with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) + if add_cls_branch: + model.adding_cls_branch(atten_dim=1536, num_class=1, time_embed_dim=0) # Check if we should use FSDP use_fsdp = training_mode or fsdp_inference @@ -340,7 +343,7 @@ def load_model_from_full_model_state_dict( unused_keys) # List of allowed parameter name patterns - ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"] # Can be extended as needed + ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l", "_cls_pred_branch", "_register_tokens", "_gan_ca_blocks"] # Can be extended as needed for new_param_name in unused_keys: if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): @@ -353,14 +356,16 @@ def load_model_from_full_model_state_dict( meta_sharded_param = meta_sd.get(new_param_name) if not hasattr(meta_sharded_param, "device_mesh"): # Initialize with zeros - sharded_tensor = torch.zeros_like(meta_sharded_param, - device=device, - dtype=param_dtype) + # sharded_tensor = torch.zeros_like(meta_sharded_param, + # device=device, + # dtype=param_dtype) + sharded_tensor = meta_sharded_param.to(device=device, dtype=param_dtype) else: # Initialize with zeros and distribute - full_tensor = torch.zeros_like(meta_sharded_param, - device=device, - dtype=param_dtype) + # full_tensor = torch.zeros_like(meta_sharded_param, + # device=device, + # dtype=param_dtype) + full_tensor = meta_sharded_param.to(device=device, dtype=param_dtype) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py index 78ef23fba..a92c4b3b4 100644 --- a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -69,6 +69,9 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): # use this scheduler for ODE trajectory timestep = timestep.unsqueeze(0) + model_output = model_output.permute(0, 2, 1, 3, 4) + sample = sample.permute(0, 2, 1, 3, 4) + self.sigmas = self.sigmas.to(model_output.device) self.timesteps = self.timesteps.to(model_output.device) timestep = timestep.to(model_output.device) @@ -81,7 +84,8 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): self.inverse_timesteps or self.reverse_sigmas) else 0 else: sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) - prev_sample = sample + model_output * (sigma_ - sigma) + prev_sample = sample.flatten(0, 1) + model_output.flatten(0, 1) * (sigma_ - sigma) + prev_sample = prev_sample.unflatten(0, model_output.shape[:2]).permute(0, 2, 1, 3, 4) if isinstance(prev_sample, torch.Tensor | float) and not return_dict: return (prev_sample, ) return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample) diff --git a/fastvideo/models/vaes/wanvae.py b/fastvideo/models/vaes/wanvae.py index 2d1502be8..a82cb831c 100644 --- a/fastvideo/models/vaes/wanvae.py +++ b/fastvideo/models/vaes/wanvae.py @@ -1219,6 +1219,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): out = torch.clamp(out, min=-1.0, max=1.0) self.clear_cache() else: + raise NotImplementedError("Feature cache is not supported for WanVAE") out = ParallelTiledVAE.decode(self, z) return out diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 8375e357d..37592ea6c 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -123,6 +123,7 @@ class ForwardBatch: # Latent tensors latents: torch.Tensor | None = None + clean_latents: torch.Tensor | None = None raw_latent_shape: tuple[int, ...] | None = None noise_pred: torch.Tensor | None = None image_latent: torch.Tensor | None = None diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 083277bc2..277e09039 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -283,8 +283,10 @@ class BasePreprocessPipeline(ComposedPipelineBase): table = pq.read_table(os.path.join(root, file)) start_idx += table.num_rows + logger.info("Start index: %s", start_idx) + # Loading dataset - train_dataset = getdataset(args) + train_dataset = getdataset(args, start_idx=start_idx) train_dataloader = DataLoader( train_dataset, @@ -301,6 +303,7 @@ class BasePreprocessPipeline(ComposedPipelineBase): for batch_idx, data in enumerate(pbar): if data is None: + logger.warning("Batch is None") continue with torch.inference_mode(): @@ -313,6 +316,7 @@ class BasePreprocessPipeline(ComposedPipelineBase): num_processed_samples += len(valid_indices) if not valid_indices: + logger.warning("No valid indices") continue # Create new batch with only valid samples diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index 171cace77..36bc45f26 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -24,8 +24,7 @@ from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, records_to_table) from fastvideo.dataset.dataloader.record_schema import ( ode_text_only_record_creator) -from fastvideo.dataset.dataloader.schema import ( - pyarrow_schema_ode_trajectory_text_only) +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( @@ -180,13 +179,14 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): batch.height = args.max_height batch.width = args.max_width batch.fps = args.train_fps - batch.guidance_scale = 6.0 + batch.guidance_scale = 3.0 batch.do_classifier_free_guidance = True result_batch = self.input_validation_stage( batch, fastvideo_args) - result_batch = self.timestep_preparation_stage( - batch, fastvideo_args) + # result_batch = self.timestep_preparation_stage( + # result_batch, fastvideo_args) + result_batch.timesteps = self.get_module("scheduler").timesteps result_batch = self.latent_preparation_stage( result_batch, fastvideo_args) result_batch = self.denoising_stage(result_batch, @@ -209,10 +209,12 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): if batch.return_trajectory_decoded: for i, decoded_frames in enumerate(trajectory_decoded): for j, decoded_frame in enumerate(decoded_frames): - save_decoded_latents_as_video( - decoded_frame, - f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", - args.train_fps) + if j in [49, 50]: + local_rank = int(os.getenv("RANK", 0)) + save_decoded_latents_as_video( + decoded_frame, + f"decoded_videos/trajectory_decoded_{local_rank}_{j}.mp4", + args.train_fps) # Prepare batch data for Parquet dataset batch_data: list[dict[str, Any]] = [] diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_tf_ode.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_tf_ode.py new file mode 100644 index 000000000..ace83ef51 --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_tf_ode.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +ODE Trajectory Data Preprocessing pipeline implementation. + +This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline +using the modular pipeline architecture. + +Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772 +""" + +import os +from collections.abc import Iterator +from typing import Any +from unittest import result + +import pyarrow as pa +import torch +from torch.utils.data import DataLoader +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm import tqdm + +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset import build_parquet_map_style_dataloader +from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, + records_to_table) +from fastvideo.dataset.dataloader.record_schema import ( + ode_text_only_record_creator) +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only, pyarrow_schema_t2v +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( + SelfForcingFlowMatchScheduler) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.preprocess.preprocess_pipeline_base import ( + BasePreprocessPipeline) +from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) +from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict +from fastvideo.distributed import get_local_torch_device + +logger = init_logger(__name__) + + +class PreprocessPipeline_TF_ODE(BasePreprocessPipeline): + """ODE Trajectory preprocessing pipeline implementation.""" + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + preprocess_dataloader: StatefulDataLoader + preprocess_loader_iter: Iterator[dict[str, Any]] + pbar: Any + num_processed_samples: int + + def get_pyarrow_schema(self) -> pa.Schema: + """Return the PyArrow schema for ODE Trajectory pipeline.""" + return pyarrow_schema_ode_trajectory_text_only + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + assert fastvideo_args.pipeline_config.flow_shift == 5 + self.modules["scheduler"] = SelfForcingFlowMatchScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True) + self.modules["scheduler"].set_timesteps(num_inference_steps=48, + denoising_strength=1.0) + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline=self, + )) + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def _process_batch(self, batch): + # Required fields from parquet (ODE trajectory schema) + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + latent = batch['vae_latent'].float() + infos = batch['info_list'][0] + + if (hasattr(self.modules["vae"], "shift_factor") + and self.modules["vae"].shift_factor is not None): + if isinstance(self.modules["vae"].shift_factor, torch.Tensor): + latent -= self.modules["vae"].shift_factor.to( + latent.device, latent.dtype) + else: + latent -= self.modules["vae"].shift_factor + else: + raise ValueError("vae.shift_factor is not set") + + if isinstance(self.modules["vae"].scaling_factor, torch.Tensor): + latent = latent * self.modules["vae"].scaling_factor.to( + latent.device, latent.dtype) + else: + latent = latent * self.modules["vae"].scaling_factor + + # Move to device + device = get_local_torch_device() + encoder_hidden_states = encoder_hidden_states.to( + device, dtype=torch.bfloat16) + encoder_attention_mask = encoder_attention_mask.to( + device, dtype=torch.bfloat16) + + return encoder_hidden_states, encoder_attention_mask, latent.to(device, dtype=torch.bfloat16), infos + + def preprocess_text_and_trajectory(self, fastvideo_args: FastVideoArgs, + args): + """Preprocess text-only data and generate trajectory information.""" + + for batch_idx, data in enumerate(self.pbar): + if data is None: + continue + + with torch.inference_mode(): + sampling_params = SamplingParam.from_pretrained(args.model_path) + + # encode negative prompt for trajectory collection + if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None: + negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text( + sampling_params.negative_prompt, + fastvideo_args, + encoder_index=[0], + return_attention_mask=True, + ) + negative_prompt_embed = negative_prompt_embeds_list[0] + negative_prompt_attention_mask = negative_prompt_masks_list[ + 0] + else: + negative_prompt_embed = None + negative_prompt_attention_mask = None + + trajectory_latents = [] + trajectory_timesteps = [] + trajectory_decoded = [] + + prompt_embeds, prompt_attention_masks, clean_latents, infos = self._process_batch(data) + assert prompt_embeds.shape[0] == 1, "Only one prompt embeds are supported" + self.num_processed_samples += prompt_embeds.shape[0] + + # Collect the trajectory data (text-to-video generation) + batch = ForwardBatch(**shallow_asdict(sampling_params), ) + batch.prompt_embeds = [prompt_embeds] + batch.prompt_attention_mask = [prompt_attention_masks] + batch.negative_prompt_embeds = [negative_prompt_embed] + batch.negative_attention_mask = [ + negative_prompt_attention_mask + ] + batch.clean_latents = clean_latents + batch.num_inference_steps = 48 + batch.return_trajectory_latents = True + # Enabling this will save the decoded trajectory videos. + # Used for debugging. + batch.return_trajectory_decoded = False + batch.height = args.max_height + batch.width = args.max_width + batch.fps = args.train_fps + batch.guidance_scale = 6.0 + batch.do_classifier_free_guidance = True + + result_batch = self.input_validation_stage( + batch, fastvideo_args) + # result_batch = self.timestep_preparation_stage( + # result_batch, fastvideo_args) + result_batch.timesteps = self.get_module("scheduler").timesteps + result_batch = self.latent_preparation_stage( + result_batch, fastvideo_args) + result_batch = self.denoising_stage(result_batch, + fastvideo_args) + if batch.return_trajectory_decoded: + result_batch = self.decoding_stage(result_batch, + fastvideo_args) + trajectory_decoded.append(result_batch.trajectory_decoded) + + result_batch.trajectory_latents = result_batch.trajectory_latents[:, [0, 12, 24, 36, -2, -1]] + result_batch.trajectory_timesteps = result_batch.trajectory_timesteps[[0, 12, 24, 36, -2, -1]] + trajectory_latents.append( + result_batch.trajectory_latents.cpu()) + trajectory_timesteps.append( + result_batch.trajectory_timesteps.cpu()) + + # Prepare extra features for text-only processing + extra_features = { + "trajectory_latents": trajectory_latents, + "trajectory_timesteps": trajectory_timesteps + } + + if batch.return_trajectory_decoded: + for i, decoded_frames in enumerate(trajectory_decoded): + for j, decoded_frame in enumerate(decoded_frames): + if j in [50, 51]: + local_rank = int(os.getenv("RANK", 0)) + save_decoded_latents_as_video( + decoded_frame, + f"decoded_videos/trajectory_decoded_{local_rank}_{j}.mp4", + args.train_fps) + + # Prepare batch data for Parquet dataset + batch_data: list[dict[str, Any]] = [] + + # Add progress bar for saving outputs + save_pbar = tqdm(enumerate([infos["file_name"]]), + desc="Saving outputs", + unit="item", + leave=False) + + for idx, video_path in save_pbar: + video_name = os.path.basename(video_path).split(".")[0] + + # Convert tensors to numpy arrays + text_embedding = prompt_embeds.float().cpu().numpy() + + # Get extra features for this sample + sample_extra_features = {} + if extra_features: + for key, value in extra_features.items(): + if isinstance(value, torch.Tensor): + sample_extra_features[key] = value[idx].cpu( + ).numpy() + else: + assert isinstance(value, list) + if isinstance(value[idx], torch.Tensor): + sample_extra_features[key] = value[idx].cpu( + ).float().numpy() + else: + sample_extra_features[key] = value[idx] + + # Create record for Parquet dataset (text-only ODE schema) + record: dict[str, Any] = ode_text_only_record_creator( + video_name=video_name, + text_embedding=text_embedding, + caption=infos["caption"], + trajectory_latents=sample_extra_features[ + "trajectory_latents"], + trajectory_timesteps=sample_extra_features[ + "trajectory_timesteps"], + ) + batch_data.append(record) + + if batch_data: + write_pbar = tqdm(total=1, + desc="Writing to Parquet dataset", + unit="batch") + table = records_to_table(batch_data, + self.get_pyarrow_schema()) + write_pbar.update(1) + write_pbar.close() + + if not hasattr(self, 'dataset_writer'): + self.dataset_writer = ParquetDatasetWriter( + out_dir=self.combined_parquet_dir, + samples_per_file=args.samples_per_file, + ) + self.dataset_writer.append_table(table) + + logger.info("Collected batch with %s samples", len(table)) + + if self.num_processed_samples >= args.flush_frequency: + written = self.dataset_writer.flush() + logger.info("Flushed %s samples to parquet", written) + self.num_processed_samples = 0 + + # Final flush for any remaining samples + if hasattr(self, 'dataset_writer'): + written = self.dataset_writer.flush(write_remainder=True) + if written: + logger.info("Final flush wrote %s samples", written) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args): + if not self.post_init_called: + self.post_init() + + self.local_rank = int(os.getenv("RANK", 0)) + os.makedirs(args.output_dir, exist_ok=True) + # Create directory for combined data + self.combined_parquet_dir = os.path.join(args.output_dir, + "combined_parquet_dataset") + os.makedirs(self.combined_parquet_dir, exist_ok=True) + + # Loading dataset + self.train_dataset, self.preprocess_dataloader = build_parquet_map_style_dataloader( + args.data_merge_path, + args.preprocess_video_batch_size, + # parquet_schema=pyarrow_schema_ode_trajectory_text_only, + parquet_schema=pyarrow_schema_t2v, + num_data_workers=args.dataloader_num_workers, + cfg_rate=0.0, + drop_last=True, + text_padding_length=fastvideo_args.pipeline_config. + text_encoder_configs[0].arch_config. + text_len, # type: ignore[attr-defined] + seed=args.seed) + + self.preprocess_loader_iter = iter(self.preprocess_dataloader) + + self.num_processed_samples = 0 + # Add progress bar for video preprocessing + self.pbar = tqdm(self.preprocess_loader_iter, + desc="Processing videos", + unit="batch", + disable=self.local_rank != 0) + + # Initialize class variables for data sharing + self.video_data: dict[str, Any] = {} # Store video metadata and paths + self.latent_data: dict[str, Any] = {} # Store latent tensors + self.preprocess_text_and_trajectory(fastvideo_args, args) + + +EntryClass = PreprocessPipeline_TF_ODE diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index d2f16a407..10ba58910 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -12,6 +12,8 @@ from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import ( PreprocessPipeline_I2V) from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import ( PreprocessPipeline_ODE_Trajectory) +from fastvideo.pipelines.preprocess.preprocess_pipeline_tf_ode import ( + PreprocessPipeline_TF_ODE) from fastvideo.pipelines.preprocess.hy15.hy15_ode_preprocess_pipeline import ( Hy15ODEPreprocessPipeline) from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import ( @@ -53,6 +55,7 @@ def main(args) -> None: vae_cpu_offload=False, text_encoder_cpu_offload=False, pipeline_config=pipeline_config, + init_weights_from_safetensors=args.init_weights_from_safetensors, ) if args.preprocess_task == "t2v": PreprocessPipeline = PreprocessPipeline_T2V @@ -64,6 +67,9 @@ def main(args) -> None: assert args.flow_shift is not None, "flow_shift is required for ode_trajectory" fastvideo_args.pipeline_config.flow_shift = args.flow_shift PreprocessPipeline = PreprocessPipeline_ODE_Trajectory + elif args.preprocess_task == "tf_ode": + fastvideo_args.pipeline_config.flow_shift = args.flow_shift + PreprocessPipeline = PreprocessPipeline_TF_ODE elif args.preprocess_task == "hy15_ode_trajectory": assert args.flow_shift is not None, "flow_shift is required for hy15_ode_trajectory" fastvideo_args.pipeline_config.flow_shift = args.flow_shift @@ -73,7 +79,7 @@ def main(args) -> None: else: raise ValueError( f"Invalid preprocess task: {args.preprocess_task}. " - f"Valid options: t2v, i2v, ode_trajectory, hy15_ode_trajectory, text_only, matrixgame") + f"Valid options: t2v, i2v, ode_trajectory, hy15_ode_trajectory, text_only, matrixgame, tf_ode") logger.info("Preprocess task: %s using %s", args.preprocess_task, PreprocessPipeline.__name__) @@ -121,7 +127,7 @@ if __name__ == "__main__": "--preprocess_task", type=str, default="t2v", - choices=["t2v", "i2v", "text_only", "ode_trajectory", "hy15_ode_trajectory", "matrixgame"], + choices=["t2v", "i2v", "text_only", "ode_trajectory", "hy15_ode_trajectory", "matrixgame", "tf_ode"], help="Type of preprocessing task to run") parser.add_argument("--train_fps", type=int, default=30) parser.add_argument("--use_image_num", type=int, default=0) @@ -144,6 +150,7 @@ if __name__ == "__main__": help= "The output directory where the model predictions and checkpoints will be written.", ) + parser.add_argument("--init_weights_from_safetensors", type=str, default=None) args = parser.parse_args() main(args) diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 1f5341506..52b0a421d 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -14,7 +14,7 @@ try: from fastvideo.attention.backends.sliding_tile_attn import ( SlidingTileAttentionBackend) st_attn_available = True -except ImportError: +except: st_attn_available = False SlidingTileAttentionBackend = None # type: ignore @@ -22,7 +22,7 @@ try: from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionBackend) vsa_available = True -except ImportError: +except: vsa_available = False VideoSparseAttentionBackend = None # type: ignore @@ -112,6 +112,13 @@ class CausalDMDDenosingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 + if batch.do_classifier_free_guidance: + raise ValueError("Classifier free guidance is not supported for causal DMD denoising") + neg_prompt_embeds = batch.negative_prompt_embeds + assert neg_prompt_embeds is not None + assert not torch.isnan( + neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan" + # Initialize or reset caches kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, @@ -218,6 +225,11 @@ class CausalDMDDenosingStage(DenoisingStage): block_sizes.pop(0) latents[:, :, :1, :, :] = first_frame_latent + # self.scheduler.set_timesteps(num_inference_steps=48, + # denoising_strength=1.0) + # timesteps = self.scheduler.timesteps.to(get_local_torch_device()) + logger.info(f"timesteps: {timesteps}") + # DMD loop in causal blocks with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: @@ -296,6 +308,30 @@ class CausalDMDDenosingStage(DenoisingStage): **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) + if batch.do_classifier_free_guidance: + batch.is_cfg_negative = True + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + pred_noise_uncond_btchw = current_model( + latent_model_input, + neg_prompt_embeds, + t_expanded_noise, + kv_cache=_get_kv_cache(t_cur), + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ).permute(0, 2, 1, 3, 4) + + pred_noise_btchw = pred_noise_uncond_btchw + batch.guidance_scale * ( + pred_noise_btchw - pred_noise_uncond_btchw) + # Convert pred noise to pred video with FM Euler scheduler utilities if boundary_timestep is not None and t_cur >= boundary_timestep: pred_video_btchw = pred_noise_to_x_bound( diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 7450ff226..f405c800d 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -54,22 +54,23 @@ class DecodingStage(PipelineStage): """Convert normalized latents into the VAE's expected latent space.""" # Some VAEs handle latent (de)normalization internally. if bool(getattr(self.vae, "handles_latent_denorm", False)): + raise NotImplementedError("handles_latent_denorm is not supported for WanVAE") return latents cfg = getattr(self.vae, "config", None) # MatrixGame-style: z = z * std + mean - if (cfg is not None and hasattr(cfg, "latents_mean") - and hasattr(cfg, "latents_std")): - latents_mean = torch.tensor(cfg.latents_mean, - device=latents.device, - dtype=latents.dtype).view( - 1, -1, 1, 1, 1) - latents_std = torch.tensor(cfg.latents_std, - device=latents.device, - dtype=latents.dtype).view( - 1, -1, 1, 1, 1) - return latents * latents_std + latents_mean + # if (cfg is not None and hasattr(cfg, "latents_mean") + # and hasattr(cfg, "latents_std")): + # latents_mean = torch.tensor(cfg.latents_mean, + # device=latents.device, + # dtype=latents.dtype).view( + # 1, -1, 1, 1, 1) + # latents_std = torch.tensor(cfg.latents_std, + # device=latents.device, + # dtype=latents.dtype).view( + # 1, -1, 1, 1, 1) + # return latents * latents_std + latents_mean # Diffusers-style: scaling_factor (+ optional shift_factor) if hasattr(self.vae, "scaling_factor"): @@ -86,6 +87,10 @@ class DecodingStage(PipelineStage): latents.device, latents.dtype) else: latents = latents + self.vae.shift_factor + else: + raise NotImplementedError("shift_factor is not supported for WanVAE") + else: + raise NotImplementedError("handles_latent_denorm is not supported for WanVAE") return latents @@ -121,8 +126,8 @@ class DecodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() + # if fastvideo_args.pipeline_config.vae_tiling: + # self.vae.enable_tiling() # if fastvideo_args.vae_sp: # self.vae.enable_parallel() if not vae_autocast_enabled: diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 4d40dec4d..558a73c6c 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -32,21 +32,21 @@ try: from fastvideo.attention.backends.sliding_tile_attn import ( SlidingTileAttentionBackend) st_attn_available = True -except ImportError: +except: st_attn_available = False try: from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend from fastvideo.utils import is_vmoba_available vmoba_attn_available = is_vmoba_available() -except ImportError: +except: vmoba_attn_available = False try: from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionBackend) vsa_available = True -except ImportError: +except: vsa_available = False logger = init_logger(__name__) @@ -176,6 +176,13 @@ class DenoisingStage(PipelineStage): }, ) + tf_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "clean_hidden_states": batch.clean_latents, + }, + ) + # Prepare STA parameters if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: self.prepare_sta_param(batch, fastvideo_args) @@ -320,9 +327,9 @@ class DenoisingStage(PipelineStage): temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep ]) timestep = temp_ts.unsqueeze(0) - t_expand = timestep.repeat(latent_model_input.shape[0], 1) + t_expand = timestep.repeat(latent_model_input.shape[0], 1).to(get_local_torch_device()) else: - t_expand = t.repeat(latent_model_input.shape[0]) + t_expand = t.repeat(latent_model_input.shape[0]).to(get_local_torch_device()) latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) @@ -415,8 +422,11 @@ class DenoisingStage(PipelineStage): **image_kwargs, **pos_cond_kwargs, **action_kwargs, + **tf_kwargs, ) + assert batch.do_classifier_free_guidance, "do_classifier_free_guidance must be True" + assert current_guidance_scale == 6.0, "guidance_scale must be 6.0" if batch.do_classifier_free_guidance: batch.is_cfg_negative = True with set_forward_context( @@ -432,6 +442,7 @@ class DenoisingStage(PipelineStage): **image_kwargs, **neg_cond_kwargs, **action_kwargs, + **tf_kwargs, ) noise_pred_text = noise_pred @@ -472,9 +483,15 @@ class DenoisingStage(PipelineStage): trajectory_tensor: torch.Tensor | None = None if trajectory_latents: trajectory_timesteps.append(torch.zeros_like(t)) + if batch.clean_latents is not None: + trajectory_latents.append(batch.clean_latents) + trajectory_timesteps.append(torch.zeros_like(t)) + trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) + # assert trajectory_tensor.shape[1] == 1, "trajectory_tensor should have 1 frame" + # assert trajectory_timesteps_tensor.shape[0] == 1, "trajectory_timesteps_tensor should have 1 frame" else: trajectory_tensor = None trajectory_timesteps_tensor = None diff --git a/fastvideo/pipelines/stages/hy15_causal_denoising.py b/fastvideo/pipelines/stages/hy15_causal_denoising.py index f4a3d0134..0dfce3f91 100644 --- a/fastvideo/pipelines/stages/hy15_causal_denoising.py +++ b/fastvideo/pipelines/stages/hy15_causal_denoising.py @@ -14,7 +14,7 @@ try: from fastvideo.attention.backends.sliding_tile_attn import ( SlidingTileAttentionBackend) st_attn_available = True -except ImportError: +except: st_attn_available = False SlidingTileAttentionBackend = None # type: ignore @@ -22,7 +22,7 @@ try: from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionBackend) vsa_available = True -except ImportError: +except: vsa_available = False VideoSparseAttentionBackend = None # type: ignore diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index 3b442b068..79f26a9f4 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -19,7 +19,7 @@ try: from fastvideo.attention.backends.sliding_tile_attn import ( SlidingTileAttentionBackend) st_attn_available = True -except ImportError: +except: st_attn_available = False SlidingTileAttentionBackend = None # type: ignore @@ -27,7 +27,7 @@ try: from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionBackend) vsa_available = True -except ImportError: +except: vsa_available = False VideoSparseAttentionBackend = None # type: ignore diff --git a/fastvideo/training/ar_diffusion_pipeline.py b/fastvideo/training/ar_diffusion_pipeline.py new file mode 100644 index 000000000..bbb4f2a5c --- /dev/null +++ b/fastvideo/training/ar_diffusion_pipeline.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any, cast + +import numpy as np +import torch +import torch.nn.functional as F + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( + SelfForcingFlowMatchScheduler) +from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import ( + WanCausalDMDPipeline) +from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.pipelines.stages.decoding import DecodingStage +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.training.training_utils import ( + clip_grad_norm_while_handling_failing_dtensor_cases) + +logger = init_logger(__name__) + + +class ARDiffusionTrainingPipeline(TrainingPipeline): + """ + Training pipeline for AR diffusion using precomputed denoising trajectories. + + Supervision: predict the next latent in the stored trajectory by + - feeding current latent at timestep t into the transformer to predict noise + """ + + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + # Match the preprocess/generation scheduler for consistent stepping + assert fastvideo_args.pipeline_config.flow_shift == 5.0, "flow_shift must be 5.0" + self.modules["scheduler"] = SelfForcingFlowMatchScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True) + self.modules["scheduler"].set_timesteps(num_inference_steps=1000, + training=True) + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_t2v + + def initialize_training_pipeline(self, training_args: TrainingArgs): + super().initialize_training_pipeline(training_args) + + self.noise_scheduler = self.get_module("scheduler") + self.vae = self.get_module("vae") + self.vae.requires_grad_(False) + + self.timestep_shift = self.training_args.pipeline_config.flow_shift + assert self.timestep_shift == 5.0, "timestep_shift must be 5.0" + self.noise_scheduler = SelfForcingFlowMatchScheduler( + shift=self.timestep_shift, sigma_min=0.0, extra_one_step=True) + self.noise_scheduler.set_timesteps(num_inference_steps=1000, + training=True) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.vae)) + + self.num_train_timestep = self.noise_scheduler.num_train_timesteps + self.manual_idx = 0 + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + args_copy.inference_mode = True + # Warm start validation with current transformer + self.validation_pipeline = WanCausalDMDPipeline.from_pretrained( + training_args.model_path, + args=args_copy, # type: ignore + inference_mode=True, + loaded_modules={ + "transformer": self.get_module("transformer"), + }, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + pin_cpu_memory=training_args.pin_cpu_memory, + dit_cpu_offload=True) + + def _get_next_batch( + self, + training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + self.train_loader_iter = iter(self.train_dataloader) + batch = next(self.train_loader_iter) + + # Required fields from parquet (ODE trajectory schema) + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + latent = batch['vae_latent'].float() + infos = batch['info_list'] + + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latent -= self.vae.shift_factor.to( + latent.device, latent.dtype) + else: + latent -= self.vae.shift_factor + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent = latent * self.vae.scaling_factor.to( + latent.device, latent.dtype) + else: + latent = latent * self.vae.scaling_factor + # [B, C, T, H, W] -> [B, T, C, H, W] + latent = latent.permute(0, 2, 1, 3, 4) + + # Move to device + device = get_local_torch_device() + training_batch.encoder_hidden_states = encoder_hidden_states.to( + device, dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + device, dtype=torch.bfloat16) + training_batch.infos = infos + + return training_batch, latent[:, :self.training_args. + num_latent_t].to( + device, + dtype=torch.bfloat16 + ) + + def _get_timestep(self, + min_timestep: int, + max_timestep: int, + batch_size: int, + num_frame: int, + num_frame_per_block: int, + uniform_timestep: bool = False) -> torch.Tensor: + if uniform_timestep: + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, 1], + device=self.device, + dtype=torch.long).repeat(1, num_frame) + return timestep + else: + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, num_frame], + device=self.device, + dtype=torch.long) + # logger.info(f"individual timestep: {timestep}") + # make the noise level the same within every block + timestep = timestep.reshape(timestep.shape[0], -1, + num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + timestep = timestep.reshape(timestep.shape[0], -1) + return timestep + + def _step_predict_next_latent( + self, latent: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + torch.Tensor]]: + latent_vis_dict: dict[str, torch.Tensor] = {} + device = get_local_torch_device() + B, num_frames, num_channels, height, width = latent.shape + + indexes = self._get_timestep( # [B, num_frames] + 0, + 1000, + B, + num_frames, + 3, + uniform_timestep=False) + timestep = self.noise_scheduler.timesteps[indexes.cpu()].to(device) + + noisy_input = self.noise_scheduler.add_noise( + latent.flatten(0, 1), + torch.randn_like(latent.flatten(0, 1)), + timestep.flatten(0, 1)).unflatten(0, (B, num_frames)) + + clean_input = latent.clone() + + # Prepare inputs for transformer + latent_vis_dict["noisy_input"] = noisy_input.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + latent_vis_dict["x0"] = latent.permute(0, 2, 1, 3, + 4).detach().clone().cpu() + + logger.info("timestep: %s", timestep) + input_kwargs = { + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep.to(device, dtype=torch.bfloat16), + "return_dict": False, + } + if self.training_args.use_tf: + raise NotImplementedError("Teacher-forcing is not implemented yet") + input_kwargs["clean_hidden_states"] = clean_input.permute(0, 2, 1, 3, 4) + # Predict noise and step the scheduler to obtain next latent + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + + from fastvideo.models.utils import pred_noise_to_pred_video + pred_video = pred_noise_to_pred_video( + pred_noise=noise_pred.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + scheduler=self.modules["scheduler"]).unflatten( + 0, noise_pred.shape[:2]) + latent_vis_dict["pred_video"] = pred_video.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + + return pred_video, latent, timestep, latent_vis_dict + + def train_one_step(self, training_batch): # type: ignore[override] + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + args = cast(TrainingArgs, self.training_args) + + # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent + + for _ in range(args.gradient_accumulation_steps): + training_batch, latent = self._get_next_batch( + training_batch) + text_embeds = training_batch.encoder_hidden_states + text_attention_mask = training_batch.encoder_attention_mask + assert latent.shape[0] == 1 + + # Forward to predict next latent by stepping scheduler with predicted noise + noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent( + latent, text_embeds) + + training_batch.latent_vis_dict.update(latent_vis_dict) + + mask = t != 0 + + # Compute loss + loss = F.mse_loss(noise_pred[mask], + target_latent[mask], + reduction="mean") + loss = loss / args.gradient_accumulation_steps + + with set_forward_context(current_timestep=t, + attn_metadata=None, + forward_batch=None): + loss.backward() + avg_loss = loss.detach().clone() + training_batch.total_loss += avg_loss.item() + + # Clip grad and step optimizers + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + [p for p in self.transformer.parameters() if p.requires_grad], + args.max_grad_norm if args.max_grad_norm is not None else 0.0) + + self.optimizer.step() + self.lr_scheduler.step() + + if grad_norm is None: + grad_value = 0.0 + else: + try: + if isinstance(grad_norm, torch.Tensor): + grad_value = float(grad_norm.detach().float().item()) + else: + grad_value = float(grad_norm) + except Exception: + grad_value = 0.0 + training_batch.grad_norm = grad_value + return training_batch + + def visualize_intermediate_latents(self, training_batch: TrainingBatch, + training_args: TrainingArgs, step: int): + tracker_loss_dict: dict[str, Any] = {} + latents_vis_dict = training_batch.latent_vis_dict + latent_log_keys = ['noisy_input', 'x0', 'pred_video'] + for latent_key in latent_log_keys: + assert latent_key in latents_vis_dict and latents_vis_dict[ + latent_key] is not None + latent = latents_vis_dict[latent_key] + pixel_latent = self.decoding_stage.decode( + latent, training_args) + + video = pixel_latent.cpu().float() + video = video.permute(0, 2, 1, 3, 4) + video = (video * 255).numpy().astype(np.uint8) + video_artifact = self.tracker.video( + video, fps=16, format="mp4") # change to 16 for Wan2.1 + if video_artifact is not None: + tracker_loss_dict[latent_key] = video_artifact + # Clean up references + del video, pixel_latent, latent + + if self.global_rank == 0 and tracker_loss_dict: + self.tracker.log_artifacts(tracker_loss_dict, step) + + +def main(args) -> None: + logger.info("Starting AR diffusion training pipeline...") + pipeline = ARDiffusionTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("AR diffusion training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 9ee0b7509..35422f43c 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -37,7 +37,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint, - shift_timestep) + shift_timestep, unshift_timestep) from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory) from fastvideo.optim.muon import get_muon_optimizer @@ -109,6 +109,8 @@ class DistillationPipeline(TrainingPipeline): if training_args.fake_score_model_path: logger.info("Loading fake score transformer from: %s", training_args.fake_score_model_path) + if training_args.use_gan_loss: + training_args._adding_cls_branch = True self.fake_score_transformer = self.load_module_from_path( training_args.fake_score_model_path, "transformer", training_args) @@ -702,8 +704,155 @@ class DistillationPipeline(TrainingPipeline): "generator_timestep"] = target_timestep.float().detach() return pred_video + def _dmd_decoupled_forward(self, generator_pred_video: torch.Tensor, + training_batch: TrainingBatch, + exit_timestep: float) -> torch.Tensor: + """Compute DMD (Diffusion Model Distillation) loss.""" + original_latent = generator_pred_video + with torch.no_grad(): + + def get_dm_training_batch(): + dm_timestep = torch.randint(0, + self.num_train_timestep, [1], + device=self.device, + dtype=torch.long) + + dm_timestep = shift_timestep( + dm_timestep, + self.timestep_shift, # type: ignore + self.num_train_timestep) + + dm_timestep = dm_timestep.clamp(self.min_timestep, self.max_timestep) + + dm_noise = torch.randn(self.video_latent_shape, + device=self.device, + dtype=generator_pred_video.dtype) + + dm_noisy_latent = self.noise_scheduler.add_noise( + generator_pred_video.flatten(0, 1), dm_noise.flatten(0, 1), + dm_timestep).detach().unflatten(0, + (1, generator_pred_video.shape[1])) + + return dm_timestep, dm_noisy_latent, dm_noise + + def get_ca_training_batch(noise): + from fastvideo.training.training_utils import unshift_timestep + exit_timestep_int = int(unshift_timestep(exit_timestep, self.timestep_shift, self.num_train_timestep)) + ca_timestep = torch.randint(0, + exit_timestep_int, [1], + device=self.device, + dtype=torch.long) + + ca_timestep = shift_timestep( + ca_timestep, + self.timestep_shift, # type: ignore + self.num_train_timestep) + + ca_timestep = ca_timestep.clamp(self.min_timestep, self.max_timestep) + + ca_noisy_latent = self.noise_scheduler.add_noise( + generator_pred_video.flatten(0, 1), noise.flatten(0, 1), + ca_timestep).detach().unflatten(0, + (1, generator_pred_video.shape[1])) + + return ca_timestep, ca_noisy_latent + + dm_timestep, dm_noisy_latent, noise = get_dm_training_batch() + ca_timestep, ca_noisy_latent = get_ca_training_batch(noise) + + # Calculate DM fake score + training_batch = self._build_distill_input_kwargs( + dm_noisy_latent.clone(), dm_timestep, training_batch.conditional_dict, + training_batch) + current_fake_score_transformer = self._get_fake_score_transformer( + dm_timestep) + dm_fake_score_pred_noise = current_fake_score_transformer( + **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + dm_faker_score_pred_video = pred_noise_to_pred_video( + pred_noise=dm_fake_score_pred_noise.flatten(0, 1), + noise_input_latent=dm_noisy_latent.flatten(0, 1), + timestep=dm_timestep, + scheduler=self.noise_scheduler).unflatten( + 0, dm_fake_score_pred_noise.shape[:2]) + + # Calculate DM real score + training_batch = self._build_distill_input_kwargs( + dm_noisy_latent.clone(), dm_timestep, training_batch.conditional_dict, + training_batch) + current_real_score_transformer = self._get_real_score_transformer( + dm_timestep) + dm_real_score_pred_noise_cond = current_real_score_transformer( + **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + dm_pred_real_video_cond = pred_noise_to_pred_video( + pred_noise=dm_real_score_pred_noise_cond.flatten(0, 1), + noise_input_latent=dm_noisy_latent.flatten(0, 1), + timestep=dm_timestep, + scheduler=self.noise_scheduler).unflatten( + 0, dm_real_score_pred_noise_cond.shape[:2]) + + # Calculate CA real score cond and uncond + training_batch = self._build_distill_input_kwargs( + ca_noisy_latent.clone(), ca_timestep, training_batch.conditional_dict, + training_batch) + ca_real_score_pred_noise_cond = current_real_score_transformer( + **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + ca_pred_real_video_cond = pred_noise_to_pred_video( + pred_noise=ca_real_score_pred_noise_cond.flatten(0, 1), + noise_input_latent=ca_noisy_latent.flatten(0, 1), + timestep=ca_timestep, + scheduler=self.noise_scheduler).unflatten( + 0, ca_real_score_pred_noise_cond.shape[:2]) + + training_batch = self._build_distill_input_kwargs( + ca_noisy_latent.clone(), ca_timestep, training_batch.unconditional_dict, + training_batch) + ca_real_score_pred_noise_uncond = current_real_score_transformer( + **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) + ca_pred_real_video_uncond = pred_noise_to_pred_video( + pred_noise=ca_real_score_pred_noise_uncond.flatten(0, 1), + noise_input_latent=ca_noisy_latent.flatten(0, 1), + timestep=ca_timestep, + scheduler=self.noise_scheduler).unflatten( + 0, ca_real_score_pred_noise_uncond.shape[:2]) + + # Calculate normalized DM loss + dm_grad = dm_faker_score_pred_video - dm_pred_real_video_cond + # dm_grad = (dm_faker_score_pred_video - dm_pred_real_video_cond) / torch.abs(original_latent - dm_pred_real_video_cond).mean() + ca_grad = (self.real_score_guidance_scale - 1) * (ca_pred_real_video_uncond - ca_pred_real_video_cond) + # ca_grad = ca_grad / torch.abs(original_latent - (self.real_score_guidance_scale - 1) * ca_pred_real_video_cond).mean() + + grad = dm_grad + ca_grad + grad = torch.nan_to_num(grad) + + if self.training_args.use_context_forcing and training_batch.trajectory_latents is not None: + context_forcing_length = training_batch.trajectory_latents.shape[1] + dmd_loss = 0.5 * F.mse_loss( + original_latent.float()[:, context_forcing_length:], + (original_latent.float()[:, context_forcing_length:] - + grad.float()[:, context_forcing_length:]).detach()) + else: + dmd_loss = 0.5 * F.mse_loss( + original_latent.float(), + (original_latent.float() - grad.float()).detach()) + + training_batch.dmd_latent_vis_dict.update({ + "training_batch_dmd_fwd_clean_latent": + training_batch.latents, + "generator_pred_video": + original_latent.detach(), + "real_score_pred_video": + dm_pred_real_video_cond.detach(), + "faker_score_pred_video": + dm_faker_score_pred_video.detach(), + "dmd_timestep": + ca_timestep.detach(), + }) + + return dmd_loss + def _dmd_forward(self, generator_pred_video: torch.Tensor, - training_batch: TrainingBatch) -> torch.Tensor: + training_batch: TrainingBatch, + exit_timestep: float) -> torch.Tensor: """Compute DMD (Diffusion Model Distillation) loss.""" original_latent = generator_pred_video with torch.no_grad(): @@ -802,6 +951,13 @@ class DistillationPipeline(TrainingPipeline): original_latent.float(), (original_latent.float() - grad.float()).detach()) + if self.training_args.use_gan_loss: + gan_generator_loss = self._gan_generator_forward( + generator_pred_video=generator_pred_video, + training_batch=training_batch, + exit_timestep=exit_timestep) + dmd_loss += gan_generator_loss + training_batch.dmd_latent_vis_dict.update({ "training_batch_dmd_fwd_clean_latent": training_batch.latents, @@ -817,6 +973,105 @@ class DistillationPipeline(TrainingPipeline): return dmd_loss + def _gan_generator_forward(self, generator_pred_video: torch.Tensor, + training_batch: TrainingBatch, + exit_timestep: float) -> torch.Tensor: + idx = torch.argmin(exit_timestep - self.denoising_step_list) + + critic_timestep = exit_timestep.clamp(self.min_timestep, self.max_timestep) * torch.ones([1, 1], device=self.device, dtype=torch.long) + + # if idx == len(self.denoising_step_list) - 1: + # start_timestep = 0 + # end_timestep = int(unshift_timestep(self.denoising_step_list[idx], self.timestep_shift, self.num_train_timestep)) + # else: + # start_timestep = int(unshift_timestep(self.denoising_step_list[idx + 1], self.timestep_shift, self.num_train_timestep)) + # end_timestep = int(unshift_timestep(self.denoising_step_list[idx], self.timestep_shift, self.num_train_timestep)) + + # critic_timestep = torch.randint(start_timestep, end_timestep, [1], device=self.device, dtype=torch.long) + # critic_timestep = shift_timestep(critic_timestep, self.timestep_shift, self.num_train_timestep) + # critic_timestep = critic_timestep.clamp(self.min_timestep, self.max_timestep) + + critic_noise = torch.randn(self.video_latent_shape, + device=self.device, + dtype=generator_pred_video.dtype) + noisy_fake_latent = self.noise_scheduler.add_noise( + generator_pred_video.flatten(0, 1), critic_noise.flatten(0, 1), + critic_timestep).detach().unflatten(0, + (1, generator_pred_video.shape[1])) + + training_batch = self._build_distill_input_kwargs( + noisy_fake_latent, critic_timestep, training_batch.conditional_dict, + training_batch) + current_fake_score_transformer = self._get_fake_score_transformer( + critic_timestep) + _, noisy_fake_logit = current_fake_score_transformer(classify_mode=True, + **training_batch.input_kwargs) + + gan_generator_loss = F.softplus(1 - noisy_fake_logit.float()).mean() + + return gan_generator_loss + + def _gan_critic_forward(self, training_batch: TrainingBatch, clean_latent: torch.Tensor) -> torch.Tensor: + with torch.no_grad(), set_forward_context( + current_timestep=training_batch.timesteps, + attn_metadata=training_batch.attn_metadata_vsa): + if self.training_args.simulate_generator_forward: + generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward( + training_batch) + else: + generator_pred_video, exit_timestep = self._generator_forward(training_batch) + + critic_timestep = exit_timestep.clamp(self.min_timestep, self.max_timestep) * torch.ones([1, 1], device=self.device, dtype=torch.long) + # idx = torch.argmin(exit_timestep - self.denoising_step_list) + # if idx == len(self.denoising_step_list) - 1: + # start_timestep = 0 + # end_timestep = int(unshift_timestep(self.denoising_step_list[idx], self.timestep_shift, self.num_train_timestep)) + # else: + # start_timestep = int(unshift_timestep(self.denoising_step_list[idx + 1], self.timestep_shift, self.num_train_timestep)) + # end_timestep = int(unshift_timestep(self.denoising_step_list[idx], self.timestep_shift, self.num_train_timestep)) + # critic_timestep = torch.randint(start_timestep, end_timestep, [1], device=self.device, dtype=torch.long) + # critic_timestep = shift_timestep(critic_timestep, self.timestep_shift, self.num_train_timestep) + # critic_timestep = critic_timestep.clamp(self.min_timestep, self.max_timestep) + + # Calculate fake latent score + critic_noise = torch.randn(self.video_latent_shape, + device=self.device, + dtype=generator_pred_video.dtype) + noisy_fake_latent = self.noise_scheduler.add_noise( + generator_pred_video.flatten(0, 1), critic_noise.flatten(0, 1), + critic_timestep).detach().unflatten(0, + (1, generator_pred_video.shape[1])) + + training_batch = self._build_distill_input_kwargs( + noisy_fake_latent, critic_timestep, training_batch.conditional_dict, + training_batch) + current_fake_score_transformer = self._get_fake_score_transformer( + critic_timestep) + _, noisy_fake_logit = current_fake_score_transformer(classify_mode=True, + **training_batch.input_kwargs) + + # Calculate real latent score + # critic_noise = torch.randn(self.video_latent_shape, + # device=self.device, + # dtype=generator_pred_video.dtype) + noisy_real_latent = self.noise_scheduler.add_noise( + clean_latent.flatten(0, 1), critic_noise.flatten(0, 1), + critic_timestep).detach().unflatten(0, + (1, clean_latent.shape[1])) + + training_batch = self._build_distill_input_kwargs( + noisy_real_latent, critic_timestep, training_batch.conditional_dict, + training_batch) + current_fake_score_transformer = self._get_fake_score_transformer( + critic_timestep) + _, noisy_real_logit = current_fake_score_transformer(classify_mode=True, + **training_batch.input_kwargs) + + gan_critic_loss = F.softplus(noisy_real_logit.float()).mean() - F.softplus(1 - noisy_fake_logit.float()).mean() + # gan_critic_loss = F.softplus(1 - noisy_real_logit.mean()) + F.softplus(1 + noisy_fake_logit.mean()) + + return gan_critic_loss + def faker_score_forward( self, training_batch: TrainingBatch ) -> tuple[TrainingBatch, torch.Tensor]: @@ -824,10 +1079,10 @@ class DistillationPipeline(TrainingPipeline): current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa): if self.training_args.simulate_generator_forward: - generator_pred_video = self._generator_multi_step_simulation_forward( + generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward( training_batch) else: - generator_pred_video = self._generator_forward(training_batch) + generator_pred_video, exit_timestep = self._generator_forward(training_batch) fake_score_timestep = torch.randint(0, self.num_train_timestep, [1], @@ -979,6 +1234,52 @@ class DistillationPipeline(TrainingPipeline): training_batch.infos = infos return training_batch + def _get_next_batch_2( + self, + training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]: + batch = next(self.train_loader_iter_2, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + self.train_loader_iter_2 = iter(self.train_dataloader_2) + batch = next(self.train_loader_iter_2) + + # Required fields from parquet (ODE trajectory schema) + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + latent = batch['vae_latent'].float() + infos = batch['info_list'] + + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latent -= self.vae.shift_factor.to( + latent.device, latent.dtype) + else: + latent -= self.vae.shift_factor + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent = latent * self.vae.scaling_factor.to( + latent.device, latent.dtype) + else: + latent = latent * self.vae.scaling_factor + # [B, C, T, H, W] -> [B, T, C, H, W] + latent = latent.permute(0, 2, 1, 3, 4) + + # Move to device + device = get_local_torch_device() + training_batch.encoder_hidden_states = encoder_hidden_states.to( + device, dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + device, dtype=torch.bfloat16) + training_batch.infos = infos + + return training_batch, latent[:, :self.training_args. + num_latent_t].to( + device, + dtype=torch.bfloat16 + ) + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: gradient_accumulation_steps = getattr(self.training_args, 'gradient_accumulation_steps', 1) @@ -1006,17 +1307,23 @@ class DistillationPipeline(TrainingPipeline): current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.attn_metadata_vsa): if self.training_args.simulate_generator_forward: - generator_pred_video = self._generator_multi_step_simulation_forward( + generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward( batch_gen) else: - generator_pred_video = self._generator_forward( + generator_pred_video, exit_timestep = self._generator_forward( batch_gen) with set_forward_context(current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.attn_metadata): - dmd_loss = self._dmd_forward( - generator_pred_video=generator_pred_video, - training_batch=batch_gen) + if self.training_args.use_decoupled_dmd: + dmd_loss = self._dmd_decoupled_forward( + generator_pred_video=generator_pred_video, + training_batch=batch_gen, + exit_timestep=exit_timestep) + else: + dmd_loss = self._dmd_forward( + generator_pred_video=generator_pred_video, + training_batch=batch_gen) with set_forward_context( current_timestep=batch_gen.timesteps, diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index 93ef1b939..fce309f6d 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/fastvideo/training/ode_causal_pipeline.py @@ -18,6 +18,7 @@ from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import ( WanCausalDMDPipeline) from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases) @@ -39,6 +40,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): def initialize_pipeline(self, fastvideo_args: FastVideoArgs): # Match the preprocess/generation scheduler for consistent stepping + assert fastvideo_args.pipeline_config.flow_shift == 5.0, "flow_shift must be 5.0" self.modules["scheduler"] = SelfForcingFlowMatchScheduler( shift=fastvideo_args.pipeline_config.flow_shift, sigma_min=0.0, @@ -57,21 +59,22 @@ class ODEInitTrainingPipeline(TrainingPipeline): self.vae.requires_grad_(False) self.timestep_shift = self.training_args.pipeline_config.flow_shift + assert self.timestep_shift == 5.0, "timestep_shift must be 5.0" self.noise_scheduler = SelfForcingFlowMatchScheduler( shift=self.timestep_shift, sigma_min=0.0, extra_one_step=True) self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True) - logger.info("dmd_denoising_steps: %s", - self.training_args.pipeline_config.dmd_denoising_steps) - self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.vae)) + + self.dmd_denoising_steps = torch.tensor(self.training_args.pipeline_config.dmd_denoising_steps, dtype=torch.long, device=get_local_torch_device()) if training_args.warp_denoising_step: # Warp the denoising step according to the scheduler time shift timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda() - logger.info("timesteps: %s", timesteps) self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps] logger.info("warped self.dmd_denoising_steps: %s", @@ -82,8 +85,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): self.dmd_denoising_steps = self.dmd_denoising_steps.to( get_local_torch_device()) - logger.info("denoising_step_list: %s", self.dmd_denoising_steps) - logger.info( "Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps)) @@ -110,6 +111,64 @@ class ODEInitTrainingPipeline(TrainingPipeline): pin_cpu_memory=training_args.pin_cpu_memory, dit_cpu_offload=True) + def _get_next_gt_batch( + self, + training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + self.train_loader_iter = iter(self.train_dataloader) + batch = next(self.train_loader_iter) + + logger.info("batch keys: %s", batch.keys()) + # Required fields from parquet (ODE trajectory schema) + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + infos = batch['info_list'] + + # Trajectory tensors may include a leading singleton batch dim per row + trajectory_latents = batch['trajectory_latents'] + if trajectory_latents.dim() == 7: + # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W] + trajectory_latents = trajectory_latents[:, 0] + elif trajectory_latents.dim() == 6: + # already [B, S, C, T, H, W] + pass + else: + raise ValueError( + f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}" + ) + + trajectory_timesteps = batch['trajectory_timesteps'] + if trajectory_timesteps.dim() == 3: + # [B, 1, S] -> [B, S] + trajectory_timesteps = trajectory_timesteps[:, 0] + elif trajectory_timesteps.dim() == 2: + # [B, S] + pass + else: + raise ValueError( + f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}" + ) + # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing + trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) + + # Move to device + device = get_local_torch_device() + training_batch.encoder_hidden_states = encoder_hidden_states.to( + device, dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + device, dtype=torch.bfloat16) + training_batch.infos = infos + + return training_batch, trajectory_latents[:, :, :self.training_args. + num_latent_t].to( + device, + dtype=torch.bfloat16 + ), trajectory_timesteps.to( + device) + def _get_next_batch( self, training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]: @@ -299,6 +358,144 @@ class ODEInitTrainingPipeline(TrainingPipeline): return pred_video, target_latent, timestep, latent_vis_dict + def _tf_step_predict_next_latent( + self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + torch.Tensor]]: + latent_vis_dict: dict[str, torch.Tensor] = {} + device = get_local_torch_device() + target_latent = traj_latents[:, -1] + del traj_latents + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, num_frames, num_channels, height, width = target_latent.shape + + indexes = self._get_timestep( # [B, num_frames] + 0, + 1000, + B, + num_frames, + 3, + uniform_timestep=False) + timestep = self.noise_scheduler.timesteps[indexes.cpu()].to(device) + + noisy_input = self.noise_scheduler.add_noise( + target_latent.flatten(0, 1), + torch.randn_like(target_latent.flatten(0, 1)), + timestep.flatten(0, 1)).unflatten(0, (B, num_frames)) + + clean_input = target_latent.clone() + + # Prepare inputs for transformer + latent_vis_dict["noisy_input"] = noisy_input.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, + 4).detach().clone().cpu() + + logger.info("timestep: %s", timestep) + input_kwargs = { + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + "clean_hidden_states": clean_input.permute(0, 2, 1, 3, 4), + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep.to(device, dtype=torch.bfloat16), + "return_dict": False, + } + # Predict noise and step the scheduler to obtain next latent + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + + from fastvideo.models.utils import pred_noise_to_pred_video + pred_video = pred_noise_to_pred_video( + pred_noise=noise_pred.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + scheduler=self.modules["scheduler"]).unflatten( + 0, noise_pred.shape[:2]) + latent_vis_dict["pred_video"] = pred_video.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + + return pred_video, target_latent, timestep, latent_vis_dict + + def _tf_ode_step_predict_next_latent( + self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + torch.Tensor]]: + latent_vis_dict: dict[str, torch.Tensor] = {} + device = get_local_torch_device() + clean_latent = traj_latents[:, -1] + target_latent = traj_latents[:, -2] + ode_latents_valid = traj_latents[:, :-1] + + # self._cached_closest_idx_per_dmd = torch.tensor( + # [0, 16, 32, 41], dtype=torch.long).cpu() + # self.dmd_denoising_steps = traj_timesteps[0, self._cached_closest_idx_per_dmd] + # logger.info(f"corresponding timesteps: {self.dmd_denoising_steps}") + + # relevant_traj_latents = torch.index_select( + # ode_latents_valid, + # dim=1, + # index=self._cached_closest_idx_per_dmd.to(device)) + relevant_traj_latents = ode_latents_valid + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, num_frames, num_channels, height, width = target_latent.shape + + indexes = self._get_timestep( # [B, num_frames] + 0, + len(self.dmd_denoising_steps), + B, + num_frames, + 3, + uniform_timestep=True) + timestep = self.dmd_denoising_steps[indexes.cpu()].to(device) + + noisy_input = torch.gather( + relevant_traj_latents, + dim=1, + index=indexes.reshape(B, 1, num_frames, 1, 1, + 1).expand(-1, -1, -1, num_channels, height, + width).to(self.device)).squeeze(1) + + # Prepare inputs for transformer + latent_vis_dict["noisy_input"] = noisy_input.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, + 4).detach().clone().cpu() + latent_vis_dict["clean_latent"] = clean_latent.permute(0, 2, 1, 3, + 4).detach().clone().cpu() + + logger.info("timestep: %s", timestep) + input_kwargs = { + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + "clean_hidden_states": clean_latent.permute(0, 2, 1, 3, 4), + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep.to(device, dtype=torch.bfloat16), + "return_dict": False, + } + # Predict noise and step the scheduler to obtain next latent + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + + from fastvideo.models.utils import pred_noise_to_pred_video + pred_video = pred_noise_to_pred_video( + pred_noise=noise_pred.flatten(0, 1), + noise_input_latent=noisy_input.flatten(0, 1), + timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + scheduler=self.modules["scheduler"]).unflatten( + 0, noise_pred.shape[:2]) + latent_vis_dict["pred_video"] = pred_video.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + + return pred_video, target_latent, timestep, latent_vis_dict + def train_one_step(self, training_batch): # type: ignore[override] self.transformer.train() self.optimizer.zero_grad() @@ -316,11 +513,11 @@ class ODEInitTrainingPipeline(TrainingPipeline): # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] _, S = traj_latents.shape[0], traj_latents.shape[1] - if S < 2: - raise ValueError("Trajectory must contain at least 2 steps") + # if S < 2: + # raise ValueError("Trajectory must contain at least 2 steps") # Forward to predict next latent by stepping scheduler with predicted noise - noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent( + noise_pred, target_latent, t, latent_vis_dict = self._tf_ode_step_predict_next_latent( traj_latents, traj_timesteps, text_embeds, text_attention_mask) training_batch.latent_vis_dict.update(latent_vis_dict) @@ -365,12 +562,12 @@ class ODEInitTrainingPipeline(TrainingPipeline): training_args: TrainingArgs, step: int): tracker_loss_dict: dict[str, Any] = {} latents_vis_dict = training_batch.latent_vis_dict - latent_log_keys = ['noisy_input', 'x0', 'pred_video'] + latent_log_keys = ['noisy_input', 'x0', 'pred_video', 'clean_latent'] for latent_key in latent_log_keys: assert latent_key in latents_vis_dict and latents_vis_dict[ latent_key] is not None latent = latents_vis_dict[latent_key] - pixel_latent = self.validation_pipeline.decoding_stage.decode( + pixel_latent = self.decoding_stage.decode( latent, training_args) video = pixel_latent.cpu().float() diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 491c39a6d..1e4ecb934 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -119,14 +119,21 @@ class SelfForcingDistillationPipeline(DistillationPipeline): with set_forward_context( current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa): - generator_pred_video = self._generator_multi_step_simulation_forward( + generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward( training_batch) with set_forward_context(current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata): - dmd_loss = self._dmd_forward( + if self.training_args.use_decoupled_dmd: + dmd_loss = self._dmd_decoupled_forward( + generator_pred_video=generator_pred_video, + training_batch=training_batch, + exit_timestep=exit_timestep) + else: + dmd_loss = self._dmd_forward( generator_pred_video=generator_pred_video, - training_batch=training_batch) + training_batch=training_batch, + exit_timestep=exit_timestep) log_dict = { "dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device) @@ -190,7 +197,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noise = training_batch.trajectory_latents.to(self.device, dtype=dtype) else: - noise = training_batch.latents.to(self.device, dtype=dtype) + noise = torch.randn_like(training_batch.latents).to(self.device, dtype=dtype) batch_size, num_frames, num_channels, height, width = noise.shape @@ -511,6 +518,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): self.denoising_step_list[exit_flags[0]], dtype=torch.float32, device=self.device) + exit_timestep = self.denoising_step_list[exit_flags[0]] # Store gradient mask information for debugging if gradient_mask is not None: @@ -528,7 +536,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): assert crossattn_cache is not None self._reset_simulation_caches(kv_cache1, crossattn_cache) - return final_output if gradient_mask is not None else pred_image_or_video + return (final_output, exit_timestep) if gradient_mask is not None else (pred_image_or_video, exit_timestep) def _initialize_simulation_caches( self, @@ -543,7 +551,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): latent_shape = self.video_latent_shape_sp _, num_frames, _, height, width = latent_shape - if isinstance(self.transformer.config.patch_size, tuple): + if isinstance(self.transformer.config.patch_size, tuple) or isinstance(self.transformer.config.patch_size, list): ph, pw = self.transformer.config.patch_size[ 1], self.transformer.config.patch_size[2] elif isinstance(self.transformer.config.patch_size, int): @@ -805,6 +813,44 @@ class SelfForcingDistillationPipeline(DistillationPipeline): self.fake_score_optimizer.step() self.fake_score_lr_scheduler.step() + if self.training_args.use_gan_loss: + # Freeze the base fake-score models and only train GAN head params. + # We must handle both experts because timestep routing can select either. + fake_score_models = [self.fake_score_transformer] + if self.fake_score_transformer_2 is not None: + fake_score_models.append(self.fake_score_transformer_2) + + for fs_model in fake_score_models: + fs_model.requires_grad_(False) + for name, param in fs_model.named_parameters(): + if "_cls_pred_branch" in name or "_gan_ca_blocks" in name or "_register_tokens" in name: + param.requires_grad_(True) + + # Prevent accidental gradient carry-over from the previous critic update. + self.fake_score_optimizer.zero_grad(set_to_none=True) + if self.fake_score_transformer_2 is not None: + self.fake_score_optimizer_2.zero_grad(set_to_none=True) + batch_gan_critic = copy.deepcopy(training_batch) + batch_gan_critic, real_latent = self._get_next_batch_2(batch_gan_critic) + with set_forward_context(current_timestep=batch_gan_critic.timesteps, + attn_metadata=batch_gan_critic.attn_metadata): + critic_gan_loss = self._gan_critic_forward(batch_gan_critic, real_latent) + (critic_gan_loss / gradient_accumulation_steps).backward() + total_critic_loss += critic_gan_loss.detach().item() + + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + self._clip_model_grad_norm_(batch_gan_critic, + self.fake_score_transformer_2) + self.fake_score_optimizer_2.step() + self.fake_score_lr_scheduler_2.step() + else: + self._clip_model_grad_norm_(batch_gan_critic, + self.fake_score_transformer) + self.fake_score_optimizer.step() + self.fake_score_lr_scheduler.step() + for fs_model in fake_score_models: + fs_model.requires_grad_(True) + avg_critic_loss = torch.tensor(total_critic_loss / gradient_accumulation_steps, device=self.device) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 417a39d7a..4a827534f 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -19,12 +19,12 @@ from torchdata.stateful_dataloader import StatefulDataLoader from tqdm.auto import tqdm import fastvideo.envs as envs -from fastvideo.attention.backends.video_sparse_attn import ( - VideoSparseAttentionMetadataBuilder) -from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder +# from fastvideo.attention.backends.video_sparse_attn import ( +# VideoSparseAttentionMetadataBuilder) +# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder from fastvideo.configs.sample import SamplingParam from fastvideo.dataset import build_parquet_map_style_dataloader -from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only, pyarrow_schema_t2v +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only, pyarrow_schema_t2v, pyarrow_schema_text_only from fastvideo.dataset.validation_dataset import ValidationDataset from fastvideo.distributed import (cleanup_dist_env_and_memory, get_local_torch_device, get_sp_group, @@ -48,8 +48,12 @@ from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict) from fastvideo.optim.muon import get_muon_optimizer -vsa_available = is_vsa_available() -vmoba_available = is_vmoba_available() +try: + vsa_available = is_vsa_available() + vmoba_available = is_vmoba_available() +except: + vsa_available = False + vmoba_available = False logger = init_logger(__name__) @@ -88,8 +92,8 @@ class TrainingPipeline(LoRAPipeline, ABC): "create_pipeline_stages should not be called for training pipeline") def set_schemas(self) -> None: - self.train_dataset_schema = pyarrow_schema_t2v - self.train_dataset_schema_2 = pyarrow_schema_ode_trajectory_text_only + self.train_dataset_schema = pyarrow_schema_text_only + self.train_dataset_schema_2 = pyarrow_schema_t2v def initialize_training_pipeline(self, training_args: TrainingArgs): logger.info("Initializing training pipeline...") @@ -394,7 +398,7 @@ class TrainingPipeline(LoRAPipeline, ABC): patch_size = self.training_args.pipeline_config.dit_config.patch_size current_vsa_sparsity = training_batch.current_vsa_sparsity assert latents_shape is not None - assert training_batch.timesteps is not None + # assert training_batch.timesteps is not None if envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": if not vsa_available: raise ImportError( diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 8eb61ec71..b9e896244 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1318,6 +1318,14 @@ def shift_timestep(timestep: torch.Tensor, shift: float, denominator = 1 + (shift - 1) * t return num_train_timestep * (shift * t / denominator) +def unshift_timestep(timestep: torch.Tensor, shift: float, + num_train_timestep: float) -> torch.Tensor: + # inverse of shift_timestep + if shift == 1: + return timestep + t = timestep / num_train_timestep + denominator = shift - (shift - 1) * t + return num_train_timestep * (t / denominator) # coding=utf-8 # Copyright 2025 The HuggingFace Inc. team. diff --git a/fastvideo/training/wan_training_pipeline.py b/fastvideo/training/wan_training_pipeline.py index 1b3785197..70151770f 100644 --- a/fastvideo/training/wan_training_pipeline.py +++ b/fastvideo/training/wan_training_pipeline.py @@ -10,7 +10,10 @@ from fastvideo.pipelines.basic.wan.wan_pipeline import WanPipeline from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.utils import is_vsa_available -vsa_available = is_vsa_available() +try: + vsa_available = is_vsa_available() +except: + vsa_available = False logger = init_logger(__name__) diff --git a/wan_ode_preprocess/launch_slurm.sh b/wan_ode_preprocess/launch_slurm.sh new file mode 100644 index 000000000..18561d6a0 --- /dev/null +++ b/wan_ode_preprocess/launch_slurm.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# Create output directory if it doesn't exist +mkdir -p preprocess_output + +# Launch 8 jobs, one for each node +# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node) +for node_id in {0..7}; do + # Calculate the starting file number for this node + start_file=$((node_id * 8 + 1)) + + echo "Launching node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt" + echo "sbatch --job-name=ode-${node_id} --output=preprocess_output/preprocess-node-${node_id}.out --error=preprocess_output/preprocess-node-${node_id}.err slurms/syn.slurm $start_file $node_id" + + sbatch --job-name=ode-${node_id} \ + --output=preprocess_output/preprocess-node-${node_id}.out \ + --error=preprocess_output/preprocess-node-${node_id}.err \ + /mnt/home/zhouw.jerry2017/FastVideo/wan_ode_preprocess/syn.slurm $start_file $node_id +done + +echo "All 8 nodes launched successfully!" \ No newline at end of file diff --git a/wan_ode_preprocess/syn.slurm b/wan_ode_preprocess/syn.slurm new file mode 100644 index 000000000..c8df7c518 --- /dev/null +++ b/wan_ode_preprocess/syn.slurm @@ -0,0 +1,99 @@ +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=1440G +#SBATCH --exclusive +#SBATCH --time=72:00:00 + +# conda init +# source ~/conda/miniconda/bin/activate +source ~/miniconda3/bin/activate +PYTHON_VIRTUAL_ENVIRONMENT=fastvideo +conda activate $PYTHON_VIRTUAL_ENVIRONMENT +nvidia-smi +nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv + +echo " " +echo " Number of nodes:= " $SLURM_JOB_NUM_NODES +echo " GPUs per node:= " $SLURM_JOB_GPUS +echo " Running on multiple nodes/GPU devices" +echo "" +echo " Run started at:- " +date + +# Accept parameters from launch script +START_FILE=${1:-1} # Starting file number for this node +NODE_ID=${2:-0} # Node identifier (0-7) + +num_gpus=1 +export MODEL_BASE=wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# Start port number - we'll increment for each job +base_port=$((29603 + NODE_ID * 100)) # Different port range per node + +# Create an array of CUDA device IDs +gpu_ids=(0 1 2 3 4 5 6 7) + +GPU_NUM=1 +MODEL_TYPE="wan" + +echo "NODE_ID: $NODE_ID" +echo "START_FILE: $START_FILE" +echo "Base port for this node: $base_port" + +# Run 8 parallel preprocessing jobs on this node +for i in {1..8}; do + # Calculate port for this job + port=$((base_port + i)) + + # Get GPU ID using modulo to cycle through available GPUs + gpu=${gpu_ids[((i-1))]} + + # Calculate which file this GPU should process + file_num=$((START_FILE + i - 1)) + # DATA_MERGE_PATH="/mnt/data/fv-ode-preprocessing-6k-mixkit-wan-1.3b/Node_${NODE_ID}_GPU_${i}_File_${file_num}" + DATA_MERGE_PATH="/mnt/data/mixkit_wan_1.3b_processed_t2v_distributed/Node_${NODE_ID}_GPU_${i}_File_${file_num}" + # INIT_WEIGHTS_FROM_SAFETENSORS="/mnt/data/wan_tf_init_3333_mixkit_shift5/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors" + INIT_WEIGHTS_FROM_SAFETENSORS="/mnt/data/wan_ar_diffusion_3333_mixkit_shift5/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors" + + # Create unique output directory based on node and GPU + OUTPUT_DIR="/mnt/data/fv-tf-ode-preprocessing-6k-mixkit-wan-1.3b/Node_${NODE_ID}_GPU_${i}_File_${file_num}" + + start_cpu=$(( (i-1)*2 )) # Reduced CPU allocation for 8 nodes + end_cpu=$(( start_cpu+1 )) + + echo "Starting GPU $gpu processing file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR" + + # Run the preprocessing command in background + CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_BASE \ + --init_weights_from_safetensors $INIT_WEIGHTS_FROM_SAFETENSORS \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --flow_shift 5.0 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "tf_ode" & +done + +# Wait for all jobs on this node to complete +wait + +echo "All processing blocks completed!" diff --git a/wan_preprocess/launch_slurm.sh b/wan_preprocess/launch_slurm.sh new file mode 100644 index 000000000..f6da828e4 --- /dev/null +++ b/wan_preprocess/launch_slurm.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# Create output directory if it doesn't exist +mkdir -p preprocess_output + +# Launch 8 jobs, one for each node +# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node) +for node_id in {0..7}; do + # Calculate the starting file number for this node + start_file=$((node_id * 8 + 1)) + + echo "Launching node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt" + echo "sbatch --job-name=ode-${node_id} --output=preprocess_output/preprocess-node-${node_id}.out --error=preprocess_output/preprocess-node-${node_id}.err slurms/syn.slurm $start_file $node_id" + + sbatch --job-name=ode-${node_id} \ + --output=preprocess_output/preprocess-node-${node_id}.out \ + --error=preprocess_output/preprocess-node-${node_id}.err \ + /mnt/home/zhouw.jerry2017/FastVideo/wan_preprocess/syn.slurm $start_file $node_id +done + +echo "All 8 nodes launched successfully!" \ No newline at end of file diff --git a/wan_preprocess/syn.slurm b/wan_preprocess/syn.slurm new file mode 100644 index 000000000..f59a39eb0 --- /dev/null +++ b/wan_preprocess/syn.slurm @@ -0,0 +1,95 @@ +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=960G +#SBATCH --exclusive +#SBATCH --time=72:00:00 + +# conda init +# source ~/conda/miniconda/bin/activate +source ~/miniconda3/bin/activate +PYTHON_VIRTUAL_ENVIRONMENT=fastvideo +conda activate $PYTHON_VIRTUAL_ENVIRONMENT +nvidia-smi +nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv + +echo " " +echo " Number of nodes:= " $SLURM_JOB_NUM_NODES +echo " GPUs per node:= " $SLURM_JOB_GPUS +echo " Running on multiple nodes/GPU devices" +echo "" +echo " Run started at:- " +date + +# Accept parameters from launch script +START_FILE=${1:-1} # Starting file number for this node +NODE_ID=${2:-0} # Node identifier (0-7) + +num_gpus=1 +export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# Start port number - we'll increment for each job +base_port=$((29603 + NODE_ID * 100)) # Different port range per node + +# Create an array of CUDA device IDs +gpu_ids=(0 1 2 3 4 5 6 7) + +GPU_NUM=1 +MODEL_TYPE="wan" + +echo "NODE_ID: $NODE_ID" +echo "START_FILE: $START_FILE" +echo "Base port for this node: $base_port" + +# Run 8 parallel preprocessing jobs on this node +for i in {1..8}; do + # Calculate port for this job + port=$((base_port + i)) + + # Get GPU ID using modulo to cycle through available GPUs + gpu=${gpu_ids[((i-1))]} + + # Calculate which file this GPU should process + file_num=$((START_FILE + i - 1)) + DATA_MERGE_PATH="/mnt/home/zhouw.jerry2017/mixkit_prompts/mixkit_prompts_${file_num}.txt" + + # Create unique output directory based on node and GPU + OUTPUT_DIR="/mnt/data/fv-ode-preprocessing-6k-mixkit-wan-1.3b/Node_${NODE_ID}_GPU_${i}_File_${file_num}" + + start_cpu=$(( (i-1)*2 )) # Reduced CPU allocation for 8 nodes + end_cpu=$(( start_cpu+1 )) + + echo "Starting GPU $gpu processing file mixkit_prompts_${file_num}.txt on port $port, output: $OUTPUT_DIR" + + # Run the preprocessing command in background + CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_BASE \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --flow_shift 5.0 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "ode_trajectory" & +done + +# Wait for all jobs on this node to complete +wait + +echo "All processing blocks completed!" diff --git a/wan_text_preprocess/launch_slurm.sh b/wan_text_preprocess/launch_slurm.sh new file mode 100644 index 000000000..8d5795d4e --- /dev/null +++ b/wan_text_preprocess/launch_slurm.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# Create output directory if it doesn't exist +mkdir -p preprocess_output + +# Launch 8 jobs, one for each node +# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node) +for node_id in {0..7}; do + # Calculate the starting file number for this node + start_file=$((node_id * 8 + 1)) + + echo "Launching node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt" + echo "sbatch --job-name=ode-${node_id} --output=preprocess_output/preprocess-node-${node_id}.out --error=preprocess_output/preprocess-node-${node_id}.err slurms/syn.slurm $start_file $node_id" + + sbatch --job-name=ode-${node_id} \ + --output=preprocess_output/preprocess-node-${node_id}.out \ + --error=preprocess_output/preprocess-node-${node_id}.err \ + /mnt/home/zhouw.jerry2017/FastVideo/wan_text_preprocess/syn.slurm $start_file $node_id +done + +echo "All 8 nodes launched successfully!" \ No newline at end of file diff --git a/wan_text_preprocess/syn.slurm b/wan_text_preprocess/syn.slurm new file mode 100644 index 000000000..6e6839749 --- /dev/null +++ b/wan_text_preprocess/syn.slurm @@ -0,0 +1,95 @@ +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=960G +#SBATCH --exclusive +#SBATCH --time=72:00:00 + +# conda init +# source ~/conda/miniconda/bin/activate +source ~/miniconda3/bin/activate +PYTHON_VIRTUAL_ENVIRONMENT=fastvideo +conda activate $PYTHON_VIRTUAL_ENVIRONMENT +nvidia-smi +nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv + +echo " " +echo " Number of nodes:= " $SLURM_JOB_NUM_NODES +echo " GPUs per node:= " $SLURM_JOB_GPUS +echo " Running on multiple nodes/GPU devices" +echo "" +echo " Run started at:- " +date + +# Accept parameters from launch script +START_FILE=${1:-1} # Starting file number for this node +NODE_ID=${2:-0} # Node identifier (0-7) + +num_gpus=1 +export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers +export HF_HOME=/mnt/data/huggingface +export CUDA_HOME=/usr/local/cuda +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/shared-bin/local/usr/lib/aarch64-linux-gnu:$CUDA_HOME/lib64" +export CUDACXX="$CUDA_HOME/bin/nvcc" +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# Start port number - we'll increment for each job +base_port=$((29603 + NODE_ID * 100)) # Different port range per node + +# Create an array of CUDA device IDs +gpu_ids=(0 1 2 3 4 5 6 7) + +GPU_NUM=1 +MODEL_TYPE="wan" + +echo "NODE_ID: $NODE_ID" +echo "START_FILE: $START_FILE" +echo "Base port for this node: $base_port" + +# Run 8 parallel preprocessing jobs on this node +for i in {1..8}; do + # Calculate port for this job + port=$((base_port + i)) + + # Get GPU ID using modulo to cycle through available GPUs + gpu=${gpu_ids[((i-1))]} + + # Calculate which file this GPU should process + file_num=$((START_FILE + i - 1)) + DATA_MERGE_PATH="/mnt/home/zhouw.jerry2017/vidprom_shards/vidprom_${file_num}.txt" + + # Create unique output directory based on node and GPU + OUTPUT_DIR="/mnt/data/vidprom_filtered_extended_umt5_text_embed/Node_${NODE_ID}_GPU_${i}_File_${file_num}" + + start_cpu=$(( (i-1)*2 )) # Reduced CPU allocation for 8 nodes + end_cpu=$(( start_cpu+1 )) + + echo "Starting GPU $gpu processing file mixkit_prompts_${file_num}.txt on port $port, output: $OUTPUT_DIR" + + # Run the preprocessing command in background + CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_BASE \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --flow_shift 5.0 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "text_only" & +done + +# Wait for all jobs on this node to complete +wait + +echo "All processing blocks completed!"