Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

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</code></pre></details>


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) -&gt; 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, &#x27;config&#x27;) and hasattr(self.transformer.config, &#x27;arch_config&#x27;):
        self.num_frame_per_block = getattr(self.transformer.config.arch_config, &#x27;num_frames_per_block&#x27;,
                                           getattr(self.transformer, &#x27;num_frame_per_block&#x27;, 1))
        self.sliding_window_num_frames = getattr(self.transformer.config.arch_config, &#x27;sliding_window_num_frames&#x27;,
                                                 15)
    else:
        self.num_frame_per_block = getattr(self.transformer, &#x27;num_frame_per_block&#x27;, 1)
        self.sliding_window_num_frames = 15

    try:
        self.local_attn_size = getattr(self.transformer, &quot;local_attn_size&quot;, -1)
    except Exception:
        self.local_attn_size = -1

    assert self.local_attn_size != -1, (f&quot;local_attn_size must be set for Matrix-Game causal inference, &quot;
                                        f&quot;got &#123;self.local_attn_size&#125;. Check MatrixGame2WanVideoArchConfig.&quot;)
    assert self.num_frame_per_block &gt; 0, (f&quot;num_frame_per_block must be positive, got &#123;self.num_frame_per_block&#125;&quot;)

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

    self.action_config = getattr(self.transformer, &#x27;action_config&#x27;, &#123;&#125;)
    self.use_action_module = len(self.action_config) &gt; 0

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

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

    latent_seq_length = batch.latents.shape&#91;-1&#93; * batch.latents.shape&#91;-2&#93;
    patch_size = self.transformer.patch_size
    patch_ratio = patch_size&#91;-1&#93; * patch_size&#91;-2&#93;
    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(&#91;0&#93;, dtype=torch.float32)))
        timesteps = scheduler_timesteps&#91;1000 - timesteps&#93;
    timesteps = timesteps.to(get_local_torch_device())

    boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, &#x27;boundary_ratio&#x27;, None)
    if boundary_ratio is not None:
        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
        high_noise_timesteps = timesteps&#91;timesteps &gt;= boundary_timestep&#93;
    else:
        boundary_timestep = None
        high_noise_timesteps = None

    image_embeds = batch.image_embeds
    if len(image_embeds) &gt; 0:
        assert torch.isnan(image_embeds&#91;0&#93;).sum() == 0
        image_embeds = &#91;image_embed.to(target_dtype) for image_embed in image_embeds&#93;

    # directly set the kwarg.
    image_kwargs = &#123;&quot;encoder_hidden_states_image&quot;: image_embeds&#125;
    pos_cond_kwargs: dict&#91;str, Any&#93; = &#123;&#125;

    assert batch.latents is not None, &quot;latents must be provided&quot;
    latents = batch.latents
    b, c, t, h, w = latents.shape
    prompt_embeds = batch.prompt_embeds
    assert torch.isnan(prompt_embeds&#91;0&#93;).sum() == 0

    kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape&#91;0&#93;, 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&#91;0&#93;,
                                              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&#91;0&#93;,
                                                                             dtype=target_dtype,
                                                                             device=latents.device)

    crossattn_cache = self._initialize_crossattn_cache(
        batch_size=latents.shape&#91;0&#93;,
        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(&quot;num_frames must be divisible by num_frame_per_block for causal denoising&quot;)
    num_blocks = t // self.num_frame_per_block
    block_sizes = &#91;self.num_frame_per_block&#93; * num_blocks
    start_index = 0

    if boundary_timestep is not None:
        block_sizes&#91;0&#93; = 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(&#91;x, cond_concat&#93;, 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, &quot;context_noise&quot;, 0),
        image_kwargs=image_kwargs,
        pos_cond_kwargs=pos_cond_kwargs,
    )

    context_noise = getattr(fastvideo_args.pipeline_config, &quot;context_noise&quot;, 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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93;

            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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93; = 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 &gt; 0:
            latents = latents&#91;:, :, :-num_frames_to_remove, :, :&#93;

    batch.latents = latents
    return batch

def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, num_frames: int) -&gt; dict&#91;str, Any&#93;:
    action_kwargs: dict&#91;str, Any&#93; = &#123;&#125;
    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, &#x27;mouse_cond&#x27;) and batch.mouse_cond is not None:
        action_kwargs&#91;&#x27;mouse_cond&#x27;&#93; = batch.mouse_cond&#91;:, :end_frame_idx&#93;
    if hasattr(batch, &#x27;keyboard_cond&#x27;) and batch.keyboard_cond is not None:
        action_kwargs&#91;&#x27;keyboard_cond&#x27;&#93; = batch.keyboard_cond&#91;:, :end_frame_idx&#93;

    # CRITICAL: Pass num_frame_per_block to model - this should be num_frames (current block size)
    action_kwargs&#91;&#x27;num_frame_per_block&#x27;&#93; = num_frames
    return action_kwargs

def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device) -&gt; list&#91;dict&#93;:
    kv_cache = &#91;&#93;
    num_attention_heads = self.transformer.num_attention_heads
    attention_head_dim = getattr(self.transformer, &#x27;attention_head_dim&#x27;,
                                 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(&#123;
            &quot;k&quot;:
            torch.zeros(&#91;batch_size, kv_cache_size, num_attention_heads, attention_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;v&quot;:
            torch.zeros(&#91;batch_size, kv_cache_size, num_attention_heads, attention_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;global_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
            &quot;local_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
        &#125;)

    return kv_cache

def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
    kv_cache_mouse = &#91;&#93;
    kv_cache_keyboard = &#91;&#93;

    action_heads = self.action_config.get(&#x27;heads_num&#x27;, 16)
    mouse_head_dim = self.action_config.get(&#x27;mouse_hidden_dim&#x27;, 1024) // action_heads
    keyboard_head_dim = self.action_config.get(&#x27;keyboard_hidden_dim&#x27;, 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(&#123;
            &quot;k&quot;:
            torch.zeros(&#91;batch_size, kv_cache_size, action_heads, keyboard_head_dim&#93;, dtype=dtype, device=device),
            &quot;v&quot;:
            torch.zeros(&#91;batch_size, kv_cache_size, action_heads, keyboard_head_dim&#93;, dtype=dtype, device=device),
            &quot;global_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
            &quot;local_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
        &#125;)
        kv_cache_mouse.append(&#123;
            &quot;k&quot;:
            torch.zeros(&#91;batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;v&quot;:
            torch.zeros(&#91;batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;global_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
            &quot;local_end_index&quot;:
            torch.tensor(&#91;0&#93;, dtype=torch.long, device=device),
        &#125;)

    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) -&gt; list&#91;dict&#93;:
    crossattn_cache = &#91;&#93;
    num_attention_heads = self.transformer.num_attention_heads
    attention_head_dim = getattr(self.transformer, &#x27;attention_head_dim&#x27;,
                                 self.transformer.hidden_size // num_attention_heads)
    for _ in range(self.num_transformer_blocks):
        crossattn_cache.append(&#123;
            &quot;k&quot;:
            torch.zeros(&#91;batch_size, max_text_len, num_attention_heads, attention_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;v&quot;:
            torch.zeros(&#91;batch_size, max_text_len, num_attention_heads, attention_head_dim&#93;,
                        dtype=dtype,
                        device=device),
            &quot;is_init&quot;:
            False,
        &#125;)
    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&#91;str, Any&#93;,
    noise_generator: Callable&#91;&#91;tuple, torch.dtype, int&#93;, torch.Tensor&#93;
    | None = None,
    progress_bar: Any | None = None,
) -&gt; 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, &quot;interrupt&quot;, False):
            break
        if ctx.boundary_timestep is not None and t_cur &lt; 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, &#x27;independent_first_frame&#x27;, False)
        if batch.image_latent is not None and independent_first_frame and start_index == 0:
            latent_model_input = torch.cat(&#91;latent_model_input, batch.image_latent.to(ctx.target_dtype)&#93;, dim=2)

        # t_expand needs to be &#91;batch * frames&#93; to match flattened pred_noise/noise_latents
        t_expand = t_cur.repeat(latent_model_input.shape&#91;0&#93; * 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&#91;-2:&#93;
                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, &quot;attn_metadata cannot be None&quot;
            else:
                attn_metadata = None
        else:
            attn_metadata = None

        with torch.autocast(device_type=&quot;cuda&quot;,
                            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 &#91;batch, num_frames&#93; for causal model
            t_expanded_noise = t_cur * torch.ones((latent_model_input.shape&#91;0&#93;, current_num_frames),
                                                  device=latent_model_input.device,
                                                  dtype=torch.long)

            model_kwargs = &#123;
                &quot;kv_cache&quot;: ctx.get_kv_cache(t_cur),
                &quot;crossattn_cache&quot;: ctx.crossattn_cache,
                &quot;current_start&quot;: start_index * self.frame_seq_length,
                &quot;start_frame&quot;: start_index,
            &#125;

            if self.use_action_module and current_model == self.transformer:
                model_kwargs.update(&#123;
                    &quot;kv_cache_mouse&quot;: ctx.kv_cache_mouse,
                    &quot;kv_cache_keyboard&quot;: ctx.kv_cache_keyboard,
                &#125;)
                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 &gt;= 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&#91;:2&#93;)
        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&#91;:2&#93;)

        if i &lt; len(timesteps) - 1:
            next_timestep = timesteps&#91;i + 1&#93; * torch.ones(&#91;1&#93;, 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&#91;0&#93; 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 &lt; 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&#91;:2&#93;)
            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&#91;:2&#93;)
            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&#91;str, Any&#93;,
    context_noise: float,
) -&gt; None:
    prompt_embeds = batch.prompt_embeds
    latents_device = current_latents.device

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

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

        context_model_kwargs = &#123;
            &quot;kv_cache&quot;: ctx.kv_cache1,
            &quot;crossattn_cache&quot;: ctx.crossattn_cache,
            &quot;current_start&quot;: start_index * self.frame_seq_length,
            &quot;start_frame&quot;: start_index,
        &#125;

        if self.use_action_module:
            context_model_kwargs.update(&#123;
                &quot;kv_cache_mouse&quot;: ctx.kv_cache_mouse,
                &quot;kv_cache_keyboard&quot;: ctx.kv_cache_keyboard,
            &#125;)
            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) -&gt; ForwardBatch:
    target_dtype = torch.bfloat16
    autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast

    latent_seq_length = batch.latents.shape&#91;-1&#93; * batch.latents.shape&#91;-2&#93;
    patch_size = self.transformer.patch_size
    patch_ratio = patch_size&#91;-1&#93; * patch_size&#91;-2&#93;
    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(&#91;0&#93;, dtype=torch.float32)))
        timesteps = scheduler_timesteps&#91;1000 - timesteps&#93;
    timesteps = timesteps.to(get_local_torch_device())

    boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, &#x27;boundary_ratio&#x27;, None)
    if boundary_ratio is not None:
        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
        high_noise_timesteps = timesteps&#91;timesteps &gt;= boundary_timestep&#93;
    else:
        boundary_timestep = None
        high_noise_timesteps = None

    image_embeds = batch.image_embeds
    if len(image_embeds) &gt; 0:
        assert torch.isnan(image_embeds&#91;0&#93;).sum() == 0
        image_embeds = &#91;image_embed.to(target_dtype) for image_embed in image_embeds&#93;

    # directly set the kwarg.
    image_kwargs = &#123;&quot;encoder_hidden_states_image&quot;: image_embeds&#125;
    pos_cond_kwargs: dict&#91;str, Any&#93; = &#123;&#125;

    assert batch.latents is not None, &quot;latents must be provided&quot;
    latents = batch.latents
    b, c, t, h, w = latents.shape
    prompt_embeds = batch.prompt_embeds
    assert torch.isnan(prompt_embeds&#91;0&#93;).sum() == 0

    # Initialize caches
    kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape&#91;0&#93;, 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&#91;0&#93;,
                                              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&#91;0&#93;,
                                                                             dtype=target_dtype,
                                                                             device=latents.device)

    crossattn_cache = self._initialize_crossattn_cache(
        batch_size=latents.shape&#91;0&#93;,
        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(&quot;num_frames must be divisible by num_frame_per_block for causal denoising&quot;)
    num_blocks = t // self.num_frame_per_block
    block_sizes = &#91;self.num_frame_per_block&#93; * num_blocks
    if boundary_timestep is not None:
        block_sizes&#91;0&#93; = 1

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

    # 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, &quot;context_noise&quot;, 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) -&gt; ForwardBatch:
    if not self._streaming_initialized or self._streaming_ctx is None:
        raise RuntimeError(&quot;Streaming not initialized! Call streaming_reset first.&quot;)

    ctx = self._streaming_ctx
    if ctx.block_idx &gt;= len(ctx.block_sizes):
        return ctx.batch

    batch = ctx.batch
    latents = batch.latents
    assert latents is not None, &quot;latents must be set in batch&quot;

    current_num_frames = ctx.block_sizes&#91;ctx.block_idx&#93;
    start_index = ctx.start_index

    current_latents = latents&#91;:, :, start_index:start_index + current_num_frames, :, :&#93;

    # 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&#91;1&#93;
            batch.keyboard_cond&#91;:, start_frame:start_frame + n&#93; = keyboard_action.to(batch.keyboard_cond.device)
        if mouse_action is not None:
            n = mouse_action.shape&#91;1&#93;
            batch.mouse_cond&#91;:, start_frame:start_frame + n&#93; = 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) -&gt; torch.Tensor:
        if ctx.noise_pool is not None and step_idx &lt; len(ctx.noise_pool):
            return ctx.noise_pool&#91;step_idx&#93;&#91;:, :shape&#91;1&#93;, :, :, :&#93;.to(latents.device)
        else:
            # Fallback to dynamic allocation if pool not available
            return torch.randn(
                shape,
                dtype=dtype,
                generator=(batch.generator&#91;0&#93; 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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93; = 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) -&gt; None:
    self._streaming_initialized = False
    self._streaming_ctx = None

def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; VerificationResult:
    result = VerificationResult()
    result.add_check(&quot;latents&quot;, batch.latents, &#91;V.is_tensor, V.with_dims(5)&#93;)
    result.add_check(&quot;prompt_embeds&quot;, batch.prompt_embeds, V.list_not_empty)
    result.add_check(&quot;image_embeds&quot;, batch.image_embeds, V.is_list)
    result.add_check(&quot;image_latent&quot;, batch.image_latent, V.none_or_tensor_with_dims(5))
    result.add_check(&quot;num_inference_steps&quot;, batch.num_inference_steps, V.positive_int)
    result.add_check(&quot;guidance_scale&quot;, batch.guidance_scale, V.positive_float)
    result.add_check(&quot;eta&quot;, batch.eta, V.non_negative_float)
    result.add_check(&quot;generator&quot;, batch.generator, V.generator_or_list_generators)
    result.add_check(&quot;do_classifier_free_guidance&quot;, batch.do_classifier_free_guidance, V.bool_value)
    result.add_check(&quot;negative_prompt_embeds&quot;, batch.negative_prompt_embeds,
                     lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x))
    return result</code></pre></details>


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&#91;-1&#93; * batch.latents.shape&#91;-2&#93;
patch_size = self.transformer.patch_size
patch_ratio = patch_size&#91;-1&#93; * patch_size&#91;-2&#93;
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(&#91;0&#93;, dtype=torch.float32)))
    timesteps = scheduler_timesteps&#91;1000 - timesteps&#93;
timesteps = timesteps.to(get_local_torch_device())

boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, &#x27;boundary_ratio&#x27;, None)
if boundary_ratio is not None:
    boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
    high_noise_timesteps = timesteps&#91;timesteps &gt;= boundary_timestep&#93;
else:
    boundary_timestep = None
    high_noise_timesteps = None

image_embeds = batch.image_embeds
if len(image_embeds) &gt; 0:
    assert torch.isnan(image_embeds&#91;0&#93;).sum() == 0
    image_embeds = &#91;image_embed.to(target_dtype) for image_embed in image_embeds&#93;

# directly set the kwarg.
image_kwargs = &#123;&quot;encoder_hidden_states_image&quot;: image_embeds&#125;
pos_cond_kwargs: dict&#91;str, Any&#93; = &#123;&#125;

assert batch.latents is not None, &quot;latents must be provided&quot;
latents = batch.latents
b, c, t, h, w = latents.shape
prompt_embeds = batch.prompt_embeds
assert torch.isnan(prompt_embeds&#91;0&#93;).sum() == 0

kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape&#91;0&#93;, 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&#91;0&#93;,
                                          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&#91;0&#93;,
                                                                         dtype=target_dtype,
                                                                         device=latents.device)

crossattn_cache = self._initialize_crossattn_cache(
    batch_size=latents.shape&#91;0&#93;,
    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(&quot;num_frames must be divisible by num_frame_per_block for causal denoising&quot;)
num_blocks = t // self.num_frame_per_block
block_sizes = &#91;self.num_frame_per_block&#93; * num_blocks
start_index = 0

if boundary_timestep is not None:
    block_sizes&#91;0&#93; = 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(&#91;x, cond_concat&#93;, 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, &quot;context_noise&quot;, 0),
    image_kwargs=image_kwargs,
    pos_cond_kwargs=pos_cond_kwargs,
)

context_noise = getattr(fastvideo_args.pipeline_config, &quot;context_noise&quot;, 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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93;

        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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93; = 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 &gt; 0:
        latents = latents&#91;:, :, :-num_frames_to_remove, :, :&#93;

batch.latents = latents
return batch</code></pre></details>


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&#91;-1&#93; * batch.latents.shape&#91;-2&#93;
patch_size = self.transformer.patch_size
patch_ratio = patch_size&#91;-1&#93; * patch_size&#91;-2&#93;
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(&#91;0&#93;, dtype=torch.float32)))
    timesteps = scheduler_timesteps&#91;1000 - timesteps&#93;
timesteps = timesteps.to(get_local_torch_device())

boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, &#x27;boundary_ratio&#x27;, None)
if boundary_ratio is not None:
    boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
    high_noise_timesteps = timesteps&#91;timesteps &gt;= boundary_timestep&#93;
else:
    boundary_timestep = None
    high_noise_timesteps = None

image_embeds = batch.image_embeds
if len(image_embeds) &gt; 0:
    assert torch.isnan(image_embeds&#91;0&#93;).sum() == 0
    image_embeds = &#91;image_embed.to(target_dtype) for image_embed in image_embeds&#93;

# directly set the kwarg.
image_kwargs = &#123;&quot;encoder_hidden_states_image&quot;: image_embeds&#125;
pos_cond_kwargs: dict&#91;str, Any&#93; = &#123;&#125;

assert batch.latents is not None, &quot;latents must be provided&quot;
latents = batch.latents
b, c, t, h, w = latents.shape
prompt_embeds = batch.prompt_embeds
assert torch.isnan(prompt_embeds&#91;0&#93;).sum() == 0

# Initialize caches
kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape&#91;0&#93;, 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&#91;0&#93;,
                                          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&#91;0&#93;,
                                                                         dtype=target_dtype,
                                                                         device=latents.device)

crossattn_cache = self._initialize_crossattn_cache(
    batch_size=latents.shape&#91;0&#93;,
    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(&quot;num_frames must be divisible by num_frame_per_block for causal denoising&quot;)
num_blocks = t // self.num_frame_per_block
block_sizes = &#91;self.num_frame_per_block&#93; * num_blocks
if boundary_timestep is not None:
    block_sizes&#91;0&#93; = 1

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

# 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, &quot;context_noise&quot;, 0),
    image_kwargs=image_kwargs,
    pos_cond_kwargs=pos_cond_kwargs,
)

self._streaming_initialized = True
return batch</code></pre></details>


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 &gt;= len(ctx.block_sizes):
    return ctx.batch

batch = ctx.batch
latents = batch.latents
assert latents is not None, &quot;latents must be set in batch&quot;

current_num_frames = ctx.block_sizes&#91;ctx.block_idx&#93;
start_index = ctx.start_index

current_latents = latents&#91;:, :, start_index:start_index + current_num_frames, :, :&#93;

# 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&#91;1&#93;
        batch.keyboard_cond&#91;:, start_frame:start_frame + n&#93; = keyboard_action.to(batch.keyboard_cond.device)
    if mouse_action is not None:
        n = mouse_action.shape&#91;1&#93;
        batch.mouse_cond&#91;:, start_frame:start_frame + n&#93; = 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) -&gt; torch.Tensor:
    if ctx.noise_pool is not None and step_idx &lt; len(ctx.noise_pool):
        return ctx.noise_pool&#91;step_idx&#93;&#91;:, :shape&#91;1&#93;, :, :, :&#93;.to(latents.device)
    else:
        # Fallback to dynamic allocation if pool not available
        return torch.randn(
            shape,
            dtype=dtype,
            generator=(batch.generator&#91;0&#93; 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&#91;:, :, start_index:start_index + current_num_frames, :, :&#93; = 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</code></pre></details>


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

Was this page helpful?