`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_pipeline.SamplingParam.check_sampling_param`

check_sampling_param() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SamplingParam.update`

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

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SamplingParam.from_pretrained`

from_pretrained(model_path: str) -> SamplingParam

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SamplingParam.add_cli_args`

add_cli_args(parser: Any) -> Any

Add CLI arguments for SamplingParam fields

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.get_local_torch_device`

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_pipeline.FastVideoArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.FastVideoArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.FastVideoArgs.from_kwargs`

from_kwargs(**kwargs: Any) -> FastVideoArgs

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.FastVideoArgs.check_fastvideo_args`

check_fastvideo_args() -> None

Validate inference arguments for consistency

`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_pipeline.TrainingArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> TrainingArgs

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.TrainingArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.set_timesteps`

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

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.step`

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

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high`

calculate_alpha_beta_high(sigma, sigma_bound)

`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.add_noise_high`

add_noise_high(original_samples, noise, timestep, boundary_timestep)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.training_target`

training_target(sample, noise, timestep)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.training_weight`

training_weight(timestep)

Input

timestep: the timestep with shape [B*T]

Output: the corresponding weighting [B*T]

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.scale_model_input`

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

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingFlowMatchScheduler.set_shift`

set_shift(shift: float) -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.pred_noise_to_pred_video`

pred_noise_to_pred_video(pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor

Convert predicted noise to clean latent.

pred_noise: the predicted noise with shape [B, C, H, W] where B is batch_size or batch_size * num_frames noise_input_latent: the noisy latent with shape [B, C, H, W], timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] scheduler: the scheduler

Returns

torch.Tensor: the predicted video with shape [B, C, H, W]

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase`

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

Base class for pipelines composed of multiple stages.

This class provides the framework for creating pipelines by composing multiple stages together. Each stage is responsible for a specific part of the diffusion process, and the pipeline orchestrates the execution of these stages.

Initialize the pipeline. After init, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.get_hf_download_component_dirs`

get_hf_download_component_dirs() -> tuple[str, ...] | None

Return component directories for an opt-in partial Hub download.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.get_hf_download_allow_patterns`

get_hf_download_allow_patterns() -> list[str] | None

Return Hub patterns for the manifest and selected components.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.set_trainable`

set_trainable() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.post_init`

post_init() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.initialize_training_pipeline`

initialize_training_pipeline(training_args: TrainingArgs)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.initialize_validation_pipeline`

initialize_validation_pipeline(training_args: TrainingArgs)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.from_pretrained`

from_pretrained(model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, pipeline_config: str | PipelineConfig | None = None, args: argparse.Namespace | None = None, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, **kwargs) -> ComposedPipelineBase

Load a pipeline from a pretrained model. loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, If provided, loaded_modules will be used instead of loading from config/pretrained weights.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.get_module`

get_module(module_name: str, default_value: Any = None) -> Any

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.add_module`

add_module(module_name: str, module: Any)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.required_config_modules`

required_config_modules: list[str]

List of modules that are required by the pipeline. The names should match the diffusers directory and model_index.json file. These modules will be loaded using the PipelineComponentLoader and made available in the modules dictionary. Access these modules using the get_module method.

class ConcretePipeline(ComposedPipelineBase): _required_config_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"]

@property
def required_config_modules(self):
    return self._required_config_modules

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.stages`

stages: list[PipelineStage]

List of stages in the pipeline.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Create the inference pipeline stages.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.create_training_stages`

create_training_stages(training_args: TrainingArgs)

Create the training pipeline stages.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.initialize_pipeline`

initialize_pipeline(fastvideo_args: FastVideoArgs)

Initialize the pipeline.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.load_modules`

load_modules(fastvideo_args: FastVideoArgs, loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]

Load the modules from the config. loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, If provided, loaded_modules will be used instead of loading from config/pretrained weights.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.add_stage`

add_stage(stage_name: str, stage: PipelineStage)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.forward`

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

Generate a video or image using the pipeline.

Parameters

batch: ForwardBatch: The batch to generate from.
fastvideo_args: FastVideoArgs: The inference arguments.

Returns: ForwardBatch: The batch with the generated video or image.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.train`

train() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.ComposedPipelineBase.close`

close() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2CausalDMDPipeline`

MatrixGame2CausalDMDPipeline(*args, **kwargs)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2CausalDMDPipeline.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs) -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2CausalDMDPipeline.streaming_reset`

streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs)

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2CausalDMDPipeline.streaming_step`

streaming_step(keyboard_action, mouse_action) -> ForwardBatch

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2CausalDMDPipeline.streaming_clear`

streaming_clear() -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_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_self_forcing_distillation_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_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline`

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

A self-forcing distillation pipeline that alternates between training the generator and critic based on the self-forcing methodology.

This implementation follows the self-forcing approach where: 1. Generator and critic are trained in alternating steps 2. Generator loss uses DMD-style loss with the critic as fake score 3. Critic loss trains the fake score model to distinguish real vs fake

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.initialize_training_pipeline`

initialize_training_pipeline(training_args: TrainingArgs)

Initialize the self-forcing training pipeline.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.generate_and_sync_list`

generate_and_sync_list(num_blocks: int, num_denoising_steps: int, device: torch.device) -> list[int]

Generate and synchronize random exit flags across distributed processes.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.generator_loss`

generator_loss(training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]

Compute generator loss using DMD-style approach. The generator tries to fool the critic (fake_score_transformer).

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.critic_loss`

critic_loss(training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]

Compute critic loss using flow matching between noise and generator output. The critic learns to predict the flow from noise to the generator's output.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.train_one_step`

train_one_step(training_batch: TrainingBatch) -> TrainingBatch

Self-forcing training step that alternates between generator and critic training.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.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_self_forcing_distillation_pipeline.SelfForcingDistillationPipeline.train`

train() -> None

Main training loop with self-forcing specific logging.

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.shift_timestep`

shift_timestep(timestep: torch.Tensor, shift: float, num_train_timestep: float) -> torch.Tensor

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.is_vsa_available`

is_vsa_available() -> bool

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.shallow_asdict`

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

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2SelfForcingDistillationPipeline`

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

A self-forcing distillation pipeline for Matrix-Game 2.0 that uses the self-forcing methodology with DMD for video generation.

Source

class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
    """
    A self-forcing distillation pipeline for Matrix-Game 2.0 that uses the self-forcing methodology
    with DMD for video generation.
    """
    _required_config_modules = [
        "scheduler",
        "transformer",
        "vae",
        "image_encoder",
        "image_processor",
    ]

    def set_schemas(self):
        self.train_dataset_schema = pyarrow_schema_matrixgame2

    def load_modules(
        self,
        fastvideo_args: FastVideoArgs,
        loaded_modules: dict[str, torch.nn.Module] | None = None,
    ) -> dict[str, Any]:
        modules = ComposedPipelineBase.load_modules(
            self,
            fastvideo_args,
            loaded_modules,
        )
        training_args = cast(TrainingArgs, fastvideo_args)
        old_override = training_args.override_transformer_cls_name
        training_args.override_transformer_cls_name = "MatrixGame2WanModel"
        try:
            if loaded_modules is not None and "real_score_transformer" in loaded_modules:
                self.real_score_transformer = loaded_modules["real_score_transformer"]
            elif training_args.real_score_model_path:
                logger.info(
                    "Loading real score transformer from: %s",
                    training_args.real_score_model_path,
                )
                self.real_score_transformer = self.load_module_from_path(
                    training_args.real_score_model_path,
                    "transformer",
                    training_args,
                )
            else:
                raise ValueError("real_score_model_path is required for DMD distillation pipeline")
            modules["real_score_transformer"] = self.real_score_transformer

            if loaded_modules is not None and "fake_score_transformer" in loaded_modules:
                self.fake_score_transformer = loaded_modules["fake_score_transformer"]
            elif training_args.fake_score_model_path:
                logger.info(
                    "Loading fake score transformer from: %s",
                    training_args.fake_score_model_path,
                )
                self.fake_score_transformer = self.load_module_from_path(
                    training_args.fake_score_model_path,
                    "transformer",
                    training_args,
                )
            else:
                raise ValueError("fake_score_model_path is required for DMD distillation pipeline")
            modules["fake_score_transformer"] = self.fake_score_transformer
        finally:
            training_args.override_transformer_cls_name = old_override

        self.real_score_transformer_2 = None
        self.fake_score_transformer_2 = None
        return modules

    def _build_matrixgame_cond_concat(
        self,
        image_latents: torch.Tensor,
    ) -> torch.Tensor:
        if image_latents.ndim != 5:
            raise ValueError("first_frame_latent must have shape [B, C, T, H, W], got "
                             f"{tuple(image_latents.shape)}")
        if image_latents.shape[1] != 16:
            raise ValueError("Matrix-Game 2.0 cond conversion expects a 16-channel "
                             f"first_frame_latent, got {image_latents.shape[1]} channels.")

        temporal_compression_ratio = (
            self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
        num_latent_t = image_latents.shape[2]
        num_frames = ((num_latent_t - 1) * temporal_compression_ratio + 1)
        batch_size, _, _, latent_height, latent_width = image_latents.shape

        mask_lat_size = torch.ones(
            batch_size,
            1,
            num_frames,
            latent_height,
            latent_width,
            device=image_latents.device,
            dtype=image_latents.dtype,
        )
        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).to(
            image_latents.device,
            dtype=image_latents.dtype,
        )

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

    def _initialize_simulation_caches(
        self,
        batch_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any] | None], list[dict[str, Any] | None]]:
        """Initialize KV cache and cross-attention cache for multi-step simulation."""
        num_transformer_blocks = len(self.transformer.blocks)
        latent_shape = self.video_latent_shape_sp
        _, num_frames, _, height, width = latent_shape

        _, p_h, p_w = self.transformer.patch_size
        post_patch_height = height // p_h
        post_patch_width = width // p_w

        frame_seq_length = post_patch_height * post_patch_width
        self.frame_seq_length = frame_seq_length

        # Get model configuration parameters - handle FSDP wrapping
        num_attention_heads = getattr(self.transformer, 'num_attention_heads', None)
        attention_head_dim = getattr(self.transformer, 'attention_head_dim', None)

        # 1 CLS token + 256 patch tokens = 257
        text_len = 257

        action_config = getattr(self.transformer, 'action_config', {})
        action_blocks = action_config.get('blocks', []) if action_config else []
        local_attn_size = getattr(self.transformer, "local_attn_size",
                                  action_config.get('local_attn_size', 6) if action_config else 6)

        if local_attn_size <= 0:
            raise ValueError("Matrix-Game 2.0 self-forcing requires transformer.local_attn_size > 0")
        kv_cache_size = local_attn_size * frame_seq_length

        kv_cache = []
        for _ in range(num_transformer_blocks):
            kv_cache.append({
                "k":
                torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim],
                            dtype=dtype,
                            device=device),
                "v":
                torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim],
                            dtype=dtype,
                            device=device),
                "global_end_index":
                torch.tensor([0], dtype=torch.long, device=device),
                "local_end_index":
                torch.tensor([0], dtype=torch.long, device=device)
            })

        # Initialize cross-attention cache
        crossattn_cache = []
        for _ in range(num_transformer_blocks):
            crossattn_cache.append({
                "k":
                torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype,
                            device=device),
                "v":
                torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype,
                            device=device),
                "is_init":
                False
            })

        # Initialize action module KV caches
        action_heads_num = action_config.get('heads_num', 16) if action_config else 16
        mouse_hidden_dim = action_config.get('mouse_hidden_dim', 1024) if action_config else 1024
        keyboard_hidden_dim = action_config.get('keyboard_hidden_dim', 1024) if action_config else 1024

        mouse_head_dim = mouse_hidden_dim // action_heads_num
        keyboard_head_dim = keyboard_hidden_dim // action_heads_num

        kv_cache_mouse: list[dict[str, Any] | None] = []
        kv_cache_keyboard: list[dict[str, Any] | None] = []
        for block_idx in range(num_transformer_blocks):
            if block_idx in action_blocks:
                kv_cache_mouse.append({
                    "k":
                    torch.zeros([batch_size * frame_seq_length, local_attn_size, action_heads_num, mouse_head_dim],
                                dtype=dtype,
                                device=device),
                    "v":
                    torch.zeros([batch_size * frame_seq_length, local_attn_size, action_heads_num, mouse_head_dim],
                                dtype=dtype,
                                device=device),
                    "global_end_index":
                    torch.tensor([0], dtype=torch.long, device=device),
                    "local_end_index":
                    torch.tensor([0], dtype=torch.long, device=device)
                })
                kv_cache_keyboard.append({
                    "k":
                    torch.zeros([batch_size, local_attn_size, action_heads_num, keyboard_head_dim],
                                dtype=dtype,
                                device=device),
                    "v":
                    torch.zeros([batch_size, local_attn_size, action_heads_num, keyboard_head_dim],
                                dtype=dtype,
                                device=device),
                    "global_end_index":
                    torch.tensor([0], dtype=torch.long, device=device),
                    "local_end_index":
                    torch.tensor([0], dtype=torch.long, device=device)
                })
            else:
                kv_cache_mouse.append(None)
                kv_cache_keyboard.append(None)

        return kv_cache, crossattn_cache, kv_cache_mouse, kv_cache_keyboard

    def _reset_simulation_caches(self, kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]],
                                 kv_cache_mouse: list[dict[str, Any]
                                                      | None] | None, kv_cache_keyboard: list[dict[str, Any]
                                                                                              | None] | None) -> None:
        """Reset KV cache, cross-attention cache, and action caches to clean state."""
        if kv_cache is not None:
            for cache_dict in kv_cache:
                cache_dict["global_end_index"].fill_(0)
                cache_dict["local_end_index"].fill_(0)
                cache_dict["k"].zero_()
                cache_dict["v"].zero_()

        if crossattn_cache is not None:
            for cache_dict in crossattn_cache:
                cache_dict["is_init"] = False
                cache_dict["k"].zero_()
                cache_dict["v"].zero_()

        for opt_caches in (kv_cache_mouse, kv_cache_keyboard):
            if opt_caches is None:
                continue
            for opt_cache in opt_caches:
                if opt_cache is None:
                    continue
                opt_cache["global_end_index"].fill_(0)
                opt_cache["local_end_index"].fill_(0)
                opt_cache["k"].zero_()
                opt_cache["v"].zero_()

    @staticmethod
    def _snapshot_streaming_kv_cache(
        kv_cache: Iterable[dict[str, Any] | None] | None, ) -> list[dict[str, Any] | None] | None:
        # Per-block index clone for checkpoint-safe recompute.
        if kv_cache is None:
            return None
        snapshot: list[dict[str, Any] | None] = []
        for block_cache in kv_cache:
            if block_cache is None:
                snapshot.append(None)
                continue
            copied = dict(block_cache)
            for key in ("global_end_index", "local_end_index"):
                tensor = block_cache.get(key)
                if isinstance(tensor, torch.Tensor):
                    copied[key] = tensor.detach().clone()
            snapshot.append(copied)
        return snapshot

    def _generator_multi_step_simulation_forward(self,
                                                 training_batch: TrainingBatch,
                                                 return_sim_steps: bool = False) -> torch.Tensor:
        """Forward pass through student transformer matching inference procedure with KV cache management.

        This function is adapted from the reference self-forcing implementation's inference_with_trajectory
        and includes gradient masking logic for dynamic frame generation.
        """
        latents = training_batch.latents
        dtype = latents.dtype
        batch_size = latents.shape[0]

        num_training_frames = getattr(self.training_args, 'num_latent_t', 21)
        min_num_frames = 20 if self.independent_first_frame else 21
        max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames
        assert max_num_frames % self.num_frame_per_block == 0
        assert min_num_frames % self.num_frame_per_block == 0
        max_num_blocks = max_num_frames // self.num_frame_per_block
        min_num_blocks = min_num_frames // self.num_frame_per_block

        # Sample number of blocks and sync across processes
        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1, ), device=self.device)
        if dist.is_initialized():
            dist.broadcast(num_generated_blocks, src=0)
        num_generated_blocks = num_generated_blocks.item()
        num_generated_frames = num_generated_blocks * self.num_frame_per_block
        if self.independent_first_frame:
            num_generated_frames += 1
            min_num_frames += 1

        noise_shape = [batch_size, num_generated_frames, *self.video_latent_shape[2:]]
        noise = torch.randn(noise_shape, device=self.device, dtype=dtype)
        if self.sp_world_size > 1:
            noise = rearrange(noise, "b (n t) c h w -> b n t c h w", n=self.sp_world_size).contiguous()
            noise = noise[:, self.rank_in_sp_group, :, :, :, :]

        batch_size, num_frames, num_channels, height, width = noise.shape

        if self.independent_first_frame:
            assert (num_frames - 1) % self.num_frame_per_block == 0
            num_blocks = (num_frames - 1) // self.num_frame_per_block
        else:
            assert num_frames % self.num_frame_per_block == 0
            num_blocks = num_frames // self.num_frame_per_block

        num_output_frames = num_frames
        output = torch.zeros([batch_size, num_output_frames, num_channels, height, width],
                             device=noise.device,
                             dtype=noise.dtype)

        # Step 1: Initialize KV cache to all zeros
        (self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse,
         self.kv_cache_keyboard) = self._initialize_simulation_caches(batch_size, dtype, self.device)

        # Step 2: Temporal denoising loop
        current_start_frame = 0
        all_num_frames = [self.num_frame_per_block] * num_blocks
        if self.independent_first_frame:
            all_num_frames = [1] + all_num_frames
        num_denoising_steps = len(self.denoising_step_list)
        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
        start_gradient_frame_index = max(0, num_output_frames - 21)

        for block_index, current_num_frames in enumerate(all_num_frames):
            noisy_input = noise[:, current_start_frame:current_start_frame + current_num_frames]

            # Step 3.1: Spatial denoising loop
            for index, current_timestep in enumerate(self.denoising_step_list):
                exit_flag = (index == exit_flags[0]) if self.same_step_across_blocks else (index
                                                                                           == exit_flags[block_index])

                timestep = torch.ones([batch_size, current_num_frames], device=noise.device,
                                      dtype=torch.int64) * current_timestep

                if self.boundary_timestep is not None and current_timestep < self.boundary_timestep and self.transformer_2 is not None:
                    current_model = self.transformer_2
                else:
                    current_model = self.transformer

                if not exit_flag:
                    with torch.no_grad():
                        # Build input kwargs
                        training_batch_temp = self._build_distill_input_kwargs(noisy_input,
                                                                               timestep,
                                                                               training_batch.conditional_dict,
                                                                               training_batch,
                                                                               frame_start=current_start_frame,
                                                                               frame_end=current_start_frame +
                                                                               current_num_frames,
                                                                               num_frame_per_block=current_num_frames)

                        pred_flow = current_model(**training_batch_temp.input_kwargs,
                                                  kv_cache=self.kv_cache1,
                                                  kv_cache_mouse=self.kv_cache_mouse,
                                                  kv_cache_keyboard=self.kv_cache_keyboard,
                                                  crossattn_cache=self.crossattn_cache,
                                                  current_start=current_start_frame * self.frame_seq_length,
                                                  start_frame=current_start_frame).permute(0, 2, 1, 3, 4)

                        denoised_pred = pred_noise_to_pred_video(pred_noise=pred_flow.flatten(0, 1),
                                                                 noise_input_latent=noisy_input.flatten(0, 1),
                                                                 timestep=timestep,
                                                                 scheduler=self.noise_scheduler).unflatten(
                                                                     0, pred_flow.shape[:2])

                        next_timestep = self.denoising_step_list[index + 1]
                        noisy_input = self.noise_scheduler.add_noise(
                            denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)),
                            next_timestep * torch.ones(
                                [batch_size * current_num_frames], device=noise.device, dtype=torch.long)).unflatten(
                                    0, denoised_pred.shape[:2])
                else:
                    # Final prediction with gradient control
                    if current_start_frame < start_gradient_frame_index:
                        with torch.no_grad():
                            training_batch_temp = self._build_distill_input_kwargs(
                                noisy_input,
                                timestep,
                                training_batch.conditional_dict,
                                training_batch,
                                frame_start=current_start_frame,
                                frame_end=current_start_frame + current_num_frames,
                                num_frame_per_block=current_num_frames)

                            pred_flow = current_model(**training_batch_temp.input_kwargs,
                                                      kv_cache=self.kv_cache1,
                                                      kv_cache_mouse=self.kv_cache_mouse,
                                                      kv_cache_keyboard=self.kv_cache_keyboard,
                                                      crossattn_cache=self.crossattn_cache,
                                                      current_start=current_start_frame * self.frame_seq_length,
                                                      start_frame=current_start_frame).permute(0, 2, 1, 3, 4)
                    else:
                        training_batch_temp = self._build_distill_input_kwargs(noisy_input,
                                                                               timestep,
                                                                               training_batch.conditional_dict,
                                                                               training_batch,
                                                                               frame_start=current_start_frame,
                                                                               frame_end=current_start_frame +
                                                                               current_num_frames,
                                                                               num_frame_per_block=current_num_frames)

                        # Snapshot streaming caches for checkpoint-safe recompute.
                        snap_kv_cache = self._snapshot_streaming_kv_cache(self.kv_cache1)
                        snap_kv_cache_mouse = self._snapshot_streaming_kv_cache(self.kv_cache_mouse)
                        snap_kv_cache_keyboard = self._snapshot_streaming_kv_cache(self.kv_cache_keyboard)

                        pred_flow = current_model(**training_batch_temp.input_kwargs,
                                                  kv_cache=snap_kv_cache,
                                                  kv_cache_mouse=snap_kv_cache_mouse,
                                                  kv_cache_keyboard=snap_kv_cache_keyboard,
                                                  crossattn_cache=self.crossattn_cache,
                                                  current_start=current_start_frame * self.frame_seq_length,
                                                  start_frame=current_start_frame).permute(0, 2, 1, 3, 4)

                    denoised_pred = pred_noise_to_pred_video(pred_noise=pred_flow.flatten(0, 1),
                                                             noise_input_latent=noisy_input.flatten(0, 1),
                                                             timestep=timestep,
                                                             scheduler=self.noise_scheduler).unflatten(
                                                                 0, pred_flow.shape[:2])
                    break

            # Step 3.2: record the model's output
            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred

            # Step 3.3: rerun with timestep zero to update the cache
            context_timestep = torch.ones_like(timestep) * self.context_noise
            denoised_pred = self.noise_scheduler.add_noise(denoised_pred.flatten(0, 1),
                                                           torch.randn_like(denoised_pred.flatten(0, 1)),
                                                           context_timestep).unflatten(0, denoised_pred.shape[:2])

            with torch.no_grad():
                training_batch_temp = self._build_distill_input_kwargs(denoised_pred,
                                                                       context_timestep,
                                                                       training_batch.conditional_dict,
                                                                       training_batch,
                                                                       frame_start=current_start_frame,
                                                                       frame_end=current_start_frame +
                                                                       current_num_frames,
                                                                       num_frame_per_block=current_num_frames)

                # context_timestep is 0 so we use transformer_2
                current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
                current_model(**training_batch_temp.input_kwargs,
                              kv_cache=self.kv_cache1,
                              kv_cache_mouse=self.kv_cache_mouse,
                              kv_cache_keyboard=self.kv_cache_keyboard,
                              crossattn_cache=self.crossattn_cache,
                              current_start=current_start_frame * self.frame_seq_length,
                              start_frame=current_start_frame)

            # Step 3.4: update the start and end frame indices
            current_start_frame += current_num_frames

        # Handle last 21 frames logic
        pred_image_or_video = output

        # Slice last 21 frames if we generated more
        gradient_mask = None
        if pred_image_or_video.shape[1] > 21:
            with torch.no_grad():
                # Re-encode to get image latent
                latent_to_decode = pred_image_or_video[:, :-20, ...]
                # Decode to video
                latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]

                # Apply VAE scaling and shift factors
                if isinstance(self.vae.scaling_factor, torch.Tensor):
                    latent_to_decode = latent_to_decode / self.vae.scaling_factor.to(
                        latent_to_decode.device, latent_to_decode.dtype)
                else:
                    latent_to_decode = latent_to_decode / self.vae.scaling_factor

                if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
                    if isinstance(self.vae.shift_factor, torch.Tensor):
                        latent_to_decode += self.vae.shift_factor.to(latent_to_decode.device, latent_to_decode.dtype)
                    else:
                        latent_to_decode += self.vae.shift_factor

                # Decode to pixels
                pixels = self.vae.decode(latent_to_decode)
                frame = pixels[:, :, -1:, :, :].to(dtype)  # Last frame [B, C, 1, H, W]

                # Encode frame back to get image latent
                image_latent = self.vae.encode(frame).to(dtype)
                image_latent = image_latent.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]

            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
        else:
            pred_image_or_video_last_21 = pred_image_or_video

        # Set up gradient mask if we generated more than minimum frames
        if num_generated_frames != min_num_frames:
            # Currently, we do not use gradient for the first chunk, since it contains image latents
            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
            if self.independent_first_frame:
                gradient_mask[:, :1] = False
            else:
                gradient_mask[:, :self.num_frame_per_block] = False

        # Apply gradient masking if needed
        final_output = pred_image_or_video_last_21.to(dtype)
        if gradient_mask is not None:
            # Apply gradient masking: detach frames that shouldn't contribute gradients
            final_output = torch.where(
                gradient_mask,
                pred_image_or_video_last_21,  # Keep original values where gradient_mask is True
                pred_image_or_video_last_21.detach()  # Detach where gradient_mask is False
            )

        # Store visualization data
        training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.as_tensor(
            self.denoising_step_list[exit_flags[0]],
            dtype=torch.float32,
            device=self.device,
        ).detach().clone()

        # Store gradient mask information for debugging
        if gradient_mask is not None:
            training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float()
            training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(num_generated_frames,
                                                                                      dtype=torch.float32,
                                                                                      device=self.device)
            training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(min_num_frames,
                                                                                dtype=torch.float32,
                                                                                device=self.device)

        # Clean up caches
        assert self.kv_cache1 is not None
        assert self.crossattn_cache is not None
        self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse, self.kv_cache_keyboard)

        return final_output if gradient_mask is not None else pred_image_or_video

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

        if not hasattr(self.validation_pipeline, "prompt_encoding_stage"):
            # validation expects a prompt stage
            def _prompt_encoding_stage(batch: ForwardBatch, _args: TrainingArgs) -> ForwardBatch:
                if not batch.prompt_embeds:
                    batch.prompt_embeds = [None]
                if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
                    batch.prompt_attention_mask = [None]
                return batch

            self.validation_pipeline.prompt_encoding_stage = _prompt_encoding_stage  # type: ignore[attr-defined]

    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
            # 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)

        clip_feature = batch['clip_feature']
        first_frame_latent = batch['first_frame_latent']
        keyboard_cond = batch.get('keyboard_cond', None)
        mouse_cond = batch.get('mouse_cond', None)
        infos = batch['info_list']

        batch_size = clip_feature.shape[0]
        vae_config = self.training_args.pipeline_config.vae_config.arch_config
        num_channels = vae_config.z_dim
        spatial_compression_ratio = vae_config.spatial_compression_ratio

        latent_height = self.training_args.num_height // spatial_compression_ratio
        latent_width = self.training_args.num_width // spatial_compression_ratio

        latents = torch.randn(batch_size, num_channels, self.training_args.num_latent_t, latent_height,
                              latent_width).to(get_local_torch_device(), dtype=torch.bfloat16)

        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.image_embeds = clip_feature.to(get_local_torch_device(), dtype=torch.bfloat16)
        training_batch.image_latents = first_frame_latent.to(get_local_torch_device(), dtype=torch.bfloat16)
        # Action conditioning
        if keyboard_cond is not None and keyboard_cond.numel() > 0:
            keyboard_cond_full = keyboard_cond.to(get_local_torch_device(), dtype=torch.bfloat16)
            training_batch.keyboard_cond = keyboard_cond_full  # For Teacher/Critic (dim=6)
        else:
            training_batch.keyboard_cond = None
        if mouse_cond is not None and mouse_cond.numel() > 0:
            training_batch.mouse_cond = mouse_cond.to(get_local_torch_device(), dtype=torch.bfloat16)
        else:
            training_batch.mouse_cond = None
        training_batch.infos = infos
        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)

        # cond_concat = [mask(4), image_latent(16)] with 20 channels.
        expected_cond_channels = 20
        if image_latents.shape[1] == 16:
            image_latents = self._build_matrixgame_cond_concat(image_latents)
        elif image_latents.shape[1] != expected_cond_channels:
            raise ValueError("Unexpected first_frame_latent channels, "
                             "Expected {expected_cond_channels} (cond_concat), got {image_latents.shape[1]}.")

        if self.sp_world_size > 1:
            total_frames = image_latents.shape[2]
            # Split cond latents to local SP shard only when tensor is still global.
            if total_frames == self.training_args.num_latent_t:
                if total_frames % self.sp_world_size != 0:
                    raise ValueError("image_latents temporal dim is not divisible by SP world size: "
                                     f"frames={total_frames}, sp_world_size={self.sp_world_size}")
                image_latents = rearrange(image_latents, "b c (n t) h w -> b c n t h w",
                                          n=self.sp_world_size).contiguous()
                image_latents = image_latents[:, :, self.rank_in_sp_group, :, :, :]

        training_batch.image_latents = image_latents

        return training_batch

    def _build_distill_input_kwargs(self,
                                    noise_input: torch.Tensor,
                                    timestep: torch.Tensor,
                                    text_dict: dict[str, torch.Tensor] | None,
                                    training_batch: TrainingBatch,
                                    frame_start: int | None = None,
                                    frame_end: int | None = None,
                                    num_frame_per_block: int | None = None) -> TrainingBatch:
        # Image Embeds for conditioning
        image_embeds = training_batch.image_embeds
        assert image_embeds is not None
        assert torch.isnan(image_embeds).sum() == 0
        image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16)
        image_embeds = [image_embeds]

        image_latents = training_batch.image_latents
        if frame_start is not None and frame_end is not None:
            image_latents = image_latents[:, :, frame_start:frame_end, :, :]

        vae_temporal_compression_ratio = 4
        if frame_end is not None:
            action_frame_end = (frame_end - 1) * vae_temporal_compression_ratio + 1
            keyboard_cond_sliced = training_batch.keyboard_cond[:, :
                                                                action_frame_end, :] if training_batch.keyboard_cond is not None else None
            mouse_cond_sliced = training_batch.mouse_cond[:, :
                                                          action_frame_end, :] if training_batch.mouse_cond is not None else None
        else:
            keyboard_cond_sliced = training_batch.keyboard_cond
            mouse_cond_sliced = training_batch.mouse_cond

        noisy_model_input = torch.cat([noise_input, image_latents.permute(0, 2, 1, 3, 4)], dim=2)

        training_batch.input_kwargs = {
            "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4),  # bs, c, t, h, w
            "encoder_hidden_states": None,
            "timestep": timestep,
            "encoder_hidden_states_image": image_embeds,
            "keyboard_cond": keyboard_cond_sliced,
            "mouse_cond": mouse_cond_sliced,
            "num_frame_per_block": num_frame_per_block if num_frame_per_block is not None else self.num_frame_per_block,
        }
        training_batch.noise_latents = noise_input

        return training_batch

    def _dmd_forward(self, generator_pred_video: torch.Tensor, training_batch: TrainingBatch) -> torch.Tensor:
        """Compute DMD (Diffusion Model Distillation) loss for Matrix-Game 2.0."""
        original_latent = generator_pred_video
        with torch.no_grad():
            timestep = torch.randint(0, self.num_train_timestep, [1], device=self.device, dtype=torch.long)

            timestep = shift_timestep(timestep, self.timestep_shift, self.num_train_timestep)

            timestep = timestep.clamp(self.min_timestep, self.max_timestep)

            noise = torch.randn(self.video_latent_shape, device=self.device, dtype=generator_pred_video.dtype)

            noisy_latent = self.noise_scheduler.add_noise(generator_pred_video.flatten(0, 1), noise.flatten(
                0, 1), timestep).detach().unflatten(0, (generator_pred_video.shape[0], generator_pred_video.shape[1]))

            # Non-causal models expect 1D timestep (batch_size,)
            critic_timestep = timestep.expand(noisy_latent.shape[0])

            self._build_distill_input_kwargs(noisy_latent, critic_timestep, None, training_batch)

            # fake_score_transformer forward
            current_fake_score_transformer = self._get_fake_score_transformer(timestep)
            fake_score_pred_noise = current_fake_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

            faker_score_pred_video = pred_noise_to_pred_video(pred_noise=fake_score_pred_noise.flatten(0, 1),
                                                              noise_input_latent=noisy_latent.flatten(0, 1),
                                                              timestep=timestep,
                                                              scheduler=self.noise_scheduler).unflatten(
                                                                  0, fake_score_pred_noise.shape[:2])

            # real_score_transformer forward
            current_real_score_transformer = self._get_real_score_transformer(timestep)
            real_score_pred_noise = current_real_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

            real_score_pred_video = pred_noise_to_pred_video(pred_noise=real_score_pred_noise.flatten(0, 1),
                                                             noise_input_latent=noisy_latent.flatten(0, 1),
                                                             timestep=timestep,
                                                             scheduler=self.noise_scheduler).unflatten(
                                                                 0, real_score_pred_noise.shape[:2])

            # No CFG for Matrix-Game 2.0 - use real_score_pred_video directly
            grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(original_latent -
                                                                                real_score_pred_video).mean()
            grad = torch.nan_to_num(grad)

        dmd_loss = 0.5 * F.mse_loss(original_latent.float(), (original_latent.float() - grad.float()).detach())

        training_batch.dmd_latent_vis_dict.update({
            "training_batch_dmd_fwd_clean_latent": training_batch.latents,
            "generator_pred_video": original_latent.detach(),
            "real_score_pred_video": real_score_pred_video.detach(),
            "faker_score_pred_video": faker_score_pred_video.detach(),
            "dmd_timestep": timestep.detach(),
        })

        return dmd_loss

    def faker_score_forward(self, training_batch: TrainingBatch) -> tuple[TrainingBatch, torch.Tensor]:
        """Forward pass for critic training with Matrix-Game 2.0 action conditioning."""
        with torch.no_grad(), set_forward_context(current_timestep=training_batch.timesteps,
                                                  attn_metadata=training_batch.attn_metadata_vsa):
            if self.training_args.simulate_generator_forward:
                generator_pred_video = self._generator_multi_step_simulation_forward(training_batch)
            else:
                generator_pred_video = self._generator_forward(training_batch)

        fake_score_timestep = torch.randint(0, self.num_train_timestep, [1], device=self.device, dtype=torch.long)

        fake_score_timestep = shift_timestep(fake_score_timestep, self.timestep_shift, self.num_train_timestep)

        fake_score_timestep = fake_score_timestep.clamp(self.min_timestep, self.max_timestep)

        fake_score_noise = torch.randn(self.video_latent_shape, device=self.device, dtype=generator_pred_video.dtype)

        noisy_generator_pred_video = self.noise_scheduler.add_noise(
            generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
            fake_score_timestep).unflatten(0, (generator_pred_video.shape[0], generator_pred_video.shape[1]))

        # Non-causal critic expects 1D timestep (batch_size,), not 2D (batch_size, num_frames).
        expanded_fake_score_timestep = fake_score_timestep.expand(noisy_generator_pred_video.shape[0])

        self._build_distill_input_kwargs(noisy_generator_pred_video, expanded_fake_score_timestep, None, training_batch)

        with set_forward_context(current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata):
            current_fake_score_transformer = self._get_fake_score_transformer(fake_score_timestep)
            fake_score_pred_noise = current_fake_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

        target = fake_score_noise - generator_pred_video
        flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)

        training_batch.fake_score_latent_vis_dict = {
            "training_batch_fakerscore_fwd_clean_latent": training_batch.latents,
            "generator_pred_video": generator_pred_video,
            "fake_score_timestep": fake_score_timestep,
        }

        return training_batch, flow_matching_loss

    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_self_forcing_distillation_pipeline.MatrixGame2SelfForcingDistillationPipeline.set_schemas`

set_schemas()

Source

def set_schemas(self):
    self.train_dataset_schema = pyarrow_schema_matrixgame2

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2SelfForcingDistillationPipeline.load_modules`

load_modules(fastvideo_args: FastVideoArgs, loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]

Source

def load_modules(
    self,
    fastvideo_args: FastVideoArgs,
    loaded_modules: dict[str, torch.nn.Module] | None = None,
) -> dict[str, Any]:
    modules = ComposedPipelineBase.load_modules(
        self,
        fastvideo_args,
        loaded_modules,
    )
    training_args = cast(TrainingArgs, fastvideo_args)
    old_override = training_args.override_transformer_cls_name
    training_args.override_transformer_cls_name = "MatrixGame2WanModel"
    try:
        if loaded_modules is not None and "real_score_transformer" in loaded_modules:
            self.real_score_transformer = loaded_modules["real_score_transformer"]
        elif training_args.real_score_model_path:
            logger.info(
                "Loading real score transformer from: %s",
                training_args.real_score_model_path,
            )
            self.real_score_transformer = self.load_module_from_path(
                training_args.real_score_model_path,
                "transformer",
                training_args,
            )
        else:
            raise ValueError("real_score_model_path is required for DMD distillation pipeline")
        modules["real_score_transformer"] = self.real_score_transformer

        if loaded_modules is not None and "fake_score_transformer" in loaded_modules:
            self.fake_score_transformer = loaded_modules["fake_score_transformer"]
        elif training_args.fake_score_model_path:
            logger.info(
                "Loading fake score transformer from: %s",
                training_args.fake_score_model_path,
            )
            self.fake_score_transformer = self.load_module_from_path(
                training_args.fake_score_model_path,
                "transformer",
                training_args,
            )
        else:
            raise ValueError("fake_score_model_path is required for DMD distillation pipeline")
        modules["fake_score_transformer"] = self.fake_score_transformer
    finally:
        training_args.override_transformer_cls_name = old_override

    self.real_score_transformer_2 = None
    self.fake_score_transformer_2 = None
    return modules

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2SelfForcingDistillationPipeline.initialize_validation_pipeline`

initialize_validation_pipeline(training_args: TrainingArgs)

Source

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

    if not hasattr(self.validation_pipeline, "prompt_encoding_stage"):
        # validation expects a prompt stage
        def _prompt_encoding_stage(batch: ForwardBatch, _args: TrainingArgs) -> ForwardBatch:
            if not batch.prompt_embeds:
                batch.prompt_embeds = [None]
            if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
                batch.prompt_attention_mask = [None]
            return batch

        self.validation_pipeline.prompt_encoding_stage = _prompt_encoding_stage  # type: ignore[attr-defined]

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.MatrixGame2SelfForcingDistillationPipeline.faker_score_forward`

faker_score_forward(training_batch: TrainingBatch) -> tuple[TrainingBatch, torch.Tensor]

Forward pass for critic training with Matrix-Game 2.0 action conditioning.

Source

def faker_score_forward(self, training_batch: TrainingBatch) -> tuple[TrainingBatch, torch.Tensor]:
    """Forward pass for critic training with Matrix-Game 2.0 action conditioning."""
    with torch.no_grad(), set_forward_context(current_timestep=training_batch.timesteps,
                                              attn_metadata=training_batch.attn_metadata_vsa):
        if self.training_args.simulate_generator_forward:
            generator_pred_video = self._generator_multi_step_simulation_forward(training_batch)
        else:
            generator_pred_video = self._generator_forward(training_batch)

    fake_score_timestep = torch.randint(0, self.num_train_timestep, [1], device=self.device, dtype=torch.long)

    fake_score_timestep = shift_timestep(fake_score_timestep, self.timestep_shift, self.num_train_timestep)

    fake_score_timestep = fake_score_timestep.clamp(self.min_timestep, self.max_timestep)

    fake_score_noise = torch.randn(self.video_latent_shape, device=self.device, dtype=generator_pred_video.dtype)

    noisy_generator_pred_video = self.noise_scheduler.add_noise(
        generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
        fake_score_timestep).unflatten(0, (generator_pred_video.shape[0], generator_pred_video.shape[1]))

    # Non-causal critic expects 1D timestep (batch_size,), not 2D (batch_size, num_frames).
    expanded_fake_score_timestep = fake_score_timestep.expand(noisy_generator_pred_video.shape[0])

    self._build_distill_input_kwargs(noisy_generator_pred_video, expanded_fake_score_timestep, None, training_batch)

    with set_forward_context(current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata):
        current_fake_score_transformer = self._get_fake_score_transformer(fake_score_timestep)
        fake_score_pred_noise = current_fake_score_transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

    target = fake_score_noise - generator_pred_video
    flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)

    training_batch.fake_score_latent_vis_dict = {
        "training_batch_fakerscore_fwd_clean_latent": training_batch.latents,
        "generator_pred_video": generator_pred_video,
        "fake_score_timestep": fake_score_timestep,
    }

    return training_batch, flow_matching_loss

`fastvideo.training.matrixgame2_self_forcing_distillation_pipeline.main`

main(args) -> None

`fastvideo.training.matrixgame2_self_forcing_distillation_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.