From 2d12867ce2dfe02bc6eab223b0d6fd8a5ebb21bb Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Fri, 3 Oct 2025 05:20:11 +0000 Subject: [PATCH] add SFWan2.2 hf string --- .../SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh | 4 +-- examples/distill/SFWan2.2-A14B/distill_dmd.sh | 13 +++++----- fastvideo/configs/pipelines/registry.py | 1 + fastvideo/configs/sample/registry.py | 3 +++ .../pipelines/stages/causal_denoising.py | 26 ++++++++++++++++--- fastvideo/training/training_pipeline.py | 6 ++--- 6 files changed, 37 insertions(+), 16 deletions(-) diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh index e35453cf0..c28f066bd 100644 --- a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh +++ b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh @@ -64,7 +64,7 @@ parallel_args=( ) model_args=( - --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --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 @@ -103,7 +103,7 @@ miscellaneous_args=( --use_ema True --ema_decay 0.99 --ema_start_step 100 - --init_weights_from_safetensors your_ode_init_weights_path + # --init_weights_from_safetensors your_ode_init_weights_path ) dmd_args=( diff --git a/examples/distill/SFWan2.2-A14B/distill_dmd.sh b/examples/distill/SFWan2.2-A14B/distill_dmd.sh index c9bbfd882..6edc9a9f3 100644 --- a/examples/distill/SFWan2.2-A14B/distill_dmd.sh +++ b/examples/distill/SFWan2.2-A14B/distill_dmd.sh @@ -32,10 +32,10 @@ export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN NUM_GPUS=8 # Model paths for Self-Forcing DMD distillation with Wan2.2: -# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2 +GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2 # REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model # FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model -GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2 +# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2 REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model @@ -58,17 +58,18 @@ training_args=( --train_sp_batch_size 1 --gradient_accumulation_steps 1 --num_latent_t 21 - --num_height 448 # Updated to match Wan2.2 config + --num_height 480 # Updated to match Wan2.2 config --num_width 832 # Updated to match Wan2.2 config --enable_gradient_checkpointing_type "full" --simulate_generator_forward - # --log_visualization + --log_visualization --num_frames 81 --num_frame_per_block 3 # Frame generation block size for self-forcing --enable_gradient_masking --gradient_mask_last_n_frames 21 - # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/ - # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/ + --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/ + --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/ + # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors ) parallel_args=( diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 62a4eefe5..690fb5a9c 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -36,6 +36,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig, "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig, "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, + "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWanT2V480PConfig, "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config, diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 7bd4c4f77..7d614dd2e 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -55,6 +55,9 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { # Causal Self-Forcing Wan2.1 "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, + + # Causal Self-Forcing Wan2.2 + "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWanT2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 3ee6415d3..74d916a39 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -363,8 +363,17 @@ 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 + # Handle FSDP-wrapped models by using getattr + num_attention_heads = getattr(self.transformer, 'num_attention_heads', + None) + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + None) + + if num_attention_heads is None or attention_head_dim is None: + raise AttributeError( + f"Could not access num_attention_heads or attention_head_dim from transformer. " + f"Model type: {type(self.transformer).__name__}") + if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: @@ -400,8 +409,17 @@ class CausalDMDDenosingStage(DenoisingStage): 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 + # Handle FSDP-wrapped models by using getattr + num_attention_heads = getattr(self.transformer, 'num_attention_heads', + None) + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + None) + + if num_attention_heads is None or attention_head_dim is None: + raise AttributeError( + f"Could not access num_attention_heads or attention_head_dim from transformer. " + f"Model type: {type(self.transformer).__name__}") + for _ in range(self.num_transformer_blocks): crossattn_cache.append({ "k": diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index ec7e126cf..e2cc38deb 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -351,10 +351,8 @@ class TrainingPipeline(LoRAPipeline, ABC): if self.train_transformer_2: u = (1 - boundary_ratio ) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1 - # elif self.transformer_2 is not None: - # u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio - # else: # patch for now to align with non-MoE timestep logic - # pass + elif self.transformer_2 is not None: + u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio indices = (u * self.noise_scheduler.config.num_train_timesteps).long() return self.noise_scheduler.timesteps[indices].to(device=device)