--- title: matrixgame2_denoising ---

fastvideo.pipelines.stages.matrixgame2_denoising.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.pipelines.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.pipelines.stages.matrixgame2_denoising.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.pipelines.stages.matrixgame2_denoising.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.pipelines.stages.matrixgame2_denoising.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.pipelines.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.pred_noise_to_pred_video

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

Convert predicted noise to clean latent.

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

Returns

torch.Tensor

the predicted video with shape [B, C, H, W]

fastvideo.pipelines.stages.matrixgame2_denoising.pred_noise_to_x_bound

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

Convert predicted noise to clean latent.

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

Returns

torch.Tensor

the predicted video with shape [B, C, H, W]

fastvideo.pipelines.stages.matrixgame2_denoising.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.stages.matrixgame2_denoising.DenoisingStage

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

Stage for running the denoising loop in diffusion pipelines.

This stage handles the iterative denoising process that transforms the initial noise into the final output.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.forward

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

Run the denoising loop.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The batch with denoised latents.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.prepare_denoising

prepare_denoising(batch, fastvideo_args, target_dtype) -> DenoisingState

Build request-local state. Family stages specialize latent conditioning.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.prepare_model_kwargs

prepare_model_kwargs(batch, fastvideo_args) -> tuple[dict[str, Any], dict[str, Any]]

Step-invariant transformer kwargs as (conditional, unconditional).

Family stages add named conditioning here (Wan-S2V's audio, reference and motion latents) without the shared loop learning about it.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.activate_transformer

activate_transformer(model, inactive_model, fastvideo_args) -> None

Keep CPU/layerwise/FSDP offload decisions independent of the sampling recipe.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.select_model

select_model(timestep, batch, fastvideo_args, state) -> tuple[Any, float]

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.prepare_model_input

prepare_model_input(latents, batch, target_dtype, state) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.prepare_timestep

prepare_timestep(timestep, latent_model_input, batch, state) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.finish_step

finish_step(latents, state) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.prepare_extra_func_kwargs

prepare_extra_func_kwargs(func, kwargs) -> dict[str, Any]

Prepare extra kwargs for the scheduler step / denoise step.

Parameters

func

The function to prepare kwargs for.

kwargs

The kwargs to prepare.

Returns

dict[str, Any]

The prepared kwargs.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.progress_bar

progress_bar(iterable: Iterable | None = None, total: int | None = None) -> tqdm

Create a progress bar for the denoising process.

Parameters

iterable: Iterable | None

The iterable to iterate over.

total: int | None

The total number of items.

Returns

tqdm

A tqdm progress bar.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.rescale_noise_cfg

rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale = 0.0) -> torch.Tensor

Rescale noise prediction according to guidance_rescale.

Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed" (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4.

Parameters

noise_cfg

The noise prediction with guidance.

noise_pred_text

The text-conditioned noise prediction.

guidance_rescale

The guidance rescale factor.

Returns

torch.Tensor

The rescaled noise prediction.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.verify_input

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

Verify denoising stage inputs.

fastvideo.pipelines.stages.matrixgame2_denoising.DenoisingStage.verify_output

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

Verify denoising stage outputs.

fastvideo.pipelines.stages.matrixgame2_denoising.V

V()

Common validators for pipeline stages.

fastvideo.pipelines.stages.matrixgame2_denoising.V.not_none

not_none(value: Any) -> bool

Check if value is not None.

fastvideo.pipelines.stages.matrixgame2_denoising.V.positive_int

positive_int(value: Any) -> bool

Check if value is a positive integer.

fastvideo.pipelines.stages.matrixgame2_denoising.V.non_negative_int

non_negative_int(value: Any) -> bool

Check if value is a non-negative integer (allows 0).

fastvideo.pipelines.stages.matrixgame2_denoising.V.positive_float

positive_float(value: Any) -> bool

Check if value is a positive float.

fastvideo.pipelines.stages.matrixgame2_denoising.V.non_negative_float

non_negative_float(value: Any) -> bool

Check if value is a non-negative float.

fastvideo.pipelines.stages.matrixgame2_denoising.V.divisible_by

divisible_by(value: Any, divisor: int) -> bool

Check if value is divisible by divisor.

fastvideo.pipelines.stages.matrixgame2_denoising.V.is_tensor

is_tensor(value: Any) -> bool

Check if value is a torch tensor and doesn't contain NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.tensor_with_dims

tensor_with_dims(value: Any, dims: int) -> bool

Check if value is a tensor with specific dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.tensor_min_dims

tensor_min_dims(value: Any, min_dims: int) -> bool

Check if value is a tensor with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.tensor_shape_matches

tensor_shape_matches(value: Any, expected_shape: tuple) -> bool

Check if tensor shape matches expected shape (None for any size) and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_not_empty

list_not_empty(value: Any) -> bool

Check if value is a non-empty list.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_length

list_length(value: Any, length: int) -> bool

Check if list has specific length.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_min_length

list_min_length(value: Any, min_length: int) -> bool

Check if list has at least min_length items.

fastvideo.pipelines.stages.matrixgame2_denoising.V.string_not_empty

string_not_empty(value: Any) -> bool

Check if value is a non-empty string.

fastvideo.pipelines.stages.matrixgame2_denoising.V.string_or_list_strings

string_or_list_strings(value: Any) -> bool

Check if value is a string or list of strings.

fastvideo.pipelines.stages.matrixgame2_denoising.V.bool_value

bool_value(value: Any) -> bool

Check if value is a boolean.

fastvideo.pipelines.stages.matrixgame2_denoising.V.generator_or_list_generators

generator_or_list_generators(value: Any) -> bool

Check if value is a Generator or list of Generators.

fastvideo.pipelines.stages.matrixgame2_denoising.V.is_list

is_list(value: Any) -> bool

Check if value is a list (can be empty).

fastvideo.pipelines.stages.matrixgame2_denoising.V.is_tuple

is_tuple(value: Any) -> bool

Check if value is a tuple.

fastvideo.pipelines.stages.matrixgame2_denoising.V.none_or_tensor

none_or_tensor(value: Any) -> bool

Check if value is None or a tensor without NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_of_tensors_with_dims

list_of_tensors_with_dims(value: Any, dims: int) -> bool

Check if value is a non-empty list where all items are tensors with specific dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_of_tensors

list_of_tensors(value: Any) -> bool

Check if value is a non-empty list where all items are tensors without NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_of_tensors_with_min_dims

list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool

Check if value is a non-empty list where all items are tensors with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.none_or_tensor_with_dims

none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is None or a tensor with specific dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.none_or_list

none_or_list(value: Any) -> bool

Check if value is None or a list.

fastvideo.pipelines.stages.matrixgame2_denoising.V.none_or_positive_int

none_or_positive_int(value: Any) -> bool

Check if value is None or a positive integer.

fastvideo.pipelines.stages.matrixgame2_denoising.V.with_dims

with_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if tensor has specific dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.min_list_length

min_list_length(min_length: int) -> Callable[[Any], bool]

Return a validator that checks if a list has at least min_length items.

fastvideo.pipelines.stages.matrixgame2_denoising.V.min_dims

min_dims(min_dims: int) -> Callable[[Any], bool]

Return a validator that checks if tensor has at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.divisible

divisible(divisor: int) -> Callable[[Any], bool]

Return a validator that checks if value is divisible by divisor.

fastvideo.pipelines.stages.matrixgame2_denoising.V.positive_int_divisible

positive_int_divisible(divisor: int) -> Callable[[Any], bool]

Return a validator that checks if value is a positive integer divisible by divisor.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_of_tensors_dims

list_of_tensors_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is a list of tensors with specific dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.V.list_of_tensors_min_dims

list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is a list of tensors with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult

VerificationResult()

Wrapper class for stage verification results.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.add_check

add_check(field_name: str, value: Any, validators: Callable[[Any], bool] | list[Callable[[Any], bool]]) -> VerificationResult

Add a validation check for a field.

Parameters

field_name: str

Name of the field being checked

value: Any

The actual value to validate

validators: Callable[[Any], bool] | list[Callable[[Any], bool]]

Single validation function or list of validation functions. Each function will be called with the value as its first argument.

Returns

VerificationResult

Self for method chaining

Examples

Single validator

result.add_check("tensor", my_tensor, V.is_tensor)

Multiple validators (all must pass)

result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])

Using partial functions for parameters

result.add_check("height", batch.height, [V.not_none, V.divisible(8)])

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.is_valid

is_valid() -> bool

Check if all validations passed.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.get_failed_fields

get_failed_fields() -> list[str]

Get list of fields that failed validation.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.get_detailed_failures

get_detailed_failures() -> dict[str, list[ValidationFailure]]

Get detailed failure information for each failed field.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.get_failure_summary

get_failure_summary() -> str

Get a comprehensive summary of all validation failures.

fastvideo.pipelines.stages.matrixgame2_denoising.VerificationResult.to_dict

to_dict() -> dict

Convert to dictionary for backward compatibility.

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend

VideoSparseAttentionBackend()

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend.get_supported_head_sizes

get_supported_head_sizes() -> list[int]

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend.get_name

get_name() -> str

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend.get_impl_cls

get_impl_cls() -> type[VideoSparseAttentionImpl]

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend.get_metadata_cls

get_metadata_cls() -> type[VideoSparseAttentionMetadata]

fastvideo.pipelines.stages.matrixgame2_denoising.VideoSparseAttentionBackend.get_builder_cls

get_builder_cls() -> type[VideoSparseAttentionMetadataBuilder]

fastvideo.pipelines.stages.matrixgame2_denoising.BlockProcessingContext

BlockProcessingContext(batch: ForwardBatch, block_idx: int, start_index: int, kv_cache1: list[dict[Any, Any]], kv_cache2: list[dict[Any, Any]] | None, kv_cache_mouse: list[dict[Any, Any]] | None, kv_cache_keyboard: list[dict[Any, Any]] | None, crossattn_cache: list[dict[Any, Any]], timesteps: torch.Tensor, block_sizes: list[int], noise_pool: list[torch.Tensor] | None, fastvideo_args: FastVideoArgs, target_dtype: torch.dtype, autocast_enabled: bool, boundary_timestep: float | None, high_noise_timesteps: torch.Tensor | None, context_noise: float, image_kwargs: dict[str, Any], pos_cond_kwargs: dict[str, Any])

Dataclass contains for block processing.

Source
@dataclass
class BlockProcessingContext:
    """Dataclass contains for block processing."""

    batch: ForwardBatch

    block_idx: int
    start_index: int

    kv_cache1: list[dict[Any, Any]]
    kv_cache2: list[dict[Any, Any]] | None
    kv_cache_mouse: list[dict[Any, Any]] | None
    kv_cache_keyboard: list[dict[Any, Any]] | None
    crossattn_cache: list[dict[Any, Any]]

    timesteps: torch.Tensor
    block_sizes: list[int]
    noise_pool: list[torch.Tensor] | None

    fastvideo_args: FastVideoArgs
    target_dtype: torch.dtype
    autocast_enabled: bool
    boundary_timestep: float | None
    high_noise_timesteps: torch.Tensor | None
    context_noise: float

    image_kwargs: dict[str, Any]
    pos_cond_kwargs: dict[str, Any]

    def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]:
        if self.boundary_timestep is not None:
            if timestep_val >= self.boundary_timestep:
                return self.kv_cache1
            else:
                assert self.kv_cache2 is not None, "kv_cache2 is not initialized"
                return self.kv_cache2
        return self.kv_cache1

fastvideo.pipelines.stages.matrixgame2_denoising.BlockProcessingContext.get_kv_cache

get_kv_cache(timestep_val: float) -> list[dict[Any, Any]]
Source
def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]:
    if self.boundary_timestep is not None:
        if timestep_val >= self.boundary_timestep:
            return self.kv_cache1
        else:
            assert self.kv_cache2 is not None, "kv_cache2 is not initialized"
            return self.kv_cache2
    return self.kv_cache1

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage

MatrixGame2CausalDenoisingStage(transformer, scheduler, pipeline = None, transformer_2 = None, vae = None)
Source
class MatrixGame2CausalDenoisingStage(DenoisingStage):

    def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
        super().__init__(transformer, scheduler, pipeline, transformer_2, vae)
        self.transformer = transformer
        self.transformer_2 = transformer_2
        self.vae = vae
        self.num_transformer_blocks = len(self.transformer.blocks)

        if hasattr(self.transformer, 'config') and hasattr(self.transformer.config, 'arch_config'):
            self.num_frame_per_block = getattr(self.transformer.config.arch_config, 'num_frames_per_block',
                                               getattr(self.transformer, 'num_frame_per_block', 1))
            self.sliding_window_num_frames = getattr(self.transformer.config.arch_config, 'sliding_window_num_frames',
                                                     15)
        else:
            self.num_frame_per_block = getattr(self.transformer, 'num_frame_per_block', 1)
            self.sliding_window_num_frames = 15

        try:
            self.local_attn_size = getattr(self.transformer, "local_attn_size", -1)
        except Exception:
            self.local_attn_size = -1

        assert self.local_attn_size != -1, (f"local_attn_size must be set for Matrix-Game causal inference, "
                                            f"got {self.local_attn_size}. Check MatrixGame2WanVideoArchConfig.")
        assert self.num_frame_per_block > 0, (f"num_frame_per_block must be positive, got {self.num_frame_per_block}")

        logger.info("Matrix-Game 2.0 causal inference initialized: "
                    "local_attn_size=%s, num_frame_per_block=%s", self.local_attn_size, self.num_frame_per_block)

        self.action_config = getattr(self.transformer, 'action_config', {})
        self.use_action_module = len(self.action_config) > 0

        self._streaming_initialized: bool = False
        self._streaming_ctx: BlockProcessingContext | None = None

    def forward(
        self,
        batch: ForwardBatch,
        fastvideo_args: FastVideoArgs,
    ) -> ForwardBatch:
        target_dtype = torch.bfloat16
        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast

        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
        patch_size = self.transformer.patch_size
        patch_ratio = patch_size[-1] * patch_size[-2]
        self.frame_seq_length = latent_seq_length // patch_ratio

        timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu()
        if fastvideo_args.pipeline_config.warp_denoising_step:
            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
            timesteps = scheduler_timesteps[1000 - timesteps]
        timesteps = timesteps.to(get_local_torch_device())

        boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None)
        if boundary_ratio is not None:
            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
        else:
            boundary_timestep = None
            high_noise_timesteps = None

        image_embeds = batch.image_embeds
        if len(image_embeds) > 0:
            assert torch.isnan(image_embeds[0]).sum() == 0
            image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]

        # directly set the kwarg.
        image_kwargs = {"encoder_hidden_states_image": image_embeds}
        pos_cond_kwargs: dict[str, Any] = {}

        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
        b, c, t, h, w = latents.shape
        prompt_embeds = batch.prompt_embeds
        assert torch.isnan(prompt_embeds[0]).sum() == 0

        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device)
        kv_cache2 = None
        if boundary_timestep is not None:
            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
                                                  dtype=target_dtype,
                                                  device=latents.device)

        kv_cache_mouse = None
        kv_cache_keyboard = None
        if self.use_action_module:
            kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0],
                                                                                 dtype=target_dtype,
                                                                                 device=latents.device)

        crossattn_cache = self._initialize_crossattn_cache(
            batch_size=latents.shape[0],
            max_text_len=257,  # 1 CLS + 256 patch tokens
            dtype=target_dtype,
            device=latents.device)

        if t % self.num_frame_per_block != 0:
            raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising")
        num_blocks = t // self.num_frame_per_block
        block_sizes = [self.num_frame_per_block] * num_blocks
        start_index = 0

        if boundary_timestep is not None:
            block_sizes[0] = 1

        # NOTE: Matrix-Game 2.0 does NOT process the first frame separately.
        # The first frame information is already encoded in batch.image_latent (cond_concat)
        # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1)

        ctx = BlockProcessingContext(
            batch=batch,
            block_idx=0,
            start_index=0,
            kv_cache1=kv_cache1,
            kv_cache2=kv_cache2,
            kv_cache_mouse=kv_cache_mouse,
            kv_cache_keyboard=kv_cache_keyboard,
            crossattn_cache=crossattn_cache,
            timesteps=timesteps,
            block_sizes=block_sizes,
            noise_pool=None,
            fastvideo_args=fastvideo_args,
            target_dtype=target_dtype,
            autocast_enabled=autocast_enabled,
            boundary_timestep=boundary_timestep,
            high_noise_timesteps=high_noise_timesteps,
            context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0),
            image_kwargs=image_kwargs,
            pos_cond_kwargs=pos_cond_kwargs,
        )

        context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0)

        with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar:
            for block_idx, current_num_frames in enumerate(block_sizes):
                ctx.block_idx = block_idx
                ctx.start_index = start_index
                current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :]

                action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames)

                current_latents = self._process_single_block(
                    current_latents=current_latents,
                    batch=batch,
                    start_index=start_index,
                    current_num_frames=current_num_frames,
                    timesteps=timesteps,
                    ctx=ctx,
                    action_kwargs=action_kwargs,
                    progress_bar=progress_bar,
                )

                latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents

                # Update KV caches with clean context
                self._update_context_cache(
                    current_latents=current_latents,
                    batch=batch,
                    start_index=start_index,
                    current_num_frames=current_num_frames,
                    ctx=ctx,
                    action_kwargs=action_kwargs,
                    context_noise=context_noise,
                )

                start_index += current_num_frames

        if boundary_timestep is not None:
            num_frames_to_remove = self.num_frame_per_block - 1
            if num_frames_to_remove > 0:
                latents = latents[:, :, :-num_frames_to_remove, :, :]

        batch.latents = latents
        return batch

    def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, num_frames: int) -> dict[str, Any]:
        action_kwargs: dict[str, Any] = {}
        if not self.use_action_module:
            return action_kwargs

        vae_time_compression_ratio = 4
        end_frame_idx = 1 + vae_time_compression_ratio * (start_index + num_frames - 1)
        if hasattr(batch, 'mouse_cond') and batch.mouse_cond is not None:
            action_kwargs['mouse_cond'] = batch.mouse_cond[:, :end_frame_idx]
        if hasattr(batch, 'keyboard_cond') and batch.keyboard_cond is not None:
            action_kwargs['keyboard_cond'] = batch.keyboard_cond[:, :end_frame_idx]

        # CRITICAL: Pass num_frame_per_block to model - this should be num_frames (current block size)
        action_kwargs['num_frame_per_block'] = num_frames
        return action_kwargs

    def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device) -> list[dict]:
        kv_cache = []
        num_attention_heads = self.transformer.num_attention_heads
        attention_head_dim = getattr(self.transformer, 'attention_head_dim',
                                     self.transformer.hidden_size // num_attention_heads)
        if self.local_attn_size != -1:
            kv_cache_size = self.local_attn_size * self.frame_seq_length
        else:
            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames

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

        return kv_cache

    def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
        kv_cache_mouse = []
        kv_cache_keyboard = []

        action_heads = self.action_config.get('heads_num', 16)
        mouse_head_dim = self.action_config.get('mouse_hidden_dim', 1024) // action_heads
        keyboard_head_dim = self.action_config.get('keyboard_hidden_dim', 1024) // action_heads

        kv_cache_size = self.local_attn_size if self.local_attn_size != -1 else 15

        for _ in range(self.num_transformer_blocks):
            kv_cache_keyboard.append({
                "k":
                torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_head_dim], dtype=dtype, device=device),
                "v":
                torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_head_dim], dtype=dtype, device=device),
                "global_end_index":
                torch.tensor([0], dtype=torch.long, device=device),
                "local_end_index":
                torch.tensor([0], dtype=torch.long, device=device),
            })
            kv_cache_mouse.append({
                "k":
                torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim],
                            dtype=dtype,
                            device=device),
                "v":
                torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim],
                            dtype=dtype,
                            device=device),
                "global_end_index":
                torch.tensor([0], dtype=torch.long, device=device),
                "local_end_index":
                torch.tensor([0], dtype=torch.long, device=device),
            })

        return kv_cache_mouse, kv_cache_keyboard

    def _initialize_crossattn_cache(self, batch_size: int, max_text_len: int, dtype: torch.dtype,
                                    device: torch.device) -> list[dict]:
        crossattn_cache = []
        num_attention_heads = self.transformer.num_attention_heads
        attention_head_dim = getattr(self.transformer, 'attention_head_dim',
                                     self.transformer.hidden_size // num_attention_heads)
        for _ in range(self.num_transformer_blocks):
            crossattn_cache.append({
                "k":
                torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim],
                            dtype=dtype,
                            device=device),
                "v":
                torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim],
                            dtype=dtype,
                            device=device),
                "is_init":
                False,
            })
        return crossattn_cache

    def _process_single_block(
        self,
        current_latents: torch.Tensor,
        batch: ForwardBatch,
        start_index: int,
        current_num_frames: int,
        timesteps: torch.Tensor,
        ctx: BlockProcessingContext,
        action_kwargs: dict[str, Any],
        noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor]
        | None = None,
        progress_bar: Any | None = None,
    ) -> torch.Tensor:
        prompt_embeds = batch.prompt_embeds
        noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)

        for i, t_cur in enumerate(timesteps):
            # Stop if interrupted
            if getattr(self, "interrupt", False):
                break
            if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep:
                current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
            else:
                current_model = self.transformer

            noise_latents = noise_latents_btchw.clone()
            latent_model_input = current_latents.to(ctx.target_dtype)

            independent_first_frame = getattr(self.transformer, 'independent_first_frame', False)
            if batch.image_latent is not None and independent_first_frame and start_index == 0:
                latent_model_input = torch.cat([latent_model_input, batch.image_latent.to(ctx.target_dtype)], dim=2)

            # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents
            t_expand = t_cur.repeat(latent_model_input.shape[0] * current_num_frames)

            # Build attention metadata if VSA is available
            if vsa_available and self.attn_backend == VideoSparseAttentionBackend:
                self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls()
                if self.attn_metadata_builder_cls is not None:
                    self.attn_metadata_builder = self.attn_metadata_builder_cls()
                    h, w = current_latents.shape[-2:]
                    attn_metadata = self.attn_metadata_builder.build(
                        current_timestep=i,
                        raw_latent_shape=(current_num_frames, h, w),
                        patch_size=ctx.fastvideo_args.pipeline_config.dit_config.patch_size,
                        VSA_sparsity=ctx.fastvideo_args.VSA_sparsity,
                        device=get_local_torch_device(),
                    )
                    assert attn_metadata is not None, "attn_metadata cannot be None"
                else:
                    attn_metadata = None
            else:
                attn_metadata = None

            with torch.autocast(device_type="cuda",
                                dtype=ctx.target_dtype,
                                enabled=ctx.autocast_enabled), \
                set_forward_context(current_timestep=i,
                                    attn_metadata=attn_metadata,
                                    forward_batch=batch):
                # Expand timestep to per-frame format [batch, num_frames] for causal model
                t_expanded_noise = t_cur * torch.ones((latent_model_input.shape[0], current_num_frames),
                                                      device=latent_model_input.device,
                                                      dtype=torch.long)

                model_kwargs = {
                    "kv_cache": ctx.get_kv_cache(t_cur),
                    "crossattn_cache": ctx.crossattn_cache,
                    "current_start": start_index * self.frame_seq_length,
                    "start_frame": start_index,
                }

                if self.use_action_module and current_model == self.transformer:
                    model_kwargs.update({
                        "kv_cache_mouse": ctx.kv_cache_mouse,
                        "kv_cache_keyboard": ctx.kv_cache_keyboard,
                    })
                    model_kwargs.update(action_kwargs)

                pred_noise_btchw = current_model(
                    latent_model_input,
                    prompt_embeds,
                    t_expanded_noise,
                    **ctx.image_kwargs,
                    **ctx.pos_cond_kwargs,
                    **model_kwargs,
                ).permute(0, 2, 1, 3, 4)

            if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep:
                pred_video_btchw = pred_noise_to_x_bound(
                    pred_noise=pred_noise_btchw.flatten(0, 1),
                    noise_input_latent=noise_latents.flatten(0, 1),
                    timestep=t_expand,
                    boundary_timestep=torch.ones_like(t_expand) * ctx.boundary_timestep,
                    scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2])
            else:
                pred_video_btchw = pred_noise_to_pred_video(pred_noise=pred_noise_btchw.flatten(0, 1),
                                                            noise_input_latent=noise_latents.flatten(0, 1),
                                                            timestep=t_expand,
                                                            scheduler=self.scheduler).unflatten(
                                                                0, pred_noise_btchw.shape[:2])

            if i < len(timesteps) - 1:
                next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video_btchw.device)

                # Use custom noise generator if provided (for streaming), else generate
                if noise_generator is not None:
                    noise = noise_generator(pred_video_btchw.shape, pred_video_btchw.dtype, i)
                else:
                    noise = torch.randn(
                        pred_video_btchw.shape,
                        dtype=pred_video_btchw.dtype,
                        generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator)).to(
                            pred_video_btchw.device)

                noise_btchw = noise
                if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i < len(
                        ctx.high_noise_timesteps) - 1:
                    noise_latents_btchw = self.scheduler.add_noise_high(
                        pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep,
                        torch.ones_like(next_timestep) * ctx.boundary_timestep).unflatten(
                            0, pred_video_btchw.shape[:2])
                elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i == len(
                        ctx.high_noise_timesteps) - 1:
                    noise_latents_btchw = pred_video_btchw
                else:
                    noise_latents_btchw = self.scheduler.add_noise(pred_video_btchw.flatten(0, 1),
                                                                   noise_btchw.flatten(0, 1), next_timestep).unflatten(
                                                                       0, pred_video_btchw.shape[:2])
                current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4)
            else:
                current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4)

            if progress_bar is not None:
                progress_bar.update()

        return current_latents

    def _update_context_cache(
        self,
        current_latents: torch.Tensor,
        batch: ForwardBatch,
        start_index: int,
        current_num_frames: int,
        ctx: BlockProcessingContext,
        action_kwargs: dict[str, Any],
        context_noise: float,
    ) -> None:
        prompt_embeds = batch.prompt_embeds
        latents_device = current_latents.device

        # Expand context timestep to per-frame format [batch, num_frames] for causal model
        t_context = torch.ones([current_latents.shape[0], current_num_frames], device=latents_device,
                               dtype=torch.long) * int(context_noise)
        context_bcthw = current_latents.to(ctx.target_dtype)

        with torch.autocast(device_type="cuda",
                            dtype=ctx.target_dtype,
                            enabled=ctx.autocast_enabled), \
            set_forward_context(current_timestep=0,
                                attn_metadata=None,
                                forward_batch=batch):

            context_model_kwargs = {
                "kv_cache": ctx.kv_cache1,
                "crossattn_cache": ctx.crossattn_cache,
                "current_start": start_index * self.frame_seq_length,
                "start_frame": start_index,
            }

            if self.use_action_module:
                context_model_kwargs.update({
                    "kv_cache_mouse": ctx.kv_cache_mouse,
                    "kv_cache_keyboard": ctx.kv_cache_keyboard,
                })
                context_model_kwargs.update(action_kwargs)

            if ctx.boundary_timestep is not None and self.transformer_2 is not None:
                self.transformer_2(
                    context_bcthw,
                    prompt_embeds,
                    t_context,
                    kv_cache=ctx.kv_cache2,
                    crossattn_cache=ctx.crossattn_cache,
                    current_start=start_index * self.frame_seq_length,
                    start_frame=start_index,
                    **ctx.image_kwargs,
                    **ctx.pos_cond_kwargs,
                )

            self.transformer(
                context_bcthw,
                prompt_embeds,
                t_context,
                **ctx.image_kwargs,
                **ctx.pos_cond_kwargs,
                **context_model_kwargs,
            )

    def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        target_dtype = torch.bfloat16
        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast

        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
        patch_size = self.transformer.patch_size
        patch_ratio = patch_size[-1] * patch_size[-2]
        self.frame_seq_length = latent_seq_length // patch_ratio

        timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu()
        if fastvideo_args.pipeline_config.warp_denoising_step:
            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
            timesteps = scheduler_timesteps[1000 - timesteps]
        timesteps = timesteps.to(get_local_torch_device())

        boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None)
        if boundary_ratio is not None:
            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
        else:
            boundary_timestep = None
            high_noise_timesteps = None

        image_embeds = batch.image_embeds
        if len(image_embeds) > 0:
            assert torch.isnan(image_embeds[0]).sum() == 0
            image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]

        # directly set the kwarg.
        image_kwargs = {"encoder_hidden_states_image": image_embeds}
        pos_cond_kwargs: dict[str, Any] = {}

        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
        b, c, t, h, w = latents.shape
        prompt_embeds = batch.prompt_embeds
        assert torch.isnan(prompt_embeds[0]).sum() == 0

        # Initialize caches
        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device)
        kv_cache2 = None
        if boundary_timestep is not None:
            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
                                                  dtype=target_dtype,
                                                  device=latents.device)

        kv_cache_mouse = None
        kv_cache_keyboard = None
        if self.use_action_module:
            kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0],
                                                                                 dtype=target_dtype,
                                                                                 device=latents.device)

        crossattn_cache = self._initialize_crossattn_cache(
            batch_size=latents.shape[0],
            max_text_len=257,  # 1 CLS + 256 patch tokens
            dtype=target_dtype,
            device=latents.device)

        # Calculate block sizes
        if t % self.num_frame_per_block != 0:
            raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising")
        num_blocks = t // self.num_frame_per_block
        block_sizes = [self.num_frame_per_block] * num_blocks
        if boundary_timestep is not None:
            block_sizes[0] = 1

        # Pre-allocate noise pool
        num_denoising_steps = len(timesteps)
        noise_shape = (b, self.num_frame_per_block, c, h, w)
        noise_pool = [
            torch.randn(
                noise_shape,
                dtype=target_dtype,
                device=latents.device,
            ) for _ in range(num_denoising_steps - 1)
        ]

        # Create and store context
        self._streaming_ctx = BlockProcessingContext(
            batch=batch,
            block_idx=0,
            start_index=0,
            kv_cache1=kv_cache1,
            kv_cache2=kv_cache2,
            kv_cache_mouse=kv_cache_mouse,
            kv_cache_keyboard=kv_cache_keyboard,
            crossattn_cache=crossattn_cache,
            timesteps=timesteps,
            block_sizes=block_sizes,
            noise_pool=noise_pool,
            fastvideo_args=fastvideo_args,
            target_dtype=target_dtype,
            autocast_enabled=autocast_enabled,
            boundary_timestep=boundary_timestep,
            high_noise_timesteps=high_noise_timesteps,
            context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0),
            image_kwargs=image_kwargs,
            pos_cond_kwargs=pos_cond_kwargs,
        )

        self._streaming_initialized = True
        return batch

    def streaming_step(self,
                       keyboard_action: torch.Tensor | None = None,
                       mouse_action: torch.Tensor | None = None) -> ForwardBatch:
        if not self._streaming_initialized or self._streaming_ctx is None:
            raise RuntimeError("Streaming not initialized! Call streaming_reset first.")

        ctx = self._streaming_ctx
        if ctx.block_idx >= len(ctx.block_sizes):
            return ctx.batch

        batch = ctx.batch
        latents = batch.latents
        assert latents is not None, "latents must be set in batch"

        current_num_frames = ctx.block_sizes[ctx.block_idx]
        start_index = ctx.start_index

        current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :]

        # Update batch with new actions for this block
        if keyboard_action is not None or mouse_action is not None:
            vae_ratio = 4
            start_frame = 0 if start_index == 0 else 1 + vae_ratio * (start_index - 1)

            if keyboard_action is not None:
                n = keyboard_action.shape[1]
                batch.keyboard_cond[:, start_frame:start_frame + n] = keyboard_action.to(batch.keyboard_cond.device)
            if mouse_action is not None:
                n = mouse_action.shape[1]
                batch.mouse_cond[:, start_frame:start_frame + n] = mouse_action.to(batch.mouse_cond.device)

        action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames)

        # Create noise generator that uses pre-allocated noise pool
        def streaming_noise_generator(shape: tuple, dtype: torch.dtype, step_idx: int) -> torch.Tensor:
            if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool):
                return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(latents.device)
            else:
                # Fallback to dynamic allocation if pool not available
                return torch.randn(
                    shape,
                    dtype=dtype,
                    generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator)).to(
                        latents.device)

        current_latents = self._process_single_block(
            current_latents=current_latents,
            batch=batch,
            start_index=start_index,
            current_num_frames=current_num_frames,
            timesteps=ctx.timesteps,
            ctx=ctx,
            action_kwargs=action_kwargs,
            noise_generator=streaming_noise_generator,
        )

        latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents

        # Update KV caches with clean context
        self._update_context_cache(
            current_latents=current_latents,
            batch=batch,
            start_index=start_index,
            current_num_frames=current_num_frames,
            ctx=ctx,
            action_kwargs=action_kwargs,
            context_noise=ctx.context_noise,
        )

        # Advance streaming state
        ctx.start_index += current_num_frames
        ctx.block_idx += 1

        return batch

    def streaming_clear(self) -> None:
        self._streaming_initialized = False
        self._streaming_ctx = None

    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
        result.add_check("image_embeds", batch.image_embeds, V.is_list)
        result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5))
        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
        result.add_check("guidance_scale", batch.guidance_scale, V.positive_float)
        result.add_check("eta", batch.eta, V.non_negative_float)
        result.add_check("generator", batch.generator, V.generator_or_list_generators)
        result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value)
        result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
                         lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x))
        return result

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(
    self,
    batch: ForwardBatch,
    fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
    target_dtype = torch.bfloat16
    autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast

    latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
    patch_size = self.transformer.patch_size
    patch_ratio = patch_size[-1] * patch_size[-2]
    self.frame_seq_length = latent_seq_length // patch_ratio

    timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu()
    if fastvideo_args.pipeline_config.warp_denoising_step:
        scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
        timesteps = scheduler_timesteps[1000 - timesteps]
    timesteps = timesteps.to(get_local_torch_device())

    boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None)
    if boundary_ratio is not None:
        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
        high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
    else:
        boundary_timestep = None
        high_noise_timesteps = None

    image_embeds = batch.image_embeds
    if len(image_embeds) > 0:
        assert torch.isnan(image_embeds[0]).sum() == 0
        image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]

    # directly set the kwarg.
    image_kwargs = {"encoder_hidden_states_image": image_embeds}
    pos_cond_kwargs: dict[str, Any] = {}

    assert batch.latents is not None, "latents must be provided"
    latents = batch.latents
    b, c, t, h, w = latents.shape
    prompt_embeds = batch.prompt_embeds
    assert torch.isnan(prompt_embeds[0]).sum() == 0

    kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device)
    kv_cache2 = None
    if boundary_timestep is not None:
        kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
                                              dtype=target_dtype,
                                              device=latents.device)

    kv_cache_mouse = None
    kv_cache_keyboard = None
    if self.use_action_module:
        kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0],
                                                                             dtype=target_dtype,
                                                                             device=latents.device)

    crossattn_cache = self._initialize_crossattn_cache(
        batch_size=latents.shape[0],
        max_text_len=257,  # 1 CLS + 256 patch tokens
        dtype=target_dtype,
        device=latents.device)

    if t % self.num_frame_per_block != 0:
        raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising")
    num_blocks = t // self.num_frame_per_block
    block_sizes = [self.num_frame_per_block] * num_blocks
    start_index = 0

    if boundary_timestep is not None:
        block_sizes[0] = 1

    # NOTE: Matrix-Game 2.0 does NOT process the first frame separately.
    # The first frame information is already encoded in batch.image_latent (cond_concat)
    # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1)

    ctx = BlockProcessingContext(
        batch=batch,
        block_idx=0,
        start_index=0,
        kv_cache1=kv_cache1,
        kv_cache2=kv_cache2,
        kv_cache_mouse=kv_cache_mouse,
        kv_cache_keyboard=kv_cache_keyboard,
        crossattn_cache=crossattn_cache,
        timesteps=timesteps,
        block_sizes=block_sizes,
        noise_pool=None,
        fastvideo_args=fastvideo_args,
        target_dtype=target_dtype,
        autocast_enabled=autocast_enabled,
        boundary_timestep=boundary_timestep,
        high_noise_timesteps=high_noise_timesteps,
        context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0),
        image_kwargs=image_kwargs,
        pos_cond_kwargs=pos_cond_kwargs,
    )

    context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0)

    with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar:
        for block_idx, current_num_frames in enumerate(block_sizes):
            ctx.block_idx = block_idx
            ctx.start_index = start_index
            current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :]

            action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames)

            current_latents = self._process_single_block(
                current_latents=current_latents,
                batch=batch,
                start_index=start_index,
                current_num_frames=current_num_frames,
                timesteps=timesteps,
                ctx=ctx,
                action_kwargs=action_kwargs,
                progress_bar=progress_bar,
            )

            latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents

            # Update KV caches with clean context
            self._update_context_cache(
                current_latents=current_latents,
                batch=batch,
                start_index=start_index,
                current_num_frames=current_num_frames,
                ctx=ctx,
                action_kwargs=action_kwargs,
                context_noise=context_noise,
            )

            start_index += current_num_frames

    if boundary_timestep is not None:
        num_frames_to_remove = self.num_frame_per_block - 1
        if num_frames_to_remove > 0:
            latents = latents[:, :, :-num_frames_to_remove, :, :]

    batch.latents = latents
    return batch

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage.streaming_reset

streaming_reset(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    target_dtype = torch.bfloat16
    autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast

    latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
    patch_size = self.transformer.patch_size
    patch_ratio = patch_size[-1] * patch_size[-2]
    self.frame_seq_length = latent_seq_length // patch_ratio

    timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu()
    if fastvideo_args.pipeline_config.warp_denoising_step:
        scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
        timesteps = scheduler_timesteps[1000 - timesteps]
    timesteps = timesteps.to(get_local_torch_device())

    boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None)
    if boundary_ratio is not None:
        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
        high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
    else:
        boundary_timestep = None
        high_noise_timesteps = None

    image_embeds = batch.image_embeds
    if len(image_embeds) > 0:
        assert torch.isnan(image_embeds[0]).sum() == 0
        image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]

    # directly set the kwarg.
    image_kwargs = {"encoder_hidden_states_image": image_embeds}
    pos_cond_kwargs: dict[str, Any] = {}

    assert batch.latents is not None, "latents must be provided"
    latents = batch.latents
    b, c, t, h, w = latents.shape
    prompt_embeds = batch.prompt_embeds
    assert torch.isnan(prompt_embeds[0]).sum() == 0

    # Initialize caches
    kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device)
    kv_cache2 = None
    if boundary_timestep is not None:
        kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
                                              dtype=target_dtype,
                                              device=latents.device)

    kv_cache_mouse = None
    kv_cache_keyboard = None
    if self.use_action_module:
        kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0],
                                                                             dtype=target_dtype,
                                                                             device=latents.device)

    crossattn_cache = self._initialize_crossattn_cache(
        batch_size=latents.shape[0],
        max_text_len=257,  # 1 CLS + 256 patch tokens
        dtype=target_dtype,
        device=latents.device)

    # Calculate block sizes
    if t % self.num_frame_per_block != 0:
        raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising")
    num_blocks = t // self.num_frame_per_block
    block_sizes = [self.num_frame_per_block] * num_blocks
    if boundary_timestep is not None:
        block_sizes[0] = 1

    # Pre-allocate noise pool
    num_denoising_steps = len(timesteps)
    noise_shape = (b, self.num_frame_per_block, c, h, w)
    noise_pool = [
        torch.randn(
            noise_shape,
            dtype=target_dtype,
            device=latents.device,
        ) for _ in range(num_denoising_steps - 1)
    ]

    # Create and store context
    self._streaming_ctx = BlockProcessingContext(
        batch=batch,
        block_idx=0,
        start_index=0,
        kv_cache1=kv_cache1,
        kv_cache2=kv_cache2,
        kv_cache_mouse=kv_cache_mouse,
        kv_cache_keyboard=kv_cache_keyboard,
        crossattn_cache=crossattn_cache,
        timesteps=timesteps,
        block_sizes=block_sizes,
        noise_pool=noise_pool,
        fastvideo_args=fastvideo_args,
        target_dtype=target_dtype,
        autocast_enabled=autocast_enabled,
        boundary_timestep=boundary_timestep,
        high_noise_timesteps=high_noise_timesteps,
        context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0),
        image_kwargs=image_kwargs,
        pos_cond_kwargs=pos_cond_kwargs,
    )

    self._streaming_initialized = True
    return batch

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage.streaming_step

streaming_step(keyboard_action: torch.Tensor | None = None, mouse_action: torch.Tensor | None = None) -> ForwardBatch
Source
def streaming_step(self,
                   keyboard_action: torch.Tensor | None = None,
                   mouse_action: torch.Tensor | None = None) -> ForwardBatch:
    if not self._streaming_initialized or self._streaming_ctx is None:
        raise RuntimeError("Streaming not initialized! Call streaming_reset first.")

    ctx = self._streaming_ctx
    if ctx.block_idx >= len(ctx.block_sizes):
        return ctx.batch

    batch = ctx.batch
    latents = batch.latents
    assert latents is not None, "latents must be set in batch"

    current_num_frames = ctx.block_sizes[ctx.block_idx]
    start_index = ctx.start_index

    current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :]

    # Update batch with new actions for this block
    if keyboard_action is not None or mouse_action is not None:
        vae_ratio = 4
        start_frame = 0 if start_index == 0 else 1 + vae_ratio * (start_index - 1)

        if keyboard_action is not None:
            n = keyboard_action.shape[1]
            batch.keyboard_cond[:, start_frame:start_frame + n] = keyboard_action.to(batch.keyboard_cond.device)
        if mouse_action is not None:
            n = mouse_action.shape[1]
            batch.mouse_cond[:, start_frame:start_frame + n] = mouse_action.to(batch.mouse_cond.device)

    action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames)

    # Create noise generator that uses pre-allocated noise pool
    def streaming_noise_generator(shape: tuple, dtype: torch.dtype, step_idx: int) -> torch.Tensor:
        if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool):
            return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(latents.device)
        else:
            # Fallback to dynamic allocation if pool not available
            return torch.randn(
                shape,
                dtype=dtype,
                generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator)).to(
                    latents.device)

    current_latents = self._process_single_block(
        current_latents=current_latents,
        batch=batch,
        start_index=start_index,
        current_num_frames=current_num_frames,
        timesteps=ctx.timesteps,
        ctx=ctx,
        action_kwargs=action_kwargs,
        noise_generator=streaming_noise_generator,
    )

    latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents

    # Update KV caches with clean context
    self._update_context_cache(
        current_latents=current_latents,
        batch=batch,
        start_index=start_index,
        current_num_frames=current_num_frames,
        ctx=ctx,
        action_kwargs=action_kwargs,
        context_noise=ctx.context_noise,
    )

    # Advance streaming state
    ctx.start_index += current_num_frames
    ctx.block_idx += 1

    return batch

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage.streaming_clear

streaming_clear() -> None
Source
def streaming_clear(self) -> None:
    self._streaming_initialized = False
    self._streaming_ctx = None

fastvideo.pipelines.stages.matrixgame2_denoising.MatrixGame2CausalDenoisingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
    result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
    result.add_check("image_embeds", batch.image_embeds, V.is_list)
    result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5))
    result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
    result.add_check("guidance_scale", batch.guidance_scale, V.positive_float)
    result.add_check("eta", batch.eta, V.non_negative_float)
    result.add_check("generator", batch.generator, V.generator_or_list_generators)
    result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value)
    result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
                     lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x))
    return result