matrixgame2_causal_dmd_pipeline

Matrix-Game causal DMD pipeline implementation.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.is_output_rank

is_output_rank: bool

Whether this process may materialize user-facing outputs.

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.finalize_device_offload_policy

finalize_device_offload_policy(device_id: int = 0) -> bool

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.FastVideoArgs.disable_offload_on_unified_memory

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

Disable host offload after a worker has selected its device.

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.init_logger

init_logger(name: str) -> _FastvideoLogger

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase

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

Base class for pipelines composed of multiple stages.

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

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.get_hf_download_component_dirs

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

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.get_hf_download_allow_patterns

get_hf_download_allow_patterns() -> list[str] | None

Return Hub patterns for the manifest and selected components.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.set_trainable

set_trainable() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.post_init

post_init() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.initialize_training_pipeline

initialize_training_pipeline(training_args: TrainingArgs)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.initialize_validation_pipeline

initialize_validation_pipeline(training_args: TrainingArgs)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.from_pretrained

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

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.get_module

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.add_module

add_module(module_name: str, module: Any)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.required_config_modules

required_config_modules: list[str]

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

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

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.stages

stages: list[PipelineStage]

List of stages in the pipeline.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.create_pipeline_stages

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Create the inference pipeline stages.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.create_training_stages

create_training_stages(training_args: TrainingArgs)

Create the training pipeline stages.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.initialize_pipeline

initialize_pipeline(fastvideo_args: FastVideoArgs)

Initialize the pipeline.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.load_modules

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

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.add_stage

add_stage(stage_name: str, stage: PipelineStage)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.forward

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

Generate a video or image using the pipeline.

Parameters

batch: ForwardBatch

The batch to generate from.

fastvideo_args: FastVideoArgs

The inference arguments.

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.train

train() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ComposedPipelineBase.close

close() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ForwardBatch

ForwardBatch(data_type: str, generator: torch.Generator | list[torch.Generator] | None = None, image_path: str | None = None, image_embeds: list[torch.Tensor] = list(), pil_image: torch.Tensor | PIL.Image.Image | None = None, last_image: torch.Tensor | PIL.Image.Image | None = None, references: list[Any] | None = None, preprocessed_image: torch.Tensor | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, video_path: str | None = None, video_latent: torch.Tensor | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, pose_latents: torch.Tensor | None = None, face_pixel_values: torch.Tensor | None = None, audio_path: str | None = None, audio_embeds: torch.Tensor | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt_embeds: list[torch.Tensor] = list(), negative_prompt_embeds: list[torch.Tensor] | None = None, prompt_attention_mask: list[torch.Tensor] | None = None, negative_attention_mask: list[torch.Tensor] | None = None, clip_embedding_pos: list[torch.Tensor] | None = None, clip_embedding_neg: list[torch.Tensor] | None = None, max_sequence_length: int | None = None, prompt_template: dict[str, Any] | None = None, do_classifier_free_guidance: bool = False, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, batch_size: int | None = None, num_videos_per_prompt: int = 1, seed: int | None = None, seeds: list[int] | None = None, is_prompt_processed: bool = False, latents: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, lq_latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_pred: torch.Tensor | None = None, image_latent: torch.Tensor | None = None, first_frame_latent: torch.Tensor | None = None, helios_latent_chunks: list[torch.Tensor] | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, num_iterations: int | None = None, use_base_model: bool = False, pose: str | None = None, c2ws_plucker_emb: torch.Tensor | None = None, action_path: str | None = None, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, height_latents: list[int] | int | None = None, width_latents: list[int] | int | None = None, num_frames: list[int] | int = 1, height: list[int] | int | None = None, width: list[int] | int | None = None, height_sr: list[int] | int | None = None, width_sr: list[int] | int | None = None, fps: list[int] | int | None = None, timesteps: torch.Tensor | None = None, timestep: torch.Tensor | float | int | None = None, step_index: int | None = None, boundary_ratio: float | None = None, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, guidance_rescale: float = 0.0, eta: float = 0.0, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = list(), ltx2_stg_blocks_audio: list[int] = list(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: torch.Tensor | None = None, ltx2_conditioning_latent_stage2: torch.Tensor | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, n_tokens: int | None = None, extra_step_kwargs: dict[str, Any] = dict(), modules: dict[str, Any] = dict(), output: torch.Tensor | None = None, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, trajectory_timesteps: list[torch.Tensor] | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_decoded: list[torch.Tensor] | None = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, extra: dict[str, Any] = dict(), save_video: bool = True, return_frames: bool = False, is_cfg_negative: bool = False, VSA_sparsity: float = 0.0, logging_info: PipelineLoggingInfo = PipelineLoggingInfo(), *, return_samples: bool = False)

Complete state passed through the pipeline execution.

This dataclass contains all information needed during the diffusion pipeline execution, allowing methods to update specific components without needing to manage numerous individual parameters.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline

LoRAPipeline(*args, **kwargs)

Pipeline that supports injecting LoRA adapters into the diffusion transformer. TODO: support training.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.add_module

add_module(module_name: str, module: Any) -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.is_target_layer

is_target_layer(module_name: str) -> bool

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.convert_adapter_state_dict

convert_adapter_state_dict(state_dict: dict[str, torch.Tensor], lora_path: str) -> dict[str, torch.Tensor]

Rewrite a foreign-layout adapter into the layout the model's name mapping speaks.

The default leaves the adapter alone. Pipelines override it for model families whose published adapters use another naming or fusion scheme, and raise when the adapter is unusable.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.set_trainable

set_trainable() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.convert_to_lora_layers

convert_to_lora_layers() -> None

Unified method to convert the transformer to a LoRA transformer.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.set_lora_adapter

set_lora_adapter(lora_nickname: str, lora_path: str | None = None, strength: float = 1.0, accumulate: bool = False)

Load a LoRA adapter into the pipeline and merge it into the transformer. Args: lora_nickname: The "nick name" of the adapter when referenced in the pipeline. lora_path: The path to the adapter, either a local path or a Hugging Face repo id. strength: Scale for the low-rank adapter. Hybrid adapters must set this at construction so their dense payload receives the same scale. accumulate: Add this adapter to an already merged pure low-rank adapter.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.merge_lora_weights

merge_lora_weights() -> None

Merge LoRA weights, then repack the packed quantized buffers for the merged state.

NVFP4 layers are only supported while their BF16 weight is retained: a purged BF16 weight cannot be unmerged and remerged.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LoRAPipeline.unmerge_lora_weights

unmerge_lora_weights() -> None

Unmerge LoRA weights, then repack the packed quantized buffers for the unmerged state.

NVFP4 is rejected: its FP4 buffers cannot be rebuilt from a purged BF16 weight.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ConditioningStage

ConditioningStage()

Stage for applying conditioning to the diffusion process.

This stage handles the application of conditioning, such as classifier-free guidance, to the diffusion process.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ConditioningStage.forward

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

Apply conditioning to the diffusion process.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The batch with applied conditioning.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ConditioningStage.verify_input

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

Verify conditioning stage inputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.ConditioningStage.verify_output

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

Verify conditioning stage outputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.DecodingStage

DecodingStage(vae, pipeline = None)

Stage for decoding latent representations into pixel space.

This stage handles the decoding of latent representations into the final output format (e.g., pixel values).

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.DecodingStage.verify_input

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

Verify decoding stage inputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.DecodingStage.verify_output

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

Verify decoding stage outputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.DecodingStage.streaming_decode

streaming_decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs, cache: list[torch.Tensor | None] | None = None, is_first_chunk: bool = False) -> tuple[torch.Tensor, list[torch.Tensor | None]]

Decode latent representations into pixel space using VAE with streaming cache.

Parameters

latents: torch.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.InputValidationStage.verify_input

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

Verify input validation stage inputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.InputValidationStage.verify_output

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

Verify input validation stage outputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.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.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LatentPreparationStage.verify_input

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

Verify latent preparation stage inputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.LatentPreparationStage.verify_output

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

Verify latent preparation stage outputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage

TextEncodingStage(text_encoders, tokenizers)

Stage for encoding text prompts into embeddings for diffusion models.

This stage handles the encoding of text prompts into the embedding space expected by the diffusion model.

Initialize the prompt encoding stage.

Parameters

enable_logging

Whether to enable logging for this stage.

is_secondary

Whether this is a secondary text encoder.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage.forward

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

Encode the prompt into text encoder hidden states.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The batch with encoded prompt embeddings.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage.verify_input

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

Verify text encoding stage inputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage.encode_text

encode_text(text: str | list[str], fastvideo_args: FastVideoArgs, encoder_index: int | list[int] | None = None, return_attention_mask: bool = False, return_type: str = 'list', device: torch.device | str | None = None, dtype: torch.dtype | None = None, max_length: int | Sequence[int | None] | None = None, truncation: bool | None = None, padding: bool | str | None = None)

Encode plain text using selected text encoder(s) and return embeddings.

Parameters

text: str | list[str]

A single string or a list of strings to encode.

fastvideo_args: FastVideoArgs

The inference arguments providing pipeline config, including tokenizer and encoder settings, preprocess and postprocess functions.

encoder_index: int | list[int] | None

Encoder selector by index. Accepts an int or list of ints.

return_attention_mask: bool

If True, also return attention masks for each selected encoder.

return_type: str

"list" (default) returns a list aligned with selection; "dict" returns a dict keyed by encoder index as a string; "stack" stacks along a new first dimension (requires matching shapes).

device: torch.device | str | None

Optional device override for inputs; defaults to local torch device.

dtype: torch.dtype | None

Optional dtype to cast returned embeddings to.

max_length: int | Sequence[int | None] | None

Optional per-call tokenizer override. Either one value applied to every selected encoder (existing behavior), or a sequence with one entry per encoder in encoder_index order -- None at a position falls back to that encoder's own configured default instead of the override.

truncation: bool | None

Optional per-call tokenizer override.

padding: bool | str | None

Optional per-call tokenizer override.

Returns

Depending on return_type and return_attention_mask:

• list: List[Tensor] or (List[Tensor], List[Tensor])

• dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])

• stack: Tensor of shape [num_encoders, ...] or a tuple with stacked attention masks

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage.return_embeds

return_embeds(embeds_list: list[torch.Tensor], attn_masks_list: list[torch.Tensor], return_type: str = 'list', return_attention_mask: bool = False, indices: list[int] | None = None) -> Any

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.TextEncodingStage.verify_output

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

Verify text encoding stage outputs.

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2ImageEncodingStage

MatrixGame2ImageEncodingStage(image_encoder, image_processor)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2ImageEncodingStage.forward

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage.forward

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage.streaming_reset

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage.streaming_step

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage.streaming_clear

streaming_clear() -> None

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDenoisingStage.verify_input

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2ImageVAEEncodingStage

MatrixGame2ImageVAEEncodingStage(vae: ParallelTiledVAE)

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2ImageVAEEncodingStage.forward

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

fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline

MatrixGame2CausalDMDPipeline(*args, **kwargs)

class MatrixGame2CausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
    _required_config_modules = ["vae", "transformer", "scheduler", "image_encoder", "image_processor"]
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
    self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())

    if (self.get_module("text_encoder", None) is not None and self.get_module("tokenizer", None) is not None):
        self.add_stage(stage_name="prompt_encoding_stage",
                       stage=TextEncodingStage(
                           text_encoders=[self.get_module("text_encoder")],
                           tokenizers=[self.get_module("tokenizer")],
                       ))

    if (self.get_module("image_encoder", None) is not None
            and self.get_module("image_processor", None) is not None):
        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="conditioning_stage", stage=ConditioningStage())

    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="image_latent_preparation_stage",
                   stage=MatrixGame2ImageVAEEncodingStage(vae=self.get_module("vae")))

    self.add_stage(stage_name="denoising_stage",
                   stage=MatrixGame2CausalDenoisingStage(transformer=self.get_module("transformer"),
                                                         transformer_2=self.get_module("transformer_2", None),
                                                         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")))

    logger.info("MatrixGame2CausalDMDPipeline initialized with action support")

@torch.no_grad()
def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs):
    if not self.post_init_called:
        self.post_init()

    # 1. Run Pre-processing stages
    stages_to_run = [
        "input_validation_stage", "prompt_encoding_stage", "image_encoding_stage", "conditioning_stage",
        "latent_preparation_stage", "image_latent_preparation_stage"
    ]

    for stage_name in stages_to_run:
        if stage_name in self._stage_name_mapping:
            batch = self._stage_name_mapping[stage_name].forward(batch, fastvideo_args)

    # 2. Reset Denoising Stage
    denoiser = self._stage_name_mapping["denoising_stage"]
    denoiser.streaming_reset(batch, fastvideo_args)

    # 3. Initialize VAE cache
    self._vae_cache = None

def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch:
    denoiser = self._stage_name_mapping["denoising_stage"]
    ctx = denoiser._streaming_ctx
    assert ctx is not None, "streaming_ctx must be set"

    start_idx = ctx.start_index
    batch = denoiser.streaming_step(keyboard_action, mouse_action)
    end_idx = ctx.start_index

    # Decode only the new generated block
    if end_idx > start_idx:
        current_latents = batch.latents[:, :, start_idx:end_idx, :, :]
        args = ctx.fastvideo_args
        decoder = self._stage_name_mapping["decoding_stage"]
        decoded_frames, self._vae_cache = decoder.streaming_decode(current_latents,
                                                                   args,
                                                                   cache=self._vae_cache,
                                                                   is_first_chunk=(start_idx == 0))
        batch.output = decoded_frames
    else:
        batch.output = None

    return batch

def streaming_clear(self) -> None:
    denoiser = self._stage_name_mapping.get("denoising_stage")
    if denoiser is not None and hasattr(denoiser, "streaming_clear"):
        denoiser.streaming_clear()
    self._vae_cache = None</code></pre></details>
fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline.create_pipeline_stages
create_pipeline_stages(fastvideo_args: FastVideoArgs) -> None
Source
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
    self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
if (self.get_module(&quot;text_encoder&quot;, None) is not None and self.get_module(&quot;tokenizer&quot;, None) is not None):
    self.add_stage(stage_name=&quot;prompt_encoding_stage&quot;,
                   stage=TextEncodingStage(
                       text_encoders=&#91;self.get_module(&quot;text_encoder&quot;)&#93;,
                       tokenizers=&#91;self.get_module(&quot;tokenizer&quot;)&#93;,
                   ))

if (self.get_module(&quot;image_encoder&quot;, None) is not None
        and self.get_module(&quot;image_processor&quot;, None) is not None):
    self.add_stage(stage_name=&quot;image_encoding_stage&quot;,
                   stage=MatrixGame2ImageEncodingStage(
                       image_encoder=self.get_module(&quot;image_encoder&quot;),
                       image_processor=self.get_module(&quot;image_processor&quot;),
                   ))

self.add_stage(stage_name=&quot;conditioning_stage&quot;, stage=ConditioningStage())

self.add_stage(stage_name=&quot;latent_preparation_stage&quot;,
               stage=LatentPreparationStage(scheduler=self.get_module(&quot;scheduler&quot;),
                                            transformer=self.get_module(&quot;transformer&quot;, None)))

self.add_stage(stage_name=&quot;image_latent_preparation_stage&quot;,
               stage=MatrixGame2ImageVAEEncodingStage(vae=self.get_module(&quot;vae&quot;)))

self.add_stage(stage_name=&quot;denoising_stage&quot;,
               stage=MatrixGame2CausalDenoisingStage(transformer=self.get_module(&quot;transformer&quot;),
                                                     transformer_2=self.get_module(&quot;transformer_2&quot;, None),
                                                     scheduler=self.get_module(&quot;scheduler&quot;),
                                                     pipeline=self,
                                                     vae=self.get_module(&quot;vae&quot;)))

self.add_stage(stage_name=&quot;decoding_stage&quot;, stage=DecodingStage(vae=self.get_module(&quot;vae&quot;)))

logger.info(&quot;MatrixGame2CausalDMDPipeline initialized with action support&quot;)</code></pre></details>
fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline.streaming_reset
streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs)
Source
@torch.no_grad()
def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs):
    if not self.post_init_called:
        self.post_init()
# 1. Run Pre-processing stages
stages_to_run = &#91;
    &quot;input_validation_stage&quot;, &quot;prompt_encoding_stage&quot;, &quot;image_encoding_stage&quot;, &quot;conditioning_stage&quot;,
    &quot;latent_preparation_stage&quot;, &quot;image_latent_preparation_stage&quot;
&#93;

for stage_name in stages_to_run:
    if stage_name in self._stage_name_mapping:
        batch = self._stage_name_mapping&#91;stage_name&#93;.forward(batch, fastvideo_args)

# 2. Reset Denoising Stage
denoiser = self._stage_name_mapping&#91;&quot;denoising_stage&quot;&#93;
denoiser.streaming_reset(batch, fastvideo_args)

# 3. Initialize VAE cache
self._vae_cache = None</code></pre></details>
fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline.streaming_step
streaming_step(keyboard_action, mouse_action) -> ForwardBatch
Source
def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch:
    denoiser = self._stage_name_mapping["denoising_stage"]
    ctx = denoiser._streaming_ctx
    assert ctx is not None, "streaming_ctx must be set"
start_idx = ctx.start_index
batch = denoiser.streaming_step(keyboard_action, mouse_action)
end_idx = ctx.start_index

# Decode only the new generated block
if end_idx &gt; start_idx:
    current_latents = batch.latents&#91;:, :, start_idx:end_idx, :, :&#93;
    args = ctx.fastvideo_args
    decoder = self._stage_name_mapping&#91;&quot;decoding_stage&quot;&#93;
    decoded_frames, self._vae_cache = decoder.streaming_decode(current_latents,
                                                               args,
                                                               cache=self._vae_cache,
                                                               is_first_chunk=(start_idx == 0))
    batch.output = decoded_frames
else:
    batch.output = None

return batch</code></pre></details>
fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline.streaming_clear
streaming_clear() -> None
Source
def streaming_clear(self) -> None:
    denoiser = self._stage_name_mapping.get("denoising_stage")
    if denoiser is not None and hasattr(denoiser, "streaming_clear"):
        denoiser.streaming_clear()
    self._vae_cache = None
fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGameCausalDMDPipeline
MatrixGameCausalDMDPipeline(*args, **kwargs)
Source
class MatrixGameCausalDMDPipeline(MatrixGame2CausalDMDPipeline):
    pass