diff --git a/dmd_t2v_output/t2v_646663.err b/dmd_t2v_output/t2v_646663.err deleted file mode 100644 index 26517a559..000000000 --- a/dmd_t2v_output/t2v_646663.err +++ /dev/null @@ -1,165 +0,0 @@ -+ source /mnt/weka/home/hao.zhang/conda/miniconda/bin/activate -++ _CONDA_ROOT=/mnt/weka/home/hao.zhang/conda/miniconda -++ . /mnt/weka/home/hao.zhang/conda/miniconda/etc/profile.d/conda.sh -+++ export CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -+++ CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -+++ export _CE_M= -+++ _CE_M= -+++ export _CE_CONDA= -+++ _CE_CONDA= -+++ export CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -+++ CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -+++ '[' -z x ']' -++ conda activate -++ local cmd=activate -++ case "$cmd" in -++ __conda_activate activate -++ '[' -n '' ']' -++ local ask_conda -+++ PS1= -+++ __conda_exe shell.posix activate -+++ '[' -n '' ']' -+++ /mnt/weka/home/hao.zhang/conda/miniconda/bin/conda shell.posix activate -++ ask_conda='unset _CE_M -unset _CE_CONDA -PS1='\''(base) '\'' -export PATH='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin'\'' -export CONDA_PREFIX='\''/mnt/weka/home/hao.zhang/conda/miniconda'\'' -export CONDA_SHLVL='\''3'\'' -export CONDA_DEFAULT_ENV='\''base'\'' -export CONDA_PROMPT_MODIFIER='\''(base) '\'' -export CONDA_PREFIX_2='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv'\'' -export CONDA_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda'\'' -export CONDA_PYTHON_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/python'\''' -++ eval 'unset _CE_M -unset _CE_CONDA -PS1='\''(base) '\'' -export PATH='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin'\'' -export CONDA_PREFIX='\''/mnt/weka/home/hao.zhang/conda/miniconda'\'' -export CONDA_SHLVL='\''3'\'' -export CONDA_DEFAULT_ENV='\''base'\'' -export CONDA_PROMPT_MODIFIER='\''(base) '\'' -export CONDA_PREFIX_2='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv'\'' -export CONDA_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda'\'' -export CONDA_PYTHON_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/python'\''' -+++ unset _CE_M -+++ unset _CE_CONDA -+++ PS1='(base) ' -+++ export PATH=/mnt/weka/home/hao.zhang/conda/miniconda/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin -+++ PATH=/mnt/weka/home/hao.zhang/conda/miniconda/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin -+++ export CONDA_PREFIX=/mnt/weka/home/hao.zhang/conda/miniconda -+++ CONDA_PREFIX=/mnt/weka/home/hao.zhang/conda/miniconda -+++ export CONDA_SHLVL=3 -+++ CONDA_SHLVL=3 -+++ export CONDA_DEFAULT_ENV=base -+++ CONDA_DEFAULT_ENV=base -+++ export 'CONDA_PROMPT_MODIFIER=(base) ' -+++ CONDA_PROMPT_MODIFIER='(base) ' -+++ export CONDA_PREFIX_2=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv -+++ CONDA_PREFIX_2=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv -+++ export CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -+++ CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -+++ export CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -+++ CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -++ __conda_hashr -++ '[' -n '' ']' -++ '[' -n '' ']' -++ hash -r -+ conda activate wei-fv -+ local cmd=activate -+ case "$cmd" in -+ __conda_activate activate wei-fv -+ '[' -n '' ']' -+ local ask_conda -++ PS1='(base) ' -++ __conda_exe shell.posix activate wei-fv -++ '[' -n '' ']' -++ /mnt/weka/home/hao.zhang/conda/miniconda/bin/conda shell.posix activate wei-fv -+ ask_conda='unset _CE_M -unset _CE_CONDA -PS1='\''(wei-fv) '\'' -export PATH='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin'\'' -export CONDA_PREFIX='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv'\'' -export CONDA_SHLVL='\''4'\'' -export CONDA_DEFAULT_ENV='\''wei-fv'\'' -export CONDA_PROMPT_MODIFIER='\''(wei-fv) '\'' -export CONDA_PREFIX_3='\''/mnt/weka/home/hao.zhang/conda/miniconda'\'' -export CONDA_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda'\'' -export CONDA_PYTHON_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/python'\''' -+ eval 'unset _CE_M -unset _CE_CONDA -PS1='\''(wei-fv) '\'' -export PATH='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin'\'' -export CONDA_PREFIX='\''/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv'\'' -export CONDA_SHLVL='\''4'\'' -export CONDA_DEFAULT_ENV='\''wei-fv'\'' -export CONDA_PROMPT_MODIFIER='\''(wei-fv) '\'' -export CONDA_PREFIX_3='\''/mnt/weka/home/hao.zhang/conda/miniconda'\'' -export CONDA_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda'\'' -export CONDA_PYTHON_EXE='\''/mnt/weka/home/hao.zhang/conda/miniconda/bin/python'\''' -++ unset _CE_M -++ unset _CE_CONDA -++ PS1='(wei-fv) ' -++ export PATH=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin -++ PATH=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv/bin:/mnt/weka/home/hao.zhang/conda/miniconda/condabin:/mnt/weka/home/hao.zhang/.cursor-server/cli/servers/Stable-53b99ce608cba35127ae3a050c1738a959750860/server/bin/remote-cli:/mnt/weka/home/hao.zhang/.local/bin:/usr/local/bin/wekaenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/weka/software/weka/bin -++ export CONDA_PREFIX=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv -++ CONDA_PREFIX=/mnt/weka/home/hao.zhang/conda/miniconda/envs/wei-fv -++ export CONDA_SHLVL=4 -++ CONDA_SHLVL=4 -++ export CONDA_DEFAULT_ENV=wei-fv -++ CONDA_DEFAULT_ENV=wei-fv -++ export 'CONDA_PROMPT_MODIFIER=(wei-fv) ' -++ CONDA_PROMPT_MODIFIER='(wei-fv) ' -++ export CONDA_PREFIX_3=/mnt/weka/home/hao.zhang/conda/miniconda -++ CONDA_PREFIX_3=/mnt/weka/home/hao.zhang/conda/miniconda -++ export CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -++ CONDA_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/conda -++ export CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -++ CONDA_PYTHON_EXE=/mnt/weka/home/hao.zhang/conda/miniconda/bin/python -+ __conda_hashr -+ '[' -n '' ']' -+ '[' -n '' ']' -+ hash -r -+ export WANDB_MODE=online -+ WANDB_MODE=online -+ export NCCL_P2P_DISABLE=1 -+ NCCL_P2P_DISABLE=1 -+ export TORCH_NCCL_ENABLE_MONITORING=0 -+ TORCH_NCCL_ENABLE_MONITORING=0 -+ export TRITON_CACHE_DIR=/tmp/triton_cache_0 -+ TRITON_CACHE_DIR=/tmp/triton_cache_0 -+ export MASTER_PORT=29500 -+ MASTER_PORT=29500 -+ export NODE_RANK=0 -+ NODE_RANK=0 -+ nodes=($(scontrol show hostnames $SLURM_JOB_NODELIST)) -++ scontrol show hostnames fs-mbz-gpu-141 -+ export MASTER_ADDR=fs-mbz-gpu-141 -+ MASTER_ADDR=fs-mbz-gpu-141 -+ export CUDA_VISIBLE_DEVICES=0 -+ CUDA_VISIBLE_DEVICES=0 -+ export TOKENIZERS_PARALLELISM=false -+ TOKENIZERS_PARALLELISM=false -+ export WANDB_BASE_URL=https://api.wandb.ai -+ WANDB_BASE_URL=https://api.wandb.ai -+ export WANDB_MODE=online -+ WANDB_MODE=online -+ export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN -+ FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN -+ echo 'MASTER_ADDR: fs-mbz-gpu-141' -+ echo 'NODE_RANK: 0' -+ NUM_GPUS=1 -+ GENERATOR_MODEL_PATH=wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers -+ REAL_SCORE_MODEL_PATH=Wan-AI/Wan2.1-T2V-14B-Diffusers -+ FAKE_SCORE_MODEL_PATH=Wan-AI/Wan2.1-T2V-1.3B-Diffusers -+ DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset/ -+ VALIDATION_DATASET_FILE=examples/distill/SFWan2.1-I2V/validation.json -+ training_args=(--tracker_project_name SFwan_t2v_distill_dmd --output_dir "checkpoints/SFwan_t2v_finetune" --max_train_steps 1000 --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 --num_latent_t 21 --num_height 480 --num_width 832 --num_frames 81 --enable_gradient_checkpointing_type "full") -+ parallel_args=(--num_gpus 1 --sp_size 1 --tp_size 1 --hsdp_replicate_dim 1 --hsdp_shard_dim 1) -+ model_args=(--model_path $GENERATOR_MODEL_PATH --pretrained_model_name_or_path $GENERATOR_MODEL_PATH --generator_model_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" --dataloader_num_workers 4) -+ validation_args=(--log_validation --validation_dataset_file "$VALIDATION_DATASET_FILE" --validation_steps 200 --validation_sampling_steps "3" --validation_guidance_scale "6.0") -+ optimizer_args=(--learning_rate 2e-6 --mixed_precision "bf16" --training_state_checkpointing_steps 500 --weight_only_checkpointing_steps 500 --weight_decay 0.01 --betas '0.0,0.999' --max_grad_norm 1.0) -+ miscellaneous_args=(--inference_mode False --checkpoints_total_limit 3 --training_cfg_rate 0.0 --dit_precision "fp32" --flow_shift 5 --seed 1000 --use_ema True --ema_decay 0.99 --ema_start_step 200 --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors") -+ dmd_args=(--dmd_denoising_steps '1000,750,500,250' --min_timestep_ratio 0.02 --max_timestep_ratio 0.98 --generator_update_interval 5 --real_score_guidance_scale 3.0 --fake_score_learning_rate 4e-7 --fake_score_betas '0.0,0.999') -+ srun torchrun --nnodes 1 --nproc_per_node 1 --node_rank 0 --rdzv_backend=c10d --rdzv_endpoint=fs-mbz-gpu-141:29500 fastvideo/training/wan_distillation_pipeline.py --num_gpus 1 --sp_size 1 --tp_size 1 --hsdp_replicate_dim 1 --hsdp_shard_dim 1 --model_path wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers --pretrained_model_name_or_path wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers --generator_model_path wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers --real_score_model_path Wan-AI/Wan2.1-T2V-14B-Diffusers --fake_score_model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --data_path data/crush-smol_processed_t2v/combined_parquet_dataset/ --dataloader_num_workers 4 --tracker_project_name SFwan_t2v_distill_dmd --output_dir checkpoints/SFwan_t2v_finetune --max_train_steps 1000 --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 --num_latent_t 21 --num_height 480 --num_width 832 --num_frames 81 --enable_gradient_checkpointing_type full --learning_rate 2e-6 --mixed_precision bf16 --training_state_checkpointing_steps 500 --weight_only_checkpointing_steps 500 --weight_decay 0.01 --betas 0.0,0.999 --max_grad_norm 1.0 --log_validation --validation_dataset_file examples/distill/SFWan2.1-I2V/validation.json --validation_steps 200 --validation_sampling_steps 3 --validation_guidance_scale 6.0 --inference_mode False --checkpoints_total_limit 3 --training_cfg_rate 0.0 --dit_precision fp32 --flow_shift 5 --seed 1000 --use_ema True --ema_decay 0.99 --ema_start_step 200 --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors --dmd_denoising_steps 1000,750,500,250 --min_timestep_ratio 0.02 --max_timestep_ratio 0.98 --generator_update_interval 5 --real_score_guidance_scale 3.0 --fake_score_learning_rate 4e-7 --fake_score_betas 0.0,0.999 diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index 7a70f6e4c..7372ad0c8 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -8,7 +8,6 @@ def main(): # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. - load_start_time = time.perf_counter() model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" generator = VideoGenerator.from_pretrained( model_name, @@ -18,47 +17,15 @@ def main(): text_encoder_cpu_offload=False, dit_cpu_offload=False, ) - load_end_time = time.perf_counter() - load_time = load_end_time - load_start_time - sampling_param = SamplingParam.from_pretrained(model_name) - start_time = time.perf_counter() prompt = ( - "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table." + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." ) video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) - end_time = time.perf_counter() - gen_time1 = end_time - start_time - e2e_gen_time = end_time - load_start_time - - # Generate another video with a different prompt, without reloading the - # model! - # gen_time_list = [] - # for i in range(10): - # start_time = time.perf_counter() - # prompt2 = ( - # "A majestic lion strides across the golden savanna, its powerful frame " - # "glistening under the warm afternoon sun. The tall grass ripples gently in " - # "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - # "embodying the raw energy of the wild. Low angle, steady tracking shot, " - # "cinematic.") - # video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) - # end_time = time.perf_counter() - # gen_time2 = end_time - start_time - # gen_time_list.append(gen_time2) - # # print(f"Time taken for video {i}: {gen_time2} seconds") - - # print(f"Time taken for first video: {gen_time1} seconds") - # print(f"Average time taken for 10 videos: {sum(gen_time_list) / len(gen_time_list)} seconds") - # for i, gen_time in enumerate(gen_time_list): - # print(f"Time taken for video {i}: {gen_time} seconds") - - # # print(f"Time taken for second video: {gen_time2} seconds") - # print(f"Time taken to load model: {load_time} seconds") - # print(f"Time taken for first e2e generation: {e2e_gen_time} seconds") - if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py index a08f81942..bce30f873 100644 --- a/fastvideo/configs/models/dits/wanvideo.py +++ b/fastvideo/configs/models/dits/wanvideo.py @@ -95,6 +95,8 @@ class WanVideoArchConfig(DiTArchConfig): # Causal Wan local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention) sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 def __post_init__(self): super().__post_init__() @@ -107,4 +109,4 @@ class WanVideoArchConfig(DiTArchConfig): class WanVideoConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig) - prefix: str = "Wan" + prefix: str = "Wan" \ No newline at end of file diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 4c435fdc8..68199e27e 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -148,5 +148,4 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): # ============================================= @dataclass class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam): - num_frames_per_block: int = 3 - sliding_window_num_frames: int = 21 + pass \ No newline at end of file diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 4d305f494..5bb30c438 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -113,21 +113,6 @@ class CausalWanSelfAttention(nn.Module): dim=1 ) - # Adjust block mask to match the actual sequence length (TODO: Check if this is okay) - actual_q_len = padded_roped_query.shape[1] - actual_kv_len = padded_roped_key.shape[1] - if block_mask.shape[-1] != actual_q_len or block_mask.shape[-2] != actual_kv_len: - def dynamic_attention_mask(b, h, q_idx, kv_idx): - return kv_idx <= q_idx - - from torch.nn.attention.flex_attention import create_block_mask - block_mask = create_block_mask( - dynamic_attention_mask, - B=None, H=None, - Q_LEN=actual_q_len, KV_LEN=actual_kv_len, - _compile=False, device=q.device - ) - x = flex_attention( query=padded_roped_query.transpose(2, 1), key=padded_roped_key.transpose(2, 1), @@ -329,6 +314,7 @@ class CausalWanTransformer3DModel(BaseDiT): inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads + self.attention_head_dim = config.attention_head_dim self.in_channels = config.in_channels self.out_channels = config.out_channels self.num_channels_latents = config.num_channels_latents @@ -659,4 +645,4 @@ class CausalWanTransformer3DModel(BaseDiT): if kwargs.get('kv_cache', None) is not None: return self._forward_inference(*args, **kwargs) else: - return self._forward_train(*args, **kwargs) + return self._forward_train(*args, **kwargs) \ No newline at end of file diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 36233720b..ae098fe17 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -137,3 +137,24 @@ def modulate(x: torch.Tensor, else: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze( 1) # type: ignore[union-attr] + + +def pred_noise_to_pred_video(pred_noise: torch.Tensor, + noise_input_latent: torch.Tensor, + timestep: torch.Tensor, + scheduler: Any) -> torch.Tensor: + """ + Convert predicted noise to clean latent. + """ + timestep = timestep.expand(noise_input_latent.shape[0]) + dtype = pred_noise.dtype + device = pred_noise.device + pred_noise = pred_noise.float().to(device) + noise_input_latent = noise_input_latent.float().to(device) + sigmas = scheduler.sigmas.float().to(device) + timesteps = scheduler.timesteps.float().to(device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + pred_video = noise_input_latent - sigma_t * pred_noise + return pred_video.to(dtype) \ No newline at end of file diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index f9850f3f0..46751d4a7 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -173,10 +173,6 @@ class ForwardBatch: logging_info: PipelineLoggingInfo = field( default_factory=PipelineLoggingInfo) - # Causal Wan parameters - num_frames_per_block: int = 3 - sliding_window_num_frames: int = 21 - def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" @@ -246,4 +242,4 @@ class TrainingBatch: @dataclass class PreprocessBatch(ForwardBatch): video_loader: list["VideoDecoder"] = field(default_factory=list) - video_file_name: list[str] = field(default_factory=list) + video_file_name: list[str] = field(default_factory=list) \ No newline at end of file diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index f81072a91..4fe2074e9 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -3,8 +3,10 @@ import torch # type: ignore from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler) +from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.denoising import DenoisingStage @@ -24,6 +26,8 @@ except ImportError: vsa_available = False VideoSparseAttentionBackend = None # type: ignore +logger = init_logger(__name__) + class CausalDMDDenosingStage(DenoisingStage): """ @@ -37,8 +41,10 @@ class CausalDMDDenosingStage(DenoisingStage): self.kv_cache1: list | None = None self.crossattn_cache: list | None = None # Model-dependent constants (aligned with causal_inference.py assumptions) - self.num_transformer_blocks = 30 - self.frame_seq_length = 1560 + self.num_transformer_blocks = self.transformer.config.arch_config.num_layers + 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 + try: self.local_attn_size = getattr(self.transformer.model, "local_attn_size", @@ -55,16 +61,12 @@ class CausalDMDDenosingStage(DenoisingStage): autocast_enabled = (target_dtype != torch.float32 ) and not fastvideo_args.disable_autocast - # Args - try: - num_frames_per_block = getattr(batch, "num_frames_per_block", 1) - except Exception: - num_frames_per_block = 1 - try: - independent_first_frame = getattr(fastvideo_args.pipeline_config, - "independent_first_frame", False) - except Exception: - independent_first_frame = False + 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] + 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 # Timesteps for DMD timesteps = torch.tensor( @@ -99,9 +101,12 @@ class CausalDMDDenosingStage(DenoisingStage): self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) - self._initialize_crossattn_cache(batch_size=latents.shape[0], - dtype=target_dtype, - device=latents.device) + 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) else: assert self.crossattn_cache is not None # reset cross-attention cache @@ -156,7 +161,7 @@ class CausalDMDDenosingStage(DenoisingStage): # process remaining input frames in blocks of num_frame_per_block while remaining_frames > 0: - block = min(num_frames_per_block, remaining_frames) + block = min(self.num_frames_per_block, remaining_frames) ref_btchw = image_latent[:, :, current_start_frame: current_start_frame + block, :, :].to(target_dtype).permute( @@ -184,20 +189,20 @@ class CausalDMDDenosingStage(DenoisingStage): # Determine block sizes if not independent_first_frame or (independent_first_frame and batch.image_latent is not None): - if t % num_frames_per_block != 0: + 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 // num_frames_per_block - block_sizes = [num_frames_per_block] * num_blocks + num_blocks = t // self.num_frames_per_block + block_sizes = [self.num_frames_per_block] * num_blocks start_index = 0 else: - if (t - 1) % num_frames_per_block != 0: + if (t - 1) % self.num_frames_per_block != 0: raise ValueError( "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" ) - num_blocks = (t - 1) // num_frames_per_block - block_sizes = [1] + [num_frames_per_block] * num_blocks + num_blocks = (t - 1) // self.num_frames_per_block + block_sizes = [1] + [self.num_frames_per_block] * num_blocks start_index = 0 # DMD loop in causal blocks @@ -271,7 +276,6 @@ class CausalDMDDenosingStage(DenoisingStage): ).permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities - from fastvideo.training.training_utils import pred_noise_to_pred_video pred_video_btchw = pred_noise_to_pred_video( pred_noise=pred_noise_btchw.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), @@ -344,19 +348,27 @@ 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 if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: - kv_cache_size = max(self.frame_seq_length, getattr(self, 'num_frames', 1) * self.frame_seq_length) + kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames for _ in range(self.num_transformer_blocks): kv_cache1.append({ "k": - torch.zeros([batch_size, kv_cache_size, 12, 128], + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], dtype=dtype, device=device), "v": - torch.zeros([batch_size, kv_cache_size, 12, 128], + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], dtype=dtype, device=device), "global_end_index": @@ -367,22 +379,31 @@ class CausalDMDDenosingStage(DenoisingStage): self.kv_cache1 = kv_cache1 - def _initialize_crossattn_cache(self, batch_size, dtype, device) -> None: + def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype, + device) -> None: """ Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. """ crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = self.transformer.attention_head_dim for _ in range(self.num_transformer_blocks): crossattn_cache.append({ "k": - torch.zeros([batch_size, 512, 12, 128], + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], dtype=dtype, device=device), "v": - torch.zeros([batch_size, 512, 12, 128], + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], dtype=dtype, device=device), "is_init": False, }) - self.crossattn_cache = crossattn_cache + self.crossattn_cache = crossattn_cache \ No newline at end of file diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 8f08b22a6..20ffa4580 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -778,8 +778,7 @@ class DmdDenoisingStage(DenoisingStage): **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) - from fastvideo.training.training_utils import ( - pred_noise_to_pred_video) + from fastvideo.models.utils import pred_noise_to_pred_video pred_video = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 58c27276f..f167ddd5b 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -37,8 +37,9 @@ from fastvideo.training.activation_checkpoint import ( from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, - load_distillation_checkpoint, pred_noise_to_pred_video, + load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep, EMA_FSDP) +from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.utils import is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory import wandb # isort: skip @@ -1085,16 +1086,16 @@ class DistillationPipeline(TrainingPipeline): latents.dtype) else: latents += self.vae.shift_factor - with torch.autocast("cuda", dtype=torch.bfloat16): - video = self.vae.decode(latents) - video = (video / 2 + 0.5).clamp(0, 1) - video = video.cpu().float() - video = video.permute(0, 2, 1, 3, 4) - video = (video * 255).numpy().astype(np.uint8) - wandb_loss_dict[latent_key] = wandb.Video( - video, fps=24, format="mp4") # change to 16 for Wan2.1 - # Clean up references - del video, latents + with torch.autocast("cuda", dtype=torch.bfloat16): + video = self.vae.decode(latents) + video = (video / 2 + 0.5).clamp(0, 1) + video = video.cpu().float() + video = video.permute(0, 2, 1, 3, 4) + video = (video * 255).numpy().astype(np.uint8) + wandb_loss_dict[latent_key] = wandb.Video( + video, fps=24, format="mp4") # change to 16 for Wan2.1 + # Clean up references + del video, latents # Process DMD training data if available - use decode_stage instead of self.vae.decode if 'generator_pred_video' in dmd_latents_vis_dict: diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 260dd713d..e939891cf 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -18,10 +18,10 @@ from fastvideo.pipelines import TrainingBatch from fastvideo.training.distillation_pipeline import DistillationPipeline from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, - pred_noise_to_pred_video, EMA_FSDP, save_distillation_checkpoint, ) +from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.distributed import get_world_group import torch.distributed as dist import numpy as np diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 92cca5e38..29d620133 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -842,27 +842,6 @@ def custom_to_hf_state_dict( return new_state_dict -def pred_noise_to_pred_video(pred_noise: torch.Tensor, - noise_input_latent: torch.Tensor, - timestep: torch.Tensor, - scheduler: Any) -> torch.Tensor: - """ - Convert predicted noise to clean latent. - """ - timestep = timestep.expand(noise_input_latent.shape[0]) - dtype = pred_noise.dtype - device = pred_noise.device - pred_noise = pred_noise.float().to(device) - noise_input_latent = noise_input_latent.float().to(device) - sigmas = scheduler.sigmas.float().to(device) - timesteps = scheduler.timesteps.float().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) - pred_video = noise_input_latent - sigma_t * pred_noise - return pred_video.to(dtype) - - def shift_timestep(timestep: torch.Tensor, shift: float, num_train_timestep: float) -> torch.Tensor: if shift == 1: