fastvideo.training.distillation_pipeline

Contents

`fastvideo.training.distillation_pipeline`#

Module Contents#

Classes#

DistillationPipeline

A distillation pipeline for training a 3 step model. Inherits from TrainingPipeline to reuse training infrastructure.

Data#

`logger`
`vsa_available`

API#

class fastvideo.training.distillation_pipeline.DistillationPipeline(model_path: str, fastvideo_args: fastvideo.fastvideo_args.TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)[source]#

Bases: fastvideo.training.training_pipeline.TrainingPipeline

A distillation pipeline for training a 3 step model. Inherits from TrainingPipeline to reuse training infrastructure.

Initialization

Initialize the pipeline. After init, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state.

apply_ema_to_model(model)[source]#: Apply EMA weights to the model for validation or inference.

create_pipeline_stages(fastvideo_args: fastvideo.fastvideo_args.FastVideoArgs)[source]#

current_epoch: int[source]#: 0

current_trainstep: int[source]#: None

faker_score_forward(training_batch: fastvideo.pipelines.TrainingBatch) → tuple[fastvideo.pipelines.TrainingBatch, torch.Tensor][source]#

get_ema_2_model_copy() → torch.nn.Module | None[source]#: Get a copy of the transformer_2 model with EMA weights applied.

get_ema_model_copy() → torch.nn.Module | None[source]#: Get a copy of the model with EMA weights applied.

get_ema_stats() → dict[str, Any][source]#: Get EMA statistics for monitoring.

init_steps: int[source]#: None

initialize_training_pipeline(training_args: fastvideo.fastvideo_args.TrainingArgs)[source]#: Initialize the distillation training pipeline with multiple models.

abstract initialize_validation_pipeline(training_args: fastvideo.fastvideo_args.TrainingArgs)[source]#: Initialize validation pipeline - must be implemented by subclasses.

is_ema_ready(current_step: int | None = None)[source]#: Check if EMA is ready for use (after ema_start_step).

load_module_from_path(model_path: str, module_type: str, training_args: fastvideo.fastvideo_args.TrainingArgs)[source]#

Load a module from a specific path using the same loading logic as the pipeline.

Parameters:

model_path – Path to the model
module_type – Type of module to load (e.g., “transformer”)
training_args – Training arguments

Returns:

The loaded module

reset_ema()[source]#: Reset EMA to current model weights.

save_ema_weights(output_dir: str, step: int)[source]#: Save EMA weights separately for inference purposes.

train() → None[source]#: Main training loop with distillation-specific logging.

train_dataloader: torchdata.stateful_dataloader.StatefulDataLoader[source]#: None

train_fake_score_transformer_2: bool[source]#: False

train_loader_iter: collections.abc.Iterator[dict[str, Any]][source]#: None

train_one_step(training_batch: fastvideo.pipelines.TrainingBatch) → fastvideo.pipelines.TrainingBatch[source]#

validation_pipeline: fastvideo.pipelines.ComposedPipelineBase[source]#: None

video_latent_shape: tuple[int, ...][source]#: None

video_latent_shape_sp: tuple[int, ...][source]#: None

visualize_intermediate_latents(training_batch: fastvideo.pipelines.TrainingBatch, training_args: fastvideo.fastvideo_args.TrainingArgs, step: int)[source]#: Add visualization data to wandb logging and save frames to disk.

fastvideo.training.distillation_pipeline.logger[source]#: ‘init_logger(…)’

fastvideo.training.distillation_pipeline.vsa_available[source]#: ‘is_vsa_available(…)’