fastvideo.v1.training.wan_training_pipeline
#
Module Contents#
Classes#
A training pipeline for Wan. |
Functions#
Data#
API#
- class fastvideo.v1.training.wan_training_pipeline.WanTrainingPipeline(model_path: str, fastvideo_args: Union[fastvideo.v1.fastvideo_args.FastVideoArgs, fastvideo.v1.fastvideo_args.TrainingArgs], required_config_modules: Optional[List[str]] = None, loaded_modules: Optional[Dict[str, torch.nn.Module]] = None)[source]#
Bases:
fastvideo.v1.training.training_pipeline.TrainingPipeline
A training pipeline for Wan.
Initialization
Initialize the pipeline. After init, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state.
- create_training_stages(training_args: fastvideo.v1.fastvideo_args.TrainingArgs)[source]#
May be used in future refactors.
- forward(batch: fastvideo.v1.pipelines.pipeline_batch_info.ForwardBatch, fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs)[source]#
- initialize_pipeline(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs)[source]#
- initialize_validation_pipeline(training_args: fastvideo.v1.fastvideo_args.TrainingArgs)[source]#
- train_one_step(transformer, model_type, optimizer, lr_scheduler, loader_iter, noise_scheduler, noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale, patch_size, current_vsa_sparsity) tuple[float, float] [source]#