diff --git a/examples/datasets/crush_smol/download_dataset.sh b/examples/datasets/crush_smol/download_dataset.sh new file mode 100644 index 000000000..f2942b7da --- /dev/null +++ b/examples/datasets/crush_smol/download_dataset.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset" \ No newline at end of file diff --git a/examples/datasets/crush_smol/validation.json b/examples/datasets/crush_smol/validation.json new file mode 100644 index 000000000..a7e327a04 --- /dev/null +++ b/examples/datasets/crush_smol/validation.json @@ -0,0 +1,76 @@ +{ + "data": [ + { + "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.", + "image_path": null, + "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 new file mode 100644 index 000000000..cec24bb69 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 new file mode 100644 index 000000000..87d4aed09 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 new file mode 100644 index 000000000..a0455d1f9 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 b/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 new file mode 100644 index 000000000..15c45a810 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 b/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 new file mode 100644 index 000000000..27cab6878 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 new file mode 100644 index 000000000..8e416c6b0 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 new file mode 100644 index 000000000..c2c781b51 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 new file mode 100644 index 000000000..abb42a2a8 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 differ diff --git a/examples/datasets/crush_smol/validation_single.json b/examples/datasets/crush_smol/validation_single.json new file mode 100644 index 000000000..30c3844a5 --- /dev/null +++ b/examples/datasets/crush_smol/validation_single.json @@ -0,0 +1,13 @@ +{ + "data": [ + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index f12a4969e..b987b6fb8 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -1,6 +1,6 @@ from fastvideo import VideoGenerator -# from fastvideo.configs.sample import SamplingParam +from fastvideo.configs.sample import SamplingParam OUTPUT_PATH = "video_samples" def main(): @@ -17,30 +17,29 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + ti2v_task=True, # image_encoder_cpu_offload=False, ) - # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") # sampling_param.num_frames = 45 - # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" + sampling_param.image_path = "test.jpg" # Generate videos with the same simple API, regardless of GPU count prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + "A girl is packing a suitcase when stuff suddently starts flying around the room." ) video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) + # prompt2 = ( + # "A majestic lion strides across the golden savanna, its powerful frame " + # "glistening under the warm afternoon sun. The tall grass ripples gently in " + # "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + # "embodying the raw energy of the wild. Low angle, steady tracking shot, " + # "cinematic.") + # video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) if __name__ == "__main__": diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index e2241a94e..56ffe527f 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -19,13 +19,15 @@ def main(): ) sampling_param = SamplingParam.from_pretrained(model_name) + sampling_param.num_frames = 81 - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) - video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + prompts = [ + "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.", + "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.", + ] + + for prompt in prompts: + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) if __name__ == "__main__": main() diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh b/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh new file mode 100644 index 000000000..2bbece456 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh @@ -0,0 +1,99 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/" +DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/" +VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +NUM_GPUS=1 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_70k" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "fixed_wan_ode_init_70k_6e-6" + # --resume_from_checkpoint "ode_init_diffusers/" + --max_train_steps 6000 + --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 77 + --warp_denoising_step + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 6e-6 + --mixed_precision "bf16" + --checkpointing_steps 1000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --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 +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm new file mode 100644 index 000000000..a21e126c0 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm @@ -0,0 +1,135 @@ +#!/bin/bash +#SBATCH --job-name=1e5B2_16kFV_warp_ode_vidprom +#SBATCH --partition=main +#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=ode_vidprom16k_warp/Dode_vidprom8b16k_1e-5.out +#SBATCH --error=ode_vidprom16k_warp/Dode_vidprom8b16k_1e-5.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv2 + +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b-81/" +VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "Dwarp_vidprom_8b16k_test_warp_1e-5" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "Dwarp_vidprom_8b16k_wan_ode_init_1e-5" + # --resume_from_checkpoint "ode_init_diffusers/" + --warp_denoising_step + --log_visualization + --max_train_steps 6001 + --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 77 + --dmd_denoising_steps "1000,750,500,250" + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim $NUM_GPUS + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" + # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --checkpointing_steps 500 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + # --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[@]}" diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm new file mode 100644 index 000000000..dc5c8d51c --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm @@ -0,0 +1,131 @@ +#!/bin/bash +#SBATCH --job-name=ode_vidprom2k +#SBATCH --partition=main +#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=ode_vidprom2k_output/ode_vidprom2k.out +#SBATCH --error=ode_vidprom2k_output/ode_vidprom2k.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv2 + +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/" +VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_vidprom2k" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "vidprom2k_wan_ode_init_5e-6" + # --resume_from_checkpoint "ode_init_diffusers/" + --max_train_steps 6001 + --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 77 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 8 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-6 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --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[@]}" diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm new file mode 100644 index 000000000..0688eefd6 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm @@ -0,0 +1,132 @@ +#!/bin/bash +#SBATCH --job-name=ode_crush +#SBATCH --partition=main +#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=ode_crush_output/ode_crush.out +#SBATCH --error=ode_crush_output/ode_crush.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv2 + +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json" +NUM_GPUS=2 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_warp_2" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "2warp_fixed_wan_ode_init_5e-6" + # --resume_from_checkpoint "ode_init_diffusers/" + # --warp_denoising_step + --max_train_steps 6001 + --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 77 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim $NUM_GPUS + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 20 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-6 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --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[@]}" diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh b/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh new file mode 100644 index 000000000..02f7a33af --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh @@ -0,0 +1,98 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/weka/home/hao.zhang/wl/FastVideo2/data/crush-smol_processed_t2v_1_3b_ode_init_single" +VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +NUM_GPUS=1 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_crush_smol" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "overfitwan_ode_init_crush_smol" + # --resume_from_checkpoint "ode_init_diffusers/" + --max_train_steps 2001 + # --warp_denoising_step + --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 77 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 20 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --checkpointing_steps 500 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --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 +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh b/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh new file mode 100644 index 000000000..a1b223666 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="data/crush-smol_single/merge.txt" +OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_single/" + +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 1 \ + --flush_frequency 1 \ + --video_length_tolerance_range 5 \ + --preprocess_task "ode_trajectory" diff --git a/examples/training/consistency_finetune/ode_init/validation.json b/examples/training/consistency_finetune/ode_init/validation.json new file mode 100644 index 000000000..7b14204e5 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/validation.json @@ -0,0 +1,40 @@ +{ + "data": [ + { + "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh index c7e2c251e..379a8a063 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh @@ -6,8 +6,8 @@ export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" -VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +DATA_DIR="data/crush-smol_processed_t2v_old" +VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" NUM_GPUS=4 # export CUDA_VISIBLE_DEVICES=4,5 @@ -52,7 +52,7 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 200 + --validation_steps 50 --validation_sampling_steps "50" --validation_guidance_scale "6.0" ) diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh index 2f1203060..d622c9af0 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh @@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" MODEL_TYPE="wan" DATA_MERGE_PATH="data/crush-smol/merge.txt" -OUTPUT_DIR="data/crush-smol_processed_t2v/" +OUTPUT_DIR="data/crush-smol_processed_t2v_old/" torchrun --nproc_per_node=$GPU_NUM \ fastvideo/pipelines/preprocess/v1_preprocess.py \ diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh index 8b850f78a..dbcdef7d3 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh @@ -1,6 +1,6 @@ #!/bin/bash -GPU_NUM=2 # 2,4,8 +GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATASET_PATH="data/crush-smol/" OUTPUT_DIR="data/crush-smol_processed_t2v/" @@ -14,7 +14,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --preprocess.dataset_type merged \ --preprocess.dataset_path $DATASET_PATH \ --preprocess.dataset_output_dir $OUTPUT_DIR \ - --preprocess.preprocess_video_batch_size 2 \ + --preprocess.preprocess_video_batch_size 8 \ --preprocess.dataloader_num_workers 0 \ --preprocess.max_height 480 \ --preprocess.max_width 832 \ diff --git a/examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh b/examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh new file mode 100644 index 000000000..ae5f85fc8 --- /dev/null +++ b/examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh @@ -0,0 +1,94 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="data/crush-smol_processed_t2v_old" +VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_i2v_finetune" + --output_dir "checkpoints/wan_t2v_i2v_finetune" + --max_train_steps 5000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 2 + --num_latent_t 20 + --num_height 480 + --num_width 832 + --num_frames 77 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 4 + --tp_size 1 + --hsdp_replicate_dim 2 + --hsdp_shard_dim 4 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path $DATA_DIR + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-5 + --mixed_precision "bf16" + --checkpointing_steps 1000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" + --t2v_as_i2v_task True + # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" +) + +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/wan_t2v_i2v_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/t2v_to_i2v/wan_t2v/preprocess_wan_data_t2v.sh b/examples/training/t2v_to_i2v/wan_t2v/preprocess_wan_data_t2v.sh new file mode 100644 index 000000000..4cde222ec --- /dev/null +++ b/examples/training/t2v_to_i2v/wan_t2v/preprocess_wan_data_t2v.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="data/crush-smol/merge.txt" +OUTPUT_DIR="data/crush-smol_processed_t2v_i2v_1_3b/" + +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 2 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 77 \ + --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 "t2v_ode_trajectory" \ No newline at end of file diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 4c1816351..fec5c08d6 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -45,6 +45,8 @@ class PipelineConfig: embedded_cfg_scale: float = 6.0 flow_shift: float | None = None disable_autocast: bool = False + ti2v_task: bool = False + t2v_as_i2v_task: bool = False # Model configuration dit_config: DiTConfig = field(default_factory=DiTConfig) @@ -85,9 +87,6 @@ class PipelineConfig: # DMD parameters dmd_denoising_steps: list[int] | None = field(default=None) - # Wan2.2 TI2V parameters - ti2v_task: bool = False - # Compilation # enable_torch_compile: bool = False @@ -214,6 +213,24 @@ class PipelineConfig: "Comma-separated list of denoising steps (e.g., '1000,757,522')", ) + # TI2V task + parser.add_argument( + f"--{prefix_with_dot}ti2v-task", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}ti2v_task", + default=PipelineConfig.ti2v_task, + help="Enable TI2V", + ) + + # T2V to I2V task + parser.add_argument( + f"--{prefix_with_dot}t2v-as-i2v-task", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}t2v_as_i2v_task", + default=PipelineConfig.t2v_as_i2v_task, + help="Enable T2V to I2V task", + ) + # Add VAE configuration arguments from fastvideo.configs.models.vaes.base import VAEConfig VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") @@ -245,7 +262,9 @@ class PipelineConfig: """ from fastvideo.configs.pipelines.registry import ( get_pipeline_config_cls_from_name) + logger.info("WTF model_path: %s", model_path) pipeline_config_cls = get_pipeline_config_cls_from_name(model_path) + logger.info("pipeline_config_cls: %s", pipeline_config_cls) return cast(PipelineConfig, pipeline_config_cls(model_path=model_path)) diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index ffac8dd5b..132e1b5f9 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -13,7 +13,7 @@ from fastvideo.configs.pipelines.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, - WanT2V720PConfig) + WanT2V720PConfig, SelfForcingWanT2V480PConfig) # isort: on from fastvideo.logger import init_logger from fastvideo.utils import (maybe_download_model_index, @@ -49,6 +49,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), "stepvideo": lambda id: "stepvideo" in id.lower(), + "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(), # Add other pipeline architecture detectors } @@ -60,7 +61,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { WanT2V480PConfig, # Base Wan config as fallback for any Wan variant "wanimagetovideo": WanI2V480PConfig, "wandmdpipeline": FastWan2_1_T2V_480P_Config, - "stepvideo": StepVideoT2VConfig + "stepvideo": StepVideoT2VConfig, + "wancausaldmdpipeline": SelfForcingWanT2V480PConfig, # Other fallbacks by architecture } diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 65b129b4a..c4108613b 100644 --- a/fastvideo/configs/sample/base.py +++ b/fastvideo/configs/sample/base.py @@ -48,6 +48,8 @@ class SamplingParam: # Misc save_video: bool = True return_frames: bool = False + return_trajectory_latents: bool = False # returns all latents for each timestep + return_trajectory_decoded: bool = False # returns decoded latents for each timestep def __post_init__(self) -> None: self.data_type = "video" if self.num_frames > 1 else "image" @@ -205,6 +207,18 @@ class SamplingParam: help= "Path to a JSON file containing V-MoBA specific configurations.", ) + parser.add_argument( + "--return-trajectory-latents", + action="store_true", + default=SamplingParam.return_trajectory_latents, + help="Whether to return the trajectory", + ) + parser.add_argument( + "--return-trajectory-decoded", + action="store_true", + default=SamplingParam.return_trajectory_decoded, + help="Whether to return the decoded trajectory", + ) return parser diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 36ae58e7b..35c651b50 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -3,6 +3,9 @@ from typing import Any, cast import numpy as np import torch +from fastvideo.logger import init_logger + +logger = init_logger(__name__) def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 019ccc527..7d97b21bc 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -344,6 +344,9 @@ class VideoGenerator: "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, "logging_info": logging_info, + "trajectory": output_batch.trajectory_latents, + "trajectory_timesteps": output_batch.trajectory_timesteps, + "trajectory_decoded": output_batch.trajectory_decoded, } def set_lora_adapter(self, diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index fac2e30fb..e52041828 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -158,6 +158,7 @@ class FastVideoArgs: "transformer": True, "vae": True, }) + override_transformer_cls_name: str | None = None # # DMD parameters # dmd_denoising_steps: List[int] | None = field(default=None) @@ -396,6 +397,12 @@ class FastVideoArgs: default=FastVideoArgs.enable_stage_verification, help="Enable input/output verification for pipeline stages", ) + parser.add_argument( + "--override-transformer-cls-name", + type=str, + default=FastVideoArgs.override_transformer_cls_name, + help="Override transformer cls name", + ) # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) @@ -627,6 +634,7 @@ class TrainingArgs(FastVideoArgs): checkpoints_total_limit: int = 0 checkpointing_steps: int = 0 resume_from_checkpoint: str = "" # specify the checkpoint folder to resume from + init_weights_from_safetensors: str = "" # path to safetensors file for initial weight loading # optimizer & scheduler num_train_epochs: int = 0 @@ -688,6 +696,8 @@ class TrainingArgs(FastVideoArgs): log_visualization: bool = False # simulate generator forward to match inference simulate_generator_forward: bool = False + warp_denoising_step: bool = False + intermediate_latents_visualization: bool = False @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": @@ -859,6 +869,10 @@ class TrainingArgs(FastVideoArgs): parser.add_argument("--resume-from-checkpoint", type=str, help="Path to checkpoint to resume from") + parser.add_argument( + "--init-weights-from-safetensors", + type=str, + help="Path to safetensors file for initial weight loading") parser.add_argument("--logging-dir", type=str, help="Directory for logging") @@ -1041,6 +1055,12 @@ class TrainingArgs(FastVideoArgs): "--simulate-generator-forward", action=StoreBoolean, help="Whether to simulate generator forward to match inference") + parser.add_argument( + "--warp-denoising-step", + action=StoreBoolean, + help= + "Whether to warp denoising step according to the scheduler time shift" + ) return parser diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 1eebd9f6f..cab009f72 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -415,6 +415,10 @@ class TransformerLoader(ComponentLoader): raise ValueError( "Model config does not contain a _class_name attribute. " "Only diffusers format is supported.") + logger.info("transformer cls_name: %s", cls_name) + if fastvideo_args.override_transformer_cls_name is not None: + cls_name = fastvideo_args.override_transformer_cls_name + logger.info("Overriding transformer cls_name to %s", cls_name) fastvideo_args.model_paths["transformer"] = model_path @@ -430,6 +434,16 @@ class TransformerLoader(ComponentLoader): if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") + # Check if we should use custom initialization weights + custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None) + use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and + fastvideo_args.training_mode and + not hasattr(fastvideo_args, '_loading_teacher_critic_model')) + + if use_custom_weights: + logger.info("Using custom initialization weights from: %s", custom_weights_path) + safetensors_list = [custom_weights_path] + logger.info("Loading model from %s safetensors files in %s", len(safetensors_list), model_path) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index fd33efe55..9fad823e6 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -121,14 +121,25 @@ class ComposedPipelineBase(ABC): model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, - pipeline_config: str | PipelineConfig | None = None, + pipeline_config: PipelineConfig | None = None, args: argparse.Namespace | None = None, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, **kwargs) -> "ComposedPipelineBase": """ - Load a pipeline from a pretrained model. + Load a pipeline from a pretrained model. + Few different patterns are supported: + - Only provide model_path: + - This will load the pipeline in inference mode. + - The pipeline will be initialized with the default config. + - The pipeline will be initialized with the default modules. + - The pipeline will be initialized with the default stages. + - The pipeline will be initialized with the default stages. + - override the default config using pipeline_config or args or kwargs + - override the default modules using loaded_modules + - override the pipelineconfig + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, If provided, loaded_modules will be used instead of loading from config/pretrained weights. """ @@ -136,9 +147,18 @@ class ComposedPipelineBase(ABC): kwargs['model_path'] = model_path fastvideo_args = FastVideoArgs.from_kwargs(**kwargs) + if pipeline_config is not None: + fastvideo_args.pipeline_config = pipeline_config + if fastvideo_args.override_transformer_cls_name is not None: + pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers") + fastvideo_args.pipeline_config = pipeline_config else: assert args is not None, "args must be provided for training mode" fastvideo_args = TrainingArgs.from_cli_args(args) + if fastvideo_args.override_transformer_cls_name is not None: + pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers") + fastvideo_args.pipeline_config = pipeline_config + logger.info("in 2 Overriding transformer cls name to %s", fastvideo_args.override_transformer_cls_name) # TODO(will): fix this so that its not so ugly fastvideo_args.model_path = model_path for key, value in kwargs.items(): @@ -149,7 +169,8 @@ class ComposedPipelineBase(ABC): # model is loaded with the correct precision. Subsequently we will # use FSDP2's MixedPrecisionPolicy to set the precision for the # fwd, bwd, and other operations' precision. - assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training' + fastvideo_args.pipeline_config.dit_precision = 'fp32' + # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training' logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index fcf88c15a..79a67d4ab 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -148,7 +148,12 @@ class ForwardBatch: modules: dict[str, Any] = field(default_factory=dict) # Final output (after pipeline completion) - output: Any = None + output: torch.Tensor | None = None + return_trajectory_latents: bool = False + return_trajectory_decoded: bool = False + trajectory_timesteps: list[int] | None = None + trajectory_latents: torch.Tensor | None = None + trajectory_decoded: list[torch.Tensor] | None = None # Extra parameters that might be needed by specific pipeline implementations extra: dict[str, Any] = field(default_factory=dict) @@ -207,6 +212,10 @@ class TrainingBatch: infos: list[dict[str, Any]] | None = None mask_lat_size: torch.Tensor | None = None + # ODE trajectory supervision + trajectory_latents: torch.Tensor | None = None + trajectory_timesteps: torch.Tensor | None = None + # Transformer inputs noisy_model_input: torch.Tensor | None = None timesteps: torch.Tensor | None = None @@ -237,6 +246,7 @@ class TrainingBatch: fake_score_loss: float = 0.0 dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + latent_vis_dict: dict[str, torch.Tensor] = field(default_factory=dict) fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py new file mode 100644 index 000000000..7c4a18b9c --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -0,0 +1,654 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +ODE Trajectory Data Preprocessing pipeline implementation. + +This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline +using the modular pipeline architecture. + +Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772 +""" + +import os +from collections.abc import Iterator +from typing import Any + +import numpy as np +import pyarrow as pa +import torch +from PIL import Image +from torch.utils.data import DataLoader +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm import tqdm + +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset import getdataset +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.preprocess.preprocess_pipeline_base import ( + BasePreprocessPipeline) +from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage, + ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) +from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict + +logger = init_logger(__name__) + + +class FlowMatchScheduler: + + order = 1 + + def __init__(self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps) + + def set_timesteps(self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + device=None): + sigma_start = self.sigma_min + \ + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, + num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, + num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / \ + (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp( + -2 * ((x - num_inference_steps / 2) / num_inference_steps)**2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * \ + (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + + def step(self, + model_output, + timestep, + sample, + to_final=False, + return_dict=False, + **kwargs): + assert return_dict is False + assert kwargs == {} + self.sigmas = self.sigmas.to(model_output.device) + self.timesteps = self.timesteps.to(model_output.device) + logger.info('step timestep: %s', timestep) + logger.info('step timestep: %s', timestep.shape) + # timestep is [num_frames] + # timestep_id = torch.argmin( + # (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + # assert timestep.ndim == 1 + # assert timestep.shape[0] == 1 + timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): + sigma_ = 1 if (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) + return (prev_sample, ) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + def add_noise(self, original_samples, noise, timestep): + """ + Diffusion forward corruption process. + Input: + - clean_latent: the clean latent with shape [B, C, H, W] + - noise: the noise with shape [B, C, H, W] + - timestep: the timestep with shape [B] + Output: the corrupted latent with shape [B, C, H, W] + """ + self.sigmas = self.sigmas.to(noise.device) + self.timesteps = self.timesteps.to(noise.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin( + (self.timesteps - timestep.to(self.timesteps.device)).abs()) + weights = self.linear_timesteps_weights[timestep_id] + return weights + + +class PreprocessPipeline_ODE_Trajectory(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]] + + def get_schema_fields(self): + """Get the schema fields for ODE Trajectory pipeline.""" + return [f.name for f in pyarrow_schema_ode_trajectory] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + fastvideo_args.pipeline_config.flow_shift = 5 + logger.info('WTF flow_shift: %s', + fastvideo_args.pipeline_config.flow_shift) + + assert fastvideo_args.pipeline_config.flow_shift == 5 + # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( + # shift=fastvideo_args.pipeline_config.flow_shift) + self.modules["scheduler"] = FlowMatchScheduler( + 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) + logger.info('WTF scheduler timesteps: %s', + self.modules["scheduler"].timesteps) + + 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="vae_encoding_stage", + stage=ImageVAEEncodingStage( + vae=self.get_module("vae"), )) + 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"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self, + )) + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def preprocess_video_and_text_and_trajectory(self, + fastvideo_args: FastVideoArgs, + args): + + for batch_idx, data in enumerate(self.pbar): + if data is None: + continue + + with torch.inference_mode(): + # Filter out invalid samples (those with all zeros) + valid_indices = [] + for i, pixel_values in enumerate(data["pixel_values"]): + if not torch.all( + pixel_values == 0): # Check if all values are zero + valid_indices.append(i) + self.num_processed_samples += len(valid_indices) + + if not valid_indices: + continue + + # Create new batch with only valid samples + valid_data = { + "pixel_values": + torch.stack( + [data["pixel_values"][i] for i in valid_indices]), + "text": [data["text"][i] for i in valid_indices], + "path": [data["path"][i] for i in valid_indices], + "fps": [data["fps"][i] for i in valid_indices], + "duration": [data["duration"][i] for i in valid_indices], + } + + # VAE + with torch.autocast("cuda", dtype=torch.float32): + latents = self.get_module("vae").encode( + valid_data["pixel_values"].to( + get_local_torch_device())).mean + + # Get extra features if needed + extra_features = self.get_extra_features( + valid_data, fastvideo_args) + + batch_captions = valid_data["text"] + logger.info(f"===== batch_captions: {batch_captions}") + # Encode text using the standalone TextEncodingStage API + prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( + batch_captions, + fastvideo_args, + encoder_index=[0], + return_attention_mask=True, + ) + prompt_embeds = prompt_embeds_list[0] + prompt_attention_masks = prompt_masks_list[0] + assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0] + + # # Get sequence lengths from attention masks (number of 1s) + # seq_lens = prompt_attention_mask.sum(dim=1) + + # non_padded_embeds = [] + # non_padded_masks = [] + + # # Process each item in the batch + # for i in range(prompt_embeds.size(0)): + # seq_len = seq_lens[i].item() + # # Slice the embeddings and masks to keep only non-padding parts + # non_padded_embeds.append(prompt_embeds[i, :seq_len]) + # non_padded_masks.append(prompt_attention_mask[i, :seq_len]) + + # Update the tensors with non-padded versions + # prompt_embeds = non_padded_embeds + # prompt_attention_masks = non_padded_masks + # prompt_embeds = prompt_embeds + + # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}") + # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}") + + 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][0] + negative_prompt_attention_mask = negative_prompt_masks_list[ + 0][0] + else: + negative_prompt_embed = None + negative_prompt_attention_mask = None + + trajectory_latents = [] + trajectory_timesteps = [] + trajectory_decoded = [] + for i, (prompt_embed, prompt_attention_mask) in enumerate( + zip(prompt_embeds, prompt_attention_masks, strict=False)): + prompt_embed = prompt_embed.unsqueeze(0) + prompt_attention_mask = prompt_attention_mask.unsqueeze(0) + logger.info("what") + logger.info(f"===== prompt_embed: {prompt_embed.shape}") + logger.info( + f"===== prompt_attention_mask: {prompt_attention_mask.shape}" + ) + # Collect the trajectory data + batch = ForwardBatch( + **shallow_asdict(sampling_params), + # data_type="video", + # seed=args.seed, + # prompt=batch_captions[i], + # prompt_embeds=[prompt_embed], + # prompt_attention_mask=[prompt_attention_mask], + # height=args.max_height, + # width=args.max_width, + # num_frames=81, + # fps=args.train_fps, + # return_trajectory_latents=True, + # guidance_scale=3.0, + # do_classifier_free_guidance=True, + ) + batch.prompt_embeds = [prompt_embed] + batch.prompt_attention_mask = [prompt_attention_mask] + batch.negative_prompt_embeds = [negative_prompt_embed] + batch.negative_attention_mask = [ + negative_prompt_attention_mask + ] + batch.return_trajectory_latents = True + batch.return_trajectory_decoded = False + batch.height = args.max_height + batch.width = args.max_width + batch.num_inference_steps = 48 + # batch.num_frames = 81 + batch.fps = args.train_fps + batch.guidance_scale = 6.0 + batch.do_classifier_free_guidance = True + # fastvideo_args.pipeline_config.ti2v_task = True + + result_batch = self.input_validation_stage( + batch, fastvideo_args) + # result_batch = self.prompt_encoding_stage(result_batch, fastvideo_args) + # result_batch = self.vae_encoding_stage(result_batch, fastvideo_args) + result_batch = self.timestep_preparation_stage( + batch, fastvideo_args) + result_batch = self.latent_preparation_stage( + result_batch, fastvideo_args) + result_batch = self.denoising_stage(result_batch, + fastvideo_args) + result_batch = self.decoding_stage(result_batch, + fastvideo_args) + # trajectory_latents = result_batch.trajectory_latents + trajectory_latents.append( + result_batch.trajectory_latents.cpu()) + trajectory_timesteps.append( + result_batch.trajectory_timesteps.cpu()) + trajectory_decoded.append(result_batch.trajectory_decoded) + + extra_features["trajectory_latents"] = trajectory_latents + extra_features["trajectory_timesteps"] = trajectory_timesteps + logger.info( + f"===== trajectory_latents: {trajectory_latents[0].shape}") + logger.info( + f"===== trajectory_latents len: {len(trajectory_latents)}") + logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}") + logger.info( + f"===== trajectory_timesteps len: {len(trajectory_timesteps)}") + + if batch.return_trajectory_decoded: + logger.info("===== SAVING TRAJECTORY DECODED") + for i, decoded_frames in enumerate(trajectory_decoded): + for j, decoded_frame in enumerate(decoded_frames): + logger.info( + f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}" + ) + save_decoded_latents_as_video( + decoded_frame, + f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", + args.train_fps) + # assert False + # Prepare batch data for Parquet dataset + batch_data = [] + + # Add progress bar for saving outputs + save_pbar = tqdm(enumerate(valid_data["path"]), + desc="Saving outputs", + unit="item", + leave=False) + for idx, video_path in save_pbar: + # Get the corresponding latent and info using video name + latent = latents[idx].cpu() + video_name = os.path.basename(video_path).split(".")[0] + + # Convert tensors to numpy arrays + vae_latent = latent.cpu().numpy() + text_embedding = prompt_embeds[idx].cpu().numpy() + + # Get extra features for this sample if needed + sample_extra_features = {} + if extra_features: + for key, value in extra_features.items(): + logger.info(f"===== key: {key}") + if isinstance(value, torch.Tensor): + logger.info(f"===== value: {value[idx].shape}") + sample_extra_features[key] = value[idx].cpu().numpy( + ) + else: + assert isinstance(value, list) + if isinstance(value[idx], torch.Tensor): + logger.info( + f"===== value in list: {value[idx].shape}") + sample_extra_features[key] = value[idx].cpu( + ).float().numpy() + else: + logger.info("===== value in list: not tensor") + sample_extra_features[key] = value[idx] + # logger.info(f"===== value: not tensor") + # sample_extra_features[key] = value[idx] + + # Create record for Parquet dataset + record = self.create_record( + video_name=video_name, + vae_latent=vae_latent, + text_embedding=text_embedding, + valid_data=valid_data, + idx=idx, + extra_features=sample_extra_features) + batch_data.append(record) + + if batch_data: + # Add progress bar for writing to Parquet dataset + write_pbar = tqdm(total=1, + desc="Writing to Parquet dataset", + unit="batch") + # Convert batch data to PyArrow arrays + arrays = [] + for field in self.get_schema_fields(): + if field.endswith('_bytes'): + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.binary())) + elif field.endswith('_shape'): + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.list_(pa.int32()))) + elif field in ['width', 'height', 'num_frames']: + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.int32())) + elif field in ['duration_sec', 'fps']: + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.float32())) + else: + arrays.append( + pa.array([record[field] for record in batch_data])) + + table = pa.Table.from_arrays(arrays, + names=self.get_schema_fields()) + write_pbar.update(1) + write_pbar.close() + + # Store the table in a list for later processing + if not hasattr(self, 'all_tables'): + self.all_tables = [] + self.all_tables.append(table) + + logger.info("Collected batch with %s samples", len(table)) + + if self.num_processed_samples >= args.flush_frequency: + self._flush_tables(self.num_processed_samples, args, + self.combined_parquet_dir) + self.num_processed_samples = 0 + self.all_tables = [] + + def get_extra_features(self, valid_data: dict[str, Any], + fastvideo_args: FastVideoArgs) -> dict[str, Any]: + + # TODO(will): move these to cpu at some point + self.get_module("vae").to(get_local_torch_device()) + + # generator = torch.Generator(device=get_local_torch_device(), seed=42) + generator = torch.Generator("cpu").manual_seed(42) + + features = {} + """Get CLIP features from the first frame of each video.""" + first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute( + 0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) + _, _, num_frames, height, width = valid_data["pixel_values"].shape + # latent_height = height // self.get_module( + # "vae").spatial_compression_ratio + # latent_width = width // self.get_module("vae").spatial_compression_ratio + + unprocessed_images = [] + pil_images = [] + # Frame has values between -1 and 1 + for frame in first_frame: + frame = (frame + 1) * 127.5 + frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8)) + pil_images.append(frame_pil) + # processed_img = self.get_module("image_processor")( + # images=frame_pil, return_tensors="pt") + unprocessed_images.append(frame_pil) + """Get VAE features from the first frame of each video""" + video_conditions = [] + for frame in unprocessed_images: + + latent = self.vae_encoding_stage.encode_image( + frame, height, width, fastvideo_args, generator) + video_conditions.append(latent) + + features["image_condition_latents"] = video_conditions + features["pil_images"] = pil_images + return features + + def create_record( + self, + video_name: str, + vae_latent: np.ndarray, + text_embedding: np.ndarray, + valid_data: dict[str, Any], + idx: int, + extra_features: dict[str, Any] | None = None) -> dict[str, Any]: + """Create a record for the Parquet dataset with CLIP features.""" + record = super().create_record(video_name=video_name, + vae_latent=vae_latent, + text_embedding=text_embedding, + valid_data=valid_data, + idx=idx, + extra_features=extra_features) + + if extra_features and "image_condition_latents" in extra_features: + image_condition_latents = extra_features["image_condition_latents"] + record.update({ + "image_condition_latents_bytes": + image_condition_latents.tobytes(), + "image_condition_latents_shape": + list(image_condition_latents.shape), + "image_condition_latents_dtype": + str(image_condition_latents.dtype), + }) + else: + record.update({ + "image_condition_latents_bytes": b"", + "image_condition_latents_shape": [], + "image_condition_latents_dtype": "", + }) + + if extra_features and "trajectory_latents" in extra_features: + trajectory_latents = extra_features["trajectory_latents"] + record.update({ + "trajectory_latents_bytes": + trajectory_latents.tobytes(), + "trajectory_latents_shape": + list(trajectory_latents.shape), + "trajectory_latents_dtype": + str(trajectory_latents.dtype), + }) + else: + record.update({ + "trajectory_latents_bytes": b"", + "trajectory_latents_shape": [], + "trajectory_latents_dtype": "", + }) + + if extra_features and "trajectory_timesteps" in extra_features: + trajectory_timesteps = extra_features["trajectory_timesteps"] + record.update({ + "trajectory_timesteps_bytes": + trajectory_timesteps.tobytes(), + "trajectory_timesteps_shape": + list(trajectory_timesteps.shape), + "trajectory_timesteps_dtype": + str(trajectory_timesteps.dtype), + }) + else: + record.update({ + "trajectory_timesteps_bytes": b"", + "trajectory_timesteps_shape": [], + "trajectory_timesteps_dtype": "", + }) + + if extra_features and "pil_image" in extra_features: + pil_image = extra_features["pil_image"] + record.update({ + "pil_image_bytes": pil_image.tobytes(), + "pil_image_shape": list(pil_image.shape), + "pil_image_dtype": str(pil_image.dtype), + }) + else: + record.update({ + "pil_image_bytes": b"", + "pil_image_shape": [], + "pil_image_dtype": "", + }) + + return record + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args): + if not self.post_init_called: + self.post_init() + + self.local_rank = int(os.getenv("RANK", 0)) + os.makedirs(args.output_dir, exist_ok=True) + # Create directory for combined data + self.combined_parquet_dir = os.path.join(args.output_dir, + "combined_parquet_dataset") + os.makedirs(self.combined_parquet_dir, exist_ok=True) + + # Loading dataset + train_dataset = getdataset(args) + + self.preprocess_dataloader = DataLoader( + train_dataset, + batch_size=args.preprocess_video_batch_size, + num_workers=args.dataloader_num_workers, + ) + + self.preprocess_loader_iter = iter(self.preprocess_dataloader) + + self.num_processed_samples = 0 + # Add progress bar for video preprocessing + self.pbar = tqdm(self.preprocess_loader_iter, + desc="Processing videos", + unit="batch", + disable=self.local_rank != 0) + + # Initialize class variables for data sharing + self.video_data: dict[str, Any] = {} # Store video metadata and paths + self.latent_data: dict[str, Any] = {} # Store latent tensors + self.preprocess_video_and_text_and_trajectory(fastvideo_args, args) + + +EntryClass = PreprocessPipeline_ODE_Trajectory diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 82d0514ac..4ca42bb08 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -10,6 +10,8 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger 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_t2v import ( PreprocessPipeline_T2V) from fastvideo.pipelines.preprocess.preprocess_pipeline_text import ( @@ -48,7 +50,6 @@ def main(args) -> None: text_encoder_cpu_offload=False, pipeline_config=pipeline_config, ) - if args.preprocess_task == "t2v": PreprocessPipeline = PreprocessPipeline_T2V elif args.preprocess_task == "i2v": diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 925c570f4..d7578bb5a 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -78,6 +78,8 @@ class CausalDMDDenosingStage(DenoisingStage): torch.tensor([0], dtype=torch.float32))) timesteps = scheduler_timesteps[1000 - timesteps] + else: + assert False, "warp_denoising_step must be true" timesteps = timesteps.to(get_local_torch_device()) logger.info("Using timesteps: %s", timesteps) diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 8f226b5c4..9ed07b5c7 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -50,6 +50,50 @@ class DecodingStage(PipelineStage): result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)]) return result + @torch.no_grad() + def decode(self, latents: torch.Tensor, + fastvideo_args: FastVideoArgs) -> torch.Tensor: + """Decode latents into pixel space.""" + self.vae = self.vae.to(get_local_torch_device()) + latents = latents.to(get_local_torch_device()) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latents = latents / self.vae.scaling_factor.to( + latents.device, latents.dtype) + else: + latents = latents / self.vae.scaling_factor + + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latents += self.vae.shift_factor.to(latents.device, + latents.dtype) + else: + latents += self.vae.shift_factor + + # Decode latents + 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.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + latents = latents.to(vae_dtype) + image = self.vae.decode(latents) + + # Normalize image to [0, 1] range + image = (image / 2 + 0.5).clamp(0, 1) + return image + @torch.no_grad() def forward( self, @@ -66,6 +110,7 @@ class DecodingStage(PipelineStage): Returns: The batch with decoded outputs. """ + # load vae if not already loaded (used for memory constrained devices) pipeline = self.pipeline() if self.pipeline else None if not fastvideo_args.model_loaded["vae"]: loader = VAELoader() @@ -75,58 +120,31 @@ class DecodingStage(PipelineStage): pipeline.add_module("vae", self.vae) fastvideo_args.model_loaded["vae"] = True - self.vae = self.vae.to(get_local_torch_device()) - - latents = batch.latents - # TODO(will): remove this once we add input/output validation for stages - if latents is None: - raise ValueError("Latents must be provided") - - # Skip decoding if output type is latent if fastvideo_args.output_type == "latent": - image = latents + frames = batch.latents else: - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32 - ) and not fastvideo_args.disable_autocast + frames = self.decode(batch.latents, fastvideo_args) - if isinstance(self.vae.scaling_factor, torch.Tensor): - latents = latents / self.vae.scaling_factor.to( - latents.device, latents.dtype) - else: - latents = latents / self.vae.scaling_factor - - # Apply shifting if needed - if (hasattr(self.vae, "shift_factor") - and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - latents += self.vae.shift_factor.to(latents.device, - latents.dtype) - else: - latents += self.vae.shift_factor - - # Decode latents - 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.vae_sp: - # self.vae.enable_parallel() - if not vae_autocast_enabled: - latents = latents.to(vae_dtype) - image = self.vae.decode(latents) - - # Normalize image to [0, 1] range - image = (image / 2 + 0.5).clamp(0, 1) + # decode trajectory latents if needed + if batch.return_trajectory_decoded: + batch.trajectory_decoded = [] + logger.info(f"batch.trajectory_latents.shape: {batch.trajectory_latents.shape}") + assert batch.trajectory_latents is not None, "batch should have trajectory latents" + for idx in range(batch.trajectory_latents.shape[1]): + # bathc.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] + cur_latent = batch.trajectory_latents[:, idx, :, :, :, :] + logger.info(f"cur_latent.shape: {cur_latent.shape}") + cur_timestep = batch.trajectory_timesteps[idx] + logger.info( + f"decoding trajectory latent for timestep: {cur_timestep}") + decoded_frames = self.decode(cur_latent, fastvideo_args) + batch.trajectory_decoded.append(decoded_frames.cpu().float()) # Convert to CPU float32 for compatibility - image = image.cpu().float() + frames = frames.cpu().float() # Update batch with decoded image - batch.output = image + batch.output = frames # Offload models if needed if hasattr(self, 'maybe_free_model_hooks'): diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index b4e182270..b5b2ffe4f 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -140,11 +140,12 @@ class DenoisingStage(PipelineStage): latents = latents[:, :, rank_in_sp_group, :, :, :] batch.latents = latents if batch.image_latent is not None: - image_latent = rearrange(batch.image_latent, - "b c (n t) h w -> b c n t h w", - n=sp_world_size).contiguous() - image_latent = image_latent[:, :, rank_in_sp_group, :, :, :] - batch.image_latent = image_latent + if not fastvideo_args.pipeline_config.ti2v_task and not fastvideo_args.pipeline_config.t2v_as_i2v_task: + image_latent = rearrange(batch.image_latent, + "b c (n t) h w -> b c n t h w", + n=sp_world_size).contiguous() + image_latent = image_latent[:, :, rank_in_sp_group, :, :, :] + batch.image_latent = image_latent # Get timesteps and calculate warmup steps timesteps = batch.timesteps # TODO(will): remove this once we add input/output validation for stages @@ -253,6 +254,9 @@ class DenoisingStage(PipelineStage): patch_size[2]) seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size + trajectory_timesteps: list[int] = [] + trajectory_latents: list[torch.Tensor] = [] + # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -280,14 +284,27 @@ class DenoisingStage(PipelineStage): # Expand latents for I2V latent_model_input = latents.to(target_dtype) - if batch.image_latent is not None: + if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task: assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task" latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1).to(target_dtype) + elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task: + assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task" + if rank_in_sp_group == 0: + logger.info("latent_model_input.shape: %s", + latent_model_input.shape) + latent_model_input = torch.cat([ + batch.image_latent, + latent_model_input[:, :, 1:, :, :], + ], + dim=2).to(target_dtype) + logger.info("latent_model_input.shape: %s", + latent_model_input.shape) assert not torch.isnan( latent_model_input).any(), "latent_model_input contains nan" + if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: timestep = torch.stack([t]).to(get_local_torch_device()) temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() @@ -302,6 +319,13 @@ class DenoisingStage(PipelineStage): latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + if rank_in_sp_group == 0: + latent_model_input = torch.cat([ + batch.image_latent, + latent_model_input[:, :, 1:, :, :], + ], + dim=2).to(target_dtype) # Prepare inputs for transformer guidance_expand = ( @@ -433,6 +457,12 @@ class DenoisingStage(PipelineStage): latents = (1. - mask2[0]) * z + mask2[0] * latents # latents = latents.unsqueeze(0) + # save trajectory latents if needed + if batch.return_trajectory_latents: + trajectory_timesteps.append(t) + # trajectory_latents.append(latents.cpu()) + trajectory_latents.append(latents) + # Update progress bar if i == len(timesteps) - 1 or ( (i + 1) > num_warmup_steps and @@ -441,8 +471,32 @@ class DenoisingStage(PipelineStage): progress_bar.update() # Gather results if using sequence parallelism + trajectory_tensor: torch.Tensor | None = None + if trajectory_latents: + trajectory_tensor = torch.stack(trajectory_latents, dim=1) + else: + trajectory_tensor = None + if sp_group: latents = sequence_model_parallel_all_gather(latents, dim=2) + if batch.return_trajectory_latents: + # logger.info("before stack trajectory_latents.shape: %s", trajectory_latents[0].shape) + logger.info("after stack trajectory_latents.shape: %s", trajectory_tensor.shape) + trajectory_tensor = trajectory_tensor.to( + get_local_torch_device()) + trajectory_tensor = sequence_model_parallel_all_gather( + trajectory_tensor, dim=3) + + if trajectory_tensor is not None: + batch.trajectory_timesteps = torch.tensor(trajectory_timesteps).cpu() + batch.trajectory_latents = trajectory_tensor.cpu() + + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + latents = torch.cat([ + batch.image_latent, + latents[:, :, 1:, :, :], + ], + dim=2) # Update batch with final latents batch.latents = latents diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 4e817e4ae..5c28d3eb4 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -105,6 +105,81 @@ class ImageVAEEncodingStage(PipelineStage): def __init__(self, vae: ParallelTiledVAE) -> None: self.vae: ParallelTiledVAE = vae + def encode_image(self, + image: PIL.Image.Image, + height: int, + width: int, + fastvideo_args: FastVideoArgs, + generator: torch.Generator | None = None) -> torch.Tensor: + """ + Encode image into latent space. + """ + image = self.preprocess( + image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=height, + width=width).to(get_local_torch_device(), dtype=torch.float32) + + # (B, C, H, W) -> (B, C, 1, H, W) + print(f"image.shape: {image.shape}") + image = image.unsqueeze(2) + print(f"after unsqueeze image.shape: {image.shape}") + return self.encode_tensor(image, fastvideo_args, generator) + + def encode_tensor(self, + video_condition: torch.Tensor, + fastvideo_args: FastVideoArgs, + generator: torch.Generator | None = None) -> torch.Tensor: + """ + Encode frames into latent space. + """ + self.vae = self.vae.to(get_local_torch_device()) + video_condition = video_condition.to(device=get_local_torch_device(), + dtype=torch.float32) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + # Encode Image + 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.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + video_condition = video_condition.to(vae_dtype) + encoder_output = self.vae.encode(video_condition) + + if fastvideo_args.mode == ExecutionMode.PREPROCESS: + latent_condition = encoder_output.mean + else: + generator = generator + if generator is None: + raise ValueError("Generator must be provided") + latent_condition = self.retrieve_latents(encoder_output, generator) + + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latent_condition -= self.vae.shift_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition -= self.vae.shift_factor + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latent_condition = latent_condition * self.vae.scaling_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition = latent_condition * self.vae.scaling_factor + + return latent_condition + def forward( self, batch: ForwardBatch, @@ -157,58 +232,29 @@ class ImageVAEEncodingStage(PipelineStage): # (B, C, H, W) -> (B, C, 1, H, W) image = image.unsqueeze(2) - video_condition = torch.cat([ - image, - image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, - image.shape[3], image.shape[4]) - ], - dim=2) - video_condition = video_condition.to(device=get_local_torch_device(), - dtype=torch.float32) + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + # repeat the image self.vae.temporal_compression_ratio times + video_condition = image.repeat(1, 1, + self.vae.temporal_compression_ratio, + 1, 1) + # video_condition = image + logger.info("video_condition.shape: %s", video_condition.shape) + else: + video_condition = torch.cat([ + image, + image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, + image.shape[3], image.shape[4]) + ], + dim=2) - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = ( - vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - - # Encode Image - 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.vae_sp: - # self.vae.enable_parallel() - if not vae_autocast_enabled: - video_condition = video_condition.to(vae_dtype) - encoder_output = self.vae.encode(video_condition) - - if fastvideo_args.mode == ExecutionMode.PREPROCESS: - latent_condition = encoder_output.mean - else: - generator = batch.generator - if generator is None: - raise ValueError("Generator must be provided") - latent_condition = self.retrieve_latents(encoder_output, generator) - - # Apply shifting if needed - if (hasattr(self.vae, "shift_factor") - and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - latent_condition -= self.vae.shift_factor.to( - latent_condition.device, latent_condition.dtype) - else: - latent_condition -= self.vae.shift_factor - - if isinstance(self.vae.scaling_factor, torch.Tensor): - latent_condition = latent_condition * self.vae.scaling_factor.to( - latent_condition.device, latent_condition.dtype) - else: - latent_condition = latent_condition * self.vae.scaling_factor + latent_condition = self.encode_tensor(video_condition, fastvideo_args, + batch.generator) if fastvideo_args.mode == ExecutionMode.PREPROCESS: batch.image_latent = latent_condition + elif fastvideo_args.pipeline_config.t2v_as_i2v_task: + logger.info("latent_condition.shape: %s", latent_condition.shape) + batch.image_latent = latent_condition else: mask_lat_size = torch.ones(1, 1, num_frames, latent_height, latent_width) diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index ff37bb13c..9054b6fbf 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -35,9 +35,15 @@ class InputValidationStage(PipelineStage): """Generate seeds for the inference""" seed = batch.seed num_videos_per_prompt = batch.num_videos_per_prompt + if isinstance(batch.prompt, list): + num_prompts = len(batch.prompt) + else: + num_prompts = 1 + + total_num_videos = num_prompts * num_videos_per_prompt assert seed is not None - seeds = [seed + i for i in range(num_videos_per_prompt)] + seeds = [seed + i for i in range(total_num_videos)] batch.seeds = seeds # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... batch.generator = [ diff --git a/fastvideo/platforms/__init__.py b/fastvideo/platforms/__init__.py index 2decdbbf6..992d61c7e 100644 --- a/fastvideo/platforms/__init__.py +++ b/fastvideo/platforms/__init__.py @@ -82,8 +82,8 @@ def rocm_platform_plugin() -> str | None: logger.info("ROCm platform is available") finally: amdsmi.amdsmi_shut_down() - except Exception as e: - logger.info("ROCm platform is unavailable: %s", e) + except Exception: + pass return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py new file mode 100644 index 000000000..b13d5f475 --- /dev/null +++ b/fastvideo/training/ode_causal_pipeline.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import cast + + +import torch +import torch.nn.functional as F +import numpy as np +import wandb +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only +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.training.training_pipeline import TrainingPipeline +from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.training.training_utils import ( + clip_grad_norm_while_handling_failing_dtensor_cases) + +logger = init_logger(__name__) + + +class ODEInitTrainingPipeline(TrainingPipeline): + """ + Training pipeline for ODE-init using precomputed denoising trajectories. + + Supervision: predict the next latent in the stored trajectory by + - feeding current latent at timestep t into the transformer to predict noise + - stepping the scheduler with the predicted noise + - minimizing MSE to the stored next latent at timestep t_next + """ + + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + # Match the preprocess/generation scheduler for consistent stepping + self.modules["scheduler"] = SelfForcingFlowMatchScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True) + self.modules["scheduler"].set_timesteps(num_inference_steps=1000, + training=True) + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_ode_trajectory_text_only + + 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, "flow_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(f"ARG dmd_denoising_steps: {training_args.pipeline_config.dmd_denoising_steps}") + logger.info( + f"ARG dmd_denoising_steps: {self.training_args.pipeline_config.dmd_denoising_steps}" + ) + self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], + dtype=torch.long, + device=get_local_torch_device()) + # self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], 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(f"timesteps: {timesteps}") + self.dmd_denoising_steps = timesteps[1000 - + self.dmd_denoising_steps] + logger.info( + f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}") + # assert False, "warp_denoising_step must be false" + else: + assert False, "warp_denoising_step must be true" + logger.info("not warped") + self.dmd_denoising_steps = self.dmd_denoising_steps.to( + get_local_torch_device()) + + logger.info(f"denoising_step_list: {self.dmd_denoising_steps}") + + logger.info( + "Initialized ODE-init training pipeline with %s denoising steps", + len(self.dmd_denoising_steps)) + # Cache for nearest trajectory index per DMD step (computed lazily on first batch) + self._cached_closest_idx_per_dmd = None + self.num_train_timestep = self.noise_scheduler.num_train_timesteps + # self.min_timestep = int(self.training_args.min_timestep_ratio * + # self.num_train_timestep) + # self.max_timestep = int(self.training_args.max_timestep_ratio * + # self.num_train_timestep) + # self.real_score_guidance_scale = self.training_args.real_score_guidance_scale + + 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, + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", + 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): # type: ignore[override] + 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'] + 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.to( + device, dtype=torch.bfloat16), trajectory_timesteps.to(device) + + def _get_timestep(self, + min_timestep: int, + max_timestep: int, + batch_size: int, + num_frame: int, + num_frame_per_block: int, + uniform_timestep: bool = False) -> torch.Tensor: + if uniform_timestep: + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, 1], + device=self.device, + dtype=torch.long).repeat(1, num_frame) + return timestep + else: + timestep = torch.randint(min_timestep, + max_timestep, [batch_size, num_frame], + device=self.device, + dtype=torch.long) + # logger.info(f"individual timestep: {timestep}") + # make the noise level the same within every block + timestep = timestep.reshape(timestep.shape[0], -1, + num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + timestep = timestep.reshape(timestep.shape[0], -1) + return timestep + + def _step_predict_next_latent( + self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: + latent_vis_dict = {} + device = get_local_torch_device() + target_latent = traj_latents[:, -1] + + # logger.info(f"traj_latents: {traj_latents.shape}") + # logger.info(f"traj_timesteps: {traj_timesteps.shape}") + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, S, num_frames, num_channels, height, width = traj_latents.shape + + # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps + if self._cached_closest_idx_per_dmd is None: + # Use the first sample's trajectory timesteps; assumed identical across batches + # s_steps = traj_timesteps[0].to(torch.long) # [S] + # dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device) # [K] + # distances_ks: [K, S] = |s_steps - dmd| + # distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs() + # self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu() # [K] + self._cached_closest_idx_per_dmd = torch.tensor( + [0, 12, 24, 36], dtype=torch.long).cpu() + logger.info( + f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}" + ) + logger.info( + f"corresponding timesteps: {self.noise_scheduler.timesteps[self._cached_closest_idx_per_dmd]}" + ) + + # logger.info(f"traj_latents: {traj_latents.shape}") + # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd + # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K] + # Output: [B, K, C, T, H, W] + relevant_traj_latents = torch.index_select( + traj_latents, + dim=1, + index=self._cached_closest_idx_per_dmd.to(traj_latents.device)) + # assert relevant_traj_latents.shape[0] == 1 + + indexes = self._get_timestep( # [B, num_frames] + 0, + len(self.dmd_denoising_steps), + B, + num_frames, + 3, + uniform_timestep=False) + logger.info(f"indexes: {indexes.shape}") + logger.info(f"indexes: {indexes}") + # noisy_input = relevant_traj_latents[indexes] + 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) + # noisy_input = noisy_input.unsqueeze(0) + + # # Sample a single DMD step for the whole batch and fetch its cached nearest S-index + # K = len(self.dmd_denoising_steps) + # dmd_idx = torch.randint(0, K, (1,), device=device) + # logger.info(f"dmd_idx: {dmd_idx}") + # assert self._cached_closest_idx_per_dmd is not None + # nearest_s_idx = int(self._cached_closest_idx_per_dmd[int(dmd_idx.item())]) + # nearest_idx = torch.full((B,), nearest_s_idx, device=device, dtype=torch.long) + + # batch_indices = torch.arange(B, device=device) + # noisy_input = traj_latents[batch_indices, nearest_idx] # [B, C, T, H, W] + # target_latent = traj_latents[batch_indices, -1] # [B, C, T, H, W] + # t = traj_timesteps[batch_indices, nearest_idx] # [B] + + # Scale model input as in inference for consistency with stored trajectories + # noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t) + # logger.info(f"indexes: {indexes.shape}") + # logger.info(f"indexes: {indexes}") + timestep = self.dmd_denoising_steps[indexes] + # logger.info(f"timestep: {timestep.shape}") + # logger.info(f"timestep: {timestep}") + + # 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() + + 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), + "encoder_attention_mask": encoder_attention_mask, + "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) + # logger.info(f"noise_pred: {noise_pred.shape}") + if isinstance(noise_pred, (tuple, list)): + noise_pred = noise_pred[0] + + 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.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() + + # noisy_input = pred_noise_to_pred_video(noise_pred, noisy_input, t, self.modules["scheduler"]) + # next_latent_pred = self.modules["scheduler"].step( + # noise_pred, t, current_latents, return_dict=False)[0] + return pred_video, target_latent, timestep, latent_vis_dict + + def train_one_step(self, training_batch): # type: ignore[override] + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + args = cast(TrainingArgs, self.training_args) + + # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent + + for _ in range(args.gradient_accumulation_steps): + training_batch, traj_latents, traj_timesteps = self._get_next_batch( + training_batch) + text_embeds = training_batch.encoder_hidden_states + text_attention_mask = training_batch.encoder_attention_mask + assert traj_latents.shape[0] == 1 + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, S = traj_latents.shape[0], traj_latents.shape[1] + if S < 2: + raise ValueError("Trajectory must contain at least 2 steps") + + # Sample per-sample current step i in [0, S-2] + + # idx = torch.randint(low=0, high=S - 1, size=(B, ), + # device=traj_latents.device) + + # Gather current latents and next latents + # batch_indices = torch.arange(B, device=traj_latents.device) + # current_latents = traj_latents[batch_indices, idx] # [B, C, T,H,W] + # current_latent = traj_timesteps[:, -1, :, :, :, :] + # target_latents = traj_latents[:, -1, :, :, :, :] + + # Corresponding timesteps t (long) -> cast per sample + # t = traj_timesteps[:, -1, :, :, :, :] + # if t.dtype != torch.long: + # t = t.long() + + # Forward to predict next latent by stepping scheduler with predicted noise + noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent( + traj_latents, traj_timesteps, text_embeds, text_attention_mask) + + 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): + """Add visualization data to wandb logging and save frames to disk.""" + wandb_loss_dict = {} + 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.validation_pipeline.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) + wandb_loss_dict[latent_key] = wandb.Video( + video, fps=16, format="mp4") # change to 16 for Wan2.1 + # Clean up references + del video, pixel_latent, latent + + # Log to wandb + if self.global_rank == 0: + wandb.log(wandb_loss_dict, step=step) + + + # dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict + # fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict + # fake_score_log_keys = ['generator_pred_video'] + # dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video'] + + +def main(args) -> None: + logger.info("Starting ODE-init training pipeline...") + logger.info(f"ARG dmd_denoising_steps: {args.dmd_denoising_steps}") + pipeline = ODEInitTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("ODE-init training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 4e0b4b917..82b096a0f 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -520,6 +520,10 @@ class TrainingPipeline(LoRAPipeline, ABC): self.transformer.train() self.sp_group.barrier() if self.training_args.log_validation and step % self.training_args.validation_steps == 0: + if self.training_args.log_visualization: + self.visualize_intermediate_latents(training_batch, + self.training_args, + step) self._log_validation(self.transformer, self.training_args, step) gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 trainable_params = round( @@ -720,3 +724,8 @@ class TrainingPipeline(LoRAPipeline, ABC): # Re-enable gradients for training training_args.inference_mode = False transformer.train() + + def visualize_intermediate_latents(self, training_batch: TrainingBatch, + training_args: TrainingArgs, step: int): + """Add visualization data to wandb logging and save frames to disk.""" + raise NotImplementedError("Visualize intermediate latents is not implemented for training pipeline") diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index aaef338af..914ef6fd6 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -346,10 +346,14 @@ def load_checkpoint(transformer, """ if not os.path.exists(checkpoint_path): logger.warning("Checkpoint path %s does not exist", checkpoint_path) + assert False return 0 # Extract step number from checkpoint path - step = int(os.path.basename(checkpoint_path).split('-')[-1]) + try: + step = int(os.path.basename(checkpoint_path).split('-')[-1]) + except: + step = 1 if rank == 0: logger.info("Loading checkpoint from step %s", step) diff --git a/fastvideo/training/wan_t2v_i2v_training_pipeline.py b/fastvideo/training/wan_t2v_i2v_training_pipeline.py new file mode 100644 index 000000000..4cbca1483 --- /dev/null +++ b/fastvideo/training/wan_t2v_i2v_training_pipeline.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any + +import torch + +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.configs.sample import SamplingParam +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.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) +from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import ( + WanImageToVideoPipeline) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.utils import is_vsa_available, shallow_asdict + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class WanT2VI2VTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Wan. + """ + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_t2v + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + args_copy.dit_cpu_offload = True + # args_copy.pipeline_config.vae_config.load_encoder = False + # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + pipeline_config = PipelineConfig.from_pretrained( + training_args.model_path) + pipeline_config.vae_config.load_encoder = True + self.validation_pipeline = WanImageToVideoPipeline.from_pretrained( + training_args.model_path, + args=None, + inference_mode=True, + pipeline_config=pipeline_config, + loaded_modules={ + "transformer": self.get_module("transformer"), + }, + required_config_modules=[ + "scheduler", "transformer", "vae", "text_encoder", "tokenizer" + ], + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + dit_cpu_offload=True, + ) + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + latents = batch['vae_latent'] + latents = latents[:, :, :self.training_args.num_latent_t] + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + # clip_features = batch['clip_feature'] + # image_latents = batch['first_frame_latent'] + # image_latents = image_latents[:, :, :self.training_args.num_latent_t] + # pil_image = batch['pil_image'] + infos = batch['info_list'] + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + # training_batch.preprocessed_image = pil_image.to( + # get_local_torch_device()) + # training_batch.image_embeds = clip_features.to(get_local_torch_device()) + # training_batch.image_latents = image_latents.to( + # get_local_torch_device()) + training_batch.infos = infos + + return training_batch + + def _prepare_dit_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning.""" + + # First, call parent method to prepare noise, timesteps, etc. for video latents + training_batch = super()._prepare_dit_inputs(training_batch) + + latents = training_batch.latents + logger.info("latents.shape: %s", latents.shape) + first_frame_latent = latents[:, :, 0, :, :] + + logger.info("first_frame_latent.shape: %s", first_frame_latent.shape) + logger.info("training_batch.noisy_model_input.shape: %s", + training_batch.noisy_model_input.shape) + + training_batch.noisy_model_input = torch.cat([ + first_frame_latent.unsqueeze(2), + training_batch.noisy_model_input[:, :, 1:, :, :] + ], + dim=2) + + return training_batch + + def _build_input_kwargs(self, + training_batch: TrainingBatch) -> TrainingBatch: + + # Image Embeds for conditioning + # image_embeds = training_batch.image_embeds + # assert torch.isnan(image_embeds).sum() == 0 + # image_embeds = image_embeds.to(get_local_torch_device(), + # dtype=torch.bfloat16) + # encoder_hidden_states_image = image_embeds + + # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs + training_batch.input_kwargs = { + "hidden_states": + training_batch.noisy_model_input, + "encoder_hidden_states": + training_batch.encoder_hidden_states, + "timestep": + training_batch.timesteps.to(get_local_torch_device(), + dtype=torch.bfloat16), + "encoder_attention_mask": + training_batch.encoder_attention_mask, + # "encoder_hidden_states_image": + # encoder_hidden_states_image, + "return_dict": + False, + } + return training_batch + + def _prepare_validation_batch(self, sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: dict[str, Any], + num_inference_steps: int) -> ForwardBatch: + sampling_param.prompt = validation_batch['prompt'] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.image_path = validation_batch['video_path'] + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + assert self.seed is not None + sampling_param.seed = self.seed + + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 + sampling_param.num_frames = num_frames + batch = ForwardBatch( + **shallow_asdict(sampling_param), + latents=None, + generator=torch.Generator(device="cpu").manual_seed(self.seed), + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + + return batch + + +def main(args) -> None: + logger.info("Starting training pipeline...") + + pipeline = WanT2VI2VTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 48a1eec7d..3a9b310e3 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -2,6 +2,10 @@ # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py import argparse +from einops import rearrange +import torchvision +import numpy as np +import imageio import ctypes import hashlib import importlib @@ -886,3 +890,16 @@ def best_output_size(w, h, dw, dh, expected_area): return ow1, oh1 else: return ow2, oh2 + + +def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int): + # Process outputs + videos = rearrange(decoded_latents, "b c t h w -> t b c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + imageio.mimsave(output_path, frames, fps=fps, format="mp4") diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index ea52ed687..d2296cdb7 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -1,9 +1,18 @@ -from huggingface_hub import save_torch_state_dict, load_state_dict_from_file -# from safetensors import safetensors -from safetensors.torch import save_file -import torch +# pyright: reportMissingImports=false +from safetensors.torch import save_file, load_file as safe_load_file +import argparse +import os import re +import sys from collections import OrderedDict +from typing import Any, Dict, Mapping, Tuple +import torch + +try: + from huggingface_hub import save_torch_state_dict, load_state_dict_from_file # type: ignore +except Exception: + save_torch_state_dict = None # type: ignore[assignment] + load_state_dict_from_file = None # type: ignore[assignment] _param_names_mapping: dict = { r"^text_embedding\.0\.(.*)$": @@ -135,27 +144,270 @@ _self_forcing_to_diffusers_param_names_mapping: dict = { r"blocks.\1.norm2.\2", } -state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt") -state_dict = state_dict["generator_ema"] -new_state_dict = OrderedDict() -for k, v in state_dict.items(): - new_key = k - for pattern, replacement in _self_forcing_to_diffusers_param_names_mapping.items(): - if re.match(pattern, k): - new_key = re.sub(pattern, replacement, k) - break # Stop at the first match - else: - # print(f"No match found for {k}") - raise ValueError(f"No match found for {k}") - new_state_dict[new_key] = v - if "norm_added_k" in new_key: - dummy_key = new_key.replace("norm_added_k", "norm_added_q") - dummy_value = torch.zeros_like(v) - new_state_dict[dummy_key] = dummy_value -del state_dict +def _replacement_to_regex_template(replacement: str) -> Tuple[str, int]: + r""" + Convert a replacement template like "blocks.\1.attn2.to_q.\2" into a regex pattern + that can be used to match in the reverse direction: "^blocks\.(.*)\.attn2\.to_q\.(.*)$". -save_torch_state_dict( - new_state_dict, - "new2/", - max_shard_size="10GB" -) + Returns the regex template and the number of capture groups. + """ + # First, protect placeholders \1..\9 + placeholder_tokens: Dict[str, str] = {} + group_count = 0 + def _token_for(idx: int) -> str: + return f"__CAP_{idx}__" + + out = replacement + for i in range(1, 10): + token = _token_for(i) + if f"\\{i}" in out: + out = out.replace(f"\\{i}", token) + placeholder_tokens[token] = f"\\{i}" + group_count = max(group_count, i) + + # Escape all regex meta in the literal parts + out = re.escape(out) + # Restore placeholders as (.*) + for token in placeholder_tokens.keys(): + out = out.replace(re.escape(token), "(.*)") + return out, group_count + + +def invert_mapping(forward_mapping: Mapping[str, str]) -> OrderedDict: + """Create a reverse regex mapping by inverting pattern→replacement pairs. + + - Maintains order from the forward mapping + - If the forward pattern is anchored with '$', the reverse is anchored as well + - If forward pattern is prefix-only (no '$'), reverse is also prefix-only + """ + reversed_mapping: "OrderedDict[str, str]" = OrderedDict() + for pattern, replacement in forward_mapping.items(): + # Build reverse pattern from replacement template + reverse_pat_core, _ = _replacement_to_regex_template(replacement) + # Respect anchoring: keep '^' always; add '$' only if original had it + anchored_end = pattern.endswith('$') + reverse_pattern = f"^{reverse_pat_core}" + ("$" if anchored_end else "") + # Reverse replacement must be a literal template with backrefs (\1, \2, ...) + reverse_replacement = _pattern_to_replacement_template(pattern) + reversed_mapping[reverse_pattern] = reverse_replacement + return reversed_mapping + + +def _pattern_to_replacement_template(pattern: str) -> str: + r""" + Convert a regex pattern like "^model.blocks\.(\d+)\.self_attn\.q\.(.*)$" into a replacement + template suitable for re.sub, e.g., "model.blocks.\1.self_attn.q.\2". + Only supports simple capturing groups of the form (.*) or (\d+), which + matches the patterns used in the forward mapping. + """ + # strip anchors + core = pattern + if core.startswith('^'): + core = core[1:] + if core.endswith('$'): + core = core[:-1] + + # replace groups (.*) or (\d+) with backref tokens in increasing order + group_index = 0 + def repl(_m: "re.Match[str]") -> str: + nonlocal group_index + group_index += 1 + return f"\\{group_index}" + + core = re.sub(r"\((?:\.\*|\\d\+)\)", repl, core) + + # unescape literal dots + core = core.replace(r"\.", ".") + return core + + +def select_inner_state_dict(loaded: Mapping[str, Any], key: str = "") -> Tuple[Mapping[str, Any], str]: + if key: + if key not in loaded: + raise KeyError(f"Key '{key}' not found in loaded object. Available keys: {list(loaded.keys())[:20]}") + return loaded[key], key + + # If looks like a state dict (all tensors) + if len(loaded) > 0 and all(torch.is_tensor(v) for v in loaded.values()): + return loaded, "" + + # Common containers + for candidate in ("state_dict", "generator_ema", "model", "ema", "module"): + if candidate in loaded and isinstance(loaded[candidate], Mapping): + inner = loaded[candidate] + if len(inner) > 0 and all(torch.is_tensor(v) for v in inner.values()): + return inner, candidate + + # Fallback: first tensor-dict value + for v in loaded.values(): + if isinstance(v, Mapping) and len(v) > 0 and all(torch.is_tensor(t) for t in v.values()): + return v, "" + + raise ValueError("Could not locate a state_dict (mapping of tensor parameters) in the loaded file.") + + +def convert_state_dict(state_dict: Mapping[str, torch.Tensor], + mapping: Mapping[str, str], + *, + strict: bool = True, + add_norm_added_q_dummy: bool = False) -> Tuple[OrderedDict, Dict[str, int]]: + new_state_dict: "OrderedDict[str, torch.Tensor]" = OrderedDict() + matched_count = 0 + unmatched_count = 0 + dummy_added = 0 + examples = [] # type: ignore[var-annotated] + for k, v in state_dict.items(): + new_key = None + for pattern, replacement in mapping.items(): + if re.match(pattern, k): + new_key = re.sub(pattern, replacement, k) + break + if new_key is None: + if strict: + raise ValueError(f"No mapping rule matched for key: {k}") + else: + new_key = k # keep original + unmatched_count += 1 + else: + matched_count += 1 + new_state_dict[new_key] = v + + if len(examples) < 5: + examples.append((k, new_key)) + + if add_norm_added_q_dummy and "norm_added_k" in new_key: + dummy_key = new_key.replace("norm_added_k", "norm_added_q") + dummy_value = torch.zeros_like(v) + new_state_dict[dummy_key] = dummy_value + dummy_added += 1 + stats = {"matched": matched_count, "unmatched": unmatched_count, "dummy_added": dummy_added} + # store examples count-wise in stats by encoding as counts in print time (examples returned separately not typed) + new_state_dict.__dict__["_examples"] = examples # lightweight attach for printing + return new_state_dict, stats + + +def save_output(new_state_dict: Mapping[str, torch.Tensor], + output: str, + *, + shard: bool = True, + max_shard_size: str = "10GB", + wrapper_key: str = "", + force_pt: bool = False) -> None: + if force_pt: + out_path = coerce_pt_output_path(output) + obj: Dict[str, Any] + if wrapper_key: + obj = {wrapper_key: OrderedDict(new_state_dict)} + else: + # Save raw state_dict mapping + obj = OrderedDict(new_state_dict) # type: ignore[assignment] + torch.save(obj, out_path) + return + + if shard or output.endswith('/') or os.path.isdir(output): + if save_torch_state_dict is None: + raise RuntimeError("Saving shards requires 'huggingface_hub'. Install it or use --single-file.") + out_dir = output if output.endswith('/') else output + '/' + os.makedirs(out_dir, exist_ok=True) + save_torch_state_dict(OrderedDict(new_state_dict), out_dir, max_shard_size=max_shard_size) + else: + # Save a single safetensors file + save_file(OrderedDict(new_state_dict), output) + + +def coerce_pt_output_path(output: str) -> str: + """Ensure output path is a .pt/.pth/.bin file. If a directory or unknown ext, coerce to .pt.""" + if output.endswith('/') or os.path.isdir(output): + os.makedirs(output, exist_ok=True) + return os.path.join(output, 'converted_wan.pt') + lower = output.lower() + if lower.endswith('.pt') or lower.endswith('.pth') or lower.endswith('.bin'): + return output + return output + '.pt' + + +def load_checkpoint(input_path: str) -> Mapping[str, Any]: + if load_state_dict_from_file is not None: + return load_state_dict_from_file(input_path) + # Fallbacks by extension + lower = input_path.lower() + if lower.endswith('.safetensors'): + return safe_load_file(input_path) + # torch serialized + obj = torch.load(input_path, map_location='cpu') + if isinstance(obj, Mapping): + return obj + raise TypeError("Unsupported checkpoint format without huggingface_hub. Provide a mapping-like object.") + + +def parse_args(argv=None): + p = argparse.ArgumentParser(description="Convert WAN <-> Diffusers state_dict key names.") + p.add_argument("--input", "-i", required=True, help="Path to input checkpoint file (.pt/.bin/.safetensors)") + p.add_argument("--output", "-o", required=True, help="Output directory (for shards) or .safetensors file") + p.add_argument("--direction", "-d", choices=["wan-to-diffusers", "diffusers-to-wan"], default="wan-to-diffusers", + help="Conversion direction") + p.add_argument("--inner-key", "-k", default="", help="WAN->Diffusers: unwrap this key. Diffusers->WAN: wrap output under this key.") + p.add_argument("--max-shard-size", default="10GB", help="Shard size when saving to a directory") + p.add_argument("--keep-unmatched", action="store_true", help="Keep keys with no mapping instead of failing") + p.add_argument("--single-file", action="store_true", help="Save a single .safetensors file instead of shards") + return p.parse_args(argv) + + +def main(argv=None): + args = parse_args(argv) + + print(f"[conversion] Direction: {args.direction}") + print(f"[conversion] Input: {args.input}") + save_mode = "torch .pt (forced)" if args.direction == "diffusers-to-wan" else ("single safetensors" if args.single_file else f"sharded (max_shard_size={args.max_shard_size})") + print(f"[conversion] Output: {args.output} [{save_mode}]") + + loaded = load_checkpoint(args.input) + if not isinstance(loaded, Mapping): + raise TypeError("Loaded checkpoint is not a mapping.") + + # Behavior of --inner-key differs by direction + if args.direction == "wan-to-diffusers": + inner, inner_source = select_inner_state_dict(loaded, key=args.inner_key) + print(f"[conversion] Using inner state_dict: {inner_source}") + else: + inner, inner_source = select_inner_state_dict(loaded, key="") # do not unwrap; wrap later if -k is provided + print(f"[conversion] Using inner state_dict: {inner_source} (ignoring --inner-key for unwrap; will wrap on save)") + print(f"[conversion] Parameters found: {len(inner)}") + + if args.direction == "wan-to-diffusers": + mapping = _self_forcing_to_diffusers_param_names_mapping + add_dummy = True + else: + mapping = invert_mapping(_self_forcing_to_diffusers_param_names_mapping) + add_dummy = False + print(f"[conversion] Mapping rules: {len(mapping)}") + + new_state, stats = convert_state_dict(inner, mapping, strict=not args.keep_unmatched, add_norm_added_q_dummy=add_dummy) + examples = getattr(new_state, "_examples", []) + if examples: + print("[conversion] Sample key mappings:") + for old_k, new_k in examples[:5]: + print(f" - {old_k} -> {new_k}") + print(f"[conversion] Converted parameters: {len(new_state)} (matched={stats['matched']}, unmatched_kept={stats['unmatched']})") + if add_dummy: + print(f"[conversion] Added dummy norm_added_q tensors: {stats['dummy_added']}") + + print("[conversion] Saving...") + wrapper_key = args.inner_key if (args.direction == "diffusers-to-wan" and args.inner_key) else "" + if args.direction == "diffusers-to-wan": + if wrapper_key: + print(f"[conversion] Wrapping output under key: {wrapper_key}") + out_path = coerce_pt_output_path(args.output) + print(f"[conversion] Final output path: {out_path}") + save_output(new_state, out_path, shard=False, max_shard_size=args.max_shard_size, wrapper_key=wrapper_key, force_pt=True) + else: + save_output(new_state, args.output, shard=not args.single_file, max_shard_size=args.max_shard_size, wrapper_key=wrapper_key) + print("[conversion] Done.") + + +if __name__ == "__main__": + try: + main() + except Exception as e: + print(f"[conversion] Error: {e}", file=sys.stderr) + sys.exit(1)