Skip to content

Commit c2d2e8e

Browse files
techvervetechverve
committed
[refactor]: decouple build_parser and fix tests
1 parent fee5c49 commit c2d2e8e

20 files changed

Lines changed: 103 additions & 130 deletions

‎fastvideo/tests/training/VSA/test_training_loss_VSA.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
1818
from fastvideo.utils import FlexibleArgumentParser
1919
from fastvideo.training.runner import main
20+
from fastvideo.utils import build_parser
2021

2122
wandb_name = "test_training_loss_VSA"
2223
reference_wandb_summary_file = "fastvideo/tests/training/VSA/h100_reference_wandb_summary_VSA.json"
@@ -28,12 +29,12 @@
2829
def run_worker():
2930
"""Worker function that will be run on each GPU"""
3031
# Create and populate args
31-
parser = FlexibleArgumentParser()
32-
parser = TrainingArgs.add_cli_args(parser)
33-
parser = FastVideoArgs.add_cli_args(parser)
32+
parser = build_parser()
3433

3534
# Set the arguments as they are in finetune_v1_test.sh
3635
args = parser.parse_args([
36+
"--pipeline_class", "WanTrainingPipeline",
37+
"--pipeline_module", "fastvideo.training.wan_training_pipeline",
3738
"--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
3839
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path",
3940
"data/mini_dataset_i2v_VSA/combined_parquet_dataset", "--validation_dataset_file",
@@ -50,9 +51,6 @@ def run_worker():
5051
"--dit_precision", "fp32", "--max_grad_norm", "1.0", "--VSA_decay_rate", "0.01", "--VSA_decay_interval_steps",
5152
"1", "--VSA_sparsity", "0.9"
5253
])
53-
54-
args.pipeline_class = "WanTrainingPipeline"
55-
args.pipeline_module = "fastvideo.training.wan_training_pipeline"
5654

5755
# Call the main training function
5856
main(args)

‎fastvideo/tests/training/Vanilla/mfu_calculation.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
1919
from fastvideo.utils import FlexibleArgumentParser
2020
from fastvideo.training.runner import main
21+
from fastvideo.utils import build_parser
2122

2223
MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
2324
DATA_PATH = "data/crush-smol_processed_t2v/training_dataset/worker_1/worker_0/"
@@ -38,12 +39,12 @@
3839
def run_worker():
3940
"""Worker function that will be run on each GPU"""
4041
# Create and populate args
41-
parser = FlexibleArgumentParser()
42-
parser = TrainingArgs.add_cli_args(parser)
43-
parser = FastVideoArgs.add_cli_args(parser)
42+
parser = build_parser()
4443

4544
# Set the arguments as they are in finetune_t2v.sh
4645
args = parser.parse_args([
46+
"--pipeline_class", "WanTrainingPipeline",
47+
"--pipeline_module", "fastvideo.training.wan_training_pipeline",
4748
"--model_path",
4849
MODEL_PATH,
4950
"--inference_mode",
@@ -122,9 +123,6 @@ def run_worker():
122123
"--hsdp_shard_dim",
123124
"1"
124125
])
125-
# Call the main training function
126-
args.pipeline_class = "WanTrainingPipeline"
127-
args.pipeline_module = "fastvideo.training.wan_training_pipeline"
128126

129127
# Call the main training function
130128
main(args)

‎fastvideo/tests/training/Vanilla/test_training_loss.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
1616
from fastvideo.utils import FlexibleArgumentParser
17+
from fastvideo.utils import FlexibleArgumentParser
1718
from fastvideo.training.runner import main
19+
from fastvideo.utils import build_parser
1820

1921
wandb_name = "test_training_loss"
2022
a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json"
@@ -28,12 +30,12 @@
2830
def run_worker():
2931
"""Worker function that will be run on each GPU"""
3032
# Create and populate args
31-
parser = FlexibleArgumentParser()
32-
parser = TrainingArgs.add_cli_args(parser)
33-
parser = FastVideoArgs.add_cli_args(parser)
33+
parser = build_parser()
3434

3535
# Set the arguments as they are in finetune_v1_test.sh
3636
args = parser.parse_args([
37+
"--pipeline_class", "WanTrainingPipeline",
38+
"--pipeline_module", "fastvideo.training.wan_training_pipeline",
3739
"--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
3840
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path",
3941
"data/crush-smol_processed_t2v/combined_parquet_dataset", "--validation_dataset_file",
@@ -49,9 +51,6 @@ def run_worker():
4951
"--num_euler_timesteps", "50", "--multi_phased_distill_schedule", "4000-1", "--weight_decay", "0.01",
5052
"--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0"
5153
])
52-
# Call the main training function
53-
args.pipeline_class = "WanTrainingPipeline"
54-
args.pipeline_module = "fastvideo.training.wan_training_pipeline"
5554

5655
# Call the main training function
5756
main(args)

‎fastvideo/tests/training/distill/test_distill_dmd.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
1616
from fastvideo.utils import FlexibleArgumentParser
17+
from fastvideo.utils import FlexibleArgumentParser
1718
from fastvideo.training.runner import main
19+
from fastvideo.utils import build_parser
1820

1921
wandb_name = "test_distill_dmd"
2022

@@ -25,12 +27,12 @@
2527
def run_worker():
2628
"""Worker function that will be run on each GPU"""
2729
# Create and populate args
28-
parser = FlexibleArgumentParser()
29-
parser = TrainingArgs.add_cli_args(parser)
30-
parser = FastVideoArgs.add_cli_args(parser)
30+
parser = build_parser()
3131

3232
# Set the arguments as they are in finetune_v1_test.sh
3333
args = parser.parse_args([
34+
"--pipeline_class", "WanDistillationPipeline",
35+
"--pipeline_module", "fastvideo.training.wan_distillation_pipeline",
3436
"--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
3537
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--real_score_model_path",
3638
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--fake_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
@@ -48,9 +50,6 @@ def run_worker():
4850
"--min_timestep_ratio", "0.02", "--max_timestep_ratio", "0.98", "--real_score_guidance_scale", "3.5",
4951
"--enable_gradient_checkpointing_type", "full"
5052
])
51-
# Call the main training function
52-
args.pipeline_class = "WanDistillationPipeline"
53-
args.pipeline_module = "fastvideo.training.wan_distillation_pipeline"
5453

5554
# Call the main training function
5655
main(args)

‎fastvideo/tests/training/self-forcing/test_self_forcing.py‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
1616
from fastvideo.utils import FlexibleArgumentParser
17+
from fastvideo.utils import FlexibleArgumentParser
1718
from fastvideo.training.runner import main
19+
from fastvideo.utils import build_parser
1820

1921
wandb_name = "test_self_forcing_distill"
2022

@@ -25,12 +27,12 @@
2527
def run_worker():
2628
"""Worker function that will be run on each GPU"""
2729
# Create and populate args
28-
parser = FlexibleArgumentParser()
29-
parser = TrainingArgs.add_cli_args(parser)
30-
parser = FastVideoArgs.add_cli_args(parser)
30+
parser = build_parser()
3131

3232
# Set the arguments based on the distill_dmd_t2v_1.3B.sh script
3333
args = parser.parse_args([
34+
"--pipeline_class", "WanSelfForcingDistillationPipeline",
35+
"--pipeline_module", "fastvideo.training.wan_self_forcing_distillation_pipeline",
3436
"--model_path",
3537
"wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
3638
"--inference_mode",
@@ -147,10 +149,6 @@ def run_worker():
147149
"--ema_start_step",
148150
"100",
149151
])
150-
151-
# Call the main training function
152-
args.pipeline_class = "WanSelfForcingDistillationPipeline"
153-
args.pipeline_module = "fastvideo.training.wan_self_forcing_distillation_pipeline"
154152

155153
# Call the main training function
156154
main(args)

‎fastvideo/tests/training/test_runner.py‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from unittest.mock import patch, MagicMock
22
from fastvideo.training.runner import main
33
import argparse
4+
from fastvideo.utils import build_parser
45

56
def test_runner_invokes_correct_class():
67
with patch('fastvideo.training.runner.importlib.import_module') as mock_import_module:
@@ -27,20 +28,6 @@ def test_runner_invokes_correct_class():
2728

2829

2930
def test_runner_cli_argument_parsing():
30-
from fastvideo.training.runner import FlexibleArgumentParser
31-
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
32-
33-
def build_parser():
34-
parser = FlexibleArgumentParser()
35-
parser.add_argument("--pipeline-class",
36-
type=str,
37-
required=True)
38-
parser.add_argument("--pipeline-module",
39-
type=str,
40-
required=True)
41-
parser = TrainingArgs.add_cli_args(parser)
42-
parser = FastVideoArgs.add_cli_args(parser)
43-
return parser
4431

4532
base_required_args = [
4633
"--data_path", "test_path",

‎fastvideo/training/cosmos2_5_training_pipeline.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -143,12 +143,11 @@ def main(args) -> None:
143143

144144

145145
if __name__ == "__main__":
146-
logger.warning(
147-
"\n"
148-
"================================================================================\n"
149-
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
150-
"Please use `fastvideo/training/runner.py` or the `fastvideo.training.runner.Runner` class instead.\n"
151-
"================================================================================")
146+
logger.warning("\n"
147+
"================================================================================\n"
148+
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
149+
"Please use `fastvideo/training/runner.py` instead.\n"
150+
"================================================================================")
152151
from fastvideo.fastvideo_args import TrainingArgs
153152
from fastvideo.utils import FlexibleArgumentParser
154153

‎fastvideo/training/ltx2_training_pipeline.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -484,12 +484,11 @@ def main(args) -> None:
484484

485485

486486
if __name__ == "__main__":
487-
logger.warning(
488-
"\n"
489-
"================================================================================\n"
490-
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
491-
"Please use `fastvideo/training/runner.py` or the `fastvideo.training.runner.Runner` class instead.\n"
492-
"================================================================================")
487+
logger.warning("\n"
488+
"================================================================================\n"
489+
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
490+
"Please use `fastvideo/training/runner.py` instead.\n"
491+
"================================================================================")
493492
argv = sys.argv
494493
from fastvideo.fastvideo_args import TrainingArgs
495494
from fastvideo.utils import FlexibleArgumentParser

‎fastvideo/training/matrixgame2_ar_diffusion_pipeline.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -438,12 +438,11 @@ def main(args) -> None:
438438

439439

440440
if __name__ == "__main__":
441-
logger.warning(
442-
"\n"
443-
"================================================================================\n"
444-
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
445-
"Please use `fastvideo/training/runner.py` or the `fastvideo.training.runner.Runner` class instead.\n"
446-
"================================================================================")
441+
logger.warning("\n"
442+
"================================================================================\n"
443+
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
444+
"Please use `fastvideo/training/runner.py` instead.\n"
445+
"================================================================================")
447446
argv = sys.argv
448447
from fastvideo.fastvideo_args import TrainingArgs
449448
from fastvideo.utils import FlexibleArgumentParser

‎fastvideo/training/matrixgame2_ode_causal_pipeline.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -404,12 +404,11 @@ def main(args) -> None:
404404

405405

406406
if __name__ == "__main__":
407-
logger.warning(
408-
"\n"
409-
"================================================================================\n"
410-
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
411-
"Please use `fastvideo/training/runner.py` or the `fastvideo.training.runner.Runner` class instead.\n"
412-
"================================================================================")
407+
logger.warning("\n"
408+
"================================================================================\n"
409+
"[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
410+
"Please use `fastvideo/training/runner.py` instead.\n"
411+
"================================================================================")
413412
argv = sys.argv
414413
from fastvideo.fastvideo_args import TrainingArgs
415414
from fastvideo.utils import FlexibleArgumentParser

0 commit comments

Comments
 (0)