add SFWan2.2 hf string
This commit is contained in:
@@ -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=(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user