[misc]: Refactor tests to remove unused pipeline main imports and use runner

Author: techverve

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -13,9 +13,10 @@
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
 
 wandb_name = "test_training_loss_VSA"
 reference_wandb_summary_file = "fastvideo/tests/training/VSA/h100_reference_wandb_summary_VSA.json"
@@ -50,6 +51,9 @@ def run_worker():
         "1", "--VSA_sparsity", "0.9"
     ])
 
+    args.pipeline_class = "WanTrainingPipeline"
+    args.pipeline_module = "fastvideo.training.wan_training_pipeline"
+    
     # Call the main training function
     main(args)
 

fastvideo/tests/training/Vanilla/mfu_calculation.py

@@ -14,7 +14,7 @@
 from huggingface_hub import snapshot_download
 from fastvideo.utils import logger
 # Import the training pipeline
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.training.wan_training_pipeline import WanTrainingPipeline

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -11,7 +11,7 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.training.wan_training_pipeline import WanTrainingPipeline

fastvideo/tests/training/distill/test_distill_dmd.py

@@ -11,7 +11,7 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.training.wan_distillation_pipeline import WanDistillationPipeline