matrixgame2_ode_causal_pipeline
fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam
SamplingParam(data_type: str = 'video', image_path: str | None = None, pil_image: Any | None = None, last_image: Any | None = None, references: list[Any] | None = None, video_path: str | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, audio_path: str | None = None, latents: Any | None = None, audio_latents: Any | None = None, mouse_cond: Any | None = None, keyboard_cond: Any | None = None, grid_sizes: Any | None = None, pose: str | None = None, prompt_attention_mask: list = list(), negative_attention_mask: list = list(), c2ws_plucker_emb: Any | None = None, action_path: str | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt: str | list[str] | None = None, negative_prompt: str = 'Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards', max_sequence_length: int | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, num_videos_per_prompt: int = 1, seed: int = 1024, num_frames: int = 125, height: int = 720, width: int = 1280, height_sr: int = 1072, width_sr: int = 1920, fps: int = 24, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, guidance_rescale: float = 0.0, boundary_ratio: float | None = None, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = (lambda: [29])(), ltx2_stg_blocks_audio: list[int] = (lambda: [29])(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: Any | None = None, ltx2_conditioning_latent_stage2: Any | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, save_video: bool = True, return_frames: bool = True, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, *, return_samples: bool = False)Sampling parameters for video generation.
fastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.check_sampling_param
check_sampling_param() -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.update
update(source_dict: dict[str, Any]) -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.from_pretrained
from_pretrained(model_path: str) -> SamplingParamfastvideo.training.matrixgame2_ode_causal_pipeline.SamplingParam.add_cli_args
add_cli_args(parser: Any) -> AnyAdd CLI arguments for SamplingParam fields
fastvideo.training.matrixgame2_ode_causal_pipeline.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn the torch device for the current rank.
fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs
FastVideoArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875)fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.is_output_rank
is_output_rank: boolWhether this process may materialize user-facing outputs.
This is runtime state assigned by an SPMD executor, not a public configuration field, so dataclass serialization and typed config parsing deliberately ignore it.
fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.add_cli_args
add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParserfastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.from_cli_args
from_cli_args(args: argparse.Namespace) -> FastVideoArgsfastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.from_kwargs
from_kwargs(**kwargs: Any) -> FastVideoArgsfastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.check_fastvideo_args
check_fastvideo_args() -> NoneValidate inference arguments for consistency
fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.finalize_device_offload_policy
finalize_device_offload_policy(device_id: int = 0) -> boolApply device-local memory policy, then resolve incompatible modes.
fastvideo.training.matrixgame2_ode_causal_pipeline.FastVideoArgs.disable_offload_on_unified_memory
disable_offload_on_unified_memory(device_id: int = 0, *, offload_flag: str | None = None) -> boolDisable host offload after a worker has selected its device.
CUDA's unified-memory probe reads runtime device properties and may
initialize a CUDA context. Callers must therefore use this only inside
a device-owning process, after selecting and binding device_id.
Returning the classification lets direct component-loader callers
apply the same policy to explicit per-call overrides. When
offload_flag is given, the return value says whether this policy
covers that component role.
fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingArgs
TrainingArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875, pipeline_class: str = '', pipeline_module: str = '', data_path: str = '', dataloader_num_workers: int = 0, num_height: int = 0, num_width: int = 0, num_frames: int = 0, train_batch_size: int = 0, num_latent_t: int = 0, group_frame: bool = False, group_resolution: bool = False, pretrained_model_name_or_path: str = '', real_score_model_path: str = '', fake_score_model_path: str = '', ema_decay: float = 0.0, ema_start_step: int = 0, training_cfg_rate: float = 0.0, precondition_outputs: bool = False, validation_dataset_file: str = '', validation_preprocessed_path: str = '', validation_sampling_steps: str = '', validation_guidance_scale: str = '', validation_steps: float = 0.0, log_validation: bool = False, trackers: list[str] = list(), tracker_project_name: str = '', wandb_run_name: str = '', seed: int | None = None, output_dir: str = '', checkpoints_total_limit: int = 0, resume_from_checkpoint: str = '', num_train_epochs: int = 0, max_train_steps: int = 0, gradient_accumulation_steps: int = 0, learning_rate: float = 0.0, scale_lr: bool = False, lr_scheduler: str = 'constant', lr_warmup_steps: int = 0, max_grad_norm: float = 0.0, enable_gradient_checkpointing_type: str | None = None, selective_checkpointing: float = 0.0, mixed_precision: str = '', train_sp_batch_size: int = 0, fsdp_sharding_startegy: str = '', weighting_scheme: str = '', logit_mean: float = 0.0, logit_std: float = 1.0, mode_scale: float = 0.0, num_euler_timesteps: int = 0, lr_num_cycles: int = 0, lr_power: float = 0.0, min_lr_ratio: float = 0.5, not_apply_cfg_solver: bool = False, distill_cfg: float = 0.0, scheduler_type: str = '', linear_quadratic_threshold: float = 0.0, linear_range: float = 0.0, weight_decay: float = 0.0, betas: str = '0.9,0.999', use_ema: bool = False, multi_phased_distill_schedule: str = '', pred_decay_weight: float = 0.0, pred_decay_type: str = '', hunyuan_teacher_disable_cfg: bool = False, master_weight_type: str = '', VSA_decay_rate: float = 0.01, VSA_decay_interval_steps: int = 1, VSA_cache_tile_buf: bool = False, lora_rank: int | None = None, lora_alpha: int | None = None, lora_training: bool = False, ltx2_first_frame_conditioning_p: float = 0.1, generator_update_interval: int = 5, dfake_gen_update_ratio: int = 5, min_timestep_ratio: float = 0.2, max_timestep_ratio: float = 0.98, real_score_guidance_scale: float = 3.5, fake_score_learning_rate: float = 0.0, fake_score_lr_scheduler: str = 'constant', fake_score_betas: str = '0.9,0.999', training_state_checkpointing_steps: int = 0, weight_only_checkpointing_steps: int = 0, log_visualization: bool = False, visualization_steps: int = 0, simulate_generator_forward: bool = False, warp_denoising_step: bool = False, num_frame_per_block: int = 3, independent_first_frame: bool = False, enable_gradient_masking: bool = True, gradient_mask_last_n_frames: int = 21, same_step_across_blocks: bool = False, last_step_only: bool = False, context_noise: int = 0)Training arguments. Inherits from FastVideoArgs and adds training-specific arguments. If there are any conflicts, the training arguments will take precedence.
fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingArgs.from_cli_args
from_cli_args(args: argparse.Namespace) -> TrainingArgsfastvideo.training.matrixgame2_ode_causal_pipeline.TrainingArgs.add_cli_args
add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParserfastvideo.training.matrixgame2_ode_causal_pipeline.set_forward_context
set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.
fastvideo.training.matrixgame2_ode_causal_pipeline.init_logger
init_logger(name: str) -> _FastvideoLoggerThe main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.
fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler
SelfForcingFlowMatchScheduler(num_inference_steps = 100, num_train_timesteps = 1000, shift = 3.0, sigma_max = 1.0, sigma_min = 0.003 / 1.002, inverse_timesteps = False, extra_one_step = False, reverse_sigmas = False, training = False)fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.set_timesteps
set_timesteps(num_inference_steps = 100, denoising_strength = 1.0, training = False, return_dict = False, **kwargs)fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.step
step(model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final = False, return_dict = False, **kwargs)fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high
calculate_alpha_beta_high(sigma, sigma_bound)fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.add_noise
add_noise(original_samples, noise, timestep)Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]
fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.add_noise_high
add_noise_high(original_samples, noise, timestep, boundary_timestep)Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]
fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.training_target
training_target(sample, noise, timestep)fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.training_weight
training_weight(timestep)Input
- timestep: the timestep with shape [B*T]
Output: the corresponding weighting [B*T]
fastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.scale_model_input
scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensorfastvideo.training.matrixgame2_ode_causal_pipeline.SelfForcingFlowMatchScheduler.set_shift
set_shift(shift: float) -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline
MatrixGame2CausalDMDPipeline(*args, **kwargs)fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.create_pipeline_stages
create_pipeline_stages(fastvideo_args: FastVideoArgs) -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_reset
streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs)fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_step
streaming_step(keyboard_action, mouse_action) -> ForwardBatchfastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2CausalDMDPipeline.streaming_clear
streaming_clear() -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage
DecodingStage(vae, pipeline = None)Stage for decoding latent representations into pixel space.
This stage handles the decoding of latent representations into the final output format (e.g., pixel values).
fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify decoding stage inputs.
fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify decoding stage outputs.
fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.decode
decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.TensorDecode latent representations into pixel space using VAE.
Parameters
latents: torch.TensorInput latent tensor with shape (batch, channels, frames, height_latents, width_latents)
fastvideo_args: FastVideoArgsConfiguration containing: - disable_autocast: Whether to disable automatic mixed precision (default: False) - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency
Returns
torch.TensorDecoded video tensor with shape (batch, channels, frames, height, width),
torch.Tensornormalized to [0, 1] range and moved to CPU as float32
fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.streaming_decode
streaming_decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs, cache: list[torch.Tensor | None] | None = None, is_first_chunk: bool = False) -> tuple[torch.Tensor, list[torch.Tensor | None]]Decode latent representations into pixel space using VAE with streaming cache.
Parameters
latents: torch.TensorInput latent tensor with shape (batch, channels, frames, height_latents, width_latents)
fastvideo_args: FastVideoArgsConfiguration object.
cache: list[torch.Tensor | None] | NoneVAE cache from previous call, or None to initialize a new cache.
is_first_chunk: boolWhether this is the first chunk.
Returns
tuple[torch.Tensor, list[torch.Tensor | None]]A tuple of (decoded_frames, updated_cache).
fastvideo.training.matrixgame2_ode_causal_pipeline.DecodingStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchDecode latent representations into pixel space.
This method processes the batch through the VAE decoder, converting latent representations to pixel-space video/images. It also optionally decodes trajectory latents for visualization purposes.
Parameters
batch: ForwardBatchThe current batch containing: - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents) - return_trajectory_decoded (optional): Flag to decode trajectory latents - trajectory_latents (optional): Latents at different timesteps - trajectory_timesteps (optional): Corresponding timesteps
fastvideo_args: FastVideoArgsConfiguration containing: - output_type: "latent" to skip decoding, otherwise decode to pixels - vae_cpu_offload: Whether to offload VAE to CPU after decoding - model_loaded: Track VAE loading state - model_paths: Path to VAE model if loading needed
Returns
ForwardBatchModified batch with: - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 - trajectory_decoded (if requested): List of decoded frames per timestep
fastvideo.training.matrixgame2_ode_causal_pipeline.ForwardBatch
ForwardBatch(data_type: str, generator: torch.Generator | list[torch.Generator] | None = None, image_path: str | None = None, image_embeds: list[torch.Tensor] = list(), pil_image: torch.Tensor | PIL.Image.Image | None = None, last_image: torch.Tensor | PIL.Image.Image | None = None, references: list[Any] | None = None, preprocessed_image: torch.Tensor | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, video_path: str | None = None, video_latent: torch.Tensor | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, pose_latents: torch.Tensor | None = None, face_pixel_values: torch.Tensor | None = None, audio_path: str | None = None, audio_embeds: torch.Tensor | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt_embeds: list[torch.Tensor] = list(), negative_prompt_embeds: list[torch.Tensor] | None = None, prompt_attention_mask: list[torch.Tensor] | None = None, negative_attention_mask: list[torch.Tensor] | None = None, clip_embedding_pos: list[torch.Tensor] | None = None, clip_embedding_neg: list[torch.Tensor] | None = None, max_sequence_length: int | None = None, prompt_template: dict[str, Any] | None = None, do_classifier_free_guidance: bool = False, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, batch_size: int | None = None, num_videos_per_prompt: int = 1, seed: int | None = None, seeds: list[int] | None = None, is_prompt_processed: bool = False, latents: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, lq_latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_pred: torch.Tensor | None = None, image_latent: torch.Tensor | None = None, first_frame_latent: torch.Tensor | None = None, helios_latent_chunks: list[torch.Tensor] | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, num_iterations: int | None = None, use_base_model: bool = False, pose: str | None = None, c2ws_plucker_emb: torch.Tensor | None = None, action_path: str | None = None, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, height_latents: list[int] | int | None = None, width_latents: list[int] | int | None = None, num_frames: list[int] | int = 1, height: list[int] | int | None = None, width: list[int] | int | None = None, height_sr: list[int] | int | None = None, width_sr: list[int] | int | None = None, fps: list[int] | int | None = None, timesteps: torch.Tensor | None = None, timestep: torch.Tensor | float | int | None = None, step_index: int | None = None, boundary_ratio: float | None = None, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, guidance_rescale: float = 0.0, eta: float = 0.0, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = list(), ltx2_stg_blocks_audio: list[int] = list(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: torch.Tensor | None = None, ltx2_conditioning_latent_stage2: torch.Tensor | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, n_tokens: int | None = None, extra_step_kwargs: dict[str, Any] = dict(), modules: dict[str, Any] = dict(), output: torch.Tensor | None = None, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, trajectory_timesteps: list[torch.Tensor] | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_decoded: list[torch.Tensor] | None = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, extra: dict[str, Any] = dict(), save_video: bool = True, return_frames: bool = False, is_cfg_negative: bool = False, VSA_sparsity: float = 0.0, logging_info: PipelineLoggingInfo = PipelineLoggingInfo(), *, return_samples: bool = False)Complete state passed through the pipeline execution.
This dataclass contains all information needed during the diffusion pipeline execution, allowing methods to update specific components without needing to manage numerous individual parameters.
fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingBatch
TrainingBatch(current_timestep: int = 0, current_vsa_sparsity: float = 0.0, latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_latents: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, audio_noisy_model_input: torch.Tensor | None = None, audio_timesteps: torch.Tensor | None = None, audio_sigmas: torch.Tensor | None = None, audio_noise: torch.Tensor | None = None, audio_encoder_hidden_states: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, conditioning_mask: torch.Tensor | None = None, preprocessed_image: torch.Tensor | None = None, image_embeds: torch.Tensor | None = None, image_latents: torch.Tensor | None = None, infos: list[dict[str, Any]] | None = None, mask_lat_size: torch.Tensor | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_timesteps: torch.Tensor | None = None, noisy_model_input: torch.Tensor | None = None, timesteps: torch.Tensor | None = None, sigmas: torch.Tensor | None = None, noise: torch.Tensor | None = None, training_target: torch.Tensor | None = None, minimax_h3_layout: Any | None = None, attn_metadata_vsa: AttentionMetadata | None = None, attn_metadata: AttentionMetadata | None = None, input_kwargs: dict[str, Any] | None = None, loss: torch.Tensor | None = None, total_loss: float | None = None, grad_norm: float | None = None, encoder_hidden_states_neg: torch.Tensor | None = None, encoder_attention_mask_neg: torch.Tensor | None = None, conditional_dict: dict[str, Any] | None = None, unconditional_dict: dict[str, Any] | None = None, generator_loss: float = 0.0, fake_score_loss: float = 0.0, dmd_latent_vis_dict: dict[str, Any] = dict(), latent_vis_dict: dict[str, Any] = dict(), fake_score_latent_vis_dict: dict[str, Any] = dict(), minimax_h3_condition_video_rows: torch.Tensor | None = None, minimax_h3_condition_audio_rows: torch.Tensor | None = None, minimax_h3_dmd_layout: Any | None = None, minimax_h3_document_video_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_document_audio_latents: tuple[torch.Tensor, ...] | None = None, minimax_h3_text_attn_metadata: AttentionMetadata | None = None, sequence_packing_num_documents: int = 1, sequence_packing_num_tokens: int = 0, sequence_padded_num_tokens: int = 0, sequence_padding_ratio: float | None = None, sequence_packing_loss_scale: float = 1.0, sequence_packing_data_forced_ratio: float | None = None)fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline
TrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)A pipeline for training a model. All training pipelines should inherit from this class. All reusable components and code should be implemented in this class.
fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.create_pipeline_stages
create_pipeline_stages(fastvideo_args: FastVideoArgs)fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.set_schemas
set_schemas() -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.initialize_training_pipeline
initialize_training_pipeline(training_args: TrainingArgs)fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.initialize_validation_pipeline
initialize_validation_pipeline(training_args: TrainingArgs)fastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.train_one_step
train_one_step(training_batch: TrainingBatch) -> TrainingBatchfastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.train
train() -> Nonefastvideo.training.matrixgame2_ode_causal_pipeline.TrainingPipeline.visualize_intermediate_latents
visualize_intermediate_latents(training_batch: TrainingBatch, training_args: TrainingArgs, step: int)Add visualization data to tracker logging and save frames to disk.
fastvideo.training.matrixgame2_ode_causal_pipeline.clip_grad_norm_while_handling_failing_dtensor_cases
clip_grad_norm_while_handling_failing_dtensor_cases(parameters: torch.Tensor | list[torch.Tensor], max_norm: float, norm_type: float = 2.0, error_if_nonfinite: bool = False, foreach: bool | None = None, pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None) -> torch.Tensor | Nonefastvideo.training.matrixgame2_ode_causal_pipeline.shallow_asdict
shallow_asdict(obj) -> dict[str, Any]fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline
MatrixGame2ODEInitTrainingPipeline(model_path: str, fastvideo_args: TrainingArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None)Training pipeline for ODE-init using precomputed denoising trajectories.
Supervision: predict the next latent in the stored trajectory by - feeding current latent at timestep t into the transformer to predict noise - stepping the scheduler with the predicted noise - minimizing MSE to the stored next latent at timestep t_next
Source
class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline):
"""
Training pipeline for ODE-init using precomputed denoising trajectories.
Supervision: predict the next latent in the stored trajectory by
- feeding current latent at timestep t into the transformer to predict noise
- stepping the scheduler with the predicted noise
- minimizing MSE to the stored next latent at timestep t_next
"""
_required_config_modules = ["scheduler", "transformer", "vae"]
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
# Match the preprocess/generation scheduler for consistent stepping
self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
sigma_min=0.0,
extra_one_step=True)
self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)
def set_schemas(self):
self.train_dataset_schema = pyarrow_schema_matrixgame2_ode_trajectory
def initialize_training_pipeline(self, training_args: TrainingArgs):
super().initialize_training_pipeline(training_args)
self.noise_scheduler = self.get_module("scheduler")
self.vae = self.get_module("vae")
self.vae.requires_grad_(False)
self.timestep_shift = self.training_args.pipeline_config.flow_shift
self.noise_scheduler = SelfForcingFlowMatchScheduler(shift=self.timestep_shift,
sigma_min=0.0,
extra_one_step=True)
self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps)
self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250, 0],
dtype=torch.long,
device=get_local_torch_device())
if training_args.warp_denoising_step: # Warp the denoising step according to the scheduler time shift
timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
logger.info("timesteps: %s", timesteps)
self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
logger.info("warped self.dmd_denoising_steps: %s", self.dmd_denoising_steps)
else:
raise ValueError("warp_denoising_step must be true")
self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())
logger.info("denoising_step_list: %s", self.dmd_denoising_steps)
logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
# Cache for nearest trajectory index per DMD step (computed lazily on first batch)
self._cached_closest_idx_per_dmd = None
self.num_train_timestep = self.noise_scheduler.num_train_timesteps
self.manual_idx = 0
def initialize_validation_pipeline(self, training_args: TrainingArgs):
logger.info("Initializing validation pipeline...")
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
# Use the same flow-matching scheduler as training for consistent validation.
validation_scheduler = SelfForcingFlowMatchScheduler(shift=args_copy.pipeline_config.flow_shift,
sigma_min=0.0,
extra_one_step=True)
validation_scheduler.set_timesteps(num_inference_steps=1000, training=True)
# Warm start validation with current transformer
self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
training_args.model_path,
args=args_copy, # type: ignore
inference_mode=True,
loaded_modules={
"transformer": self.get_module("transformer"),
"vae": self.get_module("vae"),
"scheduler": validation_scheduler,
},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus,
pin_cpu_memory=training_args.pin_cpu_memory,
dit_cpu_offload=True)
def _get_next_batch(self, training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]:
batch = next(self.train_loader_iter, None) # type: ignore
if batch is None:
self.current_epoch += 1
logger.info("Starting epoch %s", self.current_epoch)
self.train_loader_iter = iter(self.train_dataloader)
batch = next(self.train_loader_iter)
# Required fields from parquet (ODE trajectory schema)
clip_feature = batch['clip_feature']
first_frame_latent = batch['first_frame_latent']
keyboard_cond = batch.get('keyboard_cond', None)
mouse_cond = batch.get('mouse_cond', None)
infos = batch['info_list']
# Trajectory tensors may include a leading singleton batch dim per row
trajectory_latents = batch['trajectory_latents']
if trajectory_latents.dim() == 7:
# [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
trajectory_latents = trajectory_latents[:, 0]
elif trajectory_latents.dim() == 6:
# already [B, S, C, T, H, W]
pass
else:
raise ValueError(f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}")
trajectory_timesteps = batch['trajectory_timesteps']
if trajectory_timesteps.dim() == 3:
# [B, 1, S] -> [B, S]
trajectory_timesteps = trajectory_timesteps[:, 0]
elif trajectory_timesteps.dim() == 2:
# [B, S]
pass
else:
raise ValueError(f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}")
# [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
# Move to device
device = get_local_torch_device()
training_batch.image_embeds = clip_feature.to(device, dtype=torch.bfloat16)
training_batch.image_latents = first_frame_latent.to(device, dtype=torch.bfloat16)
if keyboard_cond is not None and keyboard_cond.numel() > 0:
training_batch.keyboard_cond = keyboard_cond.to(device, dtype=torch.bfloat16)
else:
training_batch.keyboard_cond = None
if mouse_cond is not None and mouse_cond.numel() > 0:
training_batch.mouse_cond = mouse_cond.to(device, dtype=torch.bfloat16)
else:
training_batch.mouse_cond = None
training_batch.infos = infos
return training_batch, trajectory_latents[:, :, :self.training_args.num_latent_t].to(
device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
def _get_timestep(self,
min_timestep: int,
max_timestep: int,
batch_size: int,
num_frame: int,
num_frame_per_block: int,
uniform_timestep: bool = False) -> torch.Tensor:
if uniform_timestep:
timestep = torch.randint(min_timestep, max_timestep, [batch_size, 1], device=self.device,
dtype=torch.long).repeat(1, num_frame)
return timestep
else:
timestep = torch.randint(min_timestep,
max_timestep, [batch_size, num_frame],
device=self.device,
dtype=torch.long)
# logger.info(f"individual timestep: {timestep}")
# make the noise level the same within every block
timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block)
timestep[:, :, 1:] = timestep[:, :, 0:1]
timestep = timestep.reshape(timestep.shape[0], -1)
return timestep
def _step_predict_next_latent(
self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor, image_embeds: torch.Tensor,
image_latents: torch.Tensor, keyboard_cond: torch.Tensor | None, mouse_cond: torch.Tensor | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
latent_vis_dict: dict[str, torch.Tensor] = {}
device = get_local_torch_device()
target_latent = traj_latents[:, -1]
# Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
B, S, num_frames, num_channels, height, width = traj_latents.shape
expected_action_frames = 1 + 4 * (num_frames - 1)
if keyboard_cond is None or mouse_cond is None:
raise ValueError("keyboard_cond/mouse_cond must both be provided for action-follow training. "
f"keyboard_cond={None if keyboard_cond is None else tuple(keyboard_cond.shape)}, "
f"mouse_cond={None if mouse_cond is None else tuple(mouse_cond.shape)}")
if keyboard_cond.shape[1] < expected_action_frames:
raise ValueError("keyboard_cond length is shorter than required for latent frames. "
f"got={keyboard_cond.shape[1]}, required>={expected_action_frames}, "
f"num_latent_frames={num_frames}")
if mouse_cond.shape[1] < expected_action_frames:
raise ValueError("mouse_cond length is shorter than required for latent frames. "
f"got={mouse_cond.shape[1]}, required>={expected_action_frames}, "
f"num_latent_frames={num_frames}")
# Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
if self._cached_closest_idx_per_dmd is None:
traj_ts = traj_timesteps[0].float().cpu()
dmd_steps = self.dmd_denoising_steps.float().cpu()
closest_idx = torch.argmin(torch.abs(traj_ts.unsqueeze(0) - dmd_steps.unsqueeze(1)), dim=1)
self._cached_closest_idx_per_dmd = closest_idx.to(torch.long).cpu()
logger.info("self._cached_closest_idx_per_dmd: %s", self._cached_closest_idx_per_dmd)
logger.info("corresponding timesteps: %s", traj_ts[self._cached_closest_idx_per_dmd])
# Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
# traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
# Output: [B, K, C, T, H, W]
assert self._cached_closest_idx_per_dmd is not None
relevant_traj_latents = torch.index_select(traj_latents,
dim=1,
index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
indexes = self._get_timestep( # [B, num_frames]
0, len(self.dmd_denoising_steps), B, num_frames, 3, uniform_timestep=False)
noisy_input = torch.gather(relevant_traj_latents,
dim=1,
index=indexes.reshape(B, 1, num_frames, 1, 1,
1).expand(-1, -1, -1, num_channels, height,
width).to(self.device)).squeeze(1)
latent_model_input = noisy_input.permute(0, 2, 1, 3, 4)
if image_latents is not None:
latent_model_input = torch.cat([
latent_model_input,
image_latents.to(latent_model_input.device, latent_model_input.dtype),
],
dim=1)
timestep = self.dmd_denoising_steps[indexes]
# Prepare inputs for transformer
latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu()
latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu()
latent_model_input = latent_model_input.to(device, dtype=torch.bfloat16)
timestep = timestep.to(device, dtype=torch.bfloat16)
input_kwargs = {
"hidden_states": latent_model_input,
"encoder_hidden_states": None,
"encoder_hidden_states_image": image_embeds,
"timestep": timestep,
"mouse_cond": mouse_cond,
"keyboard_cond": keyboard_cond,
"return_dict": False,
}
# Predict noise and step the scheduler to obtain next latent
with set_forward_context(current_timestep=timestep, attn_metadata=None, forward_batch=None):
noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
from fastvideo.models.utils import pred_noise_to_pred_video
pred_video = pred_noise_to_pred_video(pred_noise=noise_pred.flatten(0, 1),
noise_input_latent=noisy_input.flatten(0, 1),
timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1),
scheduler=self.modules["scheduler"]).unflatten(0, noise_pred.shape[:2])
latent_vis_dict["pred_video"] = pred_video.permute(0, 2, 1, 3, 4).detach().clone().cpu()
return pred_video, target_latent, timestep, latent_vis_dict
def train_one_step(self, training_batch): # type: ignore[override]
self.transformer.train()
self.optimizer.zero_grad()
training_batch.total_loss = 0.0
args = cast(TrainingArgs, self.training_args)
# Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
for _ in range(args.gradient_accumulation_steps):
training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
image_embeds = training_batch.image_embeds
image_latents = training_batch.image_latents
keyboard_cond = training_batch.keyboard_cond
mouse_cond = training_batch.mouse_cond
assert traj_latents.shape[0] == 1
# Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
_, S = traj_latents.shape[0], traj_latents.shape[1]
if S < 2:
raise ValueError("Trajectory must contain at least 2 steps")
# Forward to predict next latent by stepping scheduler with predicted noise
noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
traj_latents, traj_timesteps, image_embeds, image_latents, keyboard_cond, mouse_cond)
training_batch.latent_vis_dict.update(latent_vis_dict)
mask = t != 0
# Compute loss
loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
loss = loss / args.gradient_accumulation_steps
with set_forward_context(current_timestep=t, attn_metadata=None, forward_batch=None):
loss.backward()
avg_loss = loss.detach().clone()
training_batch.total_loss += avg_loss.item()
# Clip grad and step optimizers
grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
[p for p in self.transformer.parameters() if p.requires_grad],
args.max_grad_norm if args.max_grad_norm is not None else 0.0)
self.optimizer.step()
self.lr_scheduler.step()
if grad_norm is None:
grad_value = 0.0
else:
try:
if isinstance(grad_norm, torch.Tensor):
grad_value = float(grad_norm.detach().float().item())
else:
grad_value = float(grad_norm)
except Exception:
grad_value = 0.0
training_batch.grad_norm = grad_value
B, S, T, C, H, W = traj_latents.shape
training_batch.raw_latent_shape = (B, C, T, H, W)
return training_batch
def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args: TrainingArgs,
validation_batch: dict[str, Any], num_inference_steps: int) -> ForwardBatch:
sampling_param.prompt = validation_batch['prompt']
sampling_param.height = training_args.num_height
sampling_param.width = training_args.num_width
sampling_param.image_path = validation_batch.get('image_path') or validation_batch.get('video_path')
sampling_param.num_inference_steps = num_inference_steps
sampling_param.data_type = "video"
assert self.seed is not None
sampling_param.seed = self.seed
latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1
sampling_param.num_frames = num_frames
batch = ForwardBatch(
**shallow_asdict(sampling_param),
generator=torch.Generator(device="cpu").manual_seed(self.seed),
n_tokens=n_tokens,
eta=0.0,
VSA_sparsity=training_args.VSA_sparsity,
)
if "image" in validation_batch and validation_batch["image"] is not None:
batch.pil_image = validation_batch["image"]
if "keyboard_cond" in validation_batch and validation_batch["keyboard_cond"] is not None:
keyboard_cond = validation_batch["keyboard_cond"]
keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
keyboard_cond = keyboard_cond.unsqueeze(0)
batch.keyboard_cond = keyboard_cond
if "mouse_cond" in validation_batch and validation_batch["mouse_cond"] is not None:
mouse_cond = validation_batch["mouse_cond"]
mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
mouse_cond = mouse_cond.unsqueeze(0)
batch.mouse_cond = mouse_cond
return batch
def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
tracker_loss_dict: dict[str, Any] = {}
latents_vis_dict = training_batch.latent_vis_dict
latent_log_keys = ['noisy_input', 'x0', 'pred_video']
for latent_key in latent_log_keys:
assert latent_key in latents_vis_dict and latents_vis_dict[latent_key] is not None
latent = latents_vis_dict[latent_key]
pixel_latent = self.decoding_stage.decode(latent, training_args)
video = pixel_latent.cpu().float()
video = video.permute(0, 2, 1, 3, 4)
video = (video * 255).numpy().astype(np.uint8)
video_artifact = self.tracker.video(video, fps=16, format="mp4") # change to 16 for Wan2.1
if video_artifact is not None:
tracker_loss_dict[latent_key] = video_artifact
# Clean up references
del video, pixel_latent, latent
if self.global_rank == 0 and tracker_loss_dict:
self.tracker.log_artifacts(tracker_loss_dict, step)</code></pre></details>
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_pipeline
initialize_pipeline(fastvideo_args: FastVideoArgs)
Source
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
# Match the preprocess/generation scheduler for consistent stepping
self.modules["scheduler"] = SelfForcingFlowMatchScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
sigma_min=0.0,
extra_one_step=True)
self.modules["scheduler"].set_timesteps(num_inference_steps=1000, training=True)
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.set_schemas
set_schemas()
Source
def set_schemas(self):
self.train_dataset_schema = pyarrow_schema_matrixgame2_ode_trajectory
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_training_pipeline
initialize_training_pipeline(training_args: TrainingArgs)
Source
def initialize_training_pipeline(self, training_args: TrainingArgs):
super().initialize_training_pipeline(training_args)
self.noise_scheduler = self.get_module("scheduler")
self.vae = self.get_module("vae")
self.vae.requires_grad_(False)
self.timestep_shift = self.training_args.pipeline_config.flow_shift
self.noise_scheduler = SelfForcingFlowMatchScheduler(shift=self.timestep_shift,
sigma_min=0.0,
extra_one_step=True)
self.noise_scheduler.set_timesteps(num_inference_steps=1000, training=True)
self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps)
self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250, 0],
dtype=torch.long,
device=get_local_torch_device())
if training_args.warp_denoising_step: # Warp the denoising step according to the scheduler time shift
timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
logger.info("timesteps: %s", timesteps)
self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
logger.info("warped self.dmd_denoising_steps: %s", self.dmd_denoising_steps)
else:
raise ValueError("warp_denoising_step must be true")
self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())
logger.info("denoising_step_list: %s", self.dmd_denoising_steps)
logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
# Cache for nearest trajectory index per DMD step (computed lazily on first batch)
self._cached_closest_idx_per_dmd = None
self.num_train_timestep = self.noise_scheduler.num_train_timesteps
self.manual_idx = 0</code></pre></details>
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.initialize_validation_pipeline
initialize_validation_pipeline(training_args: TrainingArgs)
Source
def initialize_validation_pipeline(self, training_args: TrainingArgs):
logger.info("Initializing validation pipeline...")
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
# Use the same flow-matching scheduler as training for consistent validation.
validation_scheduler = SelfForcingFlowMatchScheduler(shift=args_copy.pipeline_config.flow_shift,
sigma_min=0.0,
extra_one_step=True)
validation_scheduler.set_timesteps(num_inference_steps=1000, training=True)
# Warm start validation with current transformer
self.validation_pipeline = MatrixGame2CausalDMDPipeline.from_pretrained(
training_args.model_path,
args=args_copy, # type: ignore
inference_mode=True,
loaded_modules={
"transformer": self.get_module("transformer"),
"vae": self.get_module("vae"),
"scheduler": validation_scheduler,
},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus,
pin_cpu_memory=training_args.pin_cpu_memory,
dit_cpu_offload=True)
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.train_one_step
train_one_step(training_batch)
Source
def train_one_step(self, training_batch): # type: ignore[override]
self.transformer.train()
self.optimizer.zero_grad()
training_batch.total_loss = 0.0
args = cast(TrainingArgs, self.training_args)
# Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
for _ in range(args.gradient_accumulation_steps):
training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
image_embeds = training_batch.image_embeds
image_latents = training_batch.image_latents
keyboard_cond = training_batch.keyboard_cond
mouse_cond = training_batch.mouse_cond
assert traj_latents.shape[0] == 1
# Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
_, S = traj_latents.shape[0], traj_latents.shape[1]
if S < 2:
raise ValueError("Trajectory must contain at least 2 steps")
# Forward to predict next latent by stepping scheduler with predicted noise
noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
traj_latents, traj_timesteps, image_embeds, image_latents, keyboard_cond, mouse_cond)
training_batch.latent_vis_dict.update(latent_vis_dict)
mask = t != 0
# Compute loss
loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
loss = loss / args.gradient_accumulation_steps
with set_forward_context(current_timestep=t, attn_metadata=None, forward_batch=None):
loss.backward()
avg_loss = loss.detach().clone()
training_batch.total_loss += avg_loss.item()
# Clip grad and step optimizers
grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
[p for p in self.transformer.parameters() if p.requires_grad],
args.max_grad_norm if args.max_grad_norm is not None else 0.0)
self.optimizer.step()
self.lr_scheduler.step()
if grad_norm is None:
grad_value = 0.0
else:
try:
if isinstance(grad_norm, torch.Tensor):
grad_value = float(grad_norm.detach().float().item())
else:
grad_value = float(grad_norm)
except Exception:
grad_value = 0.0
training_batch.grad_norm = grad_value
B, S, T, C, H, W = traj_latents.shape
training_batch.raw_latent_shape = (B, C, T, H, W)
return training_batch</code></pre></details>
fastvideo.training.matrixgame2_ode_causal_pipeline.MatrixGame2ODEInitTrainingPipeline.visualize_intermediate_latents
visualize_intermediate_latents(training_batch: TrainingBatch, training_args: TrainingArgs, step: int)
Source
def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
tracker_loss_dict: dict[str, Any] = {}
latents_vis_dict = training_batch.latent_vis_dict
latent_log_keys = ['noisy_input', 'x0', 'pred_video']
for latent_key in latent_log_keys:
assert latent_key in latents_vis_dict and latents_vis_dict[latent_key] is not None
latent = latents_vis_dict[latent_key]
pixel_latent = self.decoding_stage.decode(latent, training_args)
video = pixel_latent.cpu().float()
video = video.permute(0, 2, 1, 3, 4)
video = (video * 255).numpy().astype(np.uint8)
video_artifact = self.tracker.video(video, fps=16, format="mp4") # change to 16 for Wan2.1
if video_artifact is not None:
tracker_loss_dict[latent_key] = video_artifact
# Clean up references
del video, pixel_latent, latent
if self.global_rank == 0 and tracker_loss_dict:
self.tracker.log_artifacts(tracker_loss_dict, step)</code></pre></details>
fastvideo.training.matrixgame2_ode_causal_pipeline.main
main(args) -> None
fastvideo.training.matrixgame2_ode_causal_pipeline.build_parser
build_parser(pipeline_class: str | None = None, pipeline_module: str | None = None) -> FlexibleArgumentParser
Build the training-runner CLI parser.
When pipeline_class/pipeline_module are given they become optional defaults,
so the deprecated per-pipeline __main__ blocks keep working without the flags.