matrixgame2_preprocess_pipeline_ode_trajectory
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() -> Nonefastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.update
update(source_dict: dict[str, Any]) -> Nonefastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.from_pretrained
from_pretrained(model_path: str) -> SamplingParamfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SamplingParam.add_cli_args
add_cli_args(parser: Any) -> AnyAdd CLI arguments for SamplingParam fields
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.getdataset
getdataset(args) -> VideoCaptionMergedDatasetfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn 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: 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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.add_cli_args
add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParserfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.from_cli_args
from_cli_args(args: argparse.Namespace) -> FastVideoArgsfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.from_kwargs
from_kwargs(**kwargs: Any) -> FastVideoArgsfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.FastVideoArgs.check_fastvideo_args
check_fastvideo_args() -> NoneValidate 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) -> boolApply 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) -> 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.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) -> _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.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)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.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.Tensorfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.SelfForcingFlowMatchScheduler.set_shift
set_shift(shift: float) -> Nonefastvideo.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.SchemaReturn 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: PreprocessBatchThe batch containing the data to extract
schema: pa.SchemaPyArrow schema defining the expected fields
strict: boolIf True, raises an exception when required fields are missing or unfilled
Returns
dict[str, Any]Dictionary record matching the schema
Raises
ValueErrorIf 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) -> VerificationResultVerify decoding stage inputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.DecodingStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify decoding stage outputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.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.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.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.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.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) -> ForwardBatchValidate and prepare inputs.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe validated batch information.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify input validation stage inputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.InputValidationStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify 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) -> ForwardBatchPrepare initial latent variables for the diffusion process.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe batch with prepared latent variables.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify latent preparation stage inputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.LatentPreparationStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify 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) -> ForwardBatchfastvideo.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) -> ForwardBatchPrepare timesteps for the diffusion process.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe batch with prepared timesteps.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify timestep preparation stage inputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.TimestepPreparationStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify 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) -> ForwardBatchfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.streaming_reset
streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchfastvideo.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) -> ForwardBatchfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.streaming_clear
streaming_clear() -> Nonefastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline_ode_trajectory.MatrixGame2CausalDenoisingStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultfastvideo.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)</code></pre></details>
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")))</code></pre></details>
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</code></pre></details>
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)</code></pre></details>
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)</code></pre></details>