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)

`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

timestep: the timestep with shape [B*T]

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 = &#91;&quot;scheduler&quot;, &quot;transformer&quot;, &quot;vae&quot;&#93;

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&#91;&quot;scheduler&quot;&#93; = scheduler

def _log_training_info(self) -&gt; None:
    self.noise_scheduler = self.modules&#91;&quot;scheduler&quot;&#93;
    super()._log_training_info()

def set_schemas(self):
    self.train_dataset_schema = pyarrow_schema_matrixgame2

def _get_temporal_compression_ratio(self) -&gt; 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) -&gt; int:
    transformer_num_frame_per_block = getattr(self.transformer, &quot;num_frame_per_block&quot;, None)
    requested_num_frame_per_block = getattr(training_args, &quot;num_frame_per_block&quot;, 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(&quot;num_frame_per_block mismatch between loaded transformer and &quot;
                         &quot;training args: &quot;
                         f&quot;transformer=&#123;transformer_num_frame_per_block&#125;, &quot;
                         f&quot;training_args=&#123;requested_num_frame_per_block&#125;&quot;)

    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(&quot;vae&quot;)
    self.vae.requires_grad_(False)

    self.num_frame_per_block = self._resolve_num_frame_per_block(training_args)

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

def initialize_validation_pipeline(self, training_args: TrainingArgs):
    logger.info(&quot;Initializing Matrix-Game 2.0 AR validation pipeline...&quot;)
    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(&quot;,&quot;)&#91;0&#93;)
    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(
        &quot;Validation: %d-step Matrix-Game 2.0 causal denoising, &quot;
        &quot;dmd_denoising_steps has %d entries&quot;,
        num_val_steps,
        len(args_copy.pipeline_config.dmd_denoising_steps),
    )

    loaded_modules = &#123;
        &quot;transformer&quot;: self.get_module(&quot;transformer&quot;),
        &quot;vae&quot;: self.get_module(&quot;vae&quot;),
        &quot;scheduler&quot;: validation_scheduler,
    &#125;
    image_encoder = self.get_module(&quot;image_encoder&quot;)
    image_processor = self.get_module(&quot;image_processor&quot;)
    if image_encoder is not None:
        loaded_modules&#91;&quot;image_encoder&quot;&#93; = image_encoder
    if image_processor is not None:
        loaded_modules&#91;&quot;image_processor&quot;&#93; = 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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Sample one schedule index per chunk and broadcast to frames.

    Returns the shifted timestep, shape &#91;B, num_frame&#93;.
    &quot;&quot;&quot;
    device = get_local_torch_device()
    num_schedule_steps = len(self.noise_scheduler.timesteps)
    chunk_size = int(num_frame_per_block)
    if chunk_size &lt;= 0:
        raise ValueError(f&quot;num_frame_per_block must be &gt; 0, got &#123;chunk_size&#125;&quot;)
    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)&#91;:, :num_frame&#93;
    schedule_t = self.noise_scheduler.timesteps.to(device)
    return schedule_t&#91;indices&#93;

def _get_next_batch(self, training_batch: TrainingBatch) -&gt; TrainingBatch:
    batch = next(self.train_loader_iter, None)  # type: ignore
    if batch is None:
        self.current_epoch += 1
        logger.info(&quot;Starting epoch %s&quot;, self.current_epoch)
        self.train_loader_iter = iter(self.train_dataloader)
        batch = next(self.train_loader_iter)

    latents = batch&#91;&quot;vae_latent&quot;&#93;
    latents = latents&#91;:, :, :self.training_args.num_latent_t&#93;
    clip_features = batch&#91;&quot;clip_feature&quot;&#93;
    image_latents = batch&#91;&quot;first_frame_latent&quot;&#93;
    image_latents = image_latents&#91;:, :, :self.training_args.num_latent_t&#93;
    pil_image = batch&#91;&quot;pil_image&quot;&#93;
    infos = batch&#91;&quot;info_list&quot;&#93;

    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 &quot;mouse_cond&quot; in batch and batch&#91;&quot;mouse_cond&quot;&#93;.numel() &gt; 0:
        training_batch.mouse_cond = batch&#91;&quot;mouse_cond&quot;&#93;.to(get_local_torch_device(), dtype=torch.bfloat16)
    else:
        training_batch.mouse_cond = None

    if &quot;keyboard_cond&quot; in batch and batch&#91;&quot;keyboard_cond&quot;&#93;.numel() &gt; 0:
        training_batch.keyboard_cond = batch&#91;&quot;keyboard_cond&quot;&#93;.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&#91;1&#93; &gt;= expected_num_frames, (
            f&quot;keyboard_cond has &#123;training_batch.keyboard_cond.shape&#91;1&#93;&#125; &quot;
            f&quot;frames but need at least &#123;expected_num_frames&#125;&quot;)
        training_batch.keyboard_cond = training_batch.keyboard_cond&#91;:, :expected_num_frames&#93;
    if training_batch.mouse_cond is not None:
        assert training_batch.mouse_cond.shape&#91;1&#93; &gt;= expected_num_frames, (
            f&quot;mouse_cond has &#123;training_batch.mouse_cond.shape&#91;1&#93;&#125; frames &quot;
            f&quot;but need at least &#123;expected_num_frames&#125;&quot;)
        training_batch.mouse_cond = training_batch.mouse_cond&#91;:, :expected_num_frames&#93;

    return training_batch

def _prepare_dit_inputs(self, training_batch: TrainingBatch) -&gt; TrainingBatch:
    &quot;&quot;&quot;Prepare diffusion-forcing inputs and Matrix-Game 2.0 I2V concat.&quot;&quot;&quot;
    assert self.training_args is not None
    latents = training_batch.latents
    assert latents is not None
    batch_size = latents.shape&#91;0&#93;
    num_latent_t = latents.shape&#91;2&#93;

    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, &quot;noise_gen_cuda&quot;, None)
    noise = torch.randn(
        latents_btchw.shape,
        generator=noise_generator,
        device=latents_btchw.device,
        dtype=latents_btchw.dtype,
    )
    if (self.training_args.sp_size &gt; 1 and self.sp_group is not None and hasattr(self.sp_group, &quot;broadcast&quot;)):
        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&#91;:, :, 1:&#93; = 0

    first_frame_mask = mask_lat_size&#91;:, :, :1&#93;
    first_frame_mask = torch.repeat_interleave(
        first_frame_mask,
        dim=2,
        repeats=temporal_compression_ratio,
    )
    mask_lat_size = torch.cat(&#91;first_frame_mask, mask_lat_size&#91;:, :, 1:&#93;&#93;, 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(&#91;noisy_model_input, mask_lat_size, image_latents&#93;, 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) -&gt; 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&quot;Expected per-frame timesteps &#91;B, T&#93;, got shape &#123;timesteps.shape&#125;&quot;)

    training_batch.input_kwargs = &#123;
        &quot;hidden_states&quot;: training_batch.noisy_model_input,
        &quot;encoder_hidden_states&quot;: None,
        &quot;timestep&quot;: timesteps.to(get_local_torch_device()),
        &quot;encoder_hidden_states_image&quot;: image_embeds,
        &quot;mouse_cond&quot;: training_batch.mouse_cond,
        &quot;keyboard_cond&quot;: training_batch.keyboard_cond,
        &quot;num_frame_per_block&quot;: self.num_frame_per_block,
        &quot;return_dict&quot;: False,
    &#125;
    return training_batch

def _transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -&gt; TrainingBatch:
    &quot;&quot;&quot;
    Run transformer forward pass and compute diffusion-forcing loss.
    &quot;&quot;&quot;
    input_kwargs = training_batch.input_kwargs

    with set_forward_context(
            current_timestep=training_batch.timesteps,
            attn_metadata=None,
            forward_batch=None,
    ):
        with torch.autocast(&quot;cuda&quot;, 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=&quot;none&quot;,
        ).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) -&gt; 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(
        &#91;p for p in self.transformer.parameters() if p.requires_grad&#93;,
        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&#91;str, Any&#93;,
    num_inference_steps: int,
) -&gt; ForwardBatch:
    sampling_param.prompt = validation_batch&#91;&quot;prompt&quot;&#93;
    sampling_param.height = training_args.num_height
    sampling_param.width = training_args.num_width
    sampling_param.image_path = validation_batch.get(&quot;image_path&quot;) or validation_batch.get(&quot;video_path&quot;)
    sampling_param.num_inference_steps = num_inference_steps
    sampling_param.data_type = &quot;video&quot;
    assert self.seed is not None
    sampling_param.seed = self.seed

    latents_size = &#91;
        (sampling_param.num_frames - 1) // 4 + 1,
        sampling_param.height // 8,
        sampling_param.width // 8,
    &#93;
    n_tokens = latents_size&#91;0&#93; * latents_size&#91;1&#93; * latents_size&#91;2&#93;
    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=&quot;cpu&quot;).manual_seed(self.seed),
        n_tokens=n_tokens,
        eta=0.0,
        VSA_sparsity=training_args.VSA_sparsity,
    )
    if &quot;image&quot; in validation_batch and validation_batch&#91;&quot;image&quot;&#93; is not None:
        batch.pil_image = validation_batch&#91;&quot;image&quot;&#93;

    if (&quot;keyboard_cond&quot; in validation_batch and validation_batch&#91;&quot;keyboard_cond&quot;&#93; is not None):
        keyboard_cond = validation_batch&#91;&quot;keyboard_cond&quot;&#93;
        keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
        keyboard_cond = keyboard_cond.unsqueeze(0)
        batch.keyboard_cond = keyboard_cond

    if (&quot;mouse_cond&quot; in validation_batch and validation_batch&#91;&quot;mouse_cond&quot;&#93; is not None):
        mouse_cond = validation_batch&#91;&quot;mouse_cond&quot;&#93;
        mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
        mouse_cond = mouse_cond.unsqueeze(0)
        batch.mouse_cond = mouse_cond

    return batch</code></pre></details>
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(&quot;vae&quot;)
self.vae.requires_grad_(False)

self.num_frame_per_block = self._resolve_num_frame_per_block(training_args)

logger.info(
    &quot;Matrix-Game 2.0 AR diffusion pipeline initialized with &quot;
    &quot;num_frame_per_block=%d, diffusion_forcing_shift=%.1f&quot;,
    self.num_frame_per_block,
    training_args.pipeline_config.flow_shift,
)</code></pre></details>
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(&quot;,&quot;)&#91;0&#93;)
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(
    &quot;Validation: %d-step Matrix-Game 2.0 causal denoising, &quot;
    &quot;dmd_denoising_steps has %d entries&quot;,
    num_val_steps,
    len(args_copy.pipeline_config.dmd_denoising_steps),
)

loaded_modules = &#123;
    &quot;transformer&quot;: self.get_module(&quot;transformer&quot;),
    &quot;vae&quot;: self.get_module(&quot;vae&quot;),
    &quot;scheduler&quot;: validation_scheduler,
&#125;
image_encoder = self.get_module(&quot;image_encoder&quot;)
image_processor = self.get_module(&quot;image_processor&quot;)
if image_encoder is not None:
    loaded_modules&#91;&quot;image_encoder&quot;&#93; = image_encoder
if image_processor is not None:
    loaded_modules&#91;&quot;image_processor&quot;&#93; = 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,
)</code></pre></details>
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(
    &#91;p for p in self.transformer.parameters() if p.requires_grad&#93;,
    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</code></pre></details>
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.