diff --git a/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm new file mode 100644 index 000000000..d764bba31 --- /dev/null +++ b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm @@ -0,0 +1,166 @@ +#!/bin/bash +#SBATCH --job-name=dmd_3333 +#SBATCH --partition=main +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=dmd_3333_output/dmd_3333_%j.out +#SBATCH --error=dmd_3333_output/dmd_3333_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +source ~/conda/miniconda/bin/activate +conda activate wei-fv-distill +export HOME="/mnt/weka/home/hao.zhang/wei" + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=32 + +# Model paths for Self-Forcing DMD distillation with Hunyuan1.5: +GENERATOR_MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v" # Updated to Wan2.2 +REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" # Teacher model +FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" +# REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled" +# FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled" # Critic model + +# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" +# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing" +DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed" +DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121" +# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset +# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/" +# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" + +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +training_args=( + --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd # Updated for Wan2.2 + --output_dir "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_3333_1e-5_1e-5_cfg6_corrected_scheduler" + --wandb_run_name "self_forcing_3333_1e-5_1e-5_cfg3.5_corrected_scheduler" + --max_train_steps 4000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 31 + --num_height 480 + --num_width 848 + --enable_gradient_checkpointing_type "full" + --simulate_generator_forward + --log_visualization + --num_frames 121 + --num_frame_per_block 3 # Frame generation block size for self-forcing + # --resume-from-checkpoint "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333/checkpoint-300" + --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_ode_init_1333_new/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors +) + +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +model_args=( + --model_path $GENERATOR_MODEL_PATH + --pretrained_model_name_or_path $GENERATOR_MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_MODEL_PATH +) + +dataset_args=( + --data_path "$DATA_DIR" + --data_path_2 "$DATA_DIR_2" + --dataloader_num_workers 4 +) + +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" # not used for dmd inference + --text-encoder-cpu-offload + # --vae_cpu_offload True +) + +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 100 + --weight_only_checkpointing_steps 500 + --weight_decay 0.01 + --betas '0.0,0.999' + --max_grad_norm 1.0 +) + +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --flow_shift 5 + --seed 1000 + --use_ema True + --ema_decay 0.99 + --ema_start_step 200 +) + +dmd_args=( + --warp_denoising_step + # --dmd_denoising_steps '1000,875,750,625,500,375,250,125' + --dmd_denoising_steps '1000,760,520,280' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --dfake_gen_update_ratio 5 + --real_score_guidance_scale 3.5 + --fake_score_learning_rate 5e-6 + --fake_score_betas '0.0,0.999' +) + +self_forcing_args=( + --independent_first_frame False # Whether to treat first frame independently + --same_step_across_blocks True # Whether to use same denoising step across all blocks + --last_step_only False # Whether to only use the last denoising step + --context_noise 0 # Amount of noise to add during context caching (0 = no noise) +) + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/hy15_self_forcing_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ + "${self_forcing_args[@]}" \ No newline at end of file diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py index 70a66db2b..dbc4e344e 100644 --- a/examples/inference/basic/basic_hy15.py +++ b/examples/inference/basic/basic_hy15.py @@ -2,14 +2,14 @@ from fastvideo import VideoGenerator import json # from fastvideo.configs.sample import SamplingParam -OUTPUT_PATH = "video_samples_hy15_ti2v" +OUTPUT_PATH = "video_samples_hy15_t2v_distilled" def main(): # FastVideo will automatically use the optimal default arguments for the # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. generator = VideoGenerator.from_pretrained( - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled", # FastVideo will automatically handle distributed setup num_gpus=1, use_fsdp_inference=False, # set to True if GPU is out of memory @@ -18,14 +18,26 @@ def main(): text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, - ti2v_task=True, + # init_weights_from_safetensors="/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333_5e-5_2e-6_cfg3.5/checkpoint-500/ema/generator_ema.safetensors" ) + # json_path = "/mnt/weka/home/hao.zhang/wei/FastVideo/data/mixkit_i2v_full_720p.json" + # with open(json_path, 'r') as f: + # data_list = json.load(f)["data"] + + # # Now you can index into data_list however you like + # # For example: data_list[0], data_list[1:3], etc. + # for data in data_list: + # prompt = data["prompt"] + # image_path = data["image_path"] + # generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path, num_frames=121, fps=24) + # return + # prompt = ( # "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot." # ) - # video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16, image_path="data/1.png") + # video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=121, fps=24, image_path="data/1.png") # return prompt = ( @@ -34,7 +46,7 @@ def main(): "natural light filtering through the petals. Mid-shot, warm and cheerful tones." ) - video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=73, fps=16) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) prompt2 = ( "A majestic lion strides across the golden savanna, its powerful frame " @@ -43,7 +55,7 @@ def main(): "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, negative_prompt="", num_frames=73, fps=16) + video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) if __name__ == "__main__": diff --git a/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm index b5fa4d395..b885f0b48 100644 --- a/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm +++ b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm @@ -1,5 +1,5 @@ #!/bin/bash -#SBATCH --job-name=hy15_ode_1333 +#SBATCH --job-name=hy15_ode_3333 #SBATCH --partition=main #SBATCH --nodes=8 #SBATCH --ntasks=8 @@ -7,8 +7,8 @@ #SBATCH --gres=gpu:8 #SBATCH --cpus-per-task=128 #SBATCH --mem=1440G -#SBATCH --output=hy15_ode_1333_output/hy15_ode_1333_%j.out -#SBATCH --error=hy15_ode_1333_output/hy15_ode_1333_%j.err +#SBATCH --output=hy15_ode_3333_output/hy15_ode_3333_%j.out +#SBATCH --error=hy15_ode_3333_output/hy15_ode_3333_%j.err #SBATCH --exclusive # Basic Info @@ -37,20 +37,21 @@ export HOME="/mnt/weka/home/hao.zhang/wei" NUM_GPUS=8 NUM_TOTAL_GPUS=64 -MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" +MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v" DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" # IP=[MASTER NODE IP] # Training arguments training_args=( --tracker_project_name "hy15_ode_init" - --output_dir "/mnt/weka/home/hao.zhang/wei/hy15_ode_init_1333" - --wandb_run_name "hy15_ode_init_1333" - # --resume_from_checkpoint "ode_init_diffusers/" - # --warp_denoising_step + --output_dir "/mnt/weka/home/hao.zhang/wei/hy15_ode_init_1333_new" + --wandb_run_name "hy15_ode_init_1333_new" + --warp_denoising_step + --dmd_denoising_steps '1000,760,520,280,0' --log_visualization --visualization_steps 100 - --max_train_steps 6001 + --max_train_steps 4000 --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 2 @@ -58,7 +59,6 @@ training_args=( --num_height 480 --num_width 848 --num_frames 121 - --dmd_denoising_steps "1000,750,500,250" --enable_gradient_checkpointing_type "full" ) @@ -75,6 +75,7 @@ parallel_args=( model_args=( --model_path $MODEL_PATH --pretrained_model_name_or_path $MODEL_PATH + --text-encoder-cpu-offload ) # Dataset arguments @@ -88,11 +89,19 @@ optimizer_args=( --learning_rate 1e-5 --mixed_precision "bf16" --training_state_checkpointing_steps 400 - --weight_only_checkpointing_steps 200 --weight_decay 0.01 --max_grad_norm 1.0 ) +validation_args=( + # --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 10 + --validation_sampling_steps "4" + --validation_guidance_scale "1.0" # not used for dmd inference + # --vae_cpu_offload True +) + # Miscellaneous arguments miscellaneous_args=( --inference_mode False diff --git a/fastvideo/configs/models/dits/hunyuanvideo15.py b/fastvideo/configs/models/dits/hunyuanvideo15.py index 34bf7d834..3df422a35 100644 --- a/fastvideo/configs/models/dits/hunyuanvideo15.py +++ b/fastvideo/configs/models/dits/hunyuanvideo15.py @@ -26,6 +26,9 @@ class HunyuanVideo15ArchConfig(DiTArchConfig): param_names_mapping: dict = field( default_factory=lambda: { + r"^cond_type_embed\.(.*)$": + r"cond_type_embed.\1", + # 1. context_embedder.time_text_embed submodules (specific rules, applied first): r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1", @@ -55,6 +58,17 @@ class HunyuanVideo15ArchConfig(DiTArchConfig): r"txt_in.refiner_blocks.\1.mlp.fc_out.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.self_attn_qkv\.(.*)$": + r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", + + r"^image_embedder\.linear_1\.(.*)$": + r"image_embedder.linear_1.\1", + r"^image_embedder\.linear_2\.(.*)$": + r"image_embedder.linear_2.\1", + r"^image_embedder\.norm_in\.(.*)$": + r"image_embedder.norm_in.\1", + r"^image_embedder\.norm_out\.(.*)$": + r"image_embedder.norm_out.\1", # 2. txt_in_2 mapping: r"^context_embedder_2\.(.*)$": diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index 29e1c5b45..312986438 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/fastvideo/configs/pipelines/hunyuan15.py @@ -127,9 +127,12 @@ class Hunyuan15T2V480PConfig(PipelineConfig): vae_tiling: bool = True def __post_init__(self): - self.vae_config.load_encoder = False + self.vae_config.load_encoder = True self.vae_config.load_decoder = True +@dataclass +class Hunyuan15DistilledI2V480PConfig(Hunyuan15T2V480PConfig): + flow_shift: int = 7 @dataclass class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig): @@ -137,3 +140,12 @@ class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig): # HunyuanConfig-specific parameters with defaults flow_shift: int = 9 + +@dataclass +class SelfForcingHunyuan15T2V480PConfig(Hunyuan15T2V480PConfig): + is_causal: bool = True + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 760, 520, 280]) + # dmd_denoising_steps: list[int] | None = field( + # default_factory=lambda: [1000, 875, 750, 625, 500, 375, 250, 125]) + warp_denoising_step: bool = True \ No newline at end of file diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 3228f400a..0a90a7198 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -8,9 +8,9 @@ from fastvideo.configs.pipelines.base import PipelineConfig from fastvideo.configs.pipelines.cosmos import CosmosConfig from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig -from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig +from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig, SelfForcingHunyuan15T2V480PConfig, Hunyuan15DistilledI2V480PConfig from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.configs.pipelines.turbodiffusion import ( @@ -35,6 +35,10 @@ logger = init_logger(__name__) PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig, "hunyuanvideo-community/HunyuanVideo": HunyuanConfig, + "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v": + SelfForcingHunyuan15T2V480PConfig, + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled": + Hunyuan15DistilledI2V480PConfig, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": Hunyuan15T2V480PConfig, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v": diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py index 941b9ea22..aadabe02d 100644 --- a/fastvideo/configs/sample/hunyuan15.py +++ b/fastvideo/configs/sample/hunyuan15.py @@ -17,8 +17,8 @@ class Hunyuan15_480P_SamplingParam(SamplingParam): guidance_scale: float = 6.0 prompt_attention_mask: list = field(default_factory=list) negative_attention_mask: list = field(default_factory=list) - sigmas: list[float] | None = field( - default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1])) + # sigmas: list[float] | None = field( + # default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1])) negative_prompt: str = "" @@ -27,6 +27,12 @@ class Hunyuan15_480P_SamplingParam(SamplingParam): self.sigmas = list( np.linspace(1.0, 0.0, self.num_inference_steps + 1)[:-1]) +@dataclass +class Hunyuan15_480P_Distilled_SamplingParam(Hunyuan15_480P_SamplingParam): + num_inference_steps: int = 8 + fps: int = 24 + + guidance_scale: float = 1.0 @dataclass class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam): diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 22014b06b..95aabc7be 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -5,8 +5,8 @@ from typing import Any from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam, HunyuanSamplingParam) -from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam +from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam, Hunyuan15_480P_Distilled_SamplingParam from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam @@ -44,6 +44,10 @@ logger = init_logger(__name__) SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v": + Hunyuan15_480P_SamplingParam, + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled": + Hunyuan15_480P_Distilled_SamplingParam, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": Hunyuan15_480P_SamplingParam, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v": diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py index 1bc86dd7d..fa8b8676a 100644 --- a/fastvideo/dataset/dataloader/record_schema.py +++ b/fastvideo/dataset/dataloader/record_schema.py @@ -127,7 +127,10 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: def ode_text_only_record_creator( video_name: str, text_embedding: np.ndarray, caption: str, trajectory_latents: np.ndarray, - trajectory_timesteps: np.ndarray) -> dict[str, Any]: + trajectory_timesteps: np.ndarray, + text_mask: np.ndarray | None = None, + text_embedding_2: np.ndarray | None = None, + text_mask_2: np.ndarray | None = None) -> dict[str, Any]: """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only. Args: @@ -165,6 +168,25 @@ def ode_text_only_record_creator( "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype), }) + if text_embedding_2 is not None: + record.update({ + "text_embedding_2_bytes": text_embedding_2.tobytes(), + "text_embedding_2_shape": list(text_embedding_2.shape), + "text_embedding_2_dtype": str(text_embedding_2.dtype), + }) + if text_mask is not None: + record.update({ + "text_mask_bytes": text_mask.tobytes(), + "text_mask_shape": list(text_mask.shape), + "text_mask_dtype": str(text_mask.dtype), + }) + if text_mask_2 is not None: + record.update({ + "text_mask_2_bytes": text_mask_2.tobytes(), + "text_mask_2_shape": list(text_mask_2.shape), + "text_mask_2_dtype": str(text_mask_2.dtype), + }) + return record @@ -187,4 +209,4 @@ def text_only_record_creator(text_name: str, text_embedding: np.ndarray, "text_embedding_dtype": str(text_embedding.dtype), "caption": caption, } - return record + return record \ No newline at end of file diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index d94727bd6..8661cb1b7 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -17,6 +17,7 @@ from PIL import Image from transformers import AutoTokenizer from fastvideo.logger import init_logger +from fastvideo.utils import maybe_download_model logger = init_logger(__name__) @@ -672,7 +673,7 @@ class TextDataset(torch.utils.data.IterableDataset, self.seed = seed # Initialize tokenizer - tokenizer_path = os.path.join(args.model_path, "tokenizer") + tokenizer_path = os.path.join(maybe_download_model(args.model_path), "tokenizer") tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=args.cache_dir) @@ -775,4 +776,4 @@ class TextDataset(torch.utils.data.IterableDataset, def load_state_dict(self, state_dict: dict[str, Any]) -> None: """Load state dict from checkpoint.""" - self.processed_batches = state_dict["processed_batches"] + self.processed_batches = state_dict["processed_batches"] \ No newline at end of file diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 18bc5e532..5564568d2 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -830,6 +830,7 @@ class TrainingArgs(FastVideoArgs): precedence. """ data_path: str = "" + data_path_2: str | None = None dataloader_num_workers: int = 0 num_height: int = 0 num_width: int = 0 @@ -999,6 +1000,10 @@ class TrainingArgs(FastVideoArgs): type=str, required=True, help="Path to parquet files") + parser.add_argument("--data-path-2", + type=str, + required=False, + help="Path to parquet files") parser.add_argument("--dataloader-num-workers", type=int, required=True, diff --git a/fastvideo/models/dits/causal_hunyuanvideo15.py b/fastvideo/models/dits/causal_hunyuanvideo15.py index 408b72b56..50eb1dcb7 100644 --- a/fastvideo/models/dits/causal_hunyuanvideo15.py +++ b/fastvideo/models/dits/causal_hunyuanvideo15.py @@ -16,6 +16,12 @@ from typing import Any, Dict, Optional, List import math import torch +# import torch._dynamo +# torch._dynamo.config.cache_size_limit = 128 +# try: +# torch._dynamo.config.recompile_limit = 128 +# except AttributeError: +# pass import torch.nn as nn import torch.nn.functional as F @@ -25,6 +31,8 @@ flex_attention = torch.compile( flex_attention, dynamic=False, mode="default") import torch.distributed as dist +from fastvideo.attention import LocalAttention +from fastvideo.forward_context import set_forward_context from fastvideo.configs.models.dits import HunyuanVideo15Config from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) @@ -48,6 +56,24 @@ from fastvideo.models.dits.hunyuanvideo15 import ( logger = init_logger(__name__) +class CacheAppend(torch.autograd.Function): + """ + KV cache with shape [batch, seq_len, heads, head_dim]. + """ + @staticmethod + def forward(ctx, storage, active_cache, x, start, end): + # Ensure storage has the same dtype as x + storage.data[:, start:end] = x + ctx.save_for_backward(storage.to(x.dtype)) + ctx.start = start + ctx.end = end + return storage[:, :end].to(x.dtype) # Ensure returned value has same dtype as input + + @staticmethod + def backward(ctx, grad_output): + start = ctx.start + end = ctx.end + return None, grad_output[:, :start], grad_output[:, start:end], None, None class MMDoubleStreamBlock(nn.Module): """ @@ -69,7 +95,7 @@ class MMDoubleStreamBlock(nn.Module): self.deterministic = False self.num_attention_heads = num_attention_heads - head_dim = hidden_size // num_attention_heads + self.head_dim = hidden_size // num_attention_heads mlp_hidden_dim = int(hidden_size * mlp_ratio) # Image modulation components @@ -100,8 +126,8 @@ class MMDoubleStreamBlock(nn.Module): params_dtype=dtype, prefix=f"{prefix}.img_attn_qkv") - self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) - self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + self.img_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype) + self.img_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype) self.img_attn_proj = ReplicatedLinear(hidden_size, hidden_size, @@ -143,8 +169,8 @@ class MMDoubleStreamBlock(nn.Module): params_dtype=dtype) # QK norm layers for text - self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) - self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + self.txt_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype) + self.txt_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype) self.txt_attn_proj = ReplicatedLinear(hidden_size, hidden_size, @@ -153,16 +179,90 @@ class MMDoubleStreamBlock(nn.Module): self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) - def forward( + self.max_attention_size = 31 * 1590 + self.k_cache = None + self.v_cache = None + + self.attn = LocalAttention( + num_heads=self.num_attention_heads, + head_size=self.head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn" + ) + + def forward_txt( self, - img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, - vec_txt: torch.Tensor, + cache_txt: bool = False, + ): + txt_mod_outputs = self.txt_mod(vec) + ( + txt_attn_shift, + txt_attn_scale, + txt_attn_gate, + txt_mlp_shift, + txt_mlp_scale, + txt_mlp_gate, + ) = torch.chunk(txt_mod_outputs, 6, dim=-1) + + # Prepare text for attention using fused operation + txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) + + # Get QKV for text + txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) + batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] + + # Split QKV + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, + self.num_attention_heads, -1) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, + 2] + # Apply QK-Norm if needed + txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) + + t_kv = {} + if cache_txt: + t_kv["k_txt"] = txt_k + t_kv["v_txt"] = txt_v + + txt_attn = self.attn(txt_q, txt_k, txt_v) + # Process text attention output + txt_attn_out, _ = self.txt_attn_proj( + txt_attn.reshape(batch_size, text_seq_len, -1)) + + # Use fused operation for residual connection, normalization, and modulation + txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( + txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) + + # Process text MLP + txt_mlp_out = self.txt_mlp(txt_mlp_input) + txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) + + return txt, t_kv + + def forward_vision( + self, + img: torch.Tensor, + vec: torch.Tensor, freqs_cis: tuple, block_mask: BlockMask, kv_cache: dict | None = None, + txt_kv_cache: list | None = None, + current_start: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: + if kv_cache is not None: + # Then we are in autoregressive mode without block-wise attention mask + if self.k_cache is None: + del self.k_cache + self.register_buffer("k_cache", torch.empty(1, self.max_attention_size, self.num_attention_heads, self.head_dim, device=img.device, dtype=img.dtype), persistent=False) + self.k_cache.requires_grad_(True) + if self.v_cache is None: + del self.v_cache + self.register_buffer("v_cache", torch.empty(1, self.max_attention_size, self.num_attention_heads, self.head_dim, device=img.device, dtype=img.dtype), persistent=False) + self.v_cache.requires_grad_(True) + # Process modulation vectors if vec.dim() == 3: img_mod_outputs = self.img_mod(vec).unflatten(dim=-1, sizes=(6, -1)) @@ -185,16 +285,6 @@ class MMDoubleStreamBlock(nn.Module): img_mlp_gate, ) = torch.chunk(img_mod_outputs, 6, dim=-1) - txt_mod_outputs = self.txt_mod(vec_txt) - ( - txt_attn_shift, - txt_attn_scale, - txt_attn_gate, - txt_mlp_shift, - txt_mlp_scale, - txt_mlp_gate, - ) = torch.chunk(txt_mod_outputs, 6, dim=-1) - # Prepare image for attention using fused operation img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale) # Get QKV for image @@ -215,31 +305,16 @@ class MMDoubleStreamBlock(nn.Module): cos, sin = freqs_cis img_q = _apply_rotary_emb(img_q, cos, sin, is_neox_style=False) img_k = _apply_rotary_emb(img_k, cos, sin, is_neox_style=False) - # Prepare text for attention using fused operation - txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) - - # Get QKV for text - txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) - batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] - - # Split QKV - txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, - self.num_attention_heads, -1) - txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, - 2] - # Apply QK-Norm if needed - txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) - txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) # Apply flex_attention # Does not support SP padding for now if kv_cache is None: - txt_seq_len = txt_q.shape[1] - q = torch.cat([img_q, txt_q], dim=1) - k = torch.cat([img_k, txt_k], dim=1) - v = torch.cat([img_v, txt_v], dim=1) + q = img_q + k = torch.cat([img_k, txt_kv_cache["k_txt"]], dim=1) + v = torch.cat([img_v, txt_kv_cache["v_txt"]], dim=1) # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_kv_length = math.ceil(k.shape[1] / 128) * 128 - k.shape[1] padded_roped_query = torch.cat( [q, torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], @@ -248,27 +323,44 @@ class MMDoubleStreamBlock(nn.Module): ) padded_roped_key = torch.cat( - [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + [k, torch.zeros([k.shape[0], padded_kv_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype)], dim=1 ) padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + [v, torch.zeros([v.shape[0], padded_kv_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)], dim=1 ) - attn_out = flex_attention( + img_attn = flex_attention( query=padded_roped_query.transpose(2, 1), key=padded_roped_key.transpose(2, 1), value=padded_v.transpose(2, 1), block_mask=block_mask )[:, :, :-padded_length].transpose(2, 1) - assert attn_out.shape[1] == txt_seq_len + image_seq_len - img_attn = attn_out[:, :image_seq_len, :] - txt_attn = attn_out[:, image_seq_len:, :] + assert img_attn.shape[1] == image_seq_len + else: + current_end = current_start + img_q.shape[1] + num_new_tokens = img_q.shape[1] + assert kv_cache["local_end_index"].item() == kv_cache["global_end_index"].item(), "local_end_index and global_end_index must be the same" + # Assign new keys/values directly up to current_end + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + img_local_k = CacheAppend.apply(self.k_cache, kv_cache["k"], img_k, local_start_index, local_end_index) + img_local_v = CacheAppend.apply(self.v_cache, kv_cache["v"], img_v, local_start_index, local_end_index) + + q = img_q + k = torch.cat([img_local_k, txt_kv_cache["k_txt"]], dim=1) + v = torch.cat([img_local_v, txt_kv_cache["v_txt"]], dim=1) + img_attn = self.attn(q, k, v) + + kv_cache["k"] = img_local_k + kv_cache["v"] = img_local_v + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) img_attn_out, _ = self.img_attn_proj( img_attn.view(batch_size, image_seq_len, -1)) @@ -280,19 +372,20 @@ class MMDoubleStreamBlock(nn.Module): img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) - # Process text attention output - txt_attn_out, _ = self.txt_attn_proj( - txt_attn.reshape(batch_size, text_seq_len, -1)) + return img - # Use fused operation for residual connection, normalization, and modulation - txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( - txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) - - # Process text MLP - txt_mlp_out = self.txt_mlp(txt_mlp_input) - txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) - - return img, txt + def forward( + self, + txt_inference=False, + vision_inference=False, + **kwargs + ): + if txt_inference: + return self.forward_txt(**kwargs) + elif vision_inference: + return self.forward_vision(**kwargs) + else: + raise ValueError("txt_inference and vision_inference cannot be both False") class CausalHunyuanVideo15Transformer3DModel(CachableDiT): @@ -373,155 +466,18 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): self.__post_init__() - @staticmethod - def _prepare_blockwise_causal_attn_mask( - device: torch.device | str, num_frames: int = 21, - frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, - txt_attn_mask: torch.Tensor = None - ) -> BlockMask: - """ - we will divide the token sequence into the following format - [1 latent frame] [1 latent frame] ... [1 latent frame] - We use flexattention to construct the attention mask - """ - img_seq_len = num_frames * frame_seqlen - batch_size, txt_seq_len = txt_attn_mask.shape - total_len = img_seq_len + txt_seq_len - - # we do right padding to get to a multiple of 128 - padded_total_len = math.ceil(total_len / 128) * 128 - - ends = torch.zeros(img_seq_len, - device=device, dtype=torch.long) - - # Block-wise causal mask will attend to all elements that are before the end of the current chunk - frame_indices = torch.arange( - start=frame_seqlen, - end=img_seq_len, - step=frame_seqlen * num_frame_per_block, - device=device - ) - frame_indices = torch.cat([torch.tensor([0], device=device), frame_indices]) - - for i, tmp in enumerate(frame_indices): - if i == 0: - ends[tmp:tmp + frame_seqlen] = tmp + frame_seqlen - else: - ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ - frame_seqlen * num_frame_per_block - - # Generate the 2D text attention mask - attention_mask_bool = txt_attn_mask.bool() - self_attn_mask_1 = attention_mask_bool.view(batch_size, 1, 1, txt_seq_len) - self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) - txt_attn_2d_mask = (self_attn_mask_1 & self_attn_mask_2).bool() - - # Flatten tensors for safer indexing under vmap to avoid .item() calls - txt_attn_mask_flat = attention_mask_bool.reshape(-1) - txt_attn_2d_mask_flat = txt_attn_2d_mask.reshape(-1) - ends_flat = ends.long() - - def attention_mask(b, h, q_idx, kv_idx): - # Use flattened indexing to avoid vmap issues with multi-dimensional indexing - q_idx_long = q_idx.long() - kv_idx_long = kv_idx.long() - b_long = b.long() - - q_is_img = q_idx_long < img_seq_len - kv_is_img = kv_idx_long < img_seq_len - - q_idx_img_safe = q_idx_long.clamp(0, img_seq_len - 1) - kv_idx_img_safe = kv_idx_long.clamp(0, img_seq_len - 1) - - q_idx_txt_safe = (q_idx_long - img_seq_len).clamp(0, txt_seq_len - 1) - kv_idx_txt_safe = (kv_idx_long - img_seq_len).clamp(0, txt_seq_len - 1) - - # Case 1: Image attending to image - q_end = ends_flat[q_idx_img_safe] - if local_attn_size == -1: - res_img_img = (kv_idx_img_safe < q_end) | (q_idx_img_safe == kv_idx_img_safe) - else: - res_img_img = ((kv_idx_img_safe < q_end) & (kv_idx_img_safe >= (q_end - local_attn_size * frame_seqlen))) | (q_idx_img_safe == kv_idx_img_safe) - - # Case 2: Text attending to text - # txt_attn_2d_mask_flat index: b * 1 * txt_seq_len * txt_seq_len + 0 * txt_seq_len * txt_seq_len + q * txt_seq_len + kv - idx_txt_txt = b_long * (txt_seq_len * txt_seq_len) + q_idx_txt_safe * txt_seq_len + kv_idx_txt_safe - res_txt_txt = txt_attn_2d_mask_flat[idx_txt_txt] - - # Case 3: Image attending to text - # txt_attn_mask_flat index: b * txt_seq_len + kv - idx_img_txt = b_long * txt_seq_len + kv_idx_txt_safe - res_img_txt = txt_attn_mask_flat[idx_img_txt] - - # Case 4: Text attending to image - # index: b * txt_seq_len + q - idx_txt_img = b_long * txt_seq_len + q_idx_txt_safe - res_txt_img = txt_attn_mask_flat[idx_txt_img] - - mask = torch.where( - q_is_img, - torch.where(kv_is_img, res_img_img, res_img_txt), - torch.where(kv_is_img, res_txt_img, res_txt_txt) - ) - - return mask & (q_idx_long < total_len) & (kv_idx_long < total_len) - - block_mask = create_block_mask(attention_mask, B=batch_size, H=None, Q_LEN=padded_total_len, - KV_LEN=padded_total_len, _compile=False, device=device) - - if not dist.is_initialized() or dist.get_rank() == 0: - print( - f" cache a block wise causal mask with block size of {num_frame_per_block} frames") - print(block_mask) - - return block_mask - - def forward( + def get_text_and_mask( self, - hidden_states: torch.Tensor, - encoder_hidden_states: List[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: List[torch.Tensor], - encoder_attention_mask: List[torch.Tensor], - guidance: Optional[torch.Tensor] = None, - timestep_r: Optional[torch.LongTensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_2: torch.Tensor, + encoder_attention_mask: torch.Tensor, + encoder_attention_mask_2: torch.Tensor, + encoder_hidden_states_image: torch.Tensor, + timestep: torch.Tensor, ): - encoder_hidden_states_image = encoder_hidden_states_image[0] - encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states - encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # 1. RoPE - # Get rotary embeddings - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size, - self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta) - freqs_cos = freqs_cos.to(hidden_states.device) - freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - - # 2. Conditional embeddings - if timestep.dim() == 2: - ts_seq_len = timestep.shape[1] - temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len) - else: - temb = self.time_in(timestep, timestep_r=timestep_r) - # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim] - - # 2.1 Timestep embedding just for text - temb_txt = self.time_in(torch.tensor([0], device=hidden_states.device, dtype=hidden_states.dtype), timestep_r=timestep_r) - - hidden_states = self.img_in(hidden_states) - + batch_size, txt_seq_len = encoder_hidden_states.shape[0], encoder_hidden_states.shape[1] # qwen text embedding - txt_timestep = torch.tensor([0], device=hidden_states.device, dtype=hidden_states.dtype) - encoder_hidden_states = self.txt_in(encoder_hidden_states, txt_timestep, encoder_attention_mask) + encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask) encoder_hidden_states_cond_emb = self.cond_type_embed( torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long) @@ -607,40 +563,219 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): encoder_hidden_states = torch.stack(new_encoder_hidden_states) encoder_attention_mask = torch.stack(new_encoder_attention_mask) + assert encoder_hidden_states.shape[0] == 1 + return encoder_hidden_states, encoder_attention_mask + + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, + text_seq_len: int = 0 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + total_kv_length = total_length + text_seq_len + + total_length_tensor = torch.tensor(total_length, device=device) + total_kv_length_tensor = torch.tensor(total_kv_length, device=device) + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + kv_padded_length = math.ceil(total_kv_length / 128) * 128 - total_kv_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + # start=frame_seqlen, + start=0, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + # frame_indices = torch.cat([torch.tensor([0], device=device), frame_indices]) + + for i, tmp in enumerate(frame_indices): + # if i == 0: + # ends[tmp:tmp + frame_seqlen] = tmp + frame_seqlen + # else: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor)) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor)) + # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_kv_length + kv_padded_length, _compile=False, device=device) + + # if not dist.is_initialized() or dist.get_rank() == 0: + # print( + # f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + # print(block_mask) + + # import imageio + # import numpy as np + # from torch.nn.attention.flex_attention import create_mask + + # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + # padded_length, KV_LEN=total_length + padded_length, device=device) + # import cv2 + # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + def forward_txt( + self, + encoder_hidden_states: List[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: Optional[List[torch.Tensor]] = None, + encoder_attention_mask: Optional[List[torch.Tensor]] = None, + guidance: Optional[torch.Tensor] = None, + timestep_r: Optional[torch.LongTensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + cache_txt: bool = False, + **kwargs + ): + # Check that the timestep is only consisted of 0s + assert torch.all(timestep == 0), "Timestep for txt must be only consisted of 0s" + + if cache_txt: + _kv_cache_new = [] + transformer_num_layers = len(self.double_blocks) + for _ in range(transformer_num_layers): + _kv_cache_new.append( + {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None} + ) + + encoder_hidden_states_image = encoder_hidden_states_image[0] + encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states + encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask + + # 2. Conditional embeddings + if timestep.dim() == 2: + ts_seq_len = timestep.shape[1] + temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len) + else: + temb = self.time_in(timestep, timestep_r=timestep_r) + # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim] + + encoder_hidden_states, encoder_attention_mask = self.get_text_and_mask( + encoder_hidden_states, + encoder_hidden_states_2, + encoder_attention_mask, + encoder_attention_mask_2, + encoder_hidden_states_image, + timestep + ) + + encoder_hidden_states = encoder_hidden_states[encoder_attention_mask.bool().to(encoder_hidden_states.device)].unsqueeze(0) + + # 4. Transformer blocks + for index, block in enumerate(self.double_blocks): + encoder_hidden_states, t_kv = block( + txt_inference=True, + vision_inference=False, + txt=encoder_hidden_states, + vec=temb, + cache_txt=cache_txt, + ) + + if cache_txt: + _kv_cache_new[index]["k_txt"] = t_kv["k_txt"] + _kv_cache_new[index]["v_txt"] = t_kv["v_txt"] + + if cache_txt: + return _kv_cache_new + + def forward_vision( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + guidance: Optional[torch.Tensor] = None, + timestep_r: Optional[torch.LongTensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + kv_cache: dict | None = None, + txt_kv_cache: list | None = None, + current_start: int = 0, + rope_start_idx: int = 0, + ): + assert txt_kv_cache is not None, "txt_kv_cache must be provided" + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # 1. RoPE + # Get rotary embeddings + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size, + self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta, start_frame=rope_start_idx) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + # 2. Conditional embeddings + if timestep.dim() == 2: + ts_seq_len = timestep.shape[1] + temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len) + else: + temb = self.time_in(timestep, timestep_r=timestep_r) + # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim] + + hidden_states = self.img_in(hidden_states) # Prepare block-wise causal attention mask - if self.block_mask is None: + if kv_cache is None: self.block_mask = self._prepare_blockwise_causal_attn_mask( device=hidden_states.device, num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size, - txt_attn_mask=encoder_attention_mask + text_seq_len=txt_kv_cache[0]["k_txt"].shape[1] ) # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.double_blocks: - hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + for block_index, block in enumerate(self.double_blocks): + hidden_states = self._gradient_checkpointing_func( block, - hidden_states, - encoder_hidden_states, - temb, - temb_txt, - freqs_cis, - self.block_mask + txt_inference=False, + vision_inference=True, + img=hidden_states, + vec=temb, + freqs_cis=freqs_cis, + block_mask=self.block_mask, + kv_cache=kv_cache[block_index] if kv_cache is not None else None, + txt_kv_cache=txt_kv_cache[block_index], + current_start=current_start ) else: - for block in self.double_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, - encoder_hidden_states, - temb, - temb_txt, - freqs_cis, - self.block_mask + for block_index, block in enumerate(self.double_blocks): + hidden_states = block( + txt_inference=False, + vision_inference=True, + img=hidden_states, + vec=temb, + freqs_cis=freqs_cis, + block_mask=self.block_mask, + kv_cache=kv_cache[block_index] if kv_cache is not None else None, + txt_kv_cache=txt_kv_cache[block_index], + current_start=current_start ) # Final layer processing @@ -648,4 +783,17 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT): # Unpatchify to get original shape hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels) - return hidden_states \ No newline at end of file + return hidden_states + + def forward( + self, + txt_inference=False, + vision_inference=False, + **kwargs, + ): + if txt_inference: + return self.forward_txt(**kwargs) + elif vision_inference: + return self.forward_vision(**kwargs) + else: + raise ValueError("txt_inference and vision_inference cannot be both False") \ No newline at end of file diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index cf3117cf8..c6b535b40 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -474,6 +474,7 @@ class HunyuanVideo15Transformer3DModel(CachableDiT): guidance: Optional[torch.Tensor] = None, timestep_r: Optional[torch.LongTensor] = None, attention_kwargs: Optional[Dict[str, Any]] = None, + **kwargs ): encoder_hidden_states_image = encoder_hidden_states_image[0] encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index bf53d324c..b8b1acb4a 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -391,6 +391,8 @@ class TextEncoderLoader(ComponentLoader): from fastvideo.platforms import current_platform + logger.info("Loading text encoder with cpu_offload: %s", use_cpu_offload) + if use_cpu_offload: pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() # Disable FSDP for MPS as it's not compatible @@ -558,10 +560,9 @@ class VAELoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) - class_name = config.get("_class_name") - assert class_name is not None, ( - "Model config does not contain a _class_name attribute. Only diffusers format is supported." - ) + config.pop("_name_or_path", None) + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." fastvideo_args.model_paths["vae"] = model_path from fastvideo.platforms import current_platform @@ -729,6 +730,7 @@ class TransformerLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the transformer based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) + config.pop("_name_or_path", None) hf_config = deepcopy(config) cls_name = config.pop("_class_name") if cls_name is None: @@ -829,7 +831,7 @@ class TransformerLoader(ComponentLoader): ) total_params = sum(p.numel() for p in model.parameters()) - logger.info("Loaded model with %.2fB parameters", total_params / 1e9) + logger.info("Loaded model with %.2fB parameters, with cpu_offload: %s", total_params / 1e9, fastvideo_args.dit_cpu_offload) assert next(model.parameters()).dtype == default_dtype, ( "Model dtype does not match default dtype" diff --git a/fastvideo/models/loader/utils.py b/fastvideo/models/loader/utils.py index 033cd0058..403836064 100644 --- a/fastvideo/models/loader/utils.py +++ b/fastvideo/models/loader/utils.py @@ -78,9 +78,10 @@ def hf_to_custom_state_dict( for source_param_name, full_tensor in hf_param_sd: # type: ignore target_param_name, merge_index, num_params_to_merge = param_names_mapping( source_param_name) - reverse_param_names_mapping[target_param_name] = (source_param_name, - merge_index, - num_params_to_merge) + if merge_index is None: + reverse_param_names_mapping[target_param_name] = (source_param_name, + merge_index, + num_params_to_merge) if merge_index is not None: to_merge_params[target_param_name][merge_index] = full_tensor if len(to_merge_params[target_param_name]) == num_params_to_merge: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 954b96049..67949e295 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -28,6 +28,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = { ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"), "HYWorldTransformer3DModel": ("dits", "hyworld", "HYWorldTransformer3DModel"), + "CausalHunyuanVideo15Transformer3DModel": + ("dits", "causal_hunyuanvideo15", "CausalHunyuanVideo15Transformer3DModel"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"), diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index ae8a657e6..7b96dca23 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -155,8 +155,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, self.sigmas = sigmas.to( "cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() + self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item() + self.sigma_max = sigma_max if sigma_max is not None else self.sigmas[0].item() BaseScheduler.__init__(self) @@ -289,6 +289,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, sigmas: list[float] | None = None, mu: float | None = None, timesteps: list[float] | None = None, + extra_one_step: bool = False, ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -350,7 +351,10 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, if timesteps_array is None: t_max = self._sigma_to_t(self.sigma_max) t_min = self._sigma_to_t(self.sigma_min) - timesteps_array = np.linspace(t_max, t_min, num_inference_steps) + if extra_one_step: + timesteps_array = np.linspace(t_max, t_min, num_inference_steps + 1)[:-1] + else: + timesteps_array = np.linspace(t_max, t_min, num_inference_steps) sigmas_array = timesteps_array / self.config.num_train_timesteps else: sigmas_array = np.array(sigmas).astype(np.float32) @@ -644,7 +648,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, self, clean_latent: torch.Tensor, noise: torch.Tensor, - timestep: torch.IntTensor, + timestep: torch.Tensor, ) -> torch.Tensor: """ diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 493baeb81..c67b20a13 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -5,6 +5,9 @@ from typing import Any import torch +from fastvideo.logger import init_logger + +logger = init_logger(__name__) # TODO(PY): move it elsewhere def auto_attributes(init_func): diff --git a/fastvideo/models/vaes/hunyuan15vae.py b/fastvideo/models/vaes/hunyuan15vae.py index 66201bd09..73a04a0a4 100644 --- a/fastvideo/models/vaes/hunyuan15vae.py +++ b/fastvideo/models/vaes/hunyuan15vae.py @@ -663,7 +663,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE): # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the # intermediate tiles together, the memory requirement can be lowered. - self.use_tiling = False + self.use_tiling = True # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 256 diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_causal_dmd_pipeline.py new file mode 100644 index 000000000..a59eb8b10 --- /dev/null +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_causal_dmd_pipeline.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan causal DMD pipeline implementation. + +This module wires the causal DMD denoising stage into the modular pipeline. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline + +# isort: off +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + CausalDMDDenosingStage, + InputValidationStage, + TimestepPreparationStage, + Hy15ImageEncodingStage, + LatentPreparationStage, + TextEncodingStage) +# isort: on + +logger = init_logger(__name__) + + +class Hy15CausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "transformer", "scheduler" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up pipeline stages with proper dependency injection.""" + + 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"), + self.get_module("text_encoder_2") + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2") + ], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + # self.add_stage(stage_name="timestep_preparation_stage", + # stage=TimestepPreparationStage( + # scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) + + self.add_stage(stage_name="denoising_stage", + stage=CausalDMDDenosingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = Hy15CausalDMDPipeline diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py new file mode 100644 index 000000000..297a2637f --- /dev/null +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Hunyuan video diffusion pipeline implementation. + +This module contains an implementation of the Hunyuan video diffusion pipeline +using the modular pipeline architecture. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, + Hy15ImageEncodingStage) + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "transformer", "scheduler" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2") + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2") + ], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + 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"))) + + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = HunyuanVideo15ImageToVideoPipeline diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 725a96f11..735bfc7ad 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -84,6 +84,7 @@ class ComposedPipelineBase(ABC): # Load modules directly in initialization logger.info("Loading pipeline modules...") + # fastvideo_args.dit_cpu_offload = False with self.profiler_controller.region("profiler_region_model_loading"): self.modules = self.load_modules(fastvideo_args, loaded_modules) @@ -287,6 +288,7 @@ class ComposedPipelineBase(ABC): # remove keys that are not pipeline modules model_index.pop("_class_name") model_index.pop("_diffusers_version") + model_index.pop("_name_or_path", None) model_index.pop("workload_type", None) if "boundary_ratio" in model_index and model_index[ "boundary_ratio"] is not None: diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 53e838acb..8375e357d 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -232,6 +232,7 @@ class TrainingBatch: image_latents: torch.Tensor | None = None infos: list[dict[str, Any]] | None = None mask_lat_size: torch.Tensor | None = None + video_latent: torch.Tensor | None = None # ODE trajectory supervision trajectory_latents: torch.Tensor | None = None @@ -240,6 +241,9 @@ class TrainingBatch: # Transformer inputs noisy_model_input: torch.Tensor | None = None timesteps: torch.Tensor | None = None + use_gt_trajectory: bool = False + trajectory_timesteps: torch.Tensor | None = None + start_timestep_index: int | None = None sigmas: torch.Tensor | None = None noise: torch.Tensor | None = None diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index dc00f8e42..ee9aaf544 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -29,6 +29,8 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "HunyuanVideoPipeline": "hunyuan", "HunyuanVideo15Pipeline": "hunyuan15", "HYWorldPipeline": "hyworld", + "Hy15CausalDMDPipeline": "hunyuan15", + "HunyuanVideo15ImageToVideoPipeline": "hunyuan15", "Cosmos2VideoToWorldPipeline": "cosmos", "Cosmos2_5Pipeline": "cosmos", "MatrixGamePipeline": "matrixgame", diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 083277bc2..b0ef77d67 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -303,7 +303,7 @@ class BasePreprocessPipeline(ComposedPipelineBase): if data is None: continue - with torch.inference_mode(): + with torch.no_grad(): # Filter out invalid samples (those with all zeros) valid_indices = [] for i, pixel_values in enumerate(data["pixel_values"]): @@ -435,4 +435,4 @@ class BasePreprocessPipeline(ComposedPipelineBase): if num_processed_samples >= args.flush_frequency: written = self.dataset_writer.flush() logger.info("Flushed %s samples to parquet", written) - num_processed_samples = 0 + num_processed_samples = 0 \ No newline at end of file diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index 171cace77..a826a682a 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -37,7 +37,8 @@ from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage, InputValidationStage, LatentPreparationStage, TextEncodingStage, - TimestepPreparationStage) + TimestepPreparationStage, + Hy15ImageEncodingStage) from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict logger = init_logger(__name__) @@ -47,7 +48,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): """ODE Trajectory preprocessing pipeline implementation.""" _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler" ] preprocess_dataloader: StatefulDataLoader preprocess_loader_iter: Iterator[dict[str, Any]] @@ -61,19 +62,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): """Set up pipeline stages with proper dependency injection.""" assert fastvideo_args.pipeline_config.flow_shift == 5 - self.modules["scheduler"] = SelfForcingFlowMatchScheduler( - shift=fastvideo_args.pipeline_config.flow_shift, - sigma_min=0.0, - extra_one_step=True) - self.modules["scheduler"].set_timesteps(num_inference_steps=48, - denoising_strength=1.0) self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) self.add_stage(stage_name="prompt_encoding_stage", stage=TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], + text_encoders=[self.get_module("text_encoder"), self.get_module("text_encoder_2")], + tokenizers=[self.get_module("tokenizer"), self.get_module("tokenizer_2")], )) self.add_stage(stage_name="timestep_preparation_stage", stage=TimestepPreparationStage( @@ -82,6 +77,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): stage=LatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer", None))) + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) self.add_stage(stage_name="denoising_stage", stage=DenoisingStage( transformer=self.get_module("transformer"), @@ -95,11 +93,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): args): """Preprocess text-only data and generate trajectory information.""" + num_encoders = len(self.prompt_encoding_stage.text_encoders) + for batch_idx, data in enumerate(self.pbar): if data is None: continue - with torch.inference_mode(): + with torch.no_grad(): # For text-only processing, we only need text data # Filter out samples without text valid_indices = [] @@ -130,12 +130,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( batch_captions, fastvideo_args, - encoder_index=[0], + encoder_index=list(range(num_encoders)), 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] sampling_params = SamplingParam.from_pretrained(args.model_path) @@ -144,61 +141,52 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text( sampling_params.negative_prompt, fastvideo_args, - encoder_index=[0], + encoder_index=list(range(num_encoders)), 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 + negative_prompt_embeds_list = [] + negative_prompt_masks_list = [] 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) + # Collect the trajectory data (text-to-video generation) + batch = ForwardBatch(**shallow_asdict(sampling_params), ) + batch.prompt_embeds = prompt_embeds_list + batch.prompt_attention_mask = prompt_masks_list + batch.negative_prompt_embeds = negative_prompt_embeds_list + batch.negative_attention_mask = negative_prompt_masks_list + assert batch.num_inference_steps == 50 + batch.return_trajectory_latents = True + # Enabling this will save the decoded trajectory videos. + # Used for debugging. + batch.return_trajectory_decoded = False + batch.height = args.max_height + batch.width = args.max_width + batch.num_frames = args.num_frames + batch.fps = args.train_fps + assert batch.guidance_scale == 6.0 + assert batch.do_classifier_free_guidance == True - # Collect the trajectory data (text-to-video generation) - batch = ForwardBatch(**shallow_asdict(sampling_params), ) - 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.num_inference_steps = 48 - batch.return_trajectory_latents = True - # Enabling this will save the decoded trajectory videos. - # Used for debugging. - batch.return_trajectory_decoded = False - batch.height = args.max_height - batch.width = args.max_width - batch.fps = args.train_fps - batch.guidance_scale = 6.0 - batch.do_classifier_free_guidance = True + result_batch = self.input_validation_stage( + batch, fastvideo_args) + result_batch = self.timestep_preparation_stage( + batch, fastvideo_args) + result_batch = self.latent_preparation_stage( + result_batch, fastvideo_args) + result_batch = self.image_encoding_stage(result_batch, fastvideo_args) + result_batch = self.denoising_stage(result_batch, + fastvideo_args) + result_batch = self.decoding_stage(result_batch, + fastvideo_args) - result_batch = self.input_validation_stage( - 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.append( - result_batch.trajectory_latents.cpu()) - trajectory_timesteps.append( - result_batch.trajectory_timesteps.cpu()) - trajectory_decoded.append(result_batch.trajectory_decoded) + trajectory_latents.append( + result_batch.trajectory_latents.cpu()) + trajectory_timesteps.append( + result_batch.trajectory_timesteps.cpu()) + trajectory_decoded.append(result_batch.trajectory_decoded) # Prepare extra features for text-only processing extra_features = { @@ -209,10 +197,11 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): if batch.return_trajectory_decoded: for i, decoded_frames in enumerate(trajectory_decoded): for j, decoded_frame in enumerate(decoded_frames): - save_decoded_latents_as_video( - decoded_frame, - f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", - args.train_fps) + if j in [5, 7]: + save_decoded_latents_as_video( + decoded_frame, + f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", + args.train_fps) # Prepare batch data for Parquet dataset batch_data: list[dict[str, Any]] = [] @@ -227,7 +216,10 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): video_name = os.path.basename(video_path).split(".")[0] # Convert tensors to numpy arrays - text_embedding = prompt_embeds[idx].cpu().numpy() + text_embedding = prompt_embeds_list[0].float().cpu().numpy() + text_mask = prompt_masks_list[0].cpu().numpy() + text_embedding_2 = prompt_embeds_list[1].float().cpu().numpy() + text_mask_2 = prompt_masks_list[1].cpu().numpy() # Get extra features for this sample sample_extra_features = {} @@ -253,6 +245,9 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): "trajectory_latents"], trajectory_timesteps=sample_extra_features[ "trajectory_timesteps"], + text_embedding_2=text_embedding_2, + text_mask=text_mask, + text_mask_2=text_mask_2, ) batch_data.append(record) @@ -320,4 +315,4 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): self.preprocess_text_and_trajectory(fastvideo_args, args) -EntryClass = PreprocessPipeline_ODE_Trajectory +EntryClass = PreprocessPipeline_ODE_Trajectory \ No newline at end of file diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py index edad1fa1a..3f131c6f8 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py @@ -58,7 +58,7 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): if data is None: continue - with torch.inference_mode(): + with torch.no_grad(): # For text-only processing, we only need text data # Filter out samples without text valid_indices = [] @@ -181,4 +181,4 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): self.preprocess_text_only(fastvideo_args, args) -EntryClass = PreprocessPipeline_Text +EntryClass = PreprocessPipeline_Text \ No newline at end of file diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 18455d70f..318e2c497 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -24,32 +24,33 @@ logger = init_logger(__name__) def main(args) -> None: - args.model_path = maybe_download_model(args.model_path) + # args.model_path = maybe_download_model(args.model_path) maybe_init_distributed_environment_and_model_parallel(1, 1) num_gpus = int(os.environ["WORLD_SIZE"]) assert num_gpus == 1, "Only support 1 GPU" pipeline_config = PipelineConfig.from_pretrained(args.model_path) + print(pipeline_config.__class__.__name__) - kwargs: dict[str, Any] = {} - if args.preprocess_task == "text_only": - kwargs = { - "text_encoder_cpu_offload": False, - } - else: - # Full config for video/image processing - kwargs = { - "vae_precision": "fp32", - "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True), - } - pipeline_config.update_config_from_dict(kwargs) + # kwargs: dict[str, Any] = {} + # if args.preprocess_task == "text_only": + # kwargs = { + # "text_encoder_cpu_offload": False, + # } + # else: + # # Full config for video/image processing + # kwargs = { + # "vae_precision": "fp32", + # "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True), + # } + # pipeline_config.update_config_from_dict(kwargs) fastvideo_args = FastVideoArgs( model_path=args.model_path, num_gpus=get_world_size(), - dit_cpu_offload=False, - vae_cpu_offload=False, - text_encoder_cpu_offload=False, + dit_cpu_offload=True, + vae_cpu_offload=True, + text_encoder_cpu_offload=True, pipeline_config=pipeline_config, ) if args.preprocess_task == "t2v": @@ -140,4 +141,4 @@ if __name__ == "__main__": ) args = parser.parse_args() - main(args) + main(args) \ No newline at end of file diff --git a/fastvideo/pipelines/preprocess/v1_preprocessing_new.py b/fastvideo/pipelines/preprocess/v1_preprocessing_new.py index cfd89bce5..1299b04b5 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocessing_new.py +++ b/fastvideo/pipelines/preprocess/v1_preprocessing_new.py @@ -24,4 +24,4 @@ if __name__ == "__main__": parser = FastVideoArgs.add_cli_args(parser) args = parser.parse_args() fastvideo_args = FastVideoArgs.from_cli_args(args) - main(fastvideo_args) + main(fastvideo_args) \ No newline at end of file diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 777314547..7886e4627 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -1,3 +1,4 @@ +import math import torch # type: ignore from fastvideo.distributed import get_local_torch_device @@ -9,6 +10,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.denoising import DenoisingStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.utils import PRECISION_TO_TYPE try: from fastvideo.attention.backends.sliding_tile_attn import ( @@ -45,9 +47,9 @@ class CausalDMDDenosingStage(DenoisingStage): self.transformer_2 = transformer_2 self.vae = vae # Model-dependent constants (aligned with causal_inference.py assumptions) - self.num_transformer_blocks = len(self.transformer.blocks) - self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block - self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames + self.num_transformer_blocks = self.transformer.config.num_layers + self.num_frames_per_block = self.transformer.config.num_frames_per_block + self.sliding_window_num_frames = self.transformer.config.sliding_window_num_frames try: self.local_attn_size = getattr(self.transformer.model, @@ -66,8 +68,13 @@ class CausalDMDDenosingStage(DenoisingStage): ) and not fastvideo_args.disable_autocast latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] - patch_ratio = self.transformer.config.arch_config.patch_size[ - -1] * self.transformer.config.arch_config.patch_size[-2] + if isinstance(self.transformer.config.patch_size, tuple): + patch_ratio = self.transformer.config.patch_size[1] * self.transformer.config.patch_size[2] + elif isinstance(self.transformer.config.patch_size, int): + patch_ratio = self.transformer.config.patch_size**2 + else: + raise ValueError(f"Unsupported patch size type: {type(self.transformer.config.patch_size)}") + self.frame_seq_length = latent_seq_length // patch_ratio # TODO(will): make this a parameter once we add i2v support independent_first_frame = self.transformer.independent_first_frame if hasattr( @@ -76,12 +83,14 @@ class CausalDMDDenosingStage(DenoisingStage): timesteps = torch.tensor( fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() + self.scheduler.set_timesteps(num_inference_steps=1000, extra_one_step=True, device=get_local_torch_device()) if fastvideo_args.pipeline_config.warp_denoising_step: scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) timesteps = scheduler_timesteps[1000 - timesteps] timesteps = timesteps.to(get_local_torch_device()) + logger.info("[causal_denoising] timesteps: %s", timesteps) if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None: boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps @@ -91,7 +100,20 @@ class CausalDMDDenosingStage(DenoisingStage): high_noise_timesteps = None # Image kwargs (kept empty unless caller provides compatible args) - image_kwargs: dict = {} + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert not torch.isnan( + image_embeds[0]).any(), "image_embeds contains nan" + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + image_kwargs: dict = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_image": image_embeds, + }, + ) pos_cond_kwargs = self.prepare_extra_func_kwargs( self.transformer.forward, @@ -101,6 +123,10 @@ class CausalDMDDenosingStage(DenoisingStage): }, ) + if len(batch.prompt_attention_mask) < 2: + logger.error("prompt embeds length: %s", len(batch.prompt_embeds), local_main_process_only=False) + logger.error("prompt: %s", batch.prompt, local_main_process_only=False) + # STA if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: self.prepare_sta_param(batch, fastvideo_args) @@ -132,27 +158,44 @@ class CausalDMDDenosingStage(DenoisingStage): return kv_cache2 return kv_cache1 - crossattn_cache = self._initialize_crossattn_cache( - batch_size=latents.shape[0], - max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0]. - arch_config.text_len, - dtype=target_dtype, - device=latents.device) + # crossattn_cache = self._initialize_crossattn_cache( + # batch_size=latents.shape[0], + # max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0]. + # arch_config.text_len, + # dtype=target_dtype, + # device=latents.device) pos_start_base = 0 # Determine block sizes - if t % self.num_frames_per_block != 0: - raise ValueError( - "num_frames must be divisible by num_frames_per_block for causal DMD denoising" - ) - num_blocks = t // self.num_frames_per_block + # if t % self.num_frames_per_block != 0: + # raise ValueError( + # "num_frames must be divisible by num_frames_per_block for causal DMD denoising" + # ) + num_blocks = math.ceil(t / self.num_frames_per_block) block_sizes = [self.num_frames_per_block] * num_blocks start_index = 0 # For now hardcode the first block to be 1 frame assuming the model is Wan2.2-MoE - if boundary_timestep is not None: - block_sizes[0] = 1 + # if boundary_timestep is not None: + block_sizes[0] = 1 + + # Initialize txt kv cache + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=None, + forward_batch=batch): + txt_kv_cache = self.transformer( + txt_inference=True, + vision_inference=False, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_image=image_embeds, + encoder_attention_mask=batch.prompt_attention_mask, + timestep=torch.zeros([latents.shape[0]], device=latents.device), + cache_txt=True, + ) first_frame_latent = None if batch.pil_image is not None: @@ -160,7 +203,8 @@ class CausalDMDDenosingStage(DenoisingStage): # the image latent instead of appending along the channel dim assert self.vae is not None, "VAE is not provided for causal video gen task" self.vae = self.vae.to(get_local_torch_device()) - first_frame_latent = self.vae.encode(batch.pil_image).mean.float() + vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + first_frame_latent = self.vae.encode(batch.pil_image.to(vae_dtype)).mean.float() if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): if isinstance(self.vae.shift_factor, torch.Tensor): @@ -182,6 +226,16 @@ class CausalDMDDenosingStage(DenoisingStage): t_zero = torch.zeros([latents.shape[0], 1], device=latents.device, dtype=torch.long) + if batch.video_latent is not None: + video_latent_chunk = batch.video_latent[:, :, start_index:start_index + 1, :, :] + first_frame_input = torch.cat([ + first_frame_latent, + video_latent_chunk, + torch.zeros_like(first_frame_latent), + ], + dim=1) + else: + first_frame_input = first_frame_latent.clone() with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \ @@ -189,17 +243,28 @@ class CausalDMDDenosingStage(DenoisingStage): attn_metadata=None, forward_batch=batch): self.transformer( - first_frame_latent.to(target_dtype), - prompt_embeds, - t_zero, + txt_inference=False, + vision_inference=True, + hidden_states=first_frame_input.to(target_dtype), + timestep=t_zero, kv_cache=kv_cache1, - crossattn_cache=crossattn_cache, + txt_kv_cache=txt_kv_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, + rope_start_idx=start_index, ) + # self.transformer( + # first_frame_latent.to(target_dtype), + # prompt_embeds, + # t_zero, + # kv_cache=kv_cache1, + # crossattn_cache=crossattn_cache, + # current_start=(pos_start_base + start_index) * + # self.frame_seq_length, + # start_frame=start_index, + # **image_kwargs, + # **pos_cond_kwargs, + # ) if boundary_timestep is not None: self.transformer_2( first_frame_latent.to(target_dtype), @@ -218,6 +283,14 @@ class CausalDMDDenosingStage(DenoisingStage): block_sizes.pop(0) latents[:, :, :1, :, :] = first_frame_latent + vision_input_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "vision_inference": True, + "txt_inference": False, + }, + ) + # DMD loop in causal blocks with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: @@ -237,7 +310,15 @@ class CausalDMDDenosingStage(DenoisingStage): noise_latents = noise_latents_btchw.clone() latent_model_input = current_latents.to(target_dtype) - if batch.image_latent is not None and independent_first_frame and start_index == 0: + if batch.video_latent is not None: + video_latent_chunk = batch.video_latent[:, :, start_index:start_index + current_num_frames, :, :] + latent_model_input = torch.cat([ + latent_model_input, + video_latent_chunk, + torch.zeros_like(current_latents), + ], + dim=1) + elif batch.image_latent is not None and independent_first_frame and start_index == 0: latent_model_input = torch.cat([ latent_model_input, batch.image_latent.to(target_dtype) @@ -280,23 +361,41 @@ class CausalDMDDenosingStage(DenoisingStage): forward_batch=batch): # Run transformer; follow DMD stage pattern t_expanded_noise = t_cur * torch.ones( - (latent_model_input.shape[0], 1), + (latent_model_input.shape[0], current_num_frames), device=latent_model_input.device, dtype=torch.long) pred_noise_btchw = current_model( - latent_model_input, - prompt_embeds, - t_expanded_noise, + hidden_states=latent_model_input, + timestep=t_expanded_noise, kv_cache=_get_kv_cache(t_cur), - crossattn_cache=crossattn_cache, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, + rope_start_idx=start_index, + **vision_input_kwargs, ).permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities + # if i < len(timesteps) - 1: + # next_timestep = timesteps[i + 1] * torch.ones_like(t_expand) + # pred_video_btchw = pred_noise_to_x_bound( + # pred_noise=pred_noise_btchw.flatten(0, 1), + # noise_input_latent=noise_latents.flatten(0, 1), + # timestep=t_expand, + # boundary_timestep=next_timestep, + # scheduler=self.scheduler).unflatten( + # 0, pred_noise_btchw.shape[:2]) + # else: + # pred_video_btchw = pred_noise_to_pred_video( + # pred_noise=pred_noise_btchw.flatten(0, 1), + # noise_input_latent=noise_latents.flatten(0, 1), + # timestep=t_expand, + # scheduler=self.scheduler).unflatten( + # 0, pred_noise_btchw.shape[:2]) + + # current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) + if boundary_timestep is not None and t_cur >= boundary_timestep: pred_video_btchw = pred_noise_to_x_bound( pred_noise=pred_noise_btchw.flatten(0, 1), @@ -359,43 +458,46 @@ class CausalDMDDenosingStage(DenoisingStage): # Re-run with context timestep to update KV cache using clean context context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) - t_context = torch.ones([latents.shape[0]], + t_context = torch.ones([latents.shape[0], current_num_frames], device=latents.device, dtype=torch.long) * int(context_noise) context_bcthw = current_latents.to(target_dtype) + if batch.video_latent is not None: + video_latent_chunk = batch.video_latent[:, :, start_index:start_index + current_num_frames, :, :] + context_bcthw = torch.cat([ + context_bcthw, + video_latent_chunk, + torch.zeros_like(current_latents), + ], dim=1) with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \ set_forward_context(current_timestep=0, attn_metadata=attn_metadata, forward_batch=batch): - t_expanded_context = t_context.unsqueeze(1) if boundary_timestep is not None: self.transformer_2( - context_bcthw, - prompt_embeds, - t_expanded_context, + hidden_states=context_bcthw, + timestep=t_context, kv_cache=kv_cache2, - crossattn_cache=crossattn_cache, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, + rope_start_idx=start_index, + **vision_input_kwargs, ) self.transformer( - context_bcthw, - prompt_embeds, - t_expanded_context, + hidden_states=context_bcthw, + timestep=t_context, kv_cache=kv_cache1, - crossattn_cache=crossattn_cache, + txt_kv_cache=txt_kv_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, + rope_start_idx=start_index, + **vision_input_kwargs, ) start_index += current_num_frames @@ -412,8 +514,8 @@ class CausalDMDDenosingStage(DenoisingStage): Initialize a Per-GPU KV cache aligned with the Wan model assumptions. """ kv_cache1 = [] - num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = self.transformer.attention_head_dim + num_attention_heads = self.transformer.config.num_attention_heads + attention_head_dim = self.transformer.config.attention_head_dim if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index f7564ac8c..627c08d44 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -247,8 +247,9 @@ class DenoisingStage(PipelineStage): # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] + trajectory_latents: list[torch.Tensor] = [latents] + logger.info("timesteps: %s", timesteps) # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -413,6 +414,7 @@ class DenoisingStage(PipelineStage): **action_kwargs, ) + assert batch.do_classifier_free_guidance, "do_classifier_free_guidance is not supported" if batch.do_classifier_free_guidance: batch.is_cfg_negative = True with set_forward_context( @@ -467,6 +469,7 @@ class DenoisingStage(PipelineStage): trajectory_tensor: torch.Tensor | None = None if trajectory_latents: + trajectory_timesteps.append(torch.zeros_like(t)) trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 1e0834fc2..fdef281ff 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -115,10 +115,10 @@ class Hy15ImageEncodingStage(ImageEncodingStage): """ Encode the prompt into image encoder hidden states. """ - if batch.pil_image is None: - batch.image_embeds = [ - torch.zeros(1, 729, 1152, device=get_local_torch_device()) - ] + # if batch.pil_image is None: + batch.image_embeds = [ + torch.zeros(1, 729, 1152, device=get_local_torch_device()) + ] raw_latent_shape = list(batch.raw_latent_shape) raw_latent_shape[1] = 1 diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index c9ba23588..8a9c255c9 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -42,6 +42,7 @@ from fastvideo.training.training_utils import ( shift_timestep) from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory) +from fastvideo.optim.muon import get_muon_optimizer vsa_available = is_vsa_available() @@ -85,7 +86,7 @@ class DistillationPipeline(TrainingPipeline): if training_args.real_score_model_path: logger.info("Loading real score transformer from: %s", training_args.real_score_model_path) - training_args.override_transformer_cls_name = "WanTransformer3DModel" + # training_args.override_transformer_cls_name = "WanTransformer3DModel" # TODO(will): can use deepcopy instead if the model is the same self.real_score_transformer = self.load_module_from_path( training_args.real_score_model_path, "transformer", @@ -111,7 +112,7 @@ class DistillationPipeline(TrainingPipeline): if training_args.fake_score_model_path: logger.info("Loading fake score transformer from: %s", training_args.fake_score_model_path) - training_args.override_transformer_cls_name = "WanTransformer3DModel" + # training_args.override_transformer_cls_name = "WanTransformer3DModel" self.fake_score_transformer = self.load_module_from_path( training_args.fake_score_model_path, "transformer", training_args) @@ -146,8 +147,9 @@ class DistillationPipeline(TrainingPipeline): self.vae.requires_grad_(False) self.timestep_shift = self.training_args.pipeline_config.flow_shift - self.noise_scheduler = FlowMatchEulerDiscreteScheduler( - shift=self.timestep_shift) + assert self.timestep_shift == 5.0, "flow_shift must be 5.0" + # self.noise_scheduler.set_timesteps(sigmas=list(np.linspace(1.0, 0.0, 50 + 1)[:-1]), device=get_local_torch_device()) + self.noise_scheduler.set_timesteps(num_inference_steps=1000, extra_one_step=True, device=get_local_torch_device()) if self.training_args.boundary_ratio is not None: self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps @@ -182,11 +184,6 @@ class DistillationPipeline(TrainingPipeline): checkpointing_type=training_args. enable_gradient_checkpointing_type) - # Initialize optimizers - fake_score_params = list( - filter(lambda p: p.requires_grad, - self.fake_score_transformer.parameters())) - # Use separate learning rate for fake_score_transformer if specified fake_score_lr = training_args.fake_score_learning_rate if fake_score_lr == 0.0: @@ -195,12 +192,10 @@ class DistillationPipeline(TrainingPipeline): betas_str = training_args.fake_score_betas betas = tuple(float(x.strip()) for x in betas_str.split(",")) - self.fake_score_optimizer = torch.optim.AdamW( - fake_score_params, + self.fake_score_optimizer = get_muon_optimizer( + self.fake_score_transformer, lr=fake_score_lr, - betas=betas, weight_decay=training_args.weight_decay, - eps=1e-8, ) self.fake_score_lr_scheduler = get_scheduler( @@ -269,24 +264,25 @@ class DistillationPipeline(TrainingPipeline): self.num_train_timestep) self.real_score_guidance_scale = self.training_args.real_score_guidance_scale + # assert self.real_score_guidance_scale == 3.5, "Real score guidance scale must be 3.5" self.generator_ema: EMA_FSDP | None = None self.generator_ema_2: EMA_FSDP | None = None - if (self.training_args.ema_decay - is not None) and (self.training_args.ema_decay > 0.0): - self.generator_ema = EMA_FSDP(self.transformer, - decay=self.training_args.ema_decay) - logger.info("Initialized generator EMA with decay=%s", - self.training_args.ema_decay) + # if (self.training_args.ema_decay + # is not None) and (self.training_args.ema_decay > 0.0): + # self.generator_ema = EMA_FSDP(self.transformer, + # decay=self.training_args.ema_decay) + # logger.info("Initialized generator EMA with decay=%s", + # self.training_args.ema_decay) - # Initialize EMA for transformer_2 if it exists - if self.transformer_2 is not None: - self.generator_ema_2 = EMA_FSDP( - self.transformer_2, decay=self.training_args.ema_decay) - logger.info("Initialized generator EMA_2 with decay=%s", - self.training_args.ema_decay) - else: - logger.info("Generator EMA disabled (ema_decay <= 0.0)") + # # Initialize EMA for transformer_2 if it exists + # if self.transformer_2 is not None: + # self.generator_ema_2 = EMA_FSDP( + # self.transformer_2, decay=self.training_args.ema_decay) + # logger.info("Initialized generator EMA_2 with decay=%s", + # self.training_args.ema_decay) + # else: + # logger.info("Generator EMA disabled (ema_decay <= 0.0)") def load_module_from_path(self, model_path: str, module_type: str, training_args: "TrainingArgs"): @@ -572,6 +568,7 @@ class DistillationPipeline(TrainingPipeline): training_batch.input_kwargs = { "hidden_states": noise_input.permute(0, 2, 1, 3, 4), "encoder_hidden_states": text_dict["encoder_hidden_states"], + "encoder_hidden_states_image": training_batch.image_embeds, "encoder_attention_mask": text_dict["encoder_attention_mask"], "timestep": timestep, "return_dict": False, @@ -712,6 +709,11 @@ class DistillationPipeline(TrainingPipeline): timestep, self.timestep_shift, # type: ignore self.num_train_timestep) + # timestep_idx = torch.randint(1, + # len(self.denoising_step_list), [1], + # device=self.device, + # dtype=torch.long) + # timestep = self.denoising_step_list[timestep_idx] timestep = timestep.clamp(self.min_timestep, self.max_timestep) @@ -723,10 +725,17 @@ class DistillationPipeline(TrainingPipeline): generator_pred_video.flatten(0, 1), noise.flatten(0, 1), timestep).detach().unflatten(0, (1, generator_pred_video.shape[1])) + noisy_latent_copy = noisy_latent.clone() + if training_batch.video_latent is not None: + noisy_latent_copy = torch.cat([ + noisy_latent_copy, + training_batch.video_latent, + torch.zeros_like(noisy_latent_copy), + ], dim=2) # fake_score_transformer forward training_batch = self._build_distill_input_kwargs( - noisy_latent, timestep, training_batch.conditional_dict, + noisy_latent_copy, timestep, training_batch.conditional_dict, training_batch) current_fake_score_transformer = self._get_fake_score_transformer( timestep) @@ -742,7 +751,7 @@ class DistillationPipeline(TrainingPipeline): # real_score_transformer cond forward training_batch = self._build_distill_input_kwargs( - noisy_latent, timestep, training_batch.conditional_dict, + noisy_latent_copy, timestep, training_batch.conditional_dict, training_batch) current_real_score_transformer = self._get_real_score_transformer( timestep) @@ -758,7 +767,7 @@ class DistillationPipeline(TrainingPipeline): # real_score_transformer uncond forward training_batch = self._build_distill_input_kwargs( - noisy_latent, timestep, training_batch.unconditional_dict, + noisy_latent_copy, timestep, training_batch.unconditional_dict, training_batch) # Use same transformer as conditional forward for consistency real_score_pred_noise_uncond = current_real_score_transformer( @@ -774,6 +783,7 @@ class DistillationPipeline(TrainingPipeline): real_score_pred_video = pred_real_video_cond + ( pred_real_video_cond - pred_real_video_uncond) * self.real_score_guidance_scale + # real_score_pred_video = pred_real_video_cond grad = (faker_score_pred_video - real_score_pred_video) / torch.abs( original_latent - real_score_pred_video).mean() @@ -819,6 +829,11 @@ class DistillationPipeline(TrainingPipeline): fake_score_timestep, self.timestep_shift, # type: ignore self.num_train_timestep) + # fake_score_timestep_idx = torch.randint(1, + # len(self.denoising_step_list), [1], + # device=self.device, + # dtype=torch.long) + # fake_score_timestep = self.denoising_step_list[fake_score_timestep_idx] fake_score_timestep = fake_score_timestep.clamp(self.min_timestep, self.max_timestep) @@ -831,6 +846,12 @@ class DistillationPipeline(TrainingPipeline): generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1), fake_score_timestep).unflatten(0, (1, generator_pred_video.shape[1])) + if training_batch.video_latent is not None: + noisy_generator_pred_video = torch.cat([ + noisy_generator_pred_video, + training_batch.video_latent, + torch.zeros_like(noisy_generator_pred_video), + ], dim=2) with set_forward_context(current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata): @@ -877,7 +898,7 @@ class DistillationPipeline(TrainingPipeline): def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: - super()._prepare_dit_inputs(training_batch) + # super()._prepare_dit_inputs(training_batch) conditional_dict = { "encoder_hidden_states": training_batch.encoder_hidden_states, "encoder_attention_mask": training_batch.encoder_attention_mask, @@ -892,13 +913,19 @@ class DistillationPipeline(TrainingPipeline): training_batch.dmd_latent_vis_dict = {} training_batch.fake_score_latent_vis_dict = {} + training_batch.timesteps = 0 training_batch.conditional_dict = conditional_dict training_batch.raw_latent_shape = training_batch.latents.shape training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4) self.video_latent_shape = training_batch.latents.shape self.video_latent_shape_sp = training_batch.latents.shape - + training_batch.image_embeds = [ + torch.zeros(1, 729, 1152, device=get_local_torch_device(), dtype=torch.bfloat16) + ] + raw_latent_shape = list(self.video_latent_shape) + raw_latent_shape[2] = 1 + training_batch.video_latent = torch.zeros(tuple(raw_latent_shape), device=get_local_torch_device(), dtype=torch.bfloat16) return training_batch def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: @@ -1239,8 +1266,10 @@ class DistillationPipeline(TrainingPipeline): for validation_batch in validation_dataloader: batch = self._prepare_validation_batch( sampling_param, training_args, validation_batch, steps) + batch.prompt_attention_mask = [] negative_prompt = batch.negative_prompt + assert negative_prompt == "" batch_negative = ForwardBatch( data_type="video", prompt=negative_prompt, @@ -1249,8 +1278,9 @@ class DistillationPipeline(TrainingPipeline): ) result_batch = self.validation_pipeline.prompt_encoding_stage( # type: ignore batch_negative, training_args) - self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[ - 0], result_batch.prompt_attention_mask[0] + # self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[ + # 0], result_batch.prompt_attention_mask[0] + self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask logger.info( "rank: %s: rank_in_sp_group: %s, batch.prompt: %s", @@ -1354,6 +1384,7 @@ class DistillationPipeline(TrainingPipeline): if hasattr(self, 'transformer_2') and self.transformer_2 is not None: self.transformer_2.train() gc.collect() + torch.cuda.empty_cache() def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int): diff --git a/fastvideo/training/hy15_self_forcing_distillation_pipeline.py b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py new file mode 100644 index 000000000..90a15b42b --- /dev/null +++ b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy + +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.hunyuan15.hunyuan15_causal_dmd_pipeline import ( + Hy15CausalDMDPipeline) +from fastvideo.training.self_forcing_distillation_pipeline import ( + SelfForcingDistillationPipeline) +from fastvideo.utils import is_vsa_available + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline): + """ + A self-forcing distillation pipeline for Hy15 that uses the self-forcing methodology + with DMD for video generation. + """ + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "scheduler", + "transformer", + "vae", + ] + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + validation_pipeline = Hy15CausalDMDPipeline.from_pretrained( + training_args.model_path, + args=args_copy, # type: ignore + inference_mode=True, + loaded_modules={ + "transformer": self.get_module("transformer"), + }, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + pin_cpu_memory=training_args.pin_cpu_memory, + dit_cpu_offload=False) + + self.validation_pipeline = validation_pipeline + + +def main(args) -> None: + logger.info("Starting Hy15 self-forcing distillation pipeline...") + + pipeline = Hy15SelfForcingDistillationPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + + args = pipeline.training_args + pipeline.train() + logger.info("Hy15 self-forcing distillation pipeline completed") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + main(args) diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index df2e3b120..95e413844 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/fastvideo/training/ode_causal_pipeline.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import gc import sys from copy import deepcopy from typing import Any, cast @@ -17,7 +18,7 @@ from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( SelfForcingFlowMatchScheduler) # from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import ( # WanCausalDMDPipeline) -from fastvideo.pipelines.basic.hunyuan15.hunyuan15_pipeline import HunyuanVideo15Pipeline +from fastvideo.pipelines.basic.hunyuan15.hunyuan15_causal_dmd_pipeline import Hy15CausalDMDPipeline from fastvideo.pipelines.pipeline_batch_info import TrainingBatch from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( @@ -75,11 +76,9 @@ class ODEInitTrainingPipeline(TrainingPipeline): 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(sigmas=list(np.linspace(1.0, 0.0, 50 + 1)[:-1]), device=get_local_torch_device()) + self.noise_scheduler.set_timesteps(num_inference_steps=1000, extra_one_step=True, device=get_local_torch_device()) - logger.info("dmd_denoising_steps: %s", - self.training_args.pipeline_config.dmd_denoising_steps) - self.dmd_denoising_steps = torch.tensor([0, 12, 24, 36, 50], + self.dmd_denoising_steps = torch.tensor(self.training_args.pipeline_config.dmd_denoising_steps, dtype=torch.long, device=get_local_torch_device()) if training_args.warp_denoising_step: # Warp the denoising step according to the scheduler time shift @@ -115,7 +114,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): args_copy = deepcopy(training_args) args_copy.inference_mode = True # Warm start validation with current transformer - self.validation_pipeline = HunyuanVideo15Pipeline.from_pretrained( + self.validation_pipeline = Hy15CausalDMDPipeline.from_pretrained( training_args.model_path, args=args_copy, # type: ignore inference_mode=True, @@ -126,7 +125,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): sp_size=training_args.sp_size, num_gpus=training_args.num_gpus, pin_cpu_memory=training_args.pin_cpu_memory, - dit_cpu_offload=True) + dit_cpu_offload=False) def _get_next_batch( self, @@ -263,7 +262,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps if self._cached_closest_idx_per_dmd is None: self._cached_closest_idx_per_dmd = torch.tensor( - [0, 12, 24, 36, 49], dtype=torch.long).cpu() + [0, 12, 24, 36, 50], dtype=torch.long).cpu() # [0, 1, 2, 3], dtype=torch.long).cpu() logger.info("self._cached_closest_idx_per_dmd: %s", self._cached_closest_idx_per_dmd) @@ -279,6 +278,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): traj_latents, dim=1, index=self._cached_closest_idx_per_dmd.to(traj_latents.device)) + # relevant_traj_latents = traj_latents logger.info("relevant_traj_latents: %s", relevant_traj_latents.shape) # assert relevant_traj_latents.shape[0] == 1 @@ -306,31 +306,118 @@ class ODEInitTrainingPipeline(TrainingPipeline): 4).detach().clone().cpu() logger.info("timestep: %s", timestep) - input_kwargs = { - "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + txt_input_kwargs = { + "txt_inference": True, + "vision_inference": False, "encoder_hidden_states": encoder_hidden_states, "encoder_hidden_states_image": encoder_hidden_states_image, "encoder_attention_mask": encoder_attention_mask, + "timestep": torch.zeros([latents.shape[0]], device=latents.device, dtype=torch.bfloat16), + "cache_txt": True, + } + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + txt_kv_cache = self.transformer(**txt_input_kwargs) + + vision_input_kwargs = { + "txt_inference": False, + "vision_inference": True, + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), "timestep": timestep.to(device, dtype=torch.bfloat16), + "txt_kv_cache": txt_kv_cache, } # 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) + noise_pred = self.transformer(**vision_input_kwargs).permute(0, 2, 1, 3, 4) from fastvideo.models.utils import pred_noise_to_pred_video pred_video = pred_noise_to_pred_video( pred_noise=noise_pred.flatten(0, 1), noise_input_latent=latents.flatten(0, 1), timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), - scheduler=self.modules["scheduler"]).unflatten( + scheduler=self.noise_scheduler).unflatten( 0, noise_pred.shape[:2]) latent_vis_dict["pred_video"] = pred_video.permute( 0, 2, 1, 3, 4).detach().clone().cpu() return pred_video, target_latent, timestep, latent_vis_dict + # def _step_predict_next_latent( + # self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, + # encoder_hidden_states: torch.Tensor, + # encoder_attention_mask: torch.Tensor, + # encoder_hidden_states_image: torch.Tensor + # ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + # torch.Tensor]]: + # latent_vis_dict: dict[str, torch.Tensor] = {} + # device = get_local_torch_device() + # target_latent = traj_latents[:, -1] + # del traj_latents + + # # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + # B, num_frames, num_channels, height, width = target_latent.shape + + # indexes = self._get_timestep( # [B, num_frames] + # 0, + # 1000, + # B, + # num_frames, + # 3, + # uniform_timestep=False) + # timestep = self.noise_scheduler.timesteps[indexes.cpu()].to(device) + + # latents = self.noise_scheduler.add_noise(target_latent.flatten(0, 1), torch.randn_like(target_latent.flatten(0, 1)), timestep.flatten(0, 1)).unflatten(0, (B, num_frames)) + # noisy_input = torch.cat([latents, torch.zeros_like(latents), torch.zeros_like(latents[:, :, 0:1])], dim=2) + + # # Prepare inputs for transformer + # latent_vis_dict["noisy_input"] = latents.permute( + # 0, 2, 1, 3, 4).detach().clone().cpu() + # latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, + # 4).detach().clone().cpu() + + # logger.info("timestep: %s", timestep) + # txt_input_kwargs = { + # "txt_inference": True, + # "vision_inference": False, + # "encoder_hidden_states": encoder_hidden_states, + # "encoder_hidden_states_image": encoder_hidden_states_image, + # "encoder_attention_mask": encoder_attention_mask, + # "timestep": torch.zeros([latents.shape[0]], device=latents.device, dtype=torch.bfloat16), + # "cache_txt": True, + # } + # with set_forward_context(current_timestep=timestep, + # attn_metadata=None, + # forward_batch=None): + # txt_kv_cache = self.transformer(**txt_input_kwargs) + + # vision_input_kwargs = { + # "txt_inference": False, + # "vision_inference": True, + # "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + # "timestep": timestep.to(device, dtype=torch.bfloat16), + # "txt_kv_cache": txt_kv_cache, + # } + # # 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(**vision_input_kwargs).permute(0, 2, 1, 3, 4) + + # from fastvideo.models.utils import pred_noise_to_pred_video + # pred_video = pred_noise_to_pred_video( + # pred_noise=noise_pred.flatten(0, 1), + # noise_input_latent=latents.flatten(0, 1), + # timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + # scheduler=self.noise_scheduler).unflatten( + # 0, noise_pred.shape[:2]) + # latent_vis_dict["pred_video"] = pred_video.permute( + # 0, 2, 1, 3, 4).detach().clone().cpu() + + # return pred_video, target_latent, timestep, latent_vis_dict + def train_one_step(self, training_batch): # type: ignore[override] self.transformer.train() self.optimizer.zero_grad() @@ -342,6 +429,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): for _ in range(args.gradient_accumulation_steps): training_batch, traj_latents, traj_timesteps = self._get_next_batch( training_batch) + # traj_latents = traj_latents[:, :, :21] text_embeds = training_batch.encoder_hidden_states text_attention_mask = training_batch.encoder_attention_mask image_embeds = training_batch.encoder_hidden_states_image @@ -392,6 +480,10 @@ class ODEInitTrainingPipeline(TrainingPipeline): except Exception: grad_value = 0.0 training_batch.grad_norm = grad_value + + if training_batch.current_timestep % 10 == 0: + gc.collect() + torch.cuda.empty_cache() return training_batch def visualize_intermediate_latents(self, training_batch: TrainingBatch, @@ -419,6 +511,9 @@ class ODEInitTrainingPipeline(TrainingPipeline): if self.global_rank == 0 and tracker_loss_dict: self.tracker.log_artifacts(tracker_loss_dict, step) + gc.collect() + torch.cuda.empty_cache() + def main(args) -> None: logger.info("Starting ODE-init training pipeline...") diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index fc68e64ce..6a9d835dc 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import copy import os +import gc import time from collections import deque from typing import Any @@ -8,6 +9,7 @@ from typing import Any import numpy as np import torch import torch.distributed as dist +import torch.nn.functional as F from einops import rearrange from tqdm.auto import tqdm @@ -20,7 +22,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfForcingFlowMatchScheduler -from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound from fastvideo.pipelines import TrainingBatch from fastvideo.training.distillation_pipeline import DistillationPipeline from fastvideo.training.training_utils import (EMA_FSDP, @@ -66,12 +68,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline): except Exception: logger.info("Entered initialize_training_pipeline (rank unknown)") - self.noise_scheduler = SelfForcingFlowMatchScheduler( - num_inference_steps=1000, - shift=5.0, - sigma_min=0.0, - extra_one_step=True, - training=True) + # self.noise_scheduler = SelfForcingFlowMatchScheduler( + # num_inference_steps=1000, + # shift=5.0, + # sigma_min=0.0, + # extra_one_step=True, + # training=True) self.dfake_gen_update_ratio = getattr(training_args, 'dfake_gen_update_ratio', 5) @@ -84,35 +86,22 @@ class SelfForcingDistillationPipeline(DistillationPipeline): self.last_step_only = getattr(training_args, 'last_step_only', False) self.context_noise = getattr(training_args, 'context_noise', 0) - self.kv_cache1: list[dict[str, Any]] | None = None - self.crossattn_cache: list[dict[str, Any]] | None = None - logger.info("Self-forcing generator update ratio: %s", self.dfake_gen_update_ratio) - logger.info("RANK: %s, exiting initialize_training_pipeline", - self.global_rank, - local_main_process_only=False) - def generate_and_sync_list(self, num_blocks: int, num_denoising_steps: int, + def generate_and_sync_list(self, num_blocks: int, start_timestep_index: int, end_timestep_index: int, device: torch.device) -> list[int]: """Generate and synchronize random exit flags across distributed processes.""" - logger.info( - "RANK: %s, enter generate_and_sync_list blocks=%s steps=%s device=%s", - self.global_rank, - num_blocks, - num_denoising_steps, - str(device), - local_main_process_only=False) rank = dist.get_rank() if dist.is_initialized() else 0 if rank == 0: # Generate random indices - indices = torch.randint(low=0, - high=num_denoising_steps, + indices = torch.randint(low=start_timestep_index, + high=end_timestep_index, size=(num_blocks, ), device=device) if self.last_step_only: - indices = torch.ones_like(indices) * (num_denoising_steps - 1) + indices = torch.ones_like(indices) * (end_timestep_index - 1) else: indices = torch.empty(num_blocks, dtype=torch.long, device=device) @@ -120,12 +109,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): dist.broadcast(indices, src=0) # Broadcast the random indices to all ranks flags = indices.tolist() - logger.info( - "RANK: %s, exit generate_and_sync_list flags_len=%s first=%s", - self.global_rank, - len(flags), - flags[0] if len(flags) > 0 else None, - local_main_process_only=False) return flags def generator_loss( @@ -167,6 +150,26 @@ class SelfForcingDistillationPipeline(DistillationPipeline): return flow_matching_loss, log_dict + def _prepare_gt_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: + trajectory_latents = training_batch.trajectory_latents + + # assert trajectory_latents.shape[1] == len(self.denoising_step_list) + 1 + + # start_timestep_index = torch.randint(0, len(self.denoising_step_list), (1, ), device=self.device) + # if dist.is_initialized(): + # dist.broadcast(start_timestep_index, src=0) + # start_timestep_index = start_timestep_index.cpu().item() + # training_batch.trajectory_latents = trajectory_latents[:, :, :21] + _cached_closest_idx_per_dmd = torch.tensor( + [0, 12, 24, 36, 50], dtype=torch.long).cpu() + training_batch.trajectory_latents = torch.index_select( + trajectory_latents, + dim=1, + index=_cached_closest_idx_per_dmd.to(trajectory_latents.device)) + del trajectory_latents + # training_batch.start_timestep_index = start_timestep_index + return training_batch + def _generator_multi_step_simulation_forward( self, training_batch: TrainingBatch, @@ -186,21 +189,21 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # During training, the number of generated frames should be uniformly sampled from # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block - min_num_frames = 20 if self.independent_first_frame else 21 - max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames - assert max_num_frames % self.num_frame_per_block == 0 - assert min_num_frames % self.num_frame_per_block == 0 - max_num_blocks = max_num_frames // self.num_frame_per_block - min_num_blocks = min_num_frames // self.num_frame_per_block + # min_num_frames = 20 if self.independent_first_frame else 21 + # max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames + # assert max_num_frames % self.num_frame_per_block == 0 + # assert min_num_frames % self.num_frame_per_block == 0 + # max_num_blocks = max_num_frames // self.num_frame_per_block + # min_num_blocks = min_num_frames // self.num_frame_per_block # Sample number of blocks and sync across processes - num_generated_blocks = torch.randint(min_num_blocks, - max_num_blocks + 1, (1, ), - device=self.device) - if dist.is_initialized(): - dist.broadcast(num_generated_blocks, src=0) - num_generated_blocks = num_generated_blocks.item() - num_generated_frames = num_generated_blocks * self.num_frame_per_block + # num_generated_blocks = torch.randint(min_num_blocks, + # max_num_blocks + 1, (1, ), + # device=self.device) + # if dist.is_initialized(): + # dist.broadcast(num_generated_blocks, src=0) + # num_generated_blocks = num_generated_blocks.item() + num_generated_frames = num_training_frames if self.independent_first_frame and initial_latent is None: num_generated_frames += 1 min_num_frames += 1 @@ -216,6 +219,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch_size, num_generated_frames, *self.video_latent_shape[2:] ] + # if training_batch.use_gt_trajectory: + # noise = training_batch.trajectory_latents.to(self.device, dtype=dtype) + # else: noise = torch.randn(noise_shape, device=self.device, dtype=dtype) if self.sp_world_size > 1: noise = rearrange(noise, @@ -226,13 +232,15 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch_size, num_frames, num_channels, height, width = noise.shape # Block size calculation - if not self.independent_first_frame or (self.independent_first_frame - and initial_latent is not None): - assert num_frames % self.num_frame_per_block == 0 - num_blocks = num_frames // self.num_frame_per_block - else: - assert (num_frames - 1) % self.num_frame_per_block == 0 - num_blocks = (num_frames - 1) // self.num_frame_per_block + # if not self.independent_first_frame or (self.independent_first_frame + # and initial_latent is not None): + # assert num_frames % self.num_frame_per_block == 0 + # num_blocks = num_frames // self.num_frame_per_block + # else: + # assert (num_frames - 1) % self.num_frame_per_block == 0 + # num_blocks = (num_frames - 1) // self.num_frame_per_block + import math + num_blocks = math.ceil(num_frames / self.num_frame_per_block) num_input_frames = initial_latent.shape[ 1] if initial_latent is not None else 0 @@ -252,7 +260,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 1: Initialize KV cache to all zeros cache_frames = num_generated_frames + num_input_frames - self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches( + kv_cache1, crossattn_cache = self._initialize_simulation_caches( batch_size, dtype, self.device, max_num_frames=cache_frames) # Step 2: Cache context feature @@ -284,21 +292,52 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 3: Temporal denoising loop all_num_frames = [self.num_frame_per_block] * num_blocks + all_num_frames[0] = 1 if self.independent_first_frame and initial_latent is None: all_num_frames = [1] + all_num_frames - num_denoising_steps = len(self.denoising_step_list) + + # if training_batch.use_gt_trajectory: + # start_timestep_index = training_batch.start_timestep_index + # end_timestep_index = len(self.denoising_step_list) + # else: + start_timestep_index = 0 + end_timestep_index = len(self.denoising_step_list) exit_flags = self.generate_and_sync_list(len(all_num_frames), - num_denoising_steps, + start_timestep_index, + end_timestep_index, device=noise.device) - start_gradient_frame_index = max(0, num_output_frames - 21) + start_gradient_frame_index = max(0, num_output_frames - 31) + + logger.info("distillation step list: %s", self.denoising_step_list) + + # Initialize txt kv cache + txt_kv_cache = self.transformer( + txt_inference=True, + vision_inference=False, + encoder_hidden_states=training_batch.conditional_dict["encoder_hidden_states"], + encoder_hidden_states_image=training_batch.image_embeds, + encoder_attention_mask=training_batch.conditional_dict["encoder_attention_mask"], + timestep=torch.zeros([batch_size], device=noise.device), + cache_txt=True, + ) for block_index, current_num_frames in enumerate(all_num_frames): noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames] + index = start_timestep_index + current_timestep = self.denoising_step_list[index] # Step 3.1: Spatial denoising loop - for index, current_timestep in enumerate(self.denoising_step_list): + while index < len(self.denoising_step_list): + noisy_input_latent = noisy_input.clone() + if training_batch.video_latent is not None: + noisy_input_latent = torch.cat([ + noisy_input_latent, + training_batch.video_latent[:, current_start_frame:current_start_frame + current_num_frames], + torch.zeros_like(noisy_input_latent), + ], dim=2) + if self.same_step_across_blocks: exit_flag = (index == exit_flags[0]) else: @@ -317,23 +356,28 @@ class SelfForcingDistillationPipeline(DistillationPipeline): with torch.no_grad(): # Build input kwargs training_batch_temp = self._build_distill_input_kwargs( - noisy_input, timestep, + noisy_input_latent, timestep, training_batch.conditional_dict, training_batch) pred_flow = current_model( + txt_inference=False, + vision_inference=True, hidden_states=training_batch_temp. input_kwargs['hidden_states'], - encoder_hidden_states=training_batch_temp. - input_kwargs['encoder_hidden_states'], + # encoder_hidden_states=training_batch_temp. + # input_kwargs['encoder_hidden_states'], + # encoder_attention_mask=training_batch_temp. + # input_kwargs['encoder_attention_mask'], timestep=training_batch_temp. input_kwargs['timestep'], - encoder_hidden_states_image=training_batch_temp. - input_kwargs.get('encoder_hidden_states_image'), - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, + # encoder_hidden_states_image=training_batch_temp. + # input_kwargs.get('encoder_hidden_states_image'), + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - start_frame=current_start_frame).permute( + rope_start_idx=current_start_frame).permute( 0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( @@ -344,6 +388,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 0, pred_flow.shape[:2]) next_timestep = self.denoising_step_list[index + 1] + # noisy_input = pred_noise_to_x_bound( + # pred_noise=pred_flow.flatten(0, 1), + # noise_input_latent=noisy_input.flatten(0, 1), + # timestep=timestep, + # boundary_timestep=next_timestep * torch.ones_like(timestep).flatten(0, 1), + # scheduler=self.noise_scheduler).unflatten( + # 0, pred_flow.shape[:2]) noisy_input = self.noise_scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), @@ -357,43 +408,53 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if current_start_frame < start_gradient_frame_index: with torch.no_grad(): training_batch_temp = self._build_distill_input_kwargs( - noisy_input, timestep, + noisy_input_latent, timestep, training_batch.conditional_dict, training_batch) pred_flow = current_model( + txt_inference=False, + vision_inference=True, hidden_states=training_batch_temp. input_kwargs['hidden_states'], - encoder_hidden_states=training_batch_temp. - input_kwargs['encoder_hidden_states'], + # encoder_hidden_states=training_batch_temp. + # input_kwargs['encoder_hidden_states'], + # encoder_attention_mask=training_batch_temp. + # input_kwargs['encoder_attention_mask'], timestep=training_batch_temp. input_kwargs['timestep'], - encoder_hidden_states_image=training_batch_temp. - input_kwargs.get('encoder_hidden_states_image'), - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, + # encoder_hidden_states_image=training_batch_temp. + # input_kwargs.get('encoder_hidden_states_image'), + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - start_frame=current_start_frame).permute( + rope_start_idx=current_start_frame).permute( 0, 2, 1, 3, 4) else: training_batch_temp = self._build_distill_input_kwargs( - noisy_input, timestep, + noisy_input_latent, timestep, training_batch.conditional_dict, training_batch) pred_flow = current_model( + txt_inference=False, + vision_inference=True, hidden_states=training_batch_temp. input_kwargs['hidden_states'], - encoder_hidden_states=training_batch_temp. - input_kwargs['encoder_hidden_states'], + # encoder_hidden_states=training_batch_temp. + # input_kwargs['encoder_hidden_states'], + # encoder_attention_mask=training_batch_temp. + # input_kwargs['encoder_attention_mask'], timestep=training_batch_temp. input_kwargs['timestep'], - encoder_hidden_states_image=training_batch_temp. - input_kwargs.get('encoder_hidden_states_image'), - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, + # encoder_hidden_states_image=training_batch_temp. + # input_kwargs.get('encoder_hidden_states_image'), + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - start_frame=current_start_frame).permute( + rope_start_idx=current_start_frame).permute( 0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( @@ -404,6 +465,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 0, pred_flow.shape[:2]) break + index += 1 + current_timestep = self.denoising_step_list[index] + # Step 3.2: record the model's output output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred @@ -416,6 +480,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline): context_timestep).unflatten(0, denoised_pred.shape[:2]) with torch.no_grad(): + if training_batch.video_latent is not None: + denoised_pred = torch.cat([ + denoised_pred, + training_batch.video_latent[:, current_start_frame:current_start_frame + current_num_frames], + torch.zeros_like(denoised_pred), + ], dim=2) + training_batch_temp = self._build_distill_input_kwargs( denoised_pred, context_timestep, training_batch.conditional_dict, training_batch) @@ -423,17 +494,22 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # context_timestep is 0 so we use transformer_2 current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer current_model( + txt_inference=False, + vision_inference=True, hidden_states=training_batch_temp. input_kwargs['hidden_states'], - encoder_hidden_states=training_batch_temp. - input_kwargs['encoder_hidden_states'], + # encoder_hidden_states=training_batch_temp. + # input_kwargs['encoder_hidden_states'], + # encoder_attention_mask=training_batch_temp. + # input_kwargs['encoder_attention_mask'], timestep=training_batch_temp.input_kwargs['timestep'], - encoder_hidden_states_image=training_batch_temp. - input_kwargs.get('encoder_hidden_states_image'), - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, + # encoder_hidden_states_image=training_batch_temp. + # input_kwargs.get('encoder_hidden_states_image'), + kv_cache=kv_cache1, + txt_kv_cache=txt_kv_cache, + # crossattn_cache=crossattn_cache, current_start=current_start_frame * self.frame_seq_length, - start_frame=current_start_frame) + rope_start_idx=current_start_frame) # Step 3.4: update the start and end frame indices current_start_frame += current_num_frames @@ -445,7 +521,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Slice last 21 frames if we generated more gradient_mask = None - if pred_image_or_video.shape[1] > 21: + if pred_image_or_video.shape[1] > 31: with torch.no_grad(): # Re-encode to get image latent latent_to_decode = pred_image_or_video[:, :-20, ...] @@ -485,14 +561,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline): pred_image_or_video_last_21 = pred_image_or_video # Set up gradient mask if we generated more than minimum frames - if num_generated_frames != min_num_frames: - # Currently, we do not use gradient for the first chunk, since it contains image latents - gradient_mask = torch.ones_like(pred_image_or_video_last_21, - dtype=torch.bool) - if self.independent_first_frame: - gradient_mask[:, :1] = False - else: - gradient_mask[:, :self.num_frame_per_block] = False + # if num_generated_frames != min_num_frames: + # # Currently, we do not use gradient for the first chunk, since it contains image latents + # gradient_mask = torch.ones_like(pred_image_or_video_last_21, + # dtype=torch.bool) + # if self.independent_first_frame: + # gradient_mask[:, :1] = False + # else: + # gradient_mask[:, :self.num_frame_per_block] = False # Apply gradient masking if needed final_output = pred_image_or_video_last_21.to(dtype) @@ -523,9 +599,11 @@ class SelfForcingDistillationPipeline(DistillationPipeline): min_num_frames, dtype=torch.float32, device=self.device) # Clean up caches - assert self.kv_cache1 is not None - assert self.crossattn_cache is not None - self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache) + assert kv_cache1 is not None + # assert self.crossattn_cache is not None + self._cleanup_simulation_caches(kv_cache1, crossattn_cache) + + assert pred_image_or_video.shape[1] == 31, "pred_image_or_video must have 31 frames" return final_output if gradient_mask is not None else pred_image_or_video @@ -538,23 +616,26 @@ class SelfForcingDistillationPipeline(DistillationPipeline): max_num_frames: int | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """Initialize KV cache and cross-attention cache for multi-step simulation.""" - num_transformer_blocks = len(self.transformer.blocks) + num_transformer_blocks = self.transformer.config.num_layers latent_shape = self.video_latent_shape_sp _, num_frames, _, height, width = latent_shape - _, p_h, p_w = self.transformer.patch_size - post_patch_height = height // p_h - post_patch_width = width // p_w + if isinstance(self.transformer.config.patch_size, tuple): + ph, pw = self.transformer.config.patch_size[1], self.transformer.config.patch_size[2] + elif isinstance(self.transformer.config.patch_size, int): + ph, pw = self.transformer.config.patch_size, self.transformer.config.patch_size + else: + raise ValueError(f"Unsupported patch size type: {type(self.transformer.config.patch_size)}") + post_patch_height = height // ph + post_patch_width = width // pw frame_seq_length = post_patch_height * post_patch_width self.frame_seq_length = frame_seq_length # Get model configuration parameters - handle FSDP wrapping - num_attention_heads = getattr(self.transformer, 'num_attention_heads', - None) - attention_head_dim = getattr(self.transformer, 'attention_head_dim', - None) - text_len = getattr(self.transformer, 'text_len', None) + num_attention_heads = self.transformer.config.num_attention_heads + attention_head_dim = self.transformer.config.attention_head_dim + # text_len = getattr(self.transformer, 'text_len', None) if max_num_frames is None: max_num_frames = num_frames @@ -585,29 +666,54 @@ class SelfForcingDistillationPipeline(DistillationPipeline): }) # Initialize cross-attention cache - crossattn_cache = [] - for _ in range(num_transformer_blocks): - crossattn_cache.append({ - "k": - torch.zeros([ - batch_size, text_len, num_attention_heads, - attention_head_dim - ], - dtype=dtype, - device=device), - "v": - torch.zeros([ - batch_size, text_len, num_attention_heads, - attention_head_dim - ], - dtype=dtype, - device=device), - "is_init": - False - }) + crossattn_cache = None + # crossattn_cache = [] + # for _ in range(num_transformer_blocks): + # crossattn_cache.append({ + # "k": + # torch.zeros([ + # batch_size, text_len, num_attention_heads, + # attention_head_dim + # ], + # dtype=dtype, + # device=device), + # "v": + # torch.zeros([ + # batch_size, text_len, num_attention_heads, + # attention_head_dim + # ], + # dtype=dtype, + # device=device), + # "is_init": + # False + # }) return kv_cache, crossattn_cache + def _cleanup_simulation_caches(self, kv_cache: list[dict[str, Any]], + crossattn_cache: list[dict[str, Any]]) -> None: + """Properly clean up KV cache and cross-attention cache GPU memory.""" + if kv_cache is not None: + for cache_dict in kv_cache: + # Clear tensor references to free GPU memory + if "k" in cache_dict and cache_dict["k"] is not None: + cache_dict["k"] = None + if "v" in cache_dict and cache_dict["v"] is not None: + cache_dict["v"] = None + if "global_end_index" in cache_dict and cache_dict["global_end_index"] is not None: + cache_dict["global_end_index"] = None + if "local_end_index" in cache_dict and cache_dict["local_end_index"] is not None: + cache_dict["local_end_index"] = None + + if crossattn_cache is not None: + for cache_dict in crossattn_cache: + # Clear tensor references to free GPU memory + if "k" in cache_dict and cache_dict["k"] is not None: + cache_dict["k"] = None + if "v" in cache_dict and cache_dict["v"] is not None: + cache_dict["v"] = None + cache_dict["is_init"] = False + def _reset_simulation_caches(self, kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]]) -> None: """Reset KV cache and cross-attention cache to clean state.""" @@ -624,6 +730,179 @@ class SelfForcingDistillationPipeline(DistillationPipeline): cache_dict["k"].zero_() cache_dict["v"].zero_() + def _step_predict_next_latent( + self, training_batch: TrainingBatch + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, + torch.Tensor]]: + latent_vis_dict: dict[str, torch.Tensor] = {} + device = get_local_torch_device() + target_latent = training_batch.trajectory_latents[:, -1] + + # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] + B, S, num_frames, num_channels, height, width = training_batch.trajectory_latents.shape + + dmd_denoising_steps = torch.cat([self.denoising_step_list, torch.tensor([0], device=self.device, dtype=torch.long)]) + indexes = self._get_timestep( # [B, num_frames] + 0, + len(dmd_denoising_steps), + B, + num_frames, + 3, + uniform_timestep=False) + # noisy_input = relevant_traj_latents[indexes] + latents = torch.gather( + training_batch.trajectory_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 = torch.cat([latents, torch.zeros_like(latents), torch.zeros_like(latents[:, :, 0:1])], dim=2) + timestep = dmd_denoising_steps[indexes] + + logger.info("timestep: %s", timestep) + txt_input_kwargs = { + "txt_inference": True, + "vision_inference": False, + "encoder_hidden_states": training_batch.encoder_hidden_states, + "encoder_hidden_states_image": training_batch.image_embeds, + "encoder_attention_mask": training_batch.encoder_attention_mask, + "timestep": torch.zeros([latents.shape[0]], device=latents.device, dtype=torch.bfloat16), + "cache_txt": True, + } + with set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=None): + txt_kv_cache = self.transformer(**txt_input_kwargs) + + vision_input_kwargs = { + "txt_inference": False, + "vision_inference": True, + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), + "timestep": timestep.to(device, dtype=torch.bfloat16), + "txt_kv_cache": txt_kv_cache, + } + # 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(**vision_input_kwargs).permute(0, 2, 1, 3, 4) + + from fastvideo.models.utils import pred_noise_to_pred_video + pred_video = pred_noise_to_pred_video( + pred_noise=noise_pred.flatten(0, 1), + noise_input_latent=latents.flatten(0, 1), + timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1), + scheduler=self.noise_scheduler).unflatten( + 0, noise_pred.shape[:2]) + latent_vis_dict["pred_video"] = pred_video.permute( + 0, 2, 1, 3, 4).detach().clone().cpu() + + return pred_video, target_latent, timestep + + 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: + num_pad_frames = 0 + if num_frame % num_frame_per_block != 0: + # Pad num_frame to be divisible by num_frame_per_block + num_pad_frames = num_frame_per_block - (num_frame % num_frame_per_block) + num_frame += num_pad_frames + 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) + if num_pad_frames > 0: + timestep = timestep[:, num_pad_frames:] + return timestep + + def _get_next_batch_2(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter_2, None) # type: ignore + if batch is None: + self.current_epoch += 1 + self.train_loader_iter_2 = iter(self.train_dataloader_2) + batch = next(self.train_loader_iter_2) + + # Required fields from parquet (ODE trajectory schema) + device = get_local_torch_device() + encoder_hidden_states = batch['text_embedding'].to(device, dtype=torch.bfloat16).squeeze(0) + encoder_hidden_states_2 = batch['text_embedding_2'].to(device, dtype=torch.bfloat16).squeeze(0) + encoder_attention_mask = batch['text_mask'].to(device, dtype=torch.bfloat16).squeeze(0) + encoder_attention_mask_2 = batch['text_mask_2'].to(device, dtype=torch.bfloat16).squeeze(0) + infos = batch['info_list'] + + if encoder_hidden_states.dim() < 3: + prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( + infos[0]["caption"], + self.training_args, + encoder_index=[0], + return_attention_mask=True, + ) + encoder_hidden_states = prompt_embeds_list[0].to(device, dtype=torch.bfloat16) + encoder_attention_mask = prompt_masks_list[0].to(device, dtype=torch.bfloat16) + + # 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) + + batch_size = trajectory_latents.shape[0] + vae_config = self.training_args.pipeline_config.vae_config.arch_config + num_channels = vae_config.latent_channels + spatial_compression_ratio = vae_config.spatial_compression_ratio + + latent_height = self.training_args.num_height // spatial_compression_ratio + latent_width = self.training_args.num_width // spatial_compression_ratio + training_batch.latents = torch.randn(batch_size, num_channels, + self.training_args.num_latent_t, latent_height, + latent_width).to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.trajectory_latents = trajectory_latents.to( + device, dtype=torch.bfloat16) + training_batch.trajectory_timesteps = trajectory_timesteps.to(device) + training_batch.encoder_hidden_states = [encoder_hidden_states, encoder_hidden_states_2] + training_batch.encoder_attention_mask = [encoder_attention_mask, encoder_attention_mask_2] + training_batch.infos = infos + return training_batch + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: batch = next(self.train_loader_iter, None) # type: ignore if batch is None: @@ -634,13 +913,26 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch = next(self.train_loader_iter) # latents, encoder_hidden_states, encoder_attention_mask, infos = batch - encoder_hidden_states = batch['text_embedding'] - encoder_attention_mask = batch['text_attention_mask'] + # encoder_hidden_states = batch['text_embedding'] + # encoder_attention_mask = batch['text_attention_mask'] + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch infos = batch['info_list'] + batch_prompt = ForwardBatch( + data_type="video", + prompt=infos[0]["caption"], + prompt_embeds=[], + prompt_attention_mask=[], + ) + result_batch = self.validation_pipeline.prompt_encoding_stage( # type: ignore + batch_prompt, self.training_args) + encoder_hidden_states, encoder_attention_mask = result_batch.prompt_embeds, result_batch.prompt_attention_mask + for i in range(len(encoder_hidden_states)): + encoder_hidden_states[i] = encoder_hidden_states[i].to(get_local_torch_device(), dtype=torch.bfloat16) + encoder_attention_mask[i] = encoder_attention_mask[i].to(get_local_torch_device(), dtype=torch.bfloat16) - batch_size = encoder_hidden_states.shape[0] + batch_size = encoder_hidden_states[0].shape[0] vae_config = self.training_args.pipeline_config.vae_config.arch_config - num_channels = vae_config.z_dim + num_channels = vae_config.latent_channels spatial_compression_ratio = vae_config.spatial_compression_ratio latent_height = self.training_args.num_height // spatial_compression_ratio @@ -653,10 +945,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 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.encoder_hidden_states = encoder_hidden_states + training_batch.encoder_attention_mask = encoder_attention_mask training_batch.infos = infos return training_batch @@ -669,11 +959,25 @@ class SelfForcingDistillationPipeline(DistillationPipeline): train_generator = (self.current_trainstep % self.dfake_gen_update_ratio == 0) + if getattr(self, "train_dataloader_2", None) is not None and train_generator: + # use_gt_trajectory = torch.rand(1, device=get_local_torch_device()) + # if dist.is_initialized(): + # dist.broadcast(use_gt_trajectory, src=0) + # training_batch.use_gt_trajectory = (use_gt_trajectory < 0.8).item() + training_batch.use_gt_trajectory = True + else: + training_batch.use_gt_trajectory = False + batches = [] for _ in range(gradient_accumulation_steps): + # if training_batch.use_gt_trajectory: + # batch = self._get_next_batch_2(training_batch) + # else: batch = self._get_next_batch(training_batch) batch = self._normalize_dit_input(batch) batch = self._prepare_dit_inputs(batch) + # if training_batch.use_gt_trajectory: + # batch = self._prepare_gt_inputs(batch) batch = self._build_attention_metadata(batch) batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata) if batch.attn_metadata is not None: @@ -710,6 +1014,36 @@ class SelfForcingDistillationPipeline(DistillationPipeline): setattr(batch_gen, key, copy.deepcopy(value)) generator_loss, gen_log_dict = self.generator_loss(batch_gen) + if training_batch.use_gt_trajectory: + batch_gt = self._get_next_batch_2(training_batch) + batch_gt = self._prepare_dit_inputs(batch_gt) + batch_gt = self._prepare_gt_inputs(batch_gt) + batch_gt = self._build_attention_metadata(batch_gt) + batch_gt.attn_metadata_vsa = copy.deepcopy(batch_gt.attn_metadata) + batch_gt_gen = TrainingBatch() + for key, value in batch_gt.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(batch_gt_gen, key, value.detach().clone()) + elif isinstance(value, dict): + setattr( + batch_gt_gen, key, { + k: + v.detach().clone() if isinstance( + v, torch.Tensor) else copy.deepcopy(v) + for k, v in value.items() + }) + else: + setattr(batch_gt_gen, key, copy.deepcopy(value)) + + gt_pred_video, gt_target_latent, gt_timestep = self._step_predict_next_latent(batch_gt_gen) + mask = gt_timestep != 0 + # Compute loss + gt_loss = F.mse_loss(gt_pred_video[mask], + gt_target_latent[mask], + reduction="mean") + generator_loss += 0.5 * gt_loss + + # Store visualization data from generator training with set_forward_context(current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.attn_metadata): (generator_loss / gradient_accumulation_steps).backward() @@ -805,6 +1139,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.fake_score_loss = avg_critic_loss.item() training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss + + if training_batch.current_timestep % 10 == 0: + gc.collect() + torch.cuda.empty_cache() return training_batch def _log_training_info(self) -> None: @@ -979,6 +1317,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): logger.info("Starting training from scratch") self.train_loader_iter = iter(self.train_dataloader) + if getattr(self, "train_dataloader_2", None) is not None: + self.train_loader_iter_2 = iter(self.train_dataloader_2) step_times: deque[float] = deque(maxlen=100) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index fdf482faa..df5e35919 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -2,6 +2,7 @@ from dataclasses import asdict import math import os +import gc import time from abc import ABC, abstractmethod from collections import deque @@ -25,7 +26,7 @@ from fastvideo.attention.backends.video_sparse_attn import ( from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder from fastvideo.configs.sample import SamplingParam from fastvideo.dataset import build_parquet_map_style_dataloader -from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory_text_only, pyarrow_schema_t2v from fastvideo.dataset.validation_dataset import ValidationDataset from fastvideo.distributed import (cleanup_dist_env_and_memory, get_local_torch_device, get_sp_group, @@ -90,6 +91,7 @@ class TrainingPipeline(LoRAPipeline, ABC): def set_schemas(self) -> None: self.train_dataset_schema = pyarrow_schema_t2v + self.train_dataset_schema_2 = pyarrow_schema_ode_trajectory_text_only def initialize_training_pipeline(self, training_args: TrainingArgs): logger.info("Initializing training pipeline...") @@ -193,6 +195,20 @@ class TrainingPipeline(LoRAPipeline, ABC): text_len, # type: ignore[attr-defined] seed=self.seed) + if getattr(training_args, 'data_path_2', None) is not None: + self.train_dataset_2, self.train_dataloader_2 = build_parquet_map_style_dataloader( + training_args.data_path_2, + training_args.train_batch_size, + parquet_schema=self.train_dataset_schema_2, + num_data_workers=training_args.dataloader_num_workers, + cfg_rate=training_args.training_cfg_rate, + drop_last=True, + text_padding_length=training_args.pipeline_config. + text_encoder_configs[0].arch_config. + text_len, # type: ignore[attr-defined] + seed=self.seed + ) + self.noise_scheduler = noise_scheduler if self.training_args.boundary_ratio is not None: self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps @@ -281,7 +297,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # TODO(will): support other models with self.tracker.timed("timing/normalize_input"): training_batch.latents = normalize_dit_input( - 'wan', + 'hunyuan', training_batch.latents, self.get_module("vae"), ) @@ -605,7 +621,7 @@ class TrainingPipeline(LoRAPipeline, ABC): device="cpu").manual_seed(self.seed) logger.info("Initialized random seeds with seed: %s", self.seed) - self.noise_scheduler = FlowMatchEulerDiscreteScheduler() + # self.noise_scheduler = FlowMatchEulerDiscreteScheduler() if self.training_args.resume_from_checkpoint: self._resume_from_checkpoint() @@ -711,6 +727,12 @@ class TrainingPipeline(LoRAPipeline, ABC): "profiler_region_training_validation"): self._log_validation(self.transformer, self.training_args, step) + for module in self.transformer.modules(): + for name, buf in module._buffers.items(): + if "k_cache" in name or "v_cache" in name: + module._buffers[name] = None + gc.collect() + torch.cuda.empty_cache() gpu_memory_usage = current_platform.get_torch_device( ).memory_allocated() / 1024**2 trainable_params = round( diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 321bfb339..dd79f3a17 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -192,7 +192,6 @@ def save_checkpoint(transformer, json.dump(config_dict, f, indent=4) logger.info("--> checkpoint saved at step %s to %s", step, weight_path) - def save_distillation_checkpoint( generator_transformer, fake_score_transformer, @@ -257,8 +256,6 @@ def save_distillation_checkpoint( if generator_scheduler is not None: generator_states["scheduler"] = SchedulerWrapper( generator_scheduler) - if generator_ema is not None: - generator_states["ema"] = generator_ema.state_dict() generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint", "generator") @@ -290,8 +287,6 @@ def save_distillation_checkpoint( if generator_scheduler_2 is not None: generator_2_states["scheduler"] = SchedulerWrapper( generator_scheduler_2) - if generator_ema_2 is not None: - generator_2_states["ema"] = generator_ema_2.state_dict() generator_2_dcp_dir = os.path.join(save_dir, "distributed_checkpoint", @@ -417,6 +412,67 @@ def save_distillation_checkpoint( rank, local_main_process_only=False) + # Persist EMA separately to avoid shape mismatches across ranks. + # Supports: + # - mode="rank0_full": save consolidated EMA only on rank 0 + # - mode="local_shard": save per-rank EMA shard for each rank + try: + if generator_ema is not None and getattr(generator_ema, "mode", + None) == "rank0_full": + _save_rank0_full_ema_safetensors(generator_ema, + generator_transformer, rank, + save_dir, "generator_ema") + elif generator_ema is not None and getattr(generator_ema, "mode", + None) == "local_shard": + # Save per-rank shard + ema_dir_shard = os.path.join(save_dir, "ema_local_shard") + os.makedirs(ema_dir_shard, exist_ok=True) + ema_shard_path = os.path.join(ema_dir_shard, f"generator_ema_rank{rank}.pt") + torch.save(generator_ema.state_dict(), ema_shard_path) + logger.info( + "rank: %s, saved generator EMA shard (local_shard) to %s", + rank, + ema_shard_path, + local_main_process_only=False) + + # Also consolidate EMA to a single full-state file on rank 0 by applying EMA to the model and gathering + _consolidate_local_shard_ema_and_save_safetensors( + generator_ema, generator_transformer, rank, save_dir, + "generator_ema") + + except Exception as e: + logger.warning("rank: %s, failed saving EMA separately: %s", + rank, + str(e)) + + try: + if generator_ema_2 is not None and getattr(generator_ema_2, "mode", + None) == "rank0_full": + _save_rank0_full_ema_safetensors(generator_ema_2, + generator_transformer_2, rank, + save_dir, "generator_ema_2") + elif generator_ema_2 is not None and getattr(generator_ema_2, "mode", + None) == "local_shard": + # Save per-rank shard for EMA_2 + ema_dir_shard_2 = os.path.join(save_dir, "ema_local_shard") + os.makedirs(ema_dir_shard_2, exist_ok=True) + ema2_shard_path = os.path.join(ema_dir_shard_2, f"generator_ema_2_rank{rank}.pt") + torch.save(generator_ema_2.state_dict(), ema2_shard_path) + logger.info( + "rank: %s, saved generator_2 EMA shard (local_shard) to %s", + rank, + ema2_shard_path, + local_main_process_only=False) + + # Also consolidate EMA_2 to a single full-state file on rank 0 + _consolidate_local_shard_ema_and_save_safetensors( + generator_ema_2, generator_transformer_2, rank, save_dir, + "generator_ema_2") + except Exception as e: + logger.warning("rank: %s, failed saving EMA_2 separately: %s", + rank, + str(e)) + # Save generator model weights (consolidated) for inference cpu_state = gather_state_dict_on_cpu_rank0(generator_transformer, device=None) @@ -454,47 +510,46 @@ def save_distillation_checkpoint( logger.info("--> distillation checkpoint saved at step %s to %s", step, weight_path) - # Save generator_2 model weights (consolidated) for inference (MoE support) - if generator_transformer_2 is not None: - inference_save_dir_2 = os.path.join( - save_dir, "generator_2_inference_transformer") - cpu_state_2 = gather_state_dict_on_cpu_rank0( - generator_transformer_2, device=None) + # Save generator_2 model weights (consolidated) for inference (MoE support) + if generator_transformer_2 is not None: + inference_save_dir_2 = os.path.join( + save_dir, "generator_2_inference_transformer") + cpu_state_2 = gather_state_dict_on_cpu_rank0( + generator_transformer_2, device=None) - if rank == 0: - os.makedirs(inference_save_dir_2, exist_ok=True) - weight_path_2 = os.path.join( - inference_save_dir_2, "diffusion_pytorch_model.safetensors") - logger.info( - "rank: %s, saving consolidated generator_2 inference checkpoint to %s", - rank, - weight_path_2, - local_main_process_only=False) + if rank == 0: + os.makedirs(inference_save_dir_2, exist_ok=True) + weight_path_2 = os.path.join( + inference_save_dir_2, "diffusion_pytorch_model.safetensors") + logger.info( + "rank: %s, saving consolidated generator_2 inference checkpoint to %s", + rank, + weight_path_2, + local_main_process_only=False) - # Convert training format to diffusers format and save - diffusers_state_dict_2 = custom_to_hf_state_dict( - cpu_state_2, - generator_transformer_2.reverse_param_names_mapping) - save_file(diffusers_state_dict_2, weight_path_2) + # Convert training format to diffusers format and save + diffusers_state_dict_2 = custom_to_hf_state_dict( + cpu_state_2, + generator_transformer_2.reverse_param_names_mapping) + save_file(diffusers_state_dict_2, weight_path_2) - logger.info( - "rank: %s, consolidated generator_2 inference checkpoint saved to %s", - rank, - weight_path_2, - local_main_process_only=False) - - # Save model config - config_dict_2 = generator_transformer_2.hf_config - if "dtype" in config_dict_2: - del config_dict_2["dtype"] # TODO - config_path_2 = os.path.join(inference_save_dir_2, - "config.json") - with open(config_path_2, "w") as f: - json.dump(config_dict_2, f, indent=4) - logger.info( - "--> generator_2 distillation checkpoint saved at step %s to %s", - step, weight_path_2) + logger.info( + "rank: %s, consolidated generator_2 inference checkpoint saved to %s", + rank, + weight_path_2, + local_main_process_only=False) + # Save model config + config_dict_2 = generator_transformer_2.hf_config + if "dtype" in config_dict_2: + del config_dict_2["dtype"] # TODO + config_path_2 = os.path.join(inference_save_dir_2, + "config.json") + with open(config_path_2, "w") as f: + json.dump(config_dict_2, f, indent=4) + logger.info( + "--> generator_2 distillation checkpoint saved at step %s to %s", + step, weight_path_2) def load_checkpoint(transformer, rank, @@ -644,18 +699,38 @@ def load_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Load EMA state if available and generator_ema is provided + # Load EMA separately if saved in rank0_full mode if generator_ema is not None: try: - ema_state = generator_states.get("ema") - if ema_state is not None: - generator_ema.load_state_dict(ema_state) - logger.info("rank: %s, generator EMA state loaded successfully", - rank) - else: - logger.info("rank: %s, no EMA state found in checkpoint", rank) + if getattr(generator_ema, "mode", None) == "rank0_full": + ema_path = os.path.join(checkpoint_path, "ema", + "generator_ema.pt") + if rank == 0 and os.path.exists(ema_path): + ema_state = torch.load(ema_path, map_location="cpu") + generator_ema.load_state_dict(ema_state) + logger.info( + "rank: %s, generator EMA (rank0_full) loaded from %s", + rank, ema_path) + elif rank == 0: + logger.info( + "rank: %s, generator EMA file not found at %s; skipping", + rank, ema_path) + elif getattr(generator_ema, "mode", None) == "local_shard": + ema_path = os.path.join(checkpoint_path, "ema_local_shard", + f"generator_ema_rank{rank}.pt") + if os.path.exists(ema_path): + ema_state = torch.load(ema_path, map_location="cpu") + generator_ema.load_state_dict(ema_state) + logger.info( + "rank: %s, generator EMA shard (local_shard) loaded from %s", + rank, ema_path) + else: + logger.info( + "rank: %s, generator EMA shard file not found at %s; skipping", + rank, ema_path) + except Exception as e: - logger.warning("rank: %s, failed to load EMA state: %s", rank, + logger.warning("rank: %s, failed to load generator EMA: %s", rank, str(e)) # Load generator_2 distributed checkpoint (MoE support) @@ -850,7 +925,7 @@ def load_distillation_checkpoint( def normalize_dit_input(model_type, latents, vae) -> torch.Tensor: if model_type == "hunyuan_hf" or model_type == "hunyuan": - return latents * 0.476986 + return latents * vae.config.scaling_factor elif model_type == "wan": latents_mean = torch.tensor(vae.latents_mean) latents_std = 1.0 / torch.tensor(vae.latents_std) @@ -1168,6 +1243,72 @@ def custom_to_hf_state_dict( return new_state_dict +def _save_full_ema_safetensors_from_state( + state_dict: dict[str, Any], + reverse_param_names_mapping: dict[str, tuple[str, int, int]], + output_path: str, +) -> None: + """ + Convert a training-format state_dict to HF format and save as safetensors. + """ + diffusers_state_dict = custom_to_hf_state_dict(state_dict, + reverse_param_names_mapping) + save_file(diffusers_state_dict, output_path) + +def _save_rank0_full_ema_safetensors( + ema: "EMA_FSDP", + module, + rank: int, + save_dir: str, + base_name: str, +) -> None: + if rank != 0: + return + ema_dir = os.path.join(save_dir, "ema") + os.makedirs(ema_dir, exist_ok=True) + output_path = os.path.join(ema_dir, f"{base_name}.safetensors") + ema_state = ema.state_dict() + _save_full_ema_safetensors_from_state(ema_state, + module.reverse_param_names_mapping, + output_path) + logger.info( + "rank: %s, saved %s as consolidated EMA safetensors to %s", + rank, + base_name, + output_path, + local_main_process_only=False) + + +def _consolidate_local_shard_ema_and_save_safetensors( + ema: "EMA_FSDP", + module, + rank: int, + save_dir: str, + base_name: str, +) -> None: + try: + # Temporarily apply EMA to the live (sharded) module and gather full CPU state on rank 0 + with ema.apply_to_model(module): + cpu_state_full = gather_state_dict_on_cpu_rank0(module, device=None) + if rank == 0: + ema_dir = os.path.join(save_dir, "ema") + os.makedirs(ema_dir, exist_ok=True) + output_path = os.path.join(ema_dir, f"{base_name}.safetensors") + _save_full_ema_safetensors_from_state( + cpu_state_full, module.reverse_param_names_mapping, output_path) + logger.info( + "rank: %s, saved consolidated %s EMA (from local_shard) as safetensors to %s", + rank, + base_name, + output_path, + local_main_process_only=False) + except Exception as ce: + logger.warning( + "rank: %s, failed consolidating %s EMA (local_shard): %s", + rank, + base_name, + str(ce)) + def shift_timestep(timestep: torch.Tensor, shift: float, num_train_timestep: float) -> torch.Tensor: