--- title: matrixgame2_training_pipeline ---

fastvideo.training.matrixgame2_training_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_training_pipeline.SamplingParam.check_sampling_param

check_sampling_param() -> None

fastvideo.training.matrixgame2_training_pipeline.SamplingParam.update

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

fastvideo.training.matrixgame2_training_pipeline.SamplingParam.from_pretrained

from_pretrained(model_path: str) -> SamplingParam

fastvideo.training.matrixgame2_training_pipeline.SamplingParam.add_cli_args

add_cli_args(parser: Any) -> Any

Add CLI arguments for SamplingParam fields

fastvideo.training.matrixgame2_training_pipeline.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs

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

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.is_output_rank

is_output_rank: bool

Whether this process may materialize user-facing outputs.

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

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.finalize_device_offload_policy

finalize_device_offload_policy(device_id: int = 0) -> bool

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

fastvideo.training.matrixgame2_training_pipeline.FastVideoArgs.disable_offload_on_unified_memory

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

Disable host offload after a worker has selected its device.

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

fastvideo.training.matrixgame2_training_pipeline.TrainingArgs

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

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

fastvideo.training.matrixgame2_training_pipeline.TrainingArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> TrainingArgs

fastvideo.training.matrixgame2_training_pipeline.TrainingArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.training.matrixgame2_training_pipeline.init_logger

init_logger(name: str) -> _FastvideoLogger

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler

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

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

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

Parameters

num_train_timesteps: `int`, defaults to 1000

The number of diffusion steps to train the model.

solver_order: `int`, default `2`

The UniPC order which can be any positive integer. The effective order of accuracy is solver_order + 1 due to the UniC. It is recommended to use solver_order=2 for guided sampling, and solver_order=3 for unconditional sampling.

prediction_type: `str`, defaults to "flow_prediction"

Prediction type of the scheduler function; must be flow_prediction for this scheduler, which predicts the flow of the diffusion process.

thresholding: `bool`, defaults to `False`

Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion.

dynamic_thresholding_ratio: `float`, defaults to 0.995

The ratio for the dynamic thresholding method. Valid only when thresholding=True.

sample_max_value: `float`, defaults to 1.0

The threshold value for dynamic thresholding. Valid only when thresholding=True and predict_x0=True.

predict_x0: `bool`, defaults to `True`

Whether to use the updating algorithm on the predicted x0.

solver_type: `str`, default `bh2`

Solver type for UniPC. It is recommended to use bh1 for unconditional sampling when steps < 10, and bh2 otherwise.

lower_order_final: `bool`, default `True`

Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.

disable_corrector: `list`, default `[]`

Decides which step to disable the corrector to mitigate the misalignment between epsilon_theta(x_t, c) and epsilon_theta(x_t^c, c) which can influence convergence for a large guidance scale. Corrector is usually disabled during the first few steps.

solver_p: `SchedulerMixin`, default `None`

Any other scheduler that if specified, the algorithm becomes solver_p + UniC.

use_karras_sigmas: `bool`, *optional*, defaults to `False`

Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If True, the sigmas are determined according to a sequence of noise levels {σi}.

use_exponential_sigmas: `bool`, *optional*, defaults to `False`

Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.

timestep_spacing: `str`, defaults to `"linspace"`

The way the timesteps should be scaled. Refer to Table 2 of the Common Diffusion Noise Schedules and Sample Steps are Flawed for more information.

steps_offset: `int`, defaults to 0

An offset added to the inference steps, as required by some model families.

final_sigmas_type: `str`, defaults to `"zero"`

The final sigma value for the noise schedule during the sampling process. If "sigma_min", the final sigma is the same as the last sigma in the training schedule. If zero, the final sigma is set to 0.

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.step_index

step_index

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.begin_index

begin_index

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.set_begin_index

set_begin_index(begin_index: int = 0)

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

Parameters

begin_index: `int`

The begin index for the scheduler.

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.set_timesteps

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

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.time_shift

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.convert_model_output

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

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

Parameters

model_output: `torch.Tensor`

The direct output from the learned diffusion model.

timestep: `int`

The current discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

Returns

torch.Tensor

torch.Tensor: The converted model output.

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_p_bh_update

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

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

Parameters

model_output: `torch.Tensor`

The direct output from the learned diffusion model at the current timestep.

prev_timestep: `int`

The previous discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

order: `int`

The order of UniP at this timestep (corresponds to the p in UniPC-p).

Returns

torch.Tensor

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.multistep_uni_c_bh_update

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

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

Parameters

this_model_output: `torch.Tensor`

The model outputs at x_t.

this_timestep: `int`

The current timestep t.

last_sample: `torch.Tensor`

The generated sample before the last predictor x_{t-1}.

this_sample: `torch.Tensor`

The generated sample after the last predictor x_{t}.

order: `int`

The p of UniC-p at this step. The effective order of accuracy should be order + 1.

Returns

torch.Tensor

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.index_for_timestep

index_for_timestep(timestep, schedule_timesteps = None) -> int

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.step

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

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

Parameters

model_output: `torch.Tensor`

The direct output from learned diffusion model.

timestep: `int`

The current discrete timestep in the diffusion chain.

sample: `torch.Tensor`

A current instance of a sample created by the diffusion process.

return_dict: `bool`

Whether or not to return a [~schedulers.scheduling_utils.SchedulerOutput] or tuple.

Returns

SchedulerOutput | tuple

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

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.scale_model_input

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

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

Parameters

sample: `torch.Tensor`

The input sample.

Returns

torch.Tensor

torch.Tensor: A scaled input sample.

fastvideo.training.matrixgame2_training_pipeline.FlowUniPCMultistepScheduler.add_noise

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

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2I2VPipeline

MatrixGame2I2VPipeline(*args, **kwargs)

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2I2VPipeline.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2I2VPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_training_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_training_pipeline.TrainingBatch

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

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline

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

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

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.set_schemas

set_schemas() -> None

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.train_one_step

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.train

train() -> None

fastvideo.training.matrixgame2_training_pipeline.TrainingPipeline.visualize_intermediate_latents

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

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

fastvideo.training.matrixgame2_training_pipeline.is_vsa_available

is_vsa_available() -> bool

fastvideo.training.matrixgame2_training_pipeline.shallow_asdict

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

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2TrainingPipeline

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

A training pipeline for Matrix-Game-2.0.

Source
class MatrixGame2TrainingPipeline(TrainingPipeline):
    """
    A training pipeline for Matrix-Game-2.0.
    """
    _required_config_modules = ["scheduler", "transformer", "vae"]

    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
        self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)

    def create_training_stages(self, training_args: TrainingArgs):
        """
        May be used in future refactors.
        """
        pass

    def set_schemas(self):
        self.train_dataset_schema = pyarrow_schema_matrixgame2

    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        logger.info("Initializing validation pipeline...")
        args_copy = deepcopy(training_args)

        args_copy.inference_mode = True
        args_copy.dit_cpu_offload = True
        # args_copy.pipeline_config.vae_config.load_encoder = False
        # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
        self.validation_pipeline = MatrixGame2I2VPipeline.from_pretrained(training_args.model_path,
                                                                          args=None,
                                                                          inference_mode=True,
                                                                          loaded_modules={
                                                                              "transformer":
                                                                              self.get_module("transformer"),
                                                                          },
                                                                          tp_size=training_args.tp_size,
                                                                          sp_size=training_args.sp_size,
                                                                          num_gpus=training_args.num_gpus,
                                                                          dit_cpu_offload=True)

    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)
            # Reset iterator for next epoch
            self.train_loader_iter = iter(self.train_dataloader)
            # Get first batch of new epoch
            batch = next(self.train_loader_iter)

        latents = batch['vae_latent']
        latents = latents[:, :, :self.training_args.num_latent_t]
        # encoder_hidden_states = batch['text_embedding']
        # encoder_attention_mask = batch['text_attention_mask']
        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
        # Matrix-Game 2.0 doesn't use text encoder
        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

        # Action conditioning
        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

        return training_batch

    def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch:
        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""

        # First, call parent method to prepare noise, timesteps, etc. for video latents
        training_batch = super()._prepare_dit_inputs(training_batch)

        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.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
        num_frames = (self.training_args.num_latent_t - 1) * temporal_compression_ratio + 1
        batch_size, num_channels, _, 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).to(dtype=torch.bfloat16)

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

        return training_batch

    def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:

        # Image Embeds for conditioning
        image_embeds = training_batch.image_embeds
        assert torch.isnan(image_embeds).sum() == 0
        image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16)
        encoder_hidden_states_image = image_embeds

        # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
        training_batch.input_kwargs = {
            "hidden_states": training_batch.noisy_model_input,
            "encoder_hidden_states": training_batch.encoder_hidden_states,  # None for Matrix-Game 2.0
            "timestep": training_batch.timesteps.to(get_local_torch_device(), dtype=torch.bfloat16),
            # "encoder_attention_mask":
            # training_batch.encoder_attention_mask,
            "encoder_hidden_states_image": encoder_hidden_states_image,
            # Action conditioning
            "mouse_cond": training_batch.mouse_cond,
            "keyboard_cond": training_batch.keyboard_cond,
            "return_dict": False,
        }
        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_training_pipeline.MatrixGame2TrainingPipeline.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)
Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2TrainingPipeline.create_training_stages

create_training_stages(training_args: TrainingArgs)

May be used in future refactors.

Source
def create_training_stages(self, training_args: TrainingArgs):
    """
    May be used in future refactors.
    """
    pass

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2TrainingPipeline.set_schemas

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

fastvideo.training.matrixgame2_training_pipeline.MatrixGame2TrainingPipeline.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)
Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
    logger.info("Initializing validation pipeline...")
    args_copy = deepcopy(training_args)

    args_copy.inference_mode = True
    args_copy.dit_cpu_offload = True
    # args_copy.pipeline_config.vae_config.load_encoder = False
    # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
    self.validation_pipeline = MatrixGame2I2VPipeline.from_pretrained(training_args.model_path,
                                                                      args=None,
                                                                      inference_mode=True,
                                                                      loaded_modules={
                                                                          "transformer":
                                                                          self.get_module("transformer"),
                                                                      },
                                                                      tp_size=training_args.tp_size,
                                                                      sp_size=training_args.sp_size,
                                                                      num_gpus=training_args.num_gpus,
                                                                      dit_cpu_offload=True)

fastvideo.training.matrixgame2_training_pipeline.main

main(args) -> None

fastvideo.training.matrixgame2_training_pipeline.build_parser

build_parser(pipeline_class: str | None = None, pipeline_module: str | None = None) -> FlexibleArgumentParser

Build the training-runner CLI parser.

When pipeline_class/pipeline_module are given they become optional defaults, so the deprecated per-pipeline __main__ blocks keep working without the flags.