[refactor]: remove monolithic main blocks from legacy training pipelines

Author: techverve

fastvideo/training/cosmos2_5_training_pipeline.py

@@ -133,22 +133,3 @@ def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:
 # Entry point (mirrors wan_training_pipeline.py)
 # ---------------------------------------------------------------------------
 
-
-def main(args) -> None:
-    logger.info("Starting Cosmos 2.5 training pipeline...")
-    pipeline = Cosmos25TrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/ltx2_training_pipeline.py

@@ -474,23 +474,3 @@ def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
         training_batch.grad_norm = grad_norm
         return training_batch
 
-
-def main(args) -> None:
-    logger.info("Starting LTX-2 training pipeline...")
-    pipeline = LTX2TrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/matrixgame2_ar_diffusion_pipeline.py

@@ -427,24 +427,3 @@ def _prepare_validation_batch(
 
         return batch
 
-
-def main(args) -> None:
-    logger.info("Starting Matrix-Game 2.0 AR diffusion training pipeline...")
-
-    pipeline = MatrixGame2ARDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Matrix-Game 2.0 AR diffusion training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/matrixgame2_ode_causal_pipeline.py

@@ -875,23 +875,3 @@ def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args
 
         return batch
 
-
-def main(args) -> None:
-    logger.info("Starting Matrix-Game 2.0 self-forcing distillation pipeline...")
-
-    pipeline = MatrixGame2SelfForcingDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Matrix-Game 2.0 self-forcing distillation pipeline completed")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    main(args)

fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py

@@ -194,23 +194,3 @@ def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args
 
         return batch
 
-
-def main(args) -> None:
-    logger.info("Starting training pipeline...")
-
-    pipeline = MatrixGame2TrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/matrixgame2_training_pipeline.py

@@ -309,22 +309,3 @@ def visualize_intermediate_latents(self, training_batch: TrainingBatch, training
         if self.global_rank == 0 and tracker_loss_dict:
             self.tracker.log_artifacts(tracker_loss_dict, step)
 
-
-def main(args) -> None:
-    logger.info("Starting ODE-init training pipeline...")
-    pipeline = ODEInitTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("ODE-init training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/ode_causal_pipeline.py

@@ -52,25 +52,3 @@ def initialize_validation_pipeline(self, training_args: TrainingArgs):
 
         self.validation_pipeline = validation_pipeline
 
-
-def main(args) -> None:
-    logger.info("Starting Wan distillation pipeline...")
-
-    # Create pipeline with original args
-    pipeline = WanDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-
-    args = pipeline.training_args
-    # Start training
-    pipeline.train()
-    logger.info("Wan distillation pipeline completed")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    main(args)

fastvideo/training/wan_distillation_pipeline.py

@@ -174,27 +174,3 @@ def _build_distill_input_kwargs(self, noise_input: torch.Tensor, timestep: torch
 
         return training_batch
 
-
-def main(args) -> None:
-    logger.info("Starting Wan distillation pipeline...")
-
-    # Create pipeline with original args
-    pipeline = WanI2VDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-
-    args = pipeline.training_args
-
-    # Start training
-    pipeline.train()
-    logger.info("Wan distillation pipeline completed")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)

fastvideo/training/wan_i2v_distillation_pipeline.py

@@ -165,23 +165,3 @@ def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args
 
         return batch
 
-
-def main(args) -> None:
-    logger.info("Starting training pipeline...")
-
-    pipeline = WanI2VTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    main(args)