ODE Trajectory Data Preprocessing pipeline implementation.

This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline using the modular pipeline architecture.

Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.check_sampling_param`

check_sampling_param() -> None

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.update`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.from_pretrained`

from_pretrained(model_path: str) -> SamplingParam

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.add_cli_args`

add_cli_args(parser: Any) -> Any

Add CLI arguments for SamplingParam fields

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.getdataset`

getdataset(args) -> VideoCaptionMergedDataset

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.get_local_torch_device`

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.from_kwargs`

from_kwargs(**kwargs: Any) -> FastVideoArgs

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.check_fastvideo_args`

check_fastvideo_args() -> None

Validate inference arguments for consistency

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.finalize_device_offload_policy`

finalize_device_offload_policy(device_id: int = 0) -> bool

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.set_timesteps`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.step`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high`

calculate_alpha_beta_high(sigma, sigma_bound)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.add_noise_high`

add_noise_high(original_samples, noise, timestep, boundary_timestep)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.training_target`

training_target(sample, noise, timestep)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.training_weight`

training_weight(timestep)

Input

timestep: the timestep with shape [B*T]

Output: the corresponding weighting [B*T]

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.scale_model_input`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.set_shift`

set_shift(shift: float) -> None

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline`

BasePreprocessPipeline(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 preprocessing pipelines that handles common functionality.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Set up pipeline stages with proper dependency injection.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.forward`

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs, args)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.get_extra_features`

get_extra_features(valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]

Get additional features specific to the pipeline type. Override in subclasses.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.get_pyarrow_schema`

get_pyarrow_schema() -> pa.Schema

Return the PyArrow schema for this pipeline. Must be overridden.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.get_schema_fields`

get_schema_fields() -> list[str]

Get the schema fields for the pipeline type.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.create_record_for_schema`

create_record_for_schema(preprocess_batch: PreprocessBatch, schema: pa.Schema, strict: bool = False) -> dict[str, Any]

Create a record for the Parquet dataset using a generic schema-based approach.

Parameters

preprocess_batch: PreprocessBatch: The batch containing the data to extract
schema: pa.Schema: PyArrow schema defining the expected fields
strict: bool: If True, raises an exception when required fields are missing or unfilled

Returns

dict[str, Any]: Dictionary record matching the schema

Raises

ValueError: If strict=True and required fields are missing or unfilled

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.create_record`

create_record(video_name: str, vae_latent: np.ndarray, text_embedding: np.ndarray, valid_data: dict[str, Any], idx: int, extra_features: dict[str, Any] | None = None) -> dict[str, Any]

Create a record for the Parquet dataset.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.BasePreprocessPipeline.preprocess_video_and_text`

preprocess_video_and_text(fastvideo_args: FastVideoArgs, args)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.DecodingStage.verify_input`

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage inputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.DecodingStage.verify_output`

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage outputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.DecodingStage.decode`

decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor

Decode latent representations into pixel space using VAE.

Parameters

latents: torch.Tensor: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)
fastvideo_args: FastVideoArgs: Configuration 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.Tensor: Decoded video tensor with shape (batch, channels, frames, height, width),
torch.Tensor: normalized to [0, 1] range and moved to CPU as float32

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.Tensor: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)
fastvideo_args: FastVideoArgs: Configuration object.
cache: list[torch.Tensor | None] | None: VAE cache from previous call, or None to initialize a new cache.
is_first_chunk: bool: Whether this is the first chunk.

Returns

tuple[torch.Tensor, list[torch.Tensor | None]]: A tuple of (decoded_frames, updated_cache).

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.DecodingStage.forward`

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

Decode 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: ForwardBatch: The 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: FastVideoArgs: Configuration 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

ForwardBatch: Modified batch with: - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 - trajectory_decoded (if requested): List of decoded frames per timestep

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage`

InputValidationStage()

Stage for validating and preparing inputs for diffusion pipelines.

This stage validates that all required inputs are present and properly formatted before proceeding with the diffusion process.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage.forward`

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

Validate and prepare inputs.

Parameters

batch: ForwardBatch: The current batch information.
fastvideo_args: FastVideoArgs: The inference arguments.

Returns

ForwardBatch: The validated batch information.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage.verify_input`

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify input validation stage inputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage.verify_output`

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify input validation stage outputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage`

LatentPreparationStage(scheduler, transformer, use_btchw_layout: bool = False)

Stage for preparing initial latent variables for the diffusion process.

This stage handles the preparation of the initial latent variables that will be denoised during the diffusion process.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage.forward`

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

Prepare initial latent variables for the diffusion process.

Parameters

batch: ForwardBatch: The current batch information.
fastvideo_args: FastVideoArgs: The inference arguments.

Returns

ForwardBatch: The batch with prepared latent variables.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage.verify_input`

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify latent preparation stage inputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage.verify_output`

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify latent preparation stage outputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2ImageEncodingStage`

MatrixGame2ImageEncodingStage(image_encoder, image_processor)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2ImageEncodingStage.forward`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage`

TimestepPreparationStage(scheduler)

Stage for preparing timesteps for the diffusion process.

This stage handles the preparation of the timestep sequence that will be used during the diffusion process.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage.forward`

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

Prepare timesteps for the diffusion process.

Parameters

batch: ForwardBatch: The current batch information.
fastvideo_args: FastVideoArgs: The inference arguments.

Returns

ForwardBatch: The batch with prepared timesteps.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage.verify_input`

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify timestep preparation stage inputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage.verify_output`

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify timestep preparation stage outputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage`

MatrixGame2CausalDenoisingStage(transformer, scheduler, pipeline = None, transformer_2 = None, vae = None)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.forward`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.streaming_reset`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.streaming_step`

streaming_step(keyboard_action: torch.Tensor | None = None, mouse_action: torch.Tensor | None = None) -> ForwardBatch

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.streaming_clear`

streaming_clear() -> None

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.verify_input`

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.save_decoded_latents_as_video`

save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.shallow_asdict`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory`

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

ODE Trajectory preprocessing pipeline implementation.

Source

class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline):
    """ODE Trajectory preprocessing pipeline implementation."""

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

    preprocess_dataloader: StatefulDataLoader
    preprocess_loader_iter: Iterator[dict[str, Any]]
    pbar: Any
    num_processed_samples: int

    def get_pyarrow_schema(self) -> pa.Schema:
        """Return the PyArrow schema for ODE Trajectory pipeline."""
        return pyarrow_schema_matrixgame2_ode_trajectory

    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
        """Set up pipeline stages with proper dependency injection."""
        assert fastvideo_args.pipeline_config.flow_shift == 5
        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=48, denoising_strength=1.0)

        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
        self.add_stage(stage_name="image_encoding_stage",
                       stage=MatrixGame2ImageEncodingStage(
                           image_encoder=self.get_module("image_encoder"),
                           image_processor=self.get_module("image_processor"),
                       ))
        self.add_stage(stage_name="timestep_preparation_stage",
                       stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
        self.add_stage(stage_name="latent_preparation_stage",
                       stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
                                                    transformer=self.get_module("transformer", None)))
        self.add_stage(stage_name="denoising_stage",
                       stage=MatrixGame2CausalDenoisingStage(
                           transformer=self.get_module("transformer"),
                           scheduler=self.get_module("scheduler"),
                           pipeline=self,
                           vae=self.get_module("vae"),
                       ))
        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))

    def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:

        # TODO(will): move these to cpu at some point
        self.get_module("image_encoder").to(get_local_torch_device())
        self.get_module("vae").to(get_local_torch_device())

        features = {}
        """Get CLIP features from the first frame of each video."""
        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
        _, _, num_frames, height, width = valid_data["pixel_values"].shape
        # latent_height = height // self.get_module(
        #     "vae").spatial_compression_ratio
        # latent_width = width // self.get_module("vae").spatial_compression_ratio

        processed_images = []
        # Frame has values between -1 and 1
        for frame in first_frame:
            frame = (frame + 1) * 127.5
            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
            processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt")
            processed_images.append(processed_img)

        # Get CLIP features
        pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device())
        with torch.no_grad():
            image_inputs = {'pixel_values': pixel_values}
            with set_forward_context(current_timestep=0, attn_metadata=None):
                clip_features = self.get_module("image_encoder")(**image_inputs)
            clip_features = clip_features.last_hidden_state

        features["clip_feature"] = clip_features
        features["pil_image"] = first_frame
        # Get CLIP features from the first frame of each video.
        video_conditions = []
        for frame in first_frame:
            processed_img = frame.to(device="cpu", dtype=torch.float32)
            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
            # (B, H, W, C) -> (B, C, 1, H, W)
            video_condition = torch.cat([
                processed_img,
                processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1], num_frames - 1, height, width)
            ],
                                        dim=2)
            video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32)
            video_conditions.append(video_condition)

        video_conditions = torch.cat(video_conditions, dim=0)

        with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
            encoder_outputs = self.get_module("vae").encode(video_conditions)

        # Use mode() instead of mean
        latent_condition = encoder_outputs.mode()

        # Use latents_mean/latents_std normalization to match
        vae = self.get_module("vae")
        if (hasattr(vae.config, 'latents_mean') and hasattr(vae.config, 'latents_std')):
            latents_mean = torch.tensor(vae.config.latents_mean,
                                        device=latent_condition.device,
                                        dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
            latents_std = torch.tensor(vae.config.latents_std,
                                       device=latent_condition.device,
                                       dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
            latent_condition = (latent_condition - latents_mean) / latents_std
        elif (hasattr(vae, "shift_factor") and vae.shift_factor is not None):
            if isinstance(vae.shift_factor, torch.Tensor):
                latent_condition -= vae.shift_factor.to(latent_condition.device, latent_condition.dtype)
            else:
                latent_condition -= vae.shift_factor

            if isinstance(vae.scaling_factor, torch.Tensor):
                latent_condition = latent_condition * vae.scaling_factor.to(latent_condition.device,
                                                                            latent_condition.dtype)
            else:
                latent_condition = latent_condition * vae.scaling_factor

        # Create mask_cond: ones for first frame, zeros for rest
        # Shape: (B, 16, latent_frames, latent_height, latent_width)
        mask_cond = torch.ones_like(latent_condition)
        mask_cond[:, :, 1:] = 0  # Set all frames except first to 0
        # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond
        # Shape: (B, 20, latent_frames, latent_height, latent_width)
        cond_concat = torch.cat([mask_cond[:, :4], latent_condition], dim=1)
        features["first_frame_latent"] = cond_concat

        if "action_path" in valid_data and valid_data["action_path"]:
            keyboard_cond_list = []
            mouse_cond_list = []
            keyboard_dim = self.get_module("transformer").config.arch_config.action_config["keyboard_dim_in"]
            for action_path in valid_data["action_path"]:
                if action_path:
                    action_data = np.load(action_path, allow_pickle=True)
                    if isinstance(action_data, np.ndarray) and action_data.dtype == np.dtype('O'):
                        action_dict = action_data.item()
                        if "keyboard" in action_dict:
                            keyboard_cond_list.append(action_dict["keyboard"][:, :keyboard_dim].astype(np.float32))
                        if "mouse" in action_dict:
                            mouse_cond_list.append(action_dict["mouse"])
                    else:
                        keyboard_cond_list.append(action_data[:, :keyboard_dim].astype(np.float32))
            if keyboard_cond_list:
                features["keyboard_cond"] = keyboard_cond_list
            if mouse_cond_list:
                features["mouse_cond"] = mouse_cond_list
        return features

    def preprocess_action_and_trajectory(self, fastvideo_args: FastVideoArgs, args):
        """Preprocess data and generate trajectory information."""

        for batch_idx, data in enumerate(self.pbar):
            if data is None:
                continue

            with torch.inference_mode():
                # Filter out invalid samples (those with all zeros)
                valid_indices = []
                for i, pixel_values in enumerate(data["pixel_values"]):
                    if not torch.all(pixel_values == 0):  # Check if all values are zero
                        valid_indices.append(i)
                self.num_processed_samples += len(valid_indices)

                if not valid_indices:
                    continue

                # Create new batch with only valid samples
                valid_data = {
                    "pixel_values": torch.stack([data["pixel_values"][i] for i in valid_indices]),
                    "path": [data["path"][i] for i in valid_indices],
                }

                if "fps" in data:
                    valid_data["fps"] = [data["fps"][i] for i in valid_indices]
                if "duration" in data:
                    valid_data["duration"] = [data["duration"][i] for i in valid_indices]
                if "action_path" in data:
                    valid_data["action_path"] = [data["action_path"][i] for i in valid_indices]

                pixel_values = valid_data["pixel_values"]
                if pixel_values.shape[2] == 1 and args.num_frames is not None:
                    pixel_values = pixel_values.repeat(1, 1, args.num_frames, 1, 1)
                    valid_data["pixel_values"] = pixel_values

                # Get extra features if needed
                extra_features = self.get_extra_features(valid_data, fastvideo_args)

                clip_features = extra_features['clip_feature']
                image_latents = extra_features['first_frame_latent']
                image_latents = image_latents[:, :, :args.num_latent_t]
                pil_image = extra_features['pil_image']
                keyboard_cond = extra_features.get('keyboard_cond')
                mouse_cond = extra_features.get('mouse_cond')

                sampling_params = SamplingParam.from_pretrained(args.model_path)

                trajectory_latents = []
                trajectory_timesteps = []
                trajectory_decoded = []

                device = get_local_torch_device()
                for i in range(len(valid_indices)):
                    # Collect the trajectory data
                    batch = ForwardBatch(**shallow_asdict(sampling_params), )
                    batch.image_embeds = [clip_features[i].unsqueeze(0)]
                    batch.image_latent = image_latents[i].unsqueeze(0)
                    batch.keyboard_cond = (torch.from_numpy(keyboard_cond[i]).unsqueeze(0).to(device)
                                           if keyboard_cond is not None else None)
                    batch.mouse_cond = (torch.from_numpy(mouse_cond[i]).unsqueeze(0).to(device)
                                        if mouse_cond is not None else None)
                    batch.num_inference_steps = 48
                    batch.return_trajectory_latents = True
                    # Enabling this will save the decoded trajectory videos.
                    # Used for debugging.
                    batch.return_trajectory_decoded = False
                    batch.height = args.max_height
                    batch.width = args.max_width
                    batch.fps = args.train_fps
                    batch.num_frames = valid_data["pixel_values"].shape[2]
                    batch.guidance_scale = 6.0
                    batch.do_classifier_free_guidance = False
                    batch.prompt = ""

                    result_batch = self.input_validation_stage(batch, fastvideo_args)
                    result_batch = self.timestep_preparation_stage(batch, fastvideo_args)
                    result_batch.timesteps = result_batch.timesteps.to(device)
                    result_batch = self.latent_preparation_stage(result_batch, fastvideo_args)
                    result_batch = self.denoising_stage(result_batch, fastvideo_args)
                    result_batch = self.decoding_stage(result_batch, fastvideo_args)

                    trajectory_latents.append(result_batch.trajectory_latents.cpu())
                    trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
                    trajectory_decoded.append(result_batch.trajectory_decoded)

                # Prepare extra features
                extra_features = {
                    "trajectory_latents": trajectory_latents,
                    "trajectory_timesteps": trajectory_timesteps
                }

                if batch.return_trajectory_decoded:
                    for i, decoded_frames in enumerate(trajectory_decoded):
                        for j, decoded_frame in enumerate(decoded_frames):
                            save_decoded_latents_as_video(decoded_frame,
                                                          f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
                                                          args.train_fps)

                # Prepare batch data for Parquet dataset
                batch_data: list[dict[str, Any]] = []

                # Add progress bar for saving outputs
                save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False)

                for idx, video_path in save_pbar:
                    video_name = os.path.basename(video_path).split(".")[0]

                    clip_feature_np = clip_features[idx].cpu().numpy()
                    first_frame_latent_np = image_latents[idx].cpu().numpy()
                    pil_image_np = pil_image[idx].cpu().numpy()
                    keyboard_cond_np = keyboard_cond[idx] if keyboard_cond is not None else None
                    mouse_cond_np = mouse_cond[idx] if mouse_cond is not None else None

                    # Get trajectory features for this sample
                    traj_latents = extra_features["trajectory_latents"][idx]
                    traj_timesteps = extra_features["trajectory_timesteps"][idx]
                    if isinstance(traj_latents, torch.Tensor):
                        traj_latents = traj_latents.cpu().float().numpy()
                    if isinstance(traj_timesteps, torch.Tensor):
                        traj_timesteps = traj_timesteps.cpu().float().numpy()

                    # Create record for Parquet dataset
                    record: dict[str, Any] = matrixgame2_ode_record_creator(video_name=video_name,
                                                                            clip_feature=clip_feature_np,
                                                                            first_frame_latent=first_frame_latent_np,
                                                                            trajectory_latents=traj_latents,
                                                                            trajectory_timesteps=traj_timesteps,
                                                                            pil_image=pil_image_np,
                                                                            keyboard_cond=keyboard_cond_np,
                                                                            mouse_cond=mouse_cond_np,
                                                                            caption="")
                    batch_data.append(record)

                if batch_data:
                    write_pbar = tqdm(total=1, desc="Writing to Parquet dataset", unit="batch")
                    table = records_to_table(batch_data, self.get_pyarrow_schema())
                    write_pbar.update(1)
                    write_pbar.close()

                    if not hasattr(self, 'dataset_writer'):
                        self.dataset_writer = ParquetDatasetWriter(
                            out_dir=self.combined_parquet_dir,
                            samples_per_file=args.samples_per_file,
                        )
                    self.dataset_writer.append_table(table)

                    logger.info("Collected batch with %s samples", len(table))

                if self.num_processed_samples >= args.flush_frequency:
                    written = self.dataset_writer.flush()
                    logger.info("Flushed %s samples to parquet", written)
                    self.num_processed_samples = 0

        # Final flush for any remaining samples
        if hasattr(self, 'dataset_writer'):
            written = self.dataset_writer.flush(write_remainder=True)
            if written:
                logger.info("Final flush wrote %s samples", written)

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
        if not self.post_init_called:
            self.post_init()

        self.local_rank = int(os.getenv("RANK", 0))
        os.makedirs(args.output_dir, exist_ok=True)
        # Create directory for combined data
        self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset")
        os.makedirs(self.combined_parquet_dir, exist_ok=True)

        # Loading dataset
        train_dataset = getdataset(args)

        self.preprocess_dataloader = DataLoader(
            train_dataset,
            batch_size=args.preprocess_video_batch_size,
            num_workers=args.dataloader_num_workers,
        )

        self.preprocess_loader_iter = iter(self.preprocess_dataloader)

        self.num_processed_samples = 0
        # Add progress bar for video preprocessing
        self.pbar = tqdm(self.preprocess_loader_iter,
                         desc="Processing videos",
                         unit="batch",
                         disable=self.local_rank != 0)

        # Initialize class variables for data sharing
        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
        self.latent_data: dict[str, Any] = {}  # Store latent tensors
        self.preprocess_action_and_trajectory(fastvideo_args, args)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory.get_pyarrow_schema`

get_pyarrow_schema() -> pa.Schema

Return the PyArrow schema for ODE Trajectory pipeline.

Source

def get_pyarrow_schema(self) -> pa.Schema:
    """Return the PyArrow schema for ODE Trajectory pipeline."""
    return pyarrow_schema_matrixgame2_ode_trajectory

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Set up pipeline stages with proper dependency injection.

Source

def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
    """Set up pipeline stages with proper dependency injection."""
    assert fastvideo_args.pipeline_config.flow_shift == 5
    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=48, denoising_strength=1.0)

    self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
    self.add_stage(stage_name="image_encoding_stage",
                   stage=MatrixGame2ImageEncodingStage(
                       image_encoder=self.get_module("image_encoder"),
                       image_processor=self.get_module("image_processor"),
                   ))
    self.add_stage(stage_name="timestep_preparation_stage",
                   stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
    self.add_stage(stage_name="latent_preparation_stage",
                   stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
                                                transformer=self.get_module("transformer", None)))
    self.add_stage(stage_name="denoising_stage",
                   stage=MatrixGame2CausalDenoisingStage(
                       transformer=self.get_module("transformer"),
                       scheduler=self.get_module("scheduler"),
                       pipeline=self,
                       vae=self.get_module("vae"),
                   ))
    self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory.get_extra_features`

get_extra_features(valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]

Source

def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:

    # TODO(will): move these to cpu at some point
    self.get_module("image_encoder").to(get_local_torch_device())
    self.get_module("vae").to(get_local_torch_device())

    features = {}
    """Get CLIP features from the first frame of each video."""
    first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
    _, _, num_frames, height, width = valid_data["pixel_values"].shape
    # latent_height = height // self.get_module(
    #     "vae").spatial_compression_ratio
    # latent_width = width // self.get_module("vae").spatial_compression_ratio

    processed_images = []
    # Frame has values between -1 and 1
    for frame in first_frame:
        frame = (frame + 1) * 127.5
        frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
        processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt")
        processed_images.append(processed_img)

    # Get CLIP features
    pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device())
    with torch.no_grad():
        image_inputs = {'pixel_values': pixel_values}
        with set_forward_context(current_timestep=0, attn_metadata=None):
            clip_features = self.get_module("image_encoder")(**image_inputs)
        clip_features = clip_features.last_hidden_state

    features["clip_feature"] = clip_features
    features["pil_image"] = first_frame
    # Get CLIP features from the first frame of each video.
    video_conditions = []
    for frame in first_frame:
        processed_img = frame.to(device="cpu", dtype=torch.float32)
        processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
        # (B, H, W, C) -> (B, C, 1, H, W)
        video_condition = torch.cat([
            processed_img,
            processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1], num_frames - 1, height, width)
        ],
                                    dim=2)
        video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32)
        video_conditions.append(video_condition)

    video_conditions = torch.cat(video_conditions, dim=0)

    with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
        encoder_outputs = self.get_module("vae").encode(video_conditions)

    # Use mode() instead of mean
    latent_condition = encoder_outputs.mode()

    # Use latents_mean/latents_std normalization to match
    vae = self.get_module("vae")
    if (hasattr(vae.config, 'latents_mean') and hasattr(vae.config, 'latents_std')):
        latents_mean = torch.tensor(vae.config.latents_mean,
                                    device=latent_condition.device,
                                    dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
        latents_std = torch.tensor(vae.config.latents_std,
                                   device=latent_condition.device,
                                   dtype=latent_condition.dtype).view(1, -1, 1, 1, 1)
        latent_condition = (latent_condition - latents_mean) / latents_std
    elif (hasattr(vae, "shift_factor") and vae.shift_factor is not None):
        if isinstance(vae.shift_factor, torch.Tensor):
            latent_condition -= vae.shift_factor.to(latent_condition.device, latent_condition.dtype)
        else:
            latent_condition -= vae.shift_factor

        if isinstance(vae.scaling_factor, torch.Tensor):
            latent_condition = latent_condition * vae.scaling_factor.to(latent_condition.device,
                                                                        latent_condition.dtype)
        else:
            latent_condition = latent_condition * vae.scaling_factor

    # Create mask_cond: ones for first frame, zeros for rest
    # Shape: (B, 16, latent_frames, latent_height, latent_width)
    mask_cond = torch.ones_like(latent_condition)
    mask_cond[:, :, 1:] = 0  # Set all frames except first to 0
    # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond
    # Shape: (B, 20, latent_frames, latent_height, latent_width)
    cond_concat = torch.cat([mask_cond[:, :4], latent_condition], dim=1)
    features["first_frame_latent"] = cond_concat

    if "action_path" in valid_data and valid_data["action_path"]:
        keyboard_cond_list = []
        mouse_cond_list = []
        keyboard_dim = self.get_module("transformer").config.arch_config.action_config["keyboard_dim_in"]
        for action_path in valid_data["action_path"]:
            if action_path:
                action_data = np.load(action_path, allow_pickle=True)
                if isinstance(action_data, np.ndarray) and action_data.dtype == np.dtype('O'):
                    action_dict = action_data.item()
                    if "keyboard" in action_dict:
                        keyboard_cond_list.append(action_dict["keyboard"][:, :keyboard_dim].astype(np.float32))
                    if "mouse" in action_dict:
                        mouse_cond_list.append(action_dict["mouse"])
                else:
                    keyboard_cond_list.append(action_data[:, :keyboard_dim].astype(np.float32))
        if keyboard_cond_list:
            features["keyboard_cond"] = keyboard_cond_list
        if mouse_cond_list:
            features["mouse_cond"] = mouse_cond_list
    return features

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory.preprocess_action_and_trajectory`

preprocess_action_and_trajectory(fastvideo_args: FastVideoArgs, args)

Preprocess data and generate trajectory information.

Source

def preprocess_action_and_trajectory(self, fastvideo_args: FastVideoArgs, args):
    """Preprocess data and generate trajectory information."""

    for batch_idx, data in enumerate(self.pbar):
        if data is None:
            continue

        with torch.inference_mode():
            # Filter out invalid samples (those with all zeros)
            valid_indices = []
            for i, pixel_values in enumerate(data["pixel_values"]):
                if not torch.all(pixel_values == 0):  # Check if all values are zero
                    valid_indices.append(i)
            self.num_processed_samples += len(valid_indices)

            if not valid_indices:
                continue

            # Create new batch with only valid samples
            valid_data = {
                "pixel_values": torch.stack([data["pixel_values"][i] for i in valid_indices]),
                "path": [data["path"][i] for i in valid_indices],
            }

            if "fps" in data:
                valid_data["fps"] = [data["fps"][i] for i in valid_indices]
            if "duration" in data:
                valid_data["duration"] = [data["duration"][i] for i in valid_indices]
            if "action_path" in data:
                valid_data["action_path"] = [data["action_path"][i] for i in valid_indices]

            pixel_values = valid_data["pixel_values"]
            if pixel_values.shape[2] == 1 and args.num_frames is not None:
                pixel_values = pixel_values.repeat(1, 1, args.num_frames, 1, 1)
                valid_data["pixel_values"] = pixel_values

            # Get extra features if needed
            extra_features = self.get_extra_features(valid_data, fastvideo_args)

            clip_features = extra_features['clip_feature']
            image_latents = extra_features['first_frame_latent']
            image_latents = image_latents[:, :, :args.num_latent_t]
            pil_image = extra_features['pil_image']
            keyboard_cond = extra_features.get('keyboard_cond')
            mouse_cond = extra_features.get('mouse_cond')

            sampling_params = SamplingParam.from_pretrained(args.model_path)

            trajectory_latents = []
            trajectory_timesteps = []
            trajectory_decoded = []

            device = get_local_torch_device()
            for i in range(len(valid_indices)):
                # Collect the trajectory data
                batch = ForwardBatch(**shallow_asdict(sampling_params), )
                batch.image_embeds = [clip_features[i].unsqueeze(0)]
                batch.image_latent = image_latents[i].unsqueeze(0)
                batch.keyboard_cond = (torch.from_numpy(keyboard_cond[i]).unsqueeze(0).to(device)
                                       if keyboard_cond is not None else None)
                batch.mouse_cond = (torch.from_numpy(mouse_cond[i]).unsqueeze(0).to(device)
                                    if mouse_cond is not None else None)
                batch.num_inference_steps = 48
                batch.return_trajectory_latents = True
                # Enabling this will save the decoded trajectory videos.
                # Used for debugging.
                batch.return_trajectory_decoded = False
                batch.height = args.max_height
                batch.width = args.max_width
                batch.fps = args.train_fps
                batch.num_frames = valid_data["pixel_values"].shape[2]
                batch.guidance_scale = 6.0
                batch.do_classifier_free_guidance = False
                batch.prompt = ""

                result_batch = self.input_validation_stage(batch, fastvideo_args)
                result_batch = self.timestep_preparation_stage(batch, fastvideo_args)
                result_batch.timesteps = result_batch.timesteps.to(device)
                result_batch = self.latent_preparation_stage(result_batch, fastvideo_args)
                result_batch = self.denoising_stage(result_batch, fastvideo_args)
                result_batch = self.decoding_stage(result_batch, fastvideo_args)

                trajectory_latents.append(result_batch.trajectory_latents.cpu())
                trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
                trajectory_decoded.append(result_batch.trajectory_decoded)

            # Prepare extra features
            extra_features = {
                "trajectory_latents": trajectory_latents,
                "trajectory_timesteps": trajectory_timesteps
            }

            if batch.return_trajectory_decoded:
                for i, decoded_frames in enumerate(trajectory_decoded):
                    for j, decoded_frame in enumerate(decoded_frames):
                        save_decoded_latents_as_video(decoded_frame,
                                                      f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
                                                      args.train_fps)

            # Prepare batch data for Parquet dataset
            batch_data: list[dict[str, Any]] = []

            # Add progress bar for saving outputs
            save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False)

            for idx, video_path in save_pbar:
                video_name = os.path.basename(video_path).split(".")[0]

                clip_feature_np = clip_features[idx].cpu().numpy()
                first_frame_latent_np = image_latents[idx].cpu().numpy()
                pil_image_np = pil_image[idx].cpu().numpy()
                keyboard_cond_np = keyboard_cond[idx] if keyboard_cond is not None else None
                mouse_cond_np = mouse_cond[idx] if mouse_cond is not None else None

                # Get trajectory features for this sample
                traj_latents = extra_features["trajectory_latents"][idx]
                traj_timesteps = extra_features["trajectory_timesteps"][idx]
                if isinstance(traj_latents, torch.Tensor):
                    traj_latents = traj_latents.cpu().float().numpy()
                if isinstance(traj_timesteps, torch.Tensor):
                    traj_timesteps = traj_timesteps.cpu().float().numpy()

                # Create record for Parquet dataset
                record: dict[str, Any] = matrixgame2_ode_record_creator(video_name=video_name,
                                                                        clip_feature=clip_feature_np,
                                                                        first_frame_latent=first_frame_latent_np,
                                                                        trajectory_latents=traj_latents,
                                                                        trajectory_timesteps=traj_timesteps,
                                                                        pil_image=pil_image_np,
                                                                        keyboard_cond=keyboard_cond_np,
                                                                        mouse_cond=mouse_cond_np,
                                                                        caption="")
                batch_data.append(record)

            if batch_data:
                write_pbar = tqdm(total=1, desc="Writing to Parquet dataset", unit="batch")
                table = records_to_table(batch_data, self.get_pyarrow_schema())
                write_pbar.update(1)
                write_pbar.close()

                if not hasattr(self, 'dataset_writer'):
                    self.dataset_writer = ParquetDatasetWriter(
                        out_dir=self.combined_parquet_dir,
                        samples_per_file=args.samples_per_file,
                    )
                self.dataset_writer.append_table(table)

                logger.info("Collected batch with %s samples", len(table))

            if self.num_processed_samples >= args.flush_frequency:
                written = self.dataset_writer.flush()
                logger.info("Flushed %s samples to parquet", written)
                self.num_processed_samples = 0

    # Final flush for any remaining samples
    if hasattr(self, 'dataset_writer'):
        written = self.dataset_writer.flush(write_remainder=True)
        if written:
            logger.info("Final flush wrote %s samples", written)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.PreprocessPipeline_MatrixGame2_ODE_Trajectory.forward`

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs, args)

Source

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
    if not self.post_init_called:
        self.post_init()

    self.local_rank = int(os.getenv("RANK", 0))
    os.makedirs(args.output_dir, exist_ok=True)
    # Create directory for combined data
    self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset")
    os.makedirs(self.combined_parquet_dir, exist_ok=True)

    # Loading dataset
    train_dataset = getdataset(args)

    self.preprocess_dataloader = DataLoader(
        train_dataset,
        batch_size=args.preprocess_video_batch_size,
        num_workers=args.dataloader_num_workers,
    )

    self.preprocess_loader_iter = iter(self.preprocess_dataloader)

    self.num_processed_samples = 0
    # Add progress bar for video preprocessing
    self.pbar = tqdm(self.preprocess_loader_iter,
                     desc="Processing videos",
                     unit="batch",
                     disable=self.local_rank != 0)

    # Initialize class variables for data sharing
    self.video_data: dict[str, Any] = {}  # Store video metadata and paths
    self.latent_data: dict[str, Any] = {}  # Store latent tensors
    self.preprocess_action_and_trajectory(fastvideo_args, args)