From ed1e17f37f910dde21b60d9ef1bb087d827e90ff Mon Sep 17 00:00:00 2001
From: rlsu9 <a1286225768@gmail.com>
Date: Fri, 29 Nov 2024 11:42:17 +0400
Subject: [PATCH] add ema transformer

---
 fastvideo/distill.py |  4 +++-
 t2v_debug_multi.sh   | 26 ++++++++++++++------------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/fastvideo/distill.py b/fastvideo/distill.py
index 92cff4807..db3e32f52 100644
--- a/fastvideo/distill.py
+++ b/fastvideo/distill.py
@@ -13,6 +13,7 @@ from fastvideo.utils.validation import log_validation
 import time
 from torch.utils.data import DataLoader
 import torch
+from torch.distributed.fsdp import  ShardingStrategy
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
     StateDictType,
@@ -396,7 +397,7 @@ def main(args):
         params_to_optimize,
         lr=args.learning_rate,
         betas=(0.9,0.999),
-        weight_decay=0.01,
+        weight_decay=args.weight_decay,
         eps=1e-8,
     )
 
@@ -642,5 +643,6 @@ if __name__ == "__main__":
     # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
     parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
     parser.add_argument("--linear_quadratic_threshold", type=float, default=0.025, help="Threshold for linear quadratic scheduler.")
+    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
     args = parser.parse_args()
     main(args)
\ No newline at end of file
diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh
index ac567fd6b..33801e8db 100644
--- a/t2v_debug_multi.sh
+++ b/t2v_debug_multi.sh
@@ -1,4 +1,6 @@
+
 export WANDB_MODE=online
+export WANDB_API_KEY=4f6de3765d6464f43e0506ec7d785641af645e73
 
 torchrun --nnodes 4 --nproc_per_node 4\
     --node_rank=3 \
@@ -9,9 +11,9 @@ torchrun --nnodes 4 --nproc_per_node 4\
     --seed 42\
     --pretrained_model_name_or_path data/mochi\
     --cache_dir "data/.cache"\
-    --data_json_path "data/General-Video/videos2caption.json"\
-    --validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask"\
-    --uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask"\
+    --data_json_path "data/Merge-30k-Data/video2caption.json"\
+    --validation_prompt_dir "data/validation_embeddings/validation_prompt_embed_mask"\
+    --uncond_prompt_dir "data/validation_embeddings/uncond_prompt_embed_mask"\
     --gradient_checkpointing\
     --train_batch_size=1\
     --num_latent_t 28\
@@ -19,11 +21,11 @@ torchrun --nnodes 4 --nproc_per_node 4\
     --train_sp_batch_size 2\
     --dataloader_num_workers 4\
     --gradient_accumulation_steps=1\
-    --max_train_steps=20000\
+    --max_train_steps=4000\
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
-    --checkpointing_steps=250\
-    --validation_steps 25\
+    --checkpointing_steps=500\
+    --validation_steps 250\
     --validation_sampling_steps 8 \
     --checkpoints_total_limit 3\
     --allow_tf32\
@@ -31,13 +33,13 @@ torchrun --nnodes 4 --nproc_per_node 4\
     --cfg 0.0\
     --ema_decay 0.999\
     --log_validation\
-    --output_dir="data/outputs/adv_video_distill_linear_quadratic_debugged"\
+    --output_dir="data/outputs/shift8_euler_50"\
     --tracker_project_name PCM \
-    --num_frames  67 \
-    --scheduler_type pcm_linear_quadratic \
-    --validation_guidance_scale 4.5 \
-    --num_euler_timesteps 50
-    
+    --num_frames  163 \
+    --shift 8.0 \
+    --validation_guidance_scale 4.5  \
+    --num_euler_timesteps 50 \
+    --weight_decay 0.0