cosmos2_5_training_pipeline

Cosmos 2.5 training pipeline (text-to-world, full fine-tuning + LoRA).

Follows the same structure as wan_training_pipeline.py. Key Cosmos 2.5 specifics: - Text embeddings are (B, seq, 100352) from Reason1 full-concat (28 layers × 3584 dim) - forward() takes padding_mask (B, 1, H, W) and optional fps int - Normalisation is handled by Cosmos25WanVAEWrapper (handles_latent_norm=True), so _normalize_dit_input is a no-op and stored latents are already normalised. - Timesteps are sigma values in [0, 1] (flow matching); (B,) is auto-expanded inside the model to (B, 1).

`fastvideo.training.cosmos2_5_training_pipeline.get_local_torch_device`

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs`

FastVideoArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875)

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.is_output_rank`

is_output_rank: bool

Whether this process may materialize user-facing outputs.

This is runtime state assigned by an SPMD executor, not a public configuration field, so dataclass serialization and typed config parsing deliberately ignore it.

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.from_kwargs`

from_kwargs(**kwargs: Any) -> FastVideoArgs

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.check_fastvideo_args`

check_fastvideo_args() -> None

Validate inference arguments for consistency

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.finalize_device_offload_policy`

finalize_device_offload_policy(device_id: int = 0) -> bool

Apply device-local memory policy, then resolve incompatible modes.

`fastvideo.training.cosmos2_5_training_pipeline.FastVideoArgs.disable_offload_on_unified_memory`

disable_offload_on_unified_memory(device_id: int = 0, *, offload_flag: str | None = None) -> bool

Disable host offload after a worker has selected its device.

CUDA's unified-memory probe reads runtime device properties and may initialize a CUDA context. Callers must therefore use this only inside a device-owning process, after selecting and binding device_id. Returning the classification lets direct component-loader callers apply the same policy to explicit per-call overrides. When offload_flag is given, the return value says whether this policy covers that component role.

`fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs`

TrainingArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875, pipeline_class: str = '', pipeline_module: str = '', data_path: str = '', dataloader_num_workers: int = 0, num_height: int = 0, num_width: int = 0, num_frames: int = 0, train_batch_size: int = 0, num_latent_t: int = 0, group_frame: bool = False, group_resolution: bool = False, pretrained_model_name_or_path: str = '', real_score_model_path: str = '', fake_score_model_path: str = '', ema_decay: float = 0.0, ema_start_step: int = 0, training_cfg_rate: float = 0.0, precondition_outputs: bool = False, validation_dataset_file: str = '', validation_preprocessed_path: str = '', validation_sampling_steps: str = '', validation_guidance_scale: str = '', validation_steps: float = 0.0, log_validation: bool = False, trackers: list[str] = list(), tracker_project_name: str = '', wandb_run_name: str = '', seed: int | None = None, output_dir: str = '', checkpoints_total_limit: int = 0, resume_from_checkpoint: str = '', num_train_epochs: int = 0, max_train_steps: int = 0, gradient_accumulation_steps: int = 0, learning_rate: float = 0.0, scale_lr: bool = False, lr_scheduler: str = 'constant', lr_warmup_steps: int = 0, max_grad_norm: float = 0.0, enable_gradient_checkpointing_type: str | None = None, selective_checkpointing: float = 0.0, mixed_precision: str = '', train_sp_batch_size: int = 0, fsdp_sharding_startegy: str = '', weighting_scheme: str = '', logit_mean: float = 0.0, logit_std: float = 1.0, mode_scale: float = 0.0, num_euler_timesteps: int = 0, lr_num_cycles: int = 0, lr_power: float = 0.0, min_lr_ratio: float = 0.5, not_apply_cfg_solver: bool = False, distill_cfg: float = 0.0, scheduler_type: str = '', linear_quadratic_threshold: float = 0.0, linear_range: float = 0.0, weight_decay: float = 0.0, betas: str = '0.9,0.999', use_ema: bool = False, multi_phased_distill_schedule: str = '', pred_decay_weight: float = 0.0, pred_decay_type: str = '', hunyuan_teacher_disable_cfg: bool = False, master_weight_type: str = '', VSA_decay_rate: float = 0.01, VSA_decay_interval_steps: int = 1, VSA_cache_tile_buf: bool = False, lora_rank: int | None = None, lora_alpha: int | None = None, lora_training: bool = False, ltx2_first_frame_conditioning_p: float = 0.1, generator_update_interval: int = 5, dfake_gen_update_ratio: int = 5, min_timestep_ratio: float = 0.2, max_timestep_ratio: float = 0.98, real_score_guidance_scale: float = 3.5, fake_score_learning_rate: float = 0.0, fake_score_lr_scheduler: str = 'constant', fake_score_betas: str = '0.9,0.999', training_state_checkpointing_steps: int = 0, weight_only_checkpointing_steps: int = 0, log_visualization: bool = False, visualization_steps: int = 0, simulate_generator_forward: bool = False, warp_denoising_step: bool = False, num_frame_per_block: int = 3, independent_first_frame: bool = False, enable_gradient_masking: bool = True, gradient_mask_last_n_frames: int = 21, same_step_across_blocks: bool = False, last_step_only: bool = False, context_noise: int = 0)

Training arguments. Inherits from FastVideoArgs and adds training-specific arguments. If there are any conflicts, the training arguments will take precedence.

`fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> TrainingArgs

`fastvideo.training.cosmos2_5_training_pipeline.TrainingArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.training.cosmos2_5_training_pipeline.init_logger`

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler`

FlowUniPCMultistepScheduler(num_train_timesteps: int = 1000, solver_order: int = 2, prediction_type: str = 'flow_prediction', shift: float | None = 1.0, use_dynamic_shifting = False, thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, predict_x0: bool = True, solver_type: str = 'bh2', lower_order_final: bool = True, disable_corrector: tuple = (), solver_p: SchedulerMixin = None, timestep_spacing: str = 'linspace', steps_offset: int = 0, final_sigmas_type: str | None = 'zero', **kwargs)

UniPCMultistepScheduler is a training-free framework designed for the fast sampling of diffusion models.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

Parameters

num_train_timesteps: `int`, defaults to 1000: The number of diffusion steps to train the model.
solver_order: `int`, default `2`: The UniPC order which can be any positive integer. The effective order of accuracy is solver_order + 1 due to the UniC. It is recommended to use solver_order=2 for guided sampling, and solver_order=3 for unconditional sampling.
prediction_type: `str`, defaults to "flow_prediction": Prediction type of the scheduler function; must be flow_prediction for this scheduler, which predicts the flow of the diffusion process.
thresholding: `bool`, defaults to `False`: Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion.
dynamic_thresholding_ratio: `float`, defaults to 0.995: The ratio for the dynamic thresholding method. Valid only when thresholding=True.
sample_max_value: `float`, defaults to 1.0: The threshold value for dynamic thresholding. Valid only when thresholding=True and predict_x0=True.
predict_x0: `bool`, defaults to `True`: Whether to use the updating algorithm on the predicted x0.
solver_type: `str`, default `bh2`: Solver type for UniPC. It is recommended to use bh1 for unconditional sampling when steps < 10, and bh2 otherwise.
lower_order_final: `bool`, default `True`: Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
disable_corrector: `list`, default `[]`: Decides which step to disable the corrector to mitigate the misalignment between epsilon_theta(x_t, c) and epsilon_theta(x_t^c, c) which can influence convergence for a large guidance scale. Corrector is usually disabled during the first few steps.
solver_p: `SchedulerMixin`, default `None`: Any other scheduler that if specified, the algorithm becomes solver_p + UniC.
use_karras_sigmas: `bool`, *optional*, defaults to `False`: Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If True, the sigmas are determined according to a sequence of noise levels {σi}.
use_exponential_sigmas: `bool`, *optional*, defaults to `False`: Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
timestep_spacing: `str`, defaults to `"linspace"`: The way the timesteps should be scaled. Refer to Table 2 of the Common Diffusion Noise Schedules and Sample Steps are Flawed for more information.
steps_offset: `int`, defaults to 0: An offset added to the inference steps, as required by some model families.
final_sigmas_type: `str`, defaults to `"zero"`: The final sigma value for the noise schedule during the sampling process. If "sigma_min", the final sigma is the same as the last sigma in the training schedule. If zero, the final sigma is set to 0.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.step_index`

step_index

The index counter for current timestep. It will increase 1 after each scheduler step.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.begin_index`

begin_index

The index for the first timestep. It should be set from pipeline with set_begin_index method.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_shift`

set_shift(shift: float) -> None

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_begin_index`

set_begin_index(begin_index: int = 0)

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

Parameters

begin_index: `int`: The begin index for the scheduler.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.set_timesteps`

set_timesteps(num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None | None = None, shift: float | None | None = None, use_karras_sigmas: bool | None = None, use_kerras_sigma: bool | None = None)

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (int): Total number of the spacing of the time steps. device (str or torch.device, optional): The device to which the timesteps should be moved to. If None, the timesteps are not moved.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.time_shift`

time_shift(mu: float, sigma: float, t: torch.Tensor)

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.convert_model_output`

convert_model_output(model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs) -> torch.Tensor

Convert the model output to the corresponding type the UniPC algorithm needs.

Parameters

model_output: `torch.Tensor`: The direct output from the learned diffusion model.
timestep: `int`: The current discrete timestep in the diffusion chain.
sample: `torch.Tensor`: A current instance of a sample created by the diffusion process.

Returns

torch.Tensor: torch.Tensor: The converted model output.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_p_bh_update`

multistep_uni_p_bh_update(model_output: torch.Tensor, *args, sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor

One step for the UniP (B(h) version). Alternatively, self.solver_p is used if is specified.

Parameters

model_output: `torch.Tensor`: The direct output from the learned diffusion model at the current timestep.
prev_timestep: `int`: The previous discrete timestep in the diffusion chain.
sample: `torch.Tensor`: A current instance of a sample created by the diffusion process.
order: `int`: The order of UniP at this timestep (corresponds to the p in UniPC-p).

Returns

torch.Tensor: torch.Tensor: The sample tensor at the previous timestep.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_c_bh_update`

multistep_uni_c_bh_update(this_model_output: torch.Tensor, *args, last_sample: torch.Tensor = None, this_sample: torch.Tensor = None, order: int | None = None, **kwargs) -> torch.Tensor

One step for the UniC (B(h) version).

Parameters

this_model_output: `torch.Tensor`: The model outputs at x_t.
this_timestep: `int`: The current timestep t.
last_sample: `torch.Tensor`: The generated sample before the last predictor x_{t-1}.
this_sample: `torch.Tensor`: The generated sample after the last predictor x_{t}.
order: `int`: The p of UniC-p at this step. The effective order of accuracy should be order + 1.

Returns

torch.Tensor: torch.Tensor: The corrected sample tensor at the current timestep.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.index_for_timestep`

index_for_timestep(timestep, schedule_timesteps = None) -> int

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.step`

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, generator = None) -> SchedulerOutput | tuple

Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC.

Parameters

model_output: `torch.Tensor`: The direct output from learned diffusion model.
timestep: `int`: The current discrete timestep in the diffusion chain.
sample: `torch.Tensor`: A current instance of a sample created by the diffusion process.
return_dict: `bool`: Whether or not to return a [~schedulers.scheduling_utils.SchedulerOutput] or tuple.

Returns

SchedulerOutput | tuple: [~schedulers.scheduling_utils.SchedulerOutput] or tuple: If return_dict is True, [~schedulers.scheduling_utils.SchedulerOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.scale_model_input`

scale_model_input(sample: torch.Tensor, *args, **kwargs) -> torch.Tensor

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.

Parameters

sample: `torch.Tensor`: The input sample.

Returns

torch.Tensor: torch.Tensor: A scaled input sample.

`fastvideo.training.cosmos2_5_training_pipeline.FlowUniPCMultistepScheduler.add_noise`

add_noise(original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor

`fastvideo.training.cosmos2_5_training_pipeline.Cosmos2_5Pipeline`

Cosmos2_5Pipeline(model_path: str, fastvideo_args: FastVideoArgs | TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

Cosmos 2.5 video generation pipeline.

`fastvideo.training.cosmos2_5_training_pipeline.Cosmos2_5Pipeline.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

`fastvideo.training.cosmos2_5_training_pipeline.TrainingBatch`

TrainingBatch(current_timestep: int = 0, current_vsa_sparsity: float = 0.0, latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_latents: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, audio_noisy_model_input: torch.Tensor | None = None, audio_timesteps: torch.Tensor | None = None, audio_sigmas: torch.Tensor | None = None, audio_noise: torch.Tensor | None = None, audio_encoder_hidden_states: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, conditioning_mask: torch.Tensor | None = None, preprocessed_image: torch.Tensor | None = None, image_embeds: torch.Tensor | None = None, image_latents: torch.Tensor | None = None, infos: list[dict[str, Any]] | None = None, mask_lat_size: torch.Tensor | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_timesteps: torch.Tensor | None = None, noisy_model_input: torch.Tensor | None = None, timesteps: torch.Tensor | None = None, sigmas: torch.Tensor | None = None, noise: torch.Tensor | None = None, training_target: torch.Tensor | None = None, minimax_h3_layout: Any | None = None, attn_metadata_vsa: AttentionMetadata | None = None, attn_metadata: AttentionMetadata | None = None, input_kwargs: dict[str, Any] | None = None, loss: torch.Tensor | None = None, total_loss: float | None = None, grad_norm: float | None = None, encoder_hidden_states_neg: torch.Tensor | None = None, encoder_attention_mask_neg: torch.Tensor | None = None, conditional_dict: dict[str, Any] | None = None, unconditional_dict: dict[str, Any] | None = None, generator_loss: float = 0.0, fake_score_loss: float = 0.0, dmd_latent_vis_dict: dict[str, Any] = dict(), latent_vis_dict: dict[str, Any] = dict(), fake_score_latent_vis_dict: dict[str, Any] = dict(), minimax_h3_condition_video_rows: torch.Tensor | None = None, minimax_h3_condition_audio_rows: torch.Tensor | None = None, minimax_h3_dmd_layout: Any | None = None, minimax_h3_document_video_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_document_audio_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_text_attn_metadata: AttentionMetadata | None = None, sequence_packing_num_documents: int = 1, sequence_packing_num_tokens: int = 0, sequence_padded_num_tokens: int = 0, sequence_padding_ratio: float | None = None, sequence_packing_loss_scale: float = 1.0, sequence_packing_data_forced_ratio: float | None = None)

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline`

TrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

A pipeline for training a model. All training pipelines should inherit from this class. All reusable components and code should be implemented in this class.

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.set_schemas`

set_schemas() -> None

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.initialize_training_pipeline`

initialize_training_pipeline(training_args: TrainingArgs)

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.initialize_validation_pipeline`

initialize_validation_pipeline(training_args: TrainingArgs)

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.train_one_step`

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.train`

train() -> None

`fastvideo.training.cosmos2_5_training_pipeline.TrainingPipeline.visualize_intermediate_latents`

visualize_intermediate_latents(training_batch: TrainingBatch, training_args: TrainingArgs, step: int)

Add visualization data to tracker logging and save frames to disk.

`fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline`

Cosmos25TrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)

Training pipeline for Cosmos 2.5 (text-to-world).

Supports: - Full fine-tuning (all transformer parameters) - LoRA via the inherited LoRAPipeline mechanism (lora_param_names_mapping is set on Cosmos25Transformer3DModel)

Source

class Cosmos25TrainingPipeline(TrainingPipeline):
    """Training pipeline for Cosmos 2.5 (text-to-world).
Supports:
- Full fine-tuning (all transformer parameters)
- LoRA via the inherited LoRAPipeline mechanism
  (lora_param_names_mapping is set on Cosmos25Transformer3DModel)
&quot;&quot;&quot;

_required_config_modules = &#91;&quot;scheduler&quot;, &quot;transformer&quot;, &quot;vae&quot;&#93;

def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    &quot;&quot;&quot;Create the flow-matching scheduler with Cosmos 2.5&#x27;s shift=5.0.&quot;&quot;&quot;
    self.modules&#91;&quot;scheduler&quot;&#93; = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)

def initialize_validation_pipeline(self, training_args: TrainingArgs):
    &quot;&quot;&quot;Build a full Cosmos2_5Pipeline that reuses the training transformer.&quot;&quot;&quot;
    logger.info(&quot;Initializing Cosmos 2.5 validation pipeline...&quot;)
    args_copy = deepcopy(training_args)
    args_copy.inference_mode = True

    validation_pipeline = Cosmos2_5Pipeline.from_pretrained(
        training_args.model_path,
        args=args_copy,
        inference_mode=True,
        loaded_modules=&#123;
            &quot;transformer&quot;: self.get_module(&quot;transformer&quot;),
        &#125;,
        tp_size=training_args.tp_size,
        sp_size=training_args.sp_size,
        num_gpus=training_args.num_gpus,
        pin_cpu_memory=training_args.pin_cpu_memory,
        dit_cpu_offload=True,
    )
    self.validation_pipeline = validation_pipeline

# ------------------------------------------------------------------
# Cosmos 2.5-specific overrides
# ------------------------------------------------------------------

def _normalize_dit_input(self, training_batch: TrainingBatch) -&gt; TrainingBatch:
    &quot;&quot;&quot;Skip: Cosmos25WanVAEWrapper normalises latents during encoding.

    Pre-computed latents stored in the dataset are already normalised
    (per-channel mean/std applied inside the VAE wrapper), so applying a
    second normalisation here would corrupt the inputs.
    &quot;&quot;&quot;
    return training_batch

def _build_input_kwargs(self, training_batch: TrainingBatch) -&gt; TrainingBatch:
    &quot;&quot;&quot;Build the keyword arguments for Cosmos25Transformer3DModel.forward().

    Cosmos 2.5 forward() signature (T2W training subset):
        hidden_states : (B, 16, T, H, W)
        timestep      : (B,)  sigma values in &#91;0, 1&#93; — model expands to (B, 1)
        encoder_hidden_states : (B, seq, 100352)
        padding_mask  : (B, 1, H, W)  — all-ones for T2W (no spatial padding)
        fps           : int            — used only when rope_enable_fps_modulation=True
    &quot;&quot;&quot;
    assert training_batch.noisy_model_input is not None
    assert training_batch.encoder_hidden_states is not None
    assert training_batch.sigmas is not None

    noisy = training_batch.noisy_model_input  # (B, 16, T, H, W)
    batch_size, _c, _t, height, width = noisy.shape
    device = get_local_torch_device()

    # sigmas stored as (B, 1, 1, 1, 1) after get_sigmas; flatten to (B,)
    sigma_vals = training_batch.sigmas.view(batch_size).to(dtype=torch.bfloat16, device=device)

    # T2W: no spatial padding, all-ones padding mask (16 latent + 1 padding).
    padding_mask = torch.ones(
        batch_size,
        1,
        height,
        width,
        device=device,
        dtype=torch.bfloat16,
    )

    # T2W: no inpainting, all-zeros condition mask (adds 1 channel → total 18).
    condition_mask = torch.zeros(
        batch_size,
        1,
        _t,
        height,
        width,
        device=device,
        dtype=torch.bfloat16,
    )

    training_batch.input_kwargs = &#123;
        &quot;hidden_states&quot;: noisy,
        &quot;encoder_hidden_states&quot;: training_batch.encoder_hidden_states,
        # (B,) sigma in &#91;0, 1&#93;; model auto-expands to (B, 1) at line 901
        &quot;timestep&quot;: sigma_vals,
        &quot;condition_mask&quot;: condition_mask,
        &quot;padding_mask&quot;: padding_mask,
        # rope_enable_fps_modulation=False for the 2B checkpoint, so fps
        # has no effect on the embeddings, but the argument is still forwarded.
        &quot;fps&quot;: 24,
    &#125;
    return training_batch</code></pre></details>
fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline.initialize_pipeline
initialize_pipeline(fastvideo_args: FastVideoArgs)
Create the flow-matching scheduler with Cosmos 2.5's shift=5.0.
Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    """Create the flow-matching scheduler with Cosmos 2.5's shift=5.0."""
    self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
fastvideo.training.cosmos2_5_training_pipeline.Cosmos25TrainingPipeline.initialize_validation_pipeline
initialize_validation_pipeline(training_args: TrainingArgs)
Build a full Cosmos2_5Pipeline that reuses the training transformer.
Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
    """Build a full Cosmos2_5Pipeline that reuses the training transformer."""
    logger.info("Initializing Cosmos 2.5 validation pipeline...")
    args_copy = deepcopy(training_args)
    args_copy.inference_mode = True
validation_pipeline = Cosmos2_5Pipeline.from_pretrained(
    training_args.model_path,
    args=args_copy,
    inference_mode=True,
    loaded_modules=&#123;
        &quot;transformer&quot;: self.get_module(&quot;transformer&quot;),
    &#125;,
    tp_size=training_args.tp_size,
    sp_size=training_args.sp_size,
    num_gpus=training_args.num_gpus,
    pin_cpu_memory=training_args.pin_cpu_memory,
    dit_cpu_offload=True,
)
self.validation_pipeline = validation_pipeline</code></pre></details>
fastvideo.training.cosmos2_5_training_pipeline.main
main(args) -> None
fastvideo.training.cosmos2_5_training_pipeline.build_parser
build_parser(pipeline_class: str | None = None, pipeline_module: str | None = None) -> FlexibleArgumentParser
Build the training-runner CLI parser.
When pipeline_class/pipeline_module are given they become optional defaults,
so the deprecated per-pipeline __main__ blocks keep working without the flags.