add SFWan2.2 hf string

This commit is contained in:
RandNMR73
2025-10-03 05:20:11 +00:00
parent 632a06ee6d
commit 2d12867ce2
6 changed files with 37 additions and 16 deletions
@@ -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=(
@@ -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=(
+1
View File
@@ -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,
+3
View File
@@ -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
}
+22 -4
View File
@@ -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":
+2 -4
View File
@@ -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)