update

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),

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,

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