diff --git a/examples/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm b/examples/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm index f17682d56..814594db8 100644 --- a/examples/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm +++ b/examples/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm @@ -98,7 +98,7 @@ validation_args=( --log_validation --validation_dataset_file "$VALIDATION_DATASET_FILE" --validation_steps 10 - --validation_sampling_steps "1" + --validation_sampling_steps "4" --validation_guidance_scale "6.0" # not used for dmd inference ) @@ -121,10 +121,10 @@ miscellaneous_args=( --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" + --use_ema True + --ema_decay 0.99 + --ema_start_step 100 + --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors" ) # Self-forcing DMD arguments diff --git a/examples/distill/SFWan2.1-I2V/download_single_example.sh b/examples/distill/SFWan2.1-I2V/download_single_example.sh index cf41fd648..2a3aacda0 100644 --- a/examples/distill/SFWan2.1-I2V/download_single_example.sh +++ b/examples/distill/SFWan2.1-I2V/download_single_example.sh @@ -30,7 +30,7 @@ cat > "$SINGLE_EXAMPLE_DIR/validation.json" << 'EOF' "num_inference_steps": 50, "height": 480, "width": 832, - "num_frames": 77 + "num_frames": 81 } ] } diff --git a/examples/distill/SFWan2.1-I2V/preprocess_data.sh b/examples/distill/SFWan2.1-I2V/preprocess_data.sh index 2f1203060..96ea89e78 100644 --- a/examples/distill/SFWan2.1-I2V/preprocess_data.sh +++ b/examples/distill/SFWan2.1-I2V/preprocess_data.sh @@ -14,7 +14,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --seed 42 \ --max_height 480 \ --max_width 832 \ - --num_frames 77 \ + --num_frames 81 \ --dataloader_num_workers 0 \ --output_dir=$OUTPUT_DIR \ --train_fps 16 \ diff --git a/examples/distill/SFWan2.1-I2V/preprocess_single_example.sh b/examples/distill/SFWan2.1-I2V/preprocess_single_example.sh index 7b251401a..448d8a3c6 100644 --- a/examples/distill/SFWan2.1-I2V/preprocess_single_example.sh +++ b/examples/distill/SFWan2.1-I2V/preprocess_single_example.sh @@ -14,7 +14,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --seed 42 \ --max_height 480 \ --max_width 832 \ - --num_frames 77 \ + --num_frames 81 \ --dataloader_num_workers 0 \ --output_dir=$OUTPUT_DIR \ --train_fps 16 \ diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index d06dd604d..5bb30c438 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -482,7 +482,7 @@ class CausalWanTransformer3DModel(BaseDiT): freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None - print("hidden_states tensor type: ", type(hidden_states)) + hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index c24f15add..9975c382b 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -157,7 +157,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): timestep=timestep, scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2]) - # Step 7: Clean up caches to prevent memory leaks self._reset_simulation_caches(kv_cache, crossattn_cache) return pred_video @@ -200,7 +199,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 4: Determine gradient masking for frame-level selective training num_frames = current_noise_latents.shape[1] - num_frame_per_block = getattr(self.training_args, 'num_frame_per_block', 4) + num_frame_per_block = getattr(self.training_args, 'num_frame_per_block', 3) independent_first_frame = getattr(self.training_args, 'independent_first_frame', False) enable_gradient_masking = getattr(self.training_args, 'enable_gradient_masking', True) gradient_mask_last_n_frames = getattr(self.training_args, 'gradient_mask_last_n_frames', 21) @@ -220,8 +219,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): gradient_mask[:, :num_early_blocks * num_frame_per_block] = False # Step 5: Multi-step simulation with causal block processing - # Get num_frames_per_block from transformer config like causal_denoising.py - num_frames_per_block = getattr(self.training_args, 'num_frame_per_block', 4) + num_frames_per_block = getattr(self.training_args, 'num_frame_per_block', 3) t = num_frames if not independent_first_frame or (independent_first_frame and hasattr(training_batch, 'image_latent') and training_batch.image_latent is not None): @@ -384,7 +382,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): def _initialize_simulation_caches(self, batch_size: int, dtype: torch.dtype, device: torch.device): """Initialize KV cache and cross-attention cache for multi-step simulation.""" - # Get actual model parameters instead of hardcoded values num_transformer_blocks = len(self.transformer.blocks) # Calculate frame sequence length based on input dimensions and patch size @@ -509,7 +506,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch.attn_metadata.VSA_sparsity = 0.0 batches.append(batch) - # Initialize visualization dictionaries training_batch.dmd_latent_vis_dict = {} training_batch.fake_score_latent_vis_dict = {} @@ -811,7 +807,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): progress_bar.update(1) if self.global_rank == 0: - # Prepare logging data log_data = { "train_total_loss": total_loss, "train_fake_score_loss": fake_score_loss, @@ -821,7 +816,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): "avg_step_time": avg_step_time, "grad_norm": grad_norm, } - # Only log generator loss when generator is actually trained if (step % self.dfake_gen_update_ratio == 0): log_data["train_generator_loss"] = generator_loss if use_vsa: @@ -850,12 +844,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): wandb.log(log_data, step=step) - # Add visualization during validation if self.training_args.log_validation and step % self.training_args.validation_steps == 0: if self.training_args.log_visualization: self.visualize_intermediate_latents(training_batch, self.training_args, step) - # Save training state checkpoint (for resuming training) if (self.training_args.training_state_checkpointing_steps > 0 and step % self.training_args.training_state_checkpointing_steps == 0): print("rank", self.global_rank, @@ -872,7 +864,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): self.transformer.train() self.sp_group.barrier() - # Save weight-only checkpoint if (self.training_args.weight_only_checkpointing_steps > 0 and step % self.training_args.weight_only_checkpointing_steps == 0): print("rank", self.global_rank, @@ -893,7 +884,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): wandb.finish() - # Save final training state checkpoint print("rank", self.global_rank, "save final training state checkpoint at step", self.training_args.max_train_steps)