diff --git a/fastvideo/distill/discriminator.py b/fastvideo/distill/discriminator.py index b2bc057f4..44cd144b7 100644 --- a/fastvideo/distill/discriminator.py +++ b/fastvideo/distill/discriminator.py @@ -102,6 +102,8 @@ class Discriminator(nn.Module): class DMDiscriminator(nn.Module): def __init__(self): + super().__init__() + self.cls_pred_branch = nn.Sequential( nn.Conv2d(kernel_size=4, in_channels=1280, out_channels=1280, stride=2, padding=1), # 8x8 -> 4x4 nn.GroupNorm(num_groups=32, num_channels=1280), diff --git a/fastvideo/distill_dmd.py b/fastvideo/distill_dmd.py index d6304be19..905e78da6 100644 --- a/fastvideo/distill_dmd.py +++ b/fastvideo/distill_dmd.py @@ -191,7 +191,7 @@ def predict_noise( encoder_attention_mask, return_dict=False, guidance_scale=1, -) +): do_classifier_free_guidance = guidance_scale > 1.0 latent_model_input = ( torch.cat([noisy_model_input] * 2) if do_classifier_free_guidance else noisy_model_input @@ -1035,6 +1035,12 @@ if __name__ == "__main__": default=0.025, help="The threshold of the linear quadratic scheduler.", ) + parser.add_argument( + "--linear_range", + type=float, + default=0.5, + help="Range for linear quadratic scheduler.", + ) parser.add_argument( "--master_weight_type", type=str, diff --git a/scripts/distill/distill_mochi_dmd.sh b/scripts/distill/distill_mochi_dmd.sh new file mode 100644 index 000000000..b3644751f --- /dev/null +++ b/scripts/distill/distill_mochi_dmd.sh @@ -0,0 +1,38 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online + +torchrun --nnodes 1 --nproc_per_node 4 \ + fastvideo/distill_dmd.py \ + --seed 42 \ + --pretrained_model_name_or_path data/mochi \ + --model_type "mochi" \ + --cache_dir data/.cache \ + --data_json_path data/Mochi-425-Data/video2caption.json \ + --validation_prompt_dir data/Mochi-425-Data/validation \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --num_latent_t 28 \ + --sp_size 4 \ + --train_sp_batch_size 2 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=4000 \ + --learning_rate=1e-6 \ + --mixed_precision=bf16 \ + --checkpointing_steps=64 \ + --validation_steps 1 \ + --validation_sampling_steps 8 \ + --checkpoints_total_limit 3 \ + --allow_tf32 \ + --ema_start_step 0 \ + --cfg 0.0 \ + --log_validation \ + --output_dir="data/outputs/lq_euler_50_thres0.1_lrg_0.75_phase1_lr1e-6_repro" \ + --tracker_project_name PCM \ + --num_frames 29 \ + --scheduler_type pcm_linear_quadratic \ + --validation_guidance_scale 0.5,1.5,2.5 \ + --num_euler_timesteps 50 \ + --linear_quadratic_threshold 0.1 \ + --linear_range 0.75 \ + --multi_phased_distill_schedule 4000-1 \ No newline at end of file