--- title: matrixgame2_ode_causal_pipeline ---

fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam

SamplingParam(data_type: str = 'video', image_path: str | None = None, pil_image: Any | None = None, last_image: Any | None = None, references: list[Any] | None = None, video_path: str | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, audio_path: str | None = None, latents: Any | None = None, audio_latents: Any | None = None, mouse_cond: Any | None = None, keyboard_cond: Any | None = None, grid_sizes: Any | None = None, pose: str | None = None, prompt_attention_mask: list = list(), negative_attention_mask: list = list(), c2ws_plucker_emb: Any | None = None, action_path: str | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt: str | list[str] | None = None, negative_prompt: str = 'Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards', max_sequence_length: int | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, num_videos_per_prompt: int = 1, seed: int = 1024, num_frames: int = 125, height: int = 720, width: int = 1280, height_sr: int = 1072, width_sr: int = 1920, fps: int = 24, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, guidance_rescale: float = 0.0, boundary_ratio: float | None = None, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = (lambda: [29])(), ltx2_stg_blocks_audio: list[int] = (lambda: [29])(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: Any | None = None, ltx2_conditioning_latent_stage2: Any | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, save_video: bool = True, return_frames: bool = True, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, *, return_samples: bool = False)

Sampling parameters for video generation.

fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.check_sampling_param

check_sampling_param() -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.update

update(source_dict: dict[str, Any]) -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.from_pretrained

from_pretrained(model_path: str) -> SamplingParam

fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.add_cli_args

add_cli_args(parser: Any) -> Any

Add CLI arguments for SamplingParam fields

fastvideo.training.matrixgame2_ode_causal_pipeline.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.training.matrixgame2_ode_causal_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.matrixgame2_ode_causal_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.matrixgame2_ode_causal_pipeline.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.training.matrixgame2_ode_causal_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.matrixgame2_ode_causal_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.matrixgame2_ode_causal_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.matrixgame2_ode_causal_pipeline.TrainingArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> TrainingArgs

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_ode_causal_pipeline.set_forward_context

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

fastvideo.training.matrixgame2_ode_causal_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.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler

SelfForcingFlowMatchScheduler(num_inference_steps = 100, num_train_timesteps = 1000, shift = 3.0, sigma_max = 1.0, sigma_min = 0.003 / 1.002, inverse_timesteps = False, extra_one_step = False, reverse_sigmas = False, training = False)

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.set_timesteps

set_timesteps(num_inference_steps = 100, denoising_strength = 1.0, training = False, return_dict = False, **kwargs)

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.step

step(model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final = False, return_dict = False, **kwargs)

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high

calculate_alpha_beta_high(sigma, sigma_bound)

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.add_noise

add_noise(original_samples, noise, timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.add_noise_high

add_noise_high(original_samples, noise, timestep, boundary_timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.training_target

training_target(sample, noise, timestep)

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.training_weight

training_weight(timestep)

Input

Output: the corresponding weighting [B*T]

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor

fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline

MatrixGame2CausalDMDPipeline(*args, **kwargs)

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs) -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_reset

streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_step

streaming_step(keyboard_action, mouse_action) -> ForwardBatch

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_clear

streaming_clear() -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage

DecodingStage(vae, pipeline = None)

Stage for decoding latent representations into pixel space.

This stage handles the decoding of latent representations into the final output format (e.g., pixel values).

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage inputs.

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage outputs.

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.decode

decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor

Decode latent representations into pixel space using VAE.

Parameters

latents: torch.Tensor

Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)

fastvideo_args: FastVideoArgs

Configuration containing: - disable_autocast: Whether to disable automatic mixed precision (default: False) - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency

Returns

torch.Tensor

Decoded video tensor with shape (batch, channels, frames, height, width),

torch.Tensor

normalized to [0, 1] range and moved to CPU as float32

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.streaming_decode

streaming_decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs, cache: list[torch.Tensor | None] | None = None, is_first_chunk: bool = False) -> tuple[torch.Tensor, list[torch.Tensor | None]]

Decode latent representations into pixel space using VAE with streaming cache.

Parameters

latents: torch.Tensor

Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)

fastvideo_args: FastVideoArgs

Configuration object.

cache: list[torch.Tensor | None] | None

VAE cache from previous call, or None to initialize a new cache.

is_first_chunk: bool

Whether this is the first chunk.

Returns

tuple[torch.Tensor, list[torch.Tensor | None]]

A tuple of (decoded_frames, updated_cache).

fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch

Decode latent representations into pixel space.

This method processes the batch through the VAE decoder, converting latent representations to pixel-space video/images. It also optionally decodes trajectory latents for visualization purposes.

Parameters

batch: ForwardBatch

The current batch containing: - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents) - return_trajectory_decoded (optional): Flag to decode trajectory latents - trajectory_latents (optional): Latents at different timesteps - trajectory_timesteps (optional): Corresponding timesteps

fastvideo_args: FastVideoArgs

Configuration containing: - output_type: "latent" to skip decoding, otherwise decode to pixels - vae_cpu_offload: Whether to offload VAE to CPU after decoding - model_loaded: Track VAE loading state - model_paths: Path to VAE model if loading needed

Returns

ForwardBatch

Modified batch with: - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 - trajectory_decoded (if requested): List of decoded frames per timestep

fastvideo.training.matrixgame2_ode_causal_pipeline.ForwardBatch

ForwardBatch(data_type: str, generator: torch.Generator | list[torch.Generator] | None = None, image_path: str | None = None, image_embeds: list[torch.Tensor] = list(), pil_image: torch.Tensor | PIL.Image.Image | None = None, last_image: torch.Tensor | PIL.Image.Image | None = None, references: list[Any] | None = None, preprocessed_image: torch.Tensor | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, video_path: str | None = None, video_latent: torch.Tensor | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, pose_latents: torch.Tensor | None = None, face_pixel_values: torch.Tensor | None = None, audio_path: str | None = None, audio_embeds: torch.Tensor | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt_embeds: list[torch.Tensor] = list(), negative_prompt_embeds: list[torch.Tensor] | None = None, prompt_attention_mask: list[torch.Tensor] | None = None, negative_attention_mask: list[torch.Tensor] | None = None, clip_embedding_pos: list[torch.Tensor] | None = None, clip_embedding_neg: list[torch.Tensor] | None = None, max_sequence_length: int | None = None, prompt_template: dict[str, Any] | None = None, do_classifier_free_guidance: bool = False, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, batch_size: int | None = None, num_videos_per_prompt: int = 1, seed: int | None = None, seeds: list[int] | None = None, is_prompt_processed: bool = False, latents: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, lq_latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_pred: torch.Tensor | None = None, image_latent: torch.Tensor | None = None, first_frame_latent: torch.Tensor | None = None, helios_latent_chunks: list[torch.Tensor] | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, num_iterations: int | None = None, use_base_model: bool = False, pose: str | None = None, c2ws_plucker_emb: torch.Tensor | None = None, action_path: str | None = None, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, height_latents: list[int] | int | None = None, width_latents: list[int] | int | None = None, num_frames: list[int] | int = 1, height: list[int] | int | None = None, width: list[int] | int | None = None, height_sr: list[int] | int | None = None, width_sr: list[int] | int | None = None, fps: list[int] | int | None = None, timesteps: torch.Tensor | None = None, timestep: torch.Tensor | float | int | None = None, step_index: int | None = None, boundary_ratio: float | None = None, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, guidance_rescale: float = 0.0, eta: float = 0.0, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = list(), ltx2_stg_blocks_audio: list[int] = list(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: torch.Tensor | None = None, ltx2_conditioning_latent_stage2: torch.Tensor | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, n_tokens: int | None = None, extra_step_kwargs: dict[str, Any] = dict(), modules: dict[str, Any] = dict(), output: torch.Tensor | None = None, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, trajectory_timesteps: list[torch.Tensor] | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_decoded: list[torch.Tensor] | None = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, extra: dict[str, Any] = dict(), save_video: bool = True, return_frames: bool = False, is_cfg_negative: bool = False, VSA_sparsity: float = 0.0, logging_info: PipelineLoggingInfo = PipelineLoggingInfo(), *, return_samples: bool = False)

Complete state passed through the pipeline execution.

This dataclass contains all information needed during the diffusion pipeline execution, allowing methods to update specific components without needing to manage numerous individual parameters.

fastvideo.training.matrixgame2_ode_causal_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.matrixgame2_ode_causal_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.matrixgame2_ode_causal_pipeline.TrainingPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.set_schemas

set_schemas() -> None

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.train_one_step

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.train

train() -> None

fastvideo.training.matrixgame2_ode_causal_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.matrixgame2_ode_causal_pipeline.clip_grad_norm_while_handling_failing_dtensor_cases

clip_grad_norm_while_handling_failing_dtensor_cases(parameters: torch.Tensor | list[torch.Tensor], max_norm: float, norm_type: float = 2.0, error_if_nonfinite: bool = False, foreach: bool | None = None, pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None) -> torch.Tensor | None

fastvideo.training.matrixgame2_ode_causal_pipeline.shallow_asdict

shallow_asdict(obj) -> dict[str, Any]

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline

MatrixGame2ODEInitTrainingPipeline(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 ODE-init using precomputed denoising trajectories.

Supervision: predict the next latent in the stored trajectory by - feeding current latent at timestep t into the transformer to predict noise - stepping the scheduler with the predicted noise - minimizing MSE to the stored next latent at timestep t_next

Source
class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline):
    """
    Training pipeline for ODE-init using precomputed denoising trajectories.

    Supervision: predict the next latent in the stored trajectory by
    - feeding current latent at timestep t into the transformer to predict noise
    - stepping the scheduler with the predicted noise
    - minimizing MSE to the stored next latent at timestep t_next
    """

    _required_config_modules = ["scheduler", "transformer", "vae"]

    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
        # Match the preprocess/generation scheduler for consistent stepping
        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
                                                                  sigma_min=0.0,
                                                                  extra_one_step=True)
        self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)

    def set_schemas(self):
        self.train_dataset_schema = pyarrow_schema_matrixgame2_ode_trajectory

    def initialize_training_pipeline(self, training_args: TrainingArgs):
        super().initialize_training_pipeline(training_args)

        self.noise_scheduler = self.get_module("scheduler")
        self.vae = self.get_module("vae")
        self.vae.requires_grad_(False)

        self.timestep_shift = self.training_args.pipeline_config.flow_shift
        self.noise_scheduler = SelfForcingFlowMatchScheduler(shift=self.timestep_shift,
                                                             sigma_min=0.0,
                                                             extra_one_step=True)
        self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)

        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))

        logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps)
        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250, 0],
                                                dtype=torch.long,
                                                device=get_local_torch_device())
        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
            logger.info("timesteps: %s", timesteps)
            self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
            logger.info("warped self.dmd_denoising_steps: %s", self.dmd_denoising_steps)
        else:
            raise ValueError("warp_denoising_step must be true")

        self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())

        logger.info("denoising_step_list: %s", self.dmd_denoising_steps)

        logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
        # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
        self._cached_closest_idx_per_dmd = None
        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
        self.manual_idx = 0

    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        logger.info("Initializing validation pipeline...")
        args_copy = deepcopy(training_args)
        args_copy.inference_mode = True
        # Use the same flow-matching scheduler as training for consistent validation.
        validation_scheduler = SelfForcingFlowMatchScheduler(shift=args_copy.pipeline_config.flow_shift,
                                                             sigma_min=0.0,
                                                             extra_one_step=True)
        validation_scheduler.set_timesteps(num_inference_steps=1000, training=True)
        # Warm start validation with current transformer
        self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
            training_args.model_path,
            args=args_copy,  # type: ignore
            inference_mode=True,
            loaded_modules={
                "transformer": self.get_module("transformer"),
                "vae": self.get_module("vae"),
                "scheduler": validation_scheduler,
            },
            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)

    def _get_next_batch(self, training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]:
        batch = next(self.train_loader_iter, None)  # type: ignore
        if batch is None:
            self.current_epoch += 1
            logger.info("Starting epoch %s", self.current_epoch)
            self.train_loader_iter = iter(self.train_dataloader)
            batch = next(self.train_loader_iter)

        # Required fields from parquet (ODE trajectory schema)
        clip_feature = batch['clip_feature']
        first_frame_latent = batch['first_frame_latent']
        keyboard_cond = batch.get('keyboard_cond', None)
        mouse_cond = batch.get('mouse_cond', None)
        infos = batch['info_list']

        # Trajectory tensors may include a leading singleton batch dim per row
        trajectory_latents = batch['trajectory_latents']
        if trajectory_latents.dim() == 7:
            # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
            trajectory_latents = trajectory_latents[:, 0]
        elif trajectory_latents.dim() == 6:
            # already [B, S, C, T, H, W]
            pass
        else:
            raise ValueError(f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}")

        trajectory_timesteps = batch['trajectory_timesteps']
        if trajectory_timesteps.dim() == 3:
            # [B, 1, S] -> [B, S]
            trajectory_timesteps = trajectory_timesteps[:, 0]
        elif trajectory_timesteps.dim() == 2:
            # [B, S]
            pass
        else:
            raise ValueError(f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}")
        # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
        trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)

        # Move to device
        device = get_local_torch_device()
        training_batch.image_embeds = clip_feature.to(device, dtype=torch.bfloat16)
        training_batch.image_latents = first_frame_latent.to(device, dtype=torch.bfloat16)
        if keyboard_cond is not None and keyboard_cond.numel() > 0:
            training_batch.keyboard_cond = keyboard_cond.to(device, dtype=torch.bfloat16)
        else:
            training_batch.keyboard_cond = None
        if mouse_cond is not None and mouse_cond.numel() > 0:
            training_batch.mouse_cond = mouse_cond.to(device, dtype=torch.bfloat16)
        else:
            training_batch.mouse_cond = None
        training_batch.infos = infos

        return training_batch, trajectory_latents[:, :, :self.training_args.num_latent_t].to(
            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)

    def _get_timestep(self,
                      min_timestep: int,
                      max_timestep: int,
                      batch_size: int,
                      num_frame: int,
                      num_frame_per_block: int,
                      uniform_timestep: bool = False) -> torch.Tensor:
        if uniform_timestep:
            timestep = torch.randint(min_timestep, max_timestep, [batch_size, 1], device=self.device,
                                     dtype=torch.long).repeat(1, num_frame)
            return timestep
        else:
            timestep = torch.randint(min_timestep,
                                     max_timestep, [batch_size, num_frame],
                                     device=self.device,
                                     dtype=torch.long)
            # logger.info(f"individual timestep: {timestep}")
            # make the noise level the same within every block
            timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block)
            timestep[:, :, 1:] = timestep[:, :, 0:1]
            timestep = timestep.reshape(timestep.shape[0], -1)
            return timestep

    def _step_predict_next_latent(
            self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, image_embeds: torch.Tensor,
            image_latents: torch.Tensor, keyboard_cond: torch.Tensor | None, mouse_cond: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
        latent_vis_dict: dict[str, torch.Tensor] = {}
        device = get_local_torch_device()
        target_latent = traj_latents[:, -1]

        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
        B, S, num_frames, num_channels, height, width = traj_latents.shape
        expected_action_frames = 1 + 4 * (num_frames - 1)
        if keyboard_cond is None or mouse_cond is None:
            raise ValueError("keyboard_cond/mouse_cond must both be provided for action-follow training. "
                             f"keyboard_cond={None if keyboard_cond is None else tuple(keyboard_cond.shape)}, "
                             f"mouse_cond={None if mouse_cond is None else tuple(mouse_cond.shape)}")
        if keyboard_cond.shape[1] < expected_action_frames:
            raise ValueError("keyboard_cond length is shorter than required for latent frames. "
                             f"got={keyboard_cond.shape[1]}, required>={expected_action_frames}, "
                             f"num_latent_frames={num_frames}")
        if mouse_cond.shape[1] < expected_action_frames:
            raise ValueError("mouse_cond length is shorter than required for latent frames. "
                             f"got={mouse_cond.shape[1]}, required>={expected_action_frames}, "
                             f"num_latent_frames={num_frames}")

        # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
        if self._cached_closest_idx_per_dmd is None:
            traj_ts = traj_timesteps[0].float().cpu()
            dmd_steps = self.dmd_denoising_steps.float().cpu()
            closest_idx = torch.argmin(torch.abs(traj_ts.unsqueeze(0) - dmd_steps.unsqueeze(1)), dim=1)
            self._cached_closest_idx_per_dmd = closest_idx.to(torch.long).cpu()
            logger.info("self._cached_closest_idx_per_dmd: %s", self._cached_closest_idx_per_dmd)
            logger.info("corresponding timesteps: %s", traj_ts[self._cached_closest_idx_per_dmd])

        # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
        # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
        # Output: [B, K, C, T, H, W]
        assert self._cached_closest_idx_per_dmd is not None
        relevant_traj_latents = torch.index_select(traj_latents,
                                                   dim=1,
                                                   index=self._cached_closest_idx_per_dmd.to(traj_latents.device))

        indexes = self._get_timestep(  # [B, num_frames]
            0, len(self.dmd_denoising_steps), B, num_frames, 3, uniform_timestep=False)
        noisy_input = torch.gather(relevant_traj_latents,
                                   dim=1,
                                   index=indexes.reshape(B, 1, num_frames, 1, 1,
                                                         1).expand(-1, -1, -1, num_channels, height,
                                                                   width).to(self.device)).squeeze(1)
        latent_model_input = noisy_input.permute(0, 2, 1, 3, 4)
        if image_latents is not None:
            latent_model_input = torch.cat([
                latent_model_input,
                image_latents.to(latent_model_input.device, latent_model_input.dtype),
            ],
                                           dim=1)
        timestep = self.dmd_denoising_steps[indexes]

        # Prepare inputs for transformer
        latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu()
        latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu()

        latent_model_input = latent_model_input.to(device, dtype=torch.bfloat16)
        timestep = timestep.to(device, dtype=torch.bfloat16)

        input_kwargs = {
            "hidden_states": latent_model_input,
            "encoder_hidden_states": None,
            "encoder_hidden_states_image": image_embeds,
            "timestep": timestep,
            "mouse_cond": mouse_cond,
            "keyboard_cond": keyboard_cond,
            "return_dict": False,
        }
        # Predict noise and step the scheduler to obtain next latent
        with set_forward_context(current_timestep=timestep, attn_metadata=None, forward_batch=None):
            noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)

        from fastvideo.models.utils import pred_noise_to_pred_video
        pred_video = pred_noise_to_pred_video(pred_noise=noise_pred.flatten(0, 1),
                                              noise_input_latent=noisy_input.flatten(0, 1),
                                              timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1),
                                              scheduler=self.modules["scheduler"]).unflatten(0, noise_pred.shape[:2])
        latent_vis_dict["pred_video"] = pred_video.permute(0, 2, 1, 3, 4).detach().clone().cpu()

        return pred_video, target_latent, timestep, latent_vis_dict

    def train_one_step(self, training_batch):  # type: ignore[override]
        self.transformer.train()
        self.optimizer.zero_grad()
        training_batch.total_loss = 0.0
        args = cast(TrainingArgs, self.training_args)

        # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent

        for _ in range(args.gradient_accumulation_steps):
            training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
            image_embeds = training_batch.image_embeds
            image_latents = training_batch.image_latents
            keyboard_cond = training_batch.keyboard_cond
            mouse_cond = training_batch.mouse_cond
            assert traj_latents.shape[0] == 1

            # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
            _, S = traj_latents.shape[0], traj_latents.shape[1]
            if S < 2:
                raise ValueError("Trajectory must contain at least 2 steps")

            # Forward to predict next latent by stepping scheduler with predicted noise
            noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
                traj_latents, traj_timesteps, image_embeds, image_latents, keyboard_cond, mouse_cond)

            training_batch.latent_vis_dict.update(latent_vis_dict)

            mask = t != 0

            # Compute loss
            loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
            loss = loss / args.gradient_accumulation_steps

            with set_forward_context(current_timestep=t, attn_metadata=None, forward_batch=None):
                loss.backward()
            avg_loss = loss.detach().clone()
            training_batch.total_loss += avg_loss.item()

        # Clip grad and step optimizers
        grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
            [p for p in self.transformer.parameters() if p.requires_grad],
            args.max_grad_norm if args.max_grad_norm is not None else 0.0)

        self.optimizer.step()
        self.lr_scheduler.step()

        if grad_norm is None:
            grad_value = 0.0
        else:
            try:
                if isinstance(grad_norm, torch.Tensor):
                    grad_value = float(grad_norm.detach().float().item())
                else:
                    grad_value = float(grad_norm)
            except Exception:
                grad_value = 0.0
        training_batch.grad_norm = grad_value
        B, S, T, C, H, W = traj_latents.shape
        training_batch.raw_latent_shape = (B, C, T, H, W)
        return training_batch

    def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args: TrainingArgs,
                                  validation_batch: dict[str, Any], num_inference_steps: int) -> ForwardBatch:
        sampling_param.prompt = validation_batch['prompt']
        sampling_param.height = training_args.num_height
        sampling_param.width = training_args.num_width
        sampling_param.image_path = validation_batch.get('image_path') or validation_batch.get('video_path')
        sampling_param.num_inference_steps = num_inference_steps
        sampling_param.data_type = "video"
        assert self.seed is not None
        sampling_param.seed = self.seed

        latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
        num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1
        sampling_param.num_frames = num_frames
        batch = ForwardBatch(
            **shallow_asdict(sampling_param),
            generator=torch.Generator(device="cpu").manual_seed(self.seed),
            n_tokens=n_tokens,
            eta=0.0,
            VSA_sparsity=training_args.VSA_sparsity,
        )
        if "image" in validation_batch and validation_batch["image"] is not None:
            batch.pil_image = validation_batch["image"]

        if "keyboard_cond" in validation_batch and validation_batch["keyboard_cond"] is not None:
            keyboard_cond = validation_batch["keyboard_cond"]
            keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
            keyboard_cond = keyboard_cond.unsqueeze(0)
            batch.keyboard_cond = keyboard_cond

        if "mouse_cond" in validation_batch and validation_batch["mouse_cond"] is not None:
            mouse_cond = validation_batch["mouse_cond"]
            mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
            mouse_cond = mouse_cond.unsqueeze(0)
            batch.mouse_cond = mouse_cond

        return batch

    def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
        tracker_loss_dict: dict[str, Any] = {}
        latents_vis_dict = training_batch.latent_vis_dict
        latent_log_keys = ['noisy_input', 'x0', 'pred_video']
        for latent_key in latent_log_keys:
            assert latent_key in latents_vis_dict and latents_vis_dict[latent_key] is not None
            latent = latents_vis_dict[latent_key]
            pixel_latent = self.decoding_stage.decode(latent, training_args)

            video = pixel_latent.cpu().float()
            video = video.permute(0, 2, 1, 3, 4)
            video = (video * 255).numpy().astype(np.uint8)
            video_artifact = self.tracker.video(video, fps=16, format="mp4")  # change to 16 for Wan2.1
            if video_artifact is not None:
                tracker_loss_dict[latent_key] = video_artifact
            # Clean up references
            del video, pixel_latent, latent

        if self.global_rank == 0 and tracker_loss_dict:
            self.tracker.log_artifacts(tracker_loss_dict, step)

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)
Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    # Match the preprocess/generation scheduler for consistent stepping
    self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
                                                              sigma_min=0.0,
                                                              extra_one_step=True)
    self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.set_schemas

set_schemas()
Source
def set_schemas(self):
    self.train_dataset_schema = pyarrow_schema_matrixgame2_ode_trajectory

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)
Source
def initialize_training_pipeline(self, training_args: TrainingArgs):
    super().initialize_training_pipeline(training_args)

    self.noise_scheduler = self.get_module("scheduler")
    self.vae = self.get_module("vae")
    self.vae.requires_grad_(False)

    self.timestep_shift = self.training_args.pipeline_config.flow_shift
    self.noise_scheduler = SelfForcingFlowMatchScheduler(shift=self.timestep_shift,
                                                         sigma_min=0.0,
                                                         extra_one_step=True)
    self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)

    self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))

    logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps)
    self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250, 0],
                                            dtype=torch.long,
                                            device=get_local_torch_device())
    if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
        timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
        logger.info("timesteps: %s", timesteps)
        self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
        logger.info("warped self.dmd_denoising_steps: %s", self.dmd_denoising_steps)
    else:
        raise ValueError("warp_denoising_step must be true")

    self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())

    logger.info("denoising_step_list: %s", self.dmd_denoising_steps)

    logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
    # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
    self._cached_closest_idx_per_dmd = None
    self.num_train_timestep = self.noise_scheduler.num_train_timesteps
    self.manual_idx = 0

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)
Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
    logger.info("Initializing validation pipeline...")
    args_copy = deepcopy(training_args)
    args_copy.inference_mode = True
    # Use the same flow-matching scheduler as training for consistent validation.
    validation_scheduler = SelfForcingFlowMatchScheduler(shift=args_copy.pipeline_config.flow_shift,
                                                         sigma_min=0.0,
                                                         extra_one_step=True)
    validation_scheduler.set_timesteps(num_inference_steps=1000, training=True)
    # Warm start validation with current transformer
    self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
        training_args.model_path,
        args=args_copy,  # type: ignore
        inference_mode=True,
        loaded_modules={
            "transformer": self.get_module("transformer"),
            "vae": self.get_module("vae"),
            "scheduler": validation_scheduler,
        },
        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)

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.train_one_step

train_one_step(training_batch)
Source
def train_one_step(self, training_batch):  # type: ignore[override]
    self.transformer.train()
    self.optimizer.zero_grad()
    training_batch.total_loss = 0.0
    args = cast(TrainingArgs, self.training_args)

    # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent

    for _ in range(args.gradient_accumulation_steps):
        training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
        image_embeds = training_batch.image_embeds
        image_latents = training_batch.image_latents
        keyboard_cond = training_batch.keyboard_cond
        mouse_cond = training_batch.mouse_cond
        assert traj_latents.shape[0] == 1

        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
        _, S = traj_latents.shape[0], traj_latents.shape[1]
        if S < 2:
            raise ValueError("Trajectory must contain at least 2 steps")

        # Forward to predict next latent by stepping scheduler with predicted noise
        noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
            traj_latents, traj_timesteps, image_embeds, image_latents, keyboard_cond, mouse_cond)

        training_batch.latent_vis_dict.update(latent_vis_dict)

        mask = t != 0

        # Compute loss
        loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
        loss = loss / args.gradient_accumulation_steps

        with set_forward_context(current_timestep=t, attn_metadata=None, forward_batch=None):
            loss.backward()
        avg_loss = loss.detach().clone()
        training_batch.total_loss += avg_loss.item()

    # Clip grad and step optimizers
    grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
        [p for p in self.transformer.parameters() if p.requires_grad],
        args.max_grad_norm if args.max_grad_norm is not None else 0.0)

    self.optimizer.step()
    self.lr_scheduler.step()

    if grad_norm is None:
        grad_value = 0.0
    else:
        try:
            if isinstance(grad_norm, torch.Tensor):
                grad_value = float(grad_norm.detach().float().item())
            else:
                grad_value = float(grad_norm)
        except Exception:
            grad_value = 0.0
    training_batch.grad_norm = grad_value
    B, S, T, C, H, W = traj_latents.shape
    training_batch.raw_latent_shape = (B, C, T, H, W)
    return training_batch

fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.visualize_intermediate_latents

visualize_intermediate_latents(training_batch: TrainingBatch, training_args: TrainingArgs, step: int)
Source
def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
    tracker_loss_dict: dict[str, Any] = {}
    latents_vis_dict = training_batch.latent_vis_dict
    latent_log_keys = ['noisy_input', 'x0', 'pred_video']
    for latent_key in latent_log_keys:
        assert latent_key in latents_vis_dict and latents_vis_dict[latent_key] is not None
        latent = latents_vis_dict[latent_key]
        pixel_latent = self.decoding_stage.decode(latent, training_args)

        video = pixel_latent.cpu().float()
        video = video.permute(0, 2, 1, 3, 4)
        video = (video * 255).numpy().astype(np.uint8)
        video_artifact = self.tracker.video(video, fps=16, format="mp4")  # change to 16 for Wan2.1
        if video_artifact is not None:
            tracker_loss_dict[latent_key] = video_artifact
        # Clean up references
        del video, pixel_latent, latent

    if self.global_rank == 0 and tracker_loss_dict:
        self.tracker.log_artifacts(tracker_loss_dict, step)

fastvideo.training.matrixgame2_ode_causal_pipeline.main

main(args) -> None

fastvideo.training.matrixgame2_ode_causal_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.