--- title: matrixgame2_ar_diffusion_pipeline ---

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_pipeline.SamplingParam.check_sampling_param

check_sampling_param() -> None

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SamplingParam.update

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SamplingParam.from_pretrained

from_pretrained(model_path: str) -> SamplingParam

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SamplingParam.add_cli_args

add_cli_args(parser: Any) -> Any

Add CLI arguments for SamplingParam fields

fastvideo.training.matrixgame2_ar_diffusion_pipeline.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_ar_diffusion_pipeline.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.training.matrixgame2_ar_diffusion_pipeline.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.training.matrixgame2_ar_diffusion_pipeline.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.TrainingArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> TrainingArgs

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.set_timesteps

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.step

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high

calculate_alpha_beta_high(sigma, sigma_bound)

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.training_target

training_target(sample, noise, timestep)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.training_weight

training_weight(timestep)

Input

Output: the corresponding weighting [B*T]

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.scale_model_input

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.SelfForcingFlowMatchScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2CausalDMDPipeline

MatrixGame2CausalDMDPipeline(*args, **kwargs)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2CausalDMDPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs) -> None

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2CausalDMDPipeline.streaming_reset

streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2CausalDMDPipeline.streaming_step

streaming_step(keyboard_action, mouse_action) -> ForwardBatch

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2CausalDMDPipeline.streaming_clear

streaming_clear() -> None

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.TrainingPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingPipeline.set_schemas

set_schemas() -> None

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingPipeline.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingPipeline.train_one_step

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

fastvideo.training.matrixgame2_ar_diffusion_pipeline.TrainingPipeline.train

train() -> None

fastvideo.training.matrixgame2_ar_diffusion_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_ar_diffusion_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_ar_diffusion_pipeline.shallow_asdict

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline

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

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

    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
        scheduler = SelfForcingFlowMatchScheduler(
            shift=fastvideo_args.pipeline_config.flow_shift,
            sigma_min=0.0,
            extra_one_step=True,
        )
        scheduler.set_timesteps(num_inference_steps=1000, training=True)
        self.modules["scheduler"] = scheduler

    def _log_training_info(self) -> None:
        self.noise_scheduler = self.modules["scheduler"]
        super()._log_training_info()

    def set_schemas(self):
        self.train_dataset_schema = pyarrow_schema_matrixgame2

    def _get_temporal_compression_ratio(self) -> int:
        assert self.training_args is not None
        return int(self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio)

    def _resolve_num_frame_per_block(self, training_args: TrainingArgs) -> int:
        transformer_num_frame_per_block = getattr(self.transformer, "num_frame_per_block", None)
        requested_num_frame_per_block = getattr(training_args, "num_frame_per_block", None)

        if (transformer_num_frame_per_block is not None and requested_num_frame_per_block is not None
                and transformer_num_frame_per_block != requested_num_frame_per_block):
            raise ValueError("num_frame_per_block mismatch between loaded transformer and "
                             "training args: "
                             f"transformer={transformer_num_frame_per_block}, "
                             f"training_args={requested_num_frame_per_block}")

        if transformer_num_frame_per_block is not None:
            return int(transformer_num_frame_per_block)
        if requested_num_frame_per_block is not None:
            return int(requested_num_frame_per_block)
        return 3

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

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

        self.num_frame_per_block = self._resolve_num_frame_per_block(training_args)

        logger.info(
            "Matrix-Game 2.0 AR diffusion pipeline initialized with "
            "num_frame_per_block=%d, diffusion_forcing_shift=%.1f",
            self.num_frame_per_block,
            training_args.pipeline_config.flow_shift,
        )

    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        logger.info("Initializing Matrix-Game 2.0 AR validation pipeline...")
        args_copy = deepcopy(training_args)
        args_copy.inference_mode = True

        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=False)

        num_val_steps = int(training_args.validation_sampling_steps.split(",")[0])
        step_size = 1000 // num_val_steps
        args_copy.pipeline_config.dmd_denoising_steps = list(range(1000, 0, -step_size))
        args_copy.pipeline_config.warp_denoising_step = True
        training_args.pipeline_config.dmd_denoising_steps = (args_copy.pipeline_config.dmd_denoising_steps)
        training_args.pipeline_config.warp_denoising_step = True

        logger.info(
            "Validation: %d-step Matrix-Game 2.0 causal denoising, "
            "dmd_denoising_steps has %d entries",
            num_val_steps,
            len(args_copy.pipeline_config.dmd_denoising_steps),
        )

        loaded_modules = {
            "transformer": self.get_module("transformer"),
            "vae": self.get_module("vae"),
            "scheduler": validation_scheduler,
        }
        image_encoder = self.get_module("image_encoder")
        image_processor = self.get_module("image_processor")
        if image_encoder is not None:
            loaded_modules["image_encoder"] = image_encoder
        if image_processor is not None:
            loaded_modules["image_processor"] = image_processor

        self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
            training_args.model_path,
            args=args_copy,
            inference_mode=True,
            loaded_modules=loaded_modules,
            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_timestep(
        self,
        batch_size: int,
        num_frame: int,
        num_frame_per_block: int,
    ) -> torch.Tensor:
        """Sample one schedule index per chunk and broadcast to frames.

        Returns the shifted timestep, shape [B, num_frame].
        """
        device = get_local_torch_device()
        num_schedule_steps = len(self.noise_scheduler.timesteps)
        chunk_size = int(num_frame_per_block)
        if chunk_size <= 0:
            raise ValueError(f"num_frame_per_block must be > 0, got {chunk_size}")
        num_chunks = (num_frame + chunk_size - 1) // chunk_size
        chunk_indices = torch.randint(
            0,
            num_schedule_steps,
            (batch_size, num_chunks),
            device=device,
            dtype=torch.long,
        )
        indices = chunk_indices.repeat_interleave(chunk_size, dim=1)[:, :num_frame]
        schedule_t = self.noise_scheduler.timesteps.to(device)
        return schedule_t[indices]

    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
        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)

        latents = batch["vae_latent"]
        latents = latents[:, :, :self.training_args.num_latent_t]
        clip_features = batch["clip_feature"]
        image_latents = batch["first_frame_latent"]
        image_latents = image_latents[:, :, :self.training_args.num_latent_t]
        pil_image = batch["pil_image"]
        infos = batch["info_list"]

        training_batch.latents = latents.to(get_local_torch_device(), dtype=torch.bfloat16)
        training_batch.encoder_hidden_states = None
        training_batch.encoder_attention_mask = None
        training_batch.preprocessed_image = pil_image.to(get_local_torch_device())
        training_batch.image_embeds = clip_features.to(get_local_torch_device())
        training_batch.image_latents = image_latents.to(get_local_torch_device())
        training_batch.infos = infos

        if "mouse_cond" in batch and batch["mouse_cond"].numel() > 0:
            training_batch.mouse_cond = batch["mouse_cond"].to(get_local_torch_device(), dtype=torch.bfloat16)
        else:
            training_batch.mouse_cond = None

        if "keyboard_cond" in batch and batch["keyboard_cond"].numel() > 0:
            training_batch.keyboard_cond = batch["keyboard_cond"].to(get_local_torch_device(), dtype=torch.bfloat16)
        else:
            training_batch.keyboard_cond = None

        temporal_compression_ratio = self._get_temporal_compression_ratio()
        expected_num_frames = (self.training_args.num_latent_t - 1) * temporal_compression_ratio + 1
        if training_batch.keyboard_cond is not None:
            assert training_batch.keyboard_cond.shape[1] >= expected_num_frames, (
                f"keyboard_cond has {training_batch.keyboard_cond.shape[1]} "
                f"frames but need at least {expected_num_frames}")
            training_batch.keyboard_cond = training_batch.keyboard_cond[:, :expected_num_frames]
        if training_batch.mouse_cond is not None:
            assert training_batch.mouse_cond.shape[1] >= expected_num_frames, (
                f"mouse_cond has {training_batch.mouse_cond.shape[1]} frames "
                f"but need at least {expected_num_frames}")
            training_batch.mouse_cond = training_batch.mouse_cond[:, :expected_num_frames]

        return training_batch

    def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch:
        """Prepare diffusion-forcing inputs and Matrix-Game 2.0 I2V concat."""
        assert self.training_args is not None
        latents = training_batch.latents
        assert latents is not None
        batch_size = latents.shape[0]
        num_latent_t = latents.shape[2]

        latents_btchw = latents.permute(0, 2, 1, 3, 4)

        timesteps = self._get_timestep(
            batch_size=batch_size,
            num_frame=num_latent_t,
            num_frame_per_block=self.num_frame_per_block,
        )

        noise_generator = getattr(self, "noise_gen_cuda", None)
        noise = torch.randn(
            latents_btchw.shape,
            generator=noise_generator,
            device=latents_btchw.device,
            dtype=latents_btchw.dtype,
        )
        if (self.training_args.sp_size > 1 and self.sp_group is not None and hasattr(self.sp_group, "broadcast")):
            self.sp_group.broadcast(timesteps, src=0)
            self.sp_group.broadcast(noise, src=0)

        noisy_latents = self.noise_scheduler.add_noise(
            latents_btchw.flatten(0, 1),
            noise.flatten(0, 1),
            timesteps.flatten(0, 1),
        ).unflatten(0, (batch_size, num_latent_t))

        noisy_model_input = noisy_latents.permute(0, 2, 1, 3, 4)

        assert isinstance(training_batch.image_latents, torch.Tensor)
        image_latents = training_batch.image_latents.to(get_local_torch_device(), dtype=torch.bfloat16)

        temporal_compression_ratio = self._get_temporal_compression_ratio()
        num_frames = (num_latent_t - 1) * temporal_compression_ratio + 1
        _, _, _, latent_height, latent_width = image_latents.shape
        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
        mask_lat_size[:, :, 1:] = 0

        first_frame_mask = mask_lat_size[:, :, :1]
        first_frame_mask = torch.repeat_interleave(
            first_frame_mask,
            dim=2,
            repeats=temporal_compression_ratio,
        )
        mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], dim=2)
        mask_lat_size = mask_lat_size.view(
            batch_size,
            -1,
            temporal_compression_ratio,
            latent_height,
            latent_width,
        )
        mask_lat_size = mask_lat_size.transpose(1, 2)
        mask_lat_size = mask_lat_size.to(image_latents.device, dtype=torch.bfloat16)

        noisy_model_input = torch.cat([noisy_model_input, mask_lat_size, image_latents], dim=1)

        training_target = self.noise_scheduler.training_target(
            latents_btchw.flatten(0, 1),
            noise.flatten(0, 1),
            timesteps.flatten(0, 1),
        ).unflatten(0, (batch_size, num_latent_t))

        training_batch.noisy_model_input = noisy_model_input
        training_batch.timesteps = timesteps
        training_batch.noise = noise.permute(0, 2, 1, 3, 4)
        training_batch.raw_latent_shape = latents.shape
        training_batch._ar_training_target = training_target

        return training_batch

    def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:
        image_embeds = training_batch.image_embeds
        assert isinstance(image_embeds, torch.Tensor)
        assert torch.isnan(image_embeds).sum() == 0
        image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16)

        timesteps = training_batch.timesteps
        assert isinstance(timesteps, torch.Tensor)
        assert timesteps.ndim == 2, (f"Expected per-frame timesteps [B, T], got shape {timesteps.shape}")

        training_batch.input_kwargs = {
            "hidden_states": training_batch.noisy_model_input,
            "encoder_hidden_states": None,
            "timestep": timesteps.to(get_local_torch_device()),
            "encoder_hidden_states_image": image_embeds,
            "mouse_cond": training_batch.mouse_cond,
            "keyboard_cond": training_batch.keyboard_cond,
            "num_frame_per_block": self.num_frame_per_block,
            "return_dict": False,
        }
        return training_batch

    def _transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -> TrainingBatch:
        """
        Run transformer forward pass and compute diffusion-forcing loss.
        """
        input_kwargs = training_batch.input_kwargs

        with set_forward_context(
                current_timestep=training_batch.timesteps,
                attn_metadata=None,
                forward_batch=None,
        ):
            with torch.autocast("cuda", dtype=torch.bfloat16):
                model_pred = self.transformer(**input_kwargs)
            model_pred_btchw = model_pred.permute(0, 2, 1, 3, 4)

            training_target = training_batch._ar_training_target.to(model_pred_btchw.device,
                                                                    dtype=model_pred_btchw.dtype)

            per_frame_loss = F.mse_loss(
                model_pred_btchw.float(),
                training_target.float(),
                reduction="none",
            ).mean(dim=(2, 3, 4))
            timesteps = training_batch.timesteps
            weight = self.noise_scheduler.training_weight(timesteps.flatten(0, 1)).to(per_frame_loss.dtype).reshape(
                per_frame_loss.shape)
            loss = (per_frame_loss * weight).mean()
            loss = loss / self.training_args.gradient_accumulation_steps
            loss.backward()

        avg_loss = loss.detach().clone()
        training_batch.total_loss += avg_loss.item()

        return training_batch

    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
        self.transformer.train()
        self.optimizer.zero_grad()
        training_batch.total_loss = 0.0
        args = cast(TrainingArgs, self.training_args)

        for _ in range(args.gradient_accumulation_steps):
            training_batch = self._get_next_batch(training_batch)
            training_batch = self._normalize_dit_input(training_batch)
            training_batch = self._prepare_dit_inputs(training_batch)
            training_batch = self._build_input_kwargs(training_batch)
            training_batch = self._transformer_forward_and_compute_loss(training_batch)

        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
        assert training_batch.latents is not None
        training_batch.raw_latent_shape = training_batch.latents.shape
        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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)
Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    scheduler = SelfForcingFlowMatchScheduler(
        shift=fastvideo_args.pipeline_config.flow_shift,
        sigma_min=0.0,
        extra_one_step=True,
    )
    scheduler.set_timesteps(num_inference_steps=1000, training=True)
    self.modules["scheduler"] = scheduler

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline.set_schemas

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

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline.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.vae = self.get_module("vae")
    self.vae.requires_grad_(False)

    self.num_frame_per_block = self._resolve_num_frame_per_block(training_args)

    logger.info(
        "Matrix-Game 2.0 AR diffusion pipeline initialized with "
        "num_frame_per_block=%d, diffusion_forcing_shift=%.1f",
        self.num_frame_per_block,
        training_args.pipeline_config.flow_shift,
    )

fastvideo.training.matrixgame2_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)
Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
    logger.info("Initializing Matrix-Game 2.0 AR validation pipeline...")
    args_copy = deepcopy(training_args)
    args_copy.inference_mode = True

    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=False)

    num_val_steps = int(training_args.validation_sampling_steps.split(",")[0])
    step_size = 1000 // num_val_steps
    args_copy.pipeline_config.dmd_denoising_steps = list(range(1000, 0, -step_size))
    args_copy.pipeline_config.warp_denoising_step = True
    training_args.pipeline_config.dmd_denoising_steps = (args_copy.pipeline_config.dmd_denoising_steps)
    training_args.pipeline_config.warp_denoising_step = True

    logger.info(
        "Validation: %d-step Matrix-Game 2.0 causal denoising, "
        "dmd_denoising_steps has %d entries",
        num_val_steps,
        len(args_copy.pipeline_config.dmd_denoising_steps),
    )

    loaded_modules = {
        "transformer": self.get_module("transformer"),
        "vae": self.get_module("vae"),
        "scheduler": validation_scheduler,
    }
    image_encoder = self.get_module("image_encoder")
    image_processor = self.get_module("image_processor")
    if image_encoder is not None:
        loaded_modules["image_encoder"] = image_encoder
    if image_processor is not None:
        loaded_modules["image_processor"] = image_processor

    self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
        training_args.model_path,
        args=args_copy,
        inference_mode=True,
        loaded_modules=loaded_modules,
        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_ar_diffusion_pipeline.MatrixGame2ARDiffusionPipeline.train_one_step

train_one_step(training_batch: TrainingBatch) -> TrainingBatch
Source
def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
    self.transformer.train()
    self.optimizer.zero_grad()
    training_batch.total_loss = 0.0
    args = cast(TrainingArgs, self.training_args)

    for _ in range(args.gradient_accumulation_steps):
        training_batch = self._get_next_batch(training_batch)
        training_batch = self._normalize_dit_input(training_batch)
        training_batch = self._prepare_dit_inputs(training_batch)
        training_batch = self._build_input_kwargs(training_batch)
        training_batch = self._transformer_forward_and_compute_loss(training_batch)

    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
    assert training_batch.latents is not None
    training_batch.raw_latent_shape = training_batch.latents.shape
    return training_batch

fastvideo.training.matrixgame2_ar_diffusion_pipeline.main

main(args) -> None

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