--- title: matrixgame3_denoising ---

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

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.pipelines.stages.matrixgame3_denoising.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.pipelines.stages.matrixgame3_denoising.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.pipelines.stages.matrixgame3_denoising.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.compute_relative_poses

compute_relative_poses(c2ws_mat: torch.Tensor, framewise: bool = False, normalize_trans: bool = True) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_denoising.build_extrinsics_from_actions

build_extrinsics_from_actions(keyboard_conditions: torch.Tensor, mouse_conditions: torch.Tensor) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_denoising.build_matrixgame3_action_preset

build_matrixgame3_action_preset(num_frames: int, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]

fastvideo.pipelines.stages.matrixgame3_denoising.build_plucker_from_c2ws

build_plucker_from_c2ws(c2ws_seq: torch.Tensor, src_indices: np.ndarray, tgt_indices: np.ndarray, *, target_h: int, target_w: int, latent_h: int, latent_w: int, framewise: bool = True) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_denoising.build_plucker_from_pose

build_plucker_from_pose(c2ws_pose: torch.Tensor, *, target_h: int, target_w: int, latent_h: int, latent_w: int) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_denoising.interpolate_camera_poses_handedness

interpolate_camera_poses_handedness(src_indices: np.ndarray, src_rot_mat: np.ndarray, src_trans_vec: np.ndarray, tgt_indices: np.ndarray) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_denoising.select_memory_idx_fov

select_memory_idx_fov(extrinsics_all: torch.Tensor | np.ndarray, current_start_frame_idx: int, selected_index_base: list[int], *, height: int = 720, width: int = 1280, return_confidence: bool = False) -> list[int] | tuple[list[int], list[float]]

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.DenoisingStage.prepare_denoising

prepare_denoising(batch, fastvideo_args, target_dtype) -> DenoisingState

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

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.DenoisingStage.select_model

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

fastvideo.pipelines.stages.matrixgame3_denoising.DenoisingStage.prepare_model_input

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

fastvideo.pipelines.stages.matrixgame3_denoising.DenoisingStage.prepare_timestep

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

fastvideo.pipelines.stages.matrixgame3_denoising.DenoisingStage.finish_step

finish_step(latents, state) -> torch.Tensor

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.DenoisingStage.verify_input

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

Verify denoising stage inputs.

fastvideo.pipelines.stages.matrixgame3_denoising.DenoisingStage.verify_output

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

Verify denoising stage outputs.

fastvideo.pipelines.stages.matrixgame3_denoising.V

V()

Common validators for pipeline stages.

fastvideo.pipelines.stages.matrixgame3_denoising.V.not_none

not_none(value: Any) -> bool

Check if value is not None.

fastvideo.pipelines.stages.matrixgame3_denoising.V.positive_int

positive_int(value: Any) -> bool

Check if value is a positive integer.

fastvideo.pipelines.stages.matrixgame3_denoising.V.non_negative_int

non_negative_int(value: Any) -> bool

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

fastvideo.pipelines.stages.matrixgame3_denoising.V.positive_float

positive_float(value: Any) -> bool

Check if value is a positive float.

fastvideo.pipelines.stages.matrixgame3_denoising.V.non_negative_float

non_negative_float(value: Any) -> bool

Check if value is a non-negative float.

fastvideo.pipelines.stages.matrixgame3_denoising.V.divisible_by

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

Check if value is divisible by divisor.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.V.list_not_empty

list_not_empty(value: Any) -> bool

Check if value is a non-empty list.

fastvideo.pipelines.stages.matrixgame3_denoising.V.list_length

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

Check if list has specific length.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_denoising.V.string_not_empty

string_not_empty(value: Any) -> bool

Check if value is a non-empty string.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_denoising.V.bool_value

bool_value(value: Any) -> bool

Check if value is a boolean.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_denoising.V.is_list

is_list(value: Any) -> bool

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

fastvideo.pipelines.stages.matrixgame3_denoising.V.is_tuple

is_tuple(value: Any) -> bool

Check if value is a tuple.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.V.none_or_list

none_or_list(value: Any) -> bool

Check if value is None or a list.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.V.divisible

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

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

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_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.matrixgame3_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.matrixgame3_denoising.VerificationResult

VerificationResult()

Wrapper class for stage verification results.

fastvideo.pipelines.stages.matrixgame3_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.matrixgame3_denoising.VerificationResult.is_valid

is_valid() -> bool

Check if all validations passed.

fastvideo.pipelines.stages.matrixgame3_denoising.VerificationResult.get_failed_fields

get_failed_fields() -> list[str]

Get list of fields that failed validation.

fastvideo.pipelines.stages.matrixgame3_denoising.VerificationResult.get_detailed_failures

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

Get detailed failure information for each failed field.

fastvideo.pipelines.stages.matrixgame3_denoising.VerificationResult.get_failure_summary

get_failure_summary() -> str

Get a comprehensive summary of all validation failures.

fastvideo.pipelines.stages.matrixgame3_denoising.VerificationResult.to_dict

to_dict() -> dict

Convert to dictionary for backward compatibility.

fastvideo.pipelines.stages.matrixgame3_denoising.MatrixGame3DenoisingStage

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

    def _infer_num_iterations(self, batch: ForwardBatch) -> int:
        if batch.num_iterations is not None:
            return batch.num_iterations
        if isinstance(batch.num_frames, int) and batch.num_frames > 57:
            return 1 + max(0, (batch.num_frames - 57 + 39) // 40)
        return 1

    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

    def forward(
        self,
        batch: ForwardBatch,
        fastvideo_args: FastVideoArgs,
    ) -> ForwardBatch:
        assert batch.latents is not None, "latents must be prepared before MatrixGame3 denoising"
        assert batch.image_latent is not None, "MatrixGame3 requires first-frame VAE latents"
        assert batch.prompt_embeds, "MatrixGame3 requires text embeddings"

        target_dtype = torch.bfloat16
        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
        device = batch.latents.device

        extra_step_kwargs = self.prepare_extra_func_kwargs(
            self.scheduler.step,
            {
                "generator": batch.generator,
                "eta": batch.eta,
            },
        )

        prompt_embeds = batch.prompt_embeds
        negative_prompt_embeds = batch.negative_prompt_embeds or []
        use_base_model = batch.use_base_model

        latents = batch.latents
        img_cond = batch.image_latent.to(device=device, dtype=target_dtype)
        latent_h = latents.shape[-2]
        latent_w = latents.shape[-1]
        patch_h = int(getattr(self.transformer, "patch_size", (1, 2, 2))[1])
        patch_w = int(getattr(self.transformer, "patch_size", (1, 2, 2))[2])
        latent_h_aligned = (latent_h // patch_h) * patch_h
        latent_w_aligned = (latent_w // patch_w) * patch_w
        if latent_h_aligned != latent_h or latent_w_aligned != latent_w:
            logger.warning(
                "Cropping MatrixGame3 latents to patch-aligned size: (%d, %d) -> (%d, %d), patch=(%d,%d)",
                latent_h,
                latent_w,
                latent_h_aligned,
                latent_w_aligned,
                patch_h,
                patch_w,
            )
            latents = latents[:, :, :, :latent_h_aligned, :latent_w_aligned]
            img_cond = img_cond[:, :, :, :latent_h_aligned, :latent_w_aligned]
            latent_h = latent_h_aligned
            latent_w = latent_w_aligned
        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
        target_h = latent_h * spatial_ratio
        target_w = latent_w * spatial_ratio
        num_iterations = self._infer_num_iterations(batch)
        clip_frame = 56  # hardcode for now
        first_clip_frame = clip_frame + 1
        past_frame = 16

        def align_frame_to_block(frame_idx: int) -> int:
            return (frame_idx - 1) // 4 * 4 + 1 if frame_idx > 0 else 1

        def get_latent_idx(frame_idx: int) -> int:
            return (frame_idx - 1) // 4 + 1

        total_video_frames = first_clip_frame + max(0, num_iterations - 1) * (clip_frame - past_frame)
        if batch.keyboard_cond is None or batch.mouse_cond is None:
            keyboard_cond, mouse_cond = build_matrixgame3_action_preset(total_video_frames, seed=batch.seed)
            batch.keyboard_cond = keyboard_cond.unsqueeze(0).to(device=device, dtype=target_dtype)
            batch.mouse_cond = mouse_cond.unsqueeze(0).to(device=device, dtype=target_dtype)
        else:
            batch.keyboard_cond = batch.keyboard_cond.to(device=device, dtype=target_dtype)
            batch.mouse_cond = batch.mouse_cond.to(device=device, dtype=target_dtype)

        # MG3 currently assumes batch_size=1 for camera-trajectory derivation. If a
        # future caller batches multiple action streams, build_extrinsics_from_actions
        # must be made per-batch-item; until then, fail loudly rather than silently
        # using item-0's trajectory for the whole batch.
        assert batch.keyboard_cond.shape[0] == 1, (
            "MatrixGame3 denoising stage requires batch_size=1 for action conditioning; "
            f"got batch.keyboard_cond.shape[0]={batch.keyboard_cond.shape[0]}. "
            "If batched action streams are needed, update build_extrinsics_from_actions "
            "to handle per-batch-item trajectories.")
        extrinsics_all = build_extrinsics_from_actions(batch.keyboard_cond[0], batch.mouse_cond[0]).to(device)
        all_latents: list[torch.Tensor] = []

        for clip_idx in range(num_iterations):
            try:
                self.scheduler.set_timesteps(
                    batch.num_inference_steps,
                    device=device,
                    shift=fastvideo_args.pipeline_config.flow_shift,
                )
            except TypeError:
                self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
            timesteps = self.scheduler.timesteps

            first_clip = clip_idx == 0
            current_end_frame_idx = first_clip_frame if first_clip else first_clip_frame + clip_idx * (clip_frame -
                                                                                                       past_frame)
            current_start_frame_idx = 0 if first_clip else current_end_frame_idx - clip_frame
            # 15 for first clip, 14 for later clips
            current_latent_frames = (first_clip_frame - 1) // 4 + 1 if first_clip else clip_frame // 4
            cond_frames = 1 if first_clip else 4
            latent_start_idx = get_latent_idx(current_start_frame_idx)
            latent_end_idx = get_latent_idx(current_end_frame_idx)

            clip_keyboard = batch.keyboard_cond[:, current_start_frame_idx:current_end_frame_idx]
            clip_mouse = batch.mouse_cond[:, current_start_frame_idx:current_end_frame_idx]

            cond_frames = min(cond_frames, img_cond.shape[2])
            clip_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
            current_latents = torch.randn(
                (latents.shape[0], latents.shape[1], latent_end_idx - latent_start_idx, latent_h, latent_w),
                generator=clip_generator if isinstance(clip_generator, torch.Generator) else None,
                dtype=target_dtype).to(device)
            current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames]

            c2ws_chunk = extrinsics_all[current_start_frame_idx:current_end_frame_idx]
            src_indices = np.linspace(current_start_frame_idx, current_end_frame_idx - 1,
                                      current_end_frame_idx - current_start_frame_idx)
            tgt_indices = np.linspace(0 if first_clip else current_start_frame_idx + 3, current_end_frame_idx - 1,
                                      current_latent_frames)

            if batch.c2ws_plucker_emb is not None:
                plucker_no_mem = batch.c2ws_plucker_emb[:, :, current_start_frame_idx:current_end_frame_idx]
                plucker_no_mem = plucker_no_mem.to(device=device, dtype=target_dtype)
            else:
                plucker_no_mem = build_plucker_from_c2ws(
                    c2ws_chunk,
                    src_indices,
                    tgt_indices,
                    target_h=target_h,
                    target_w=target_w,
                    latent_h=latent_h,
                    latent_w=latent_w,
                    framewise=True,
                ).to(device=device, dtype=target_dtype)

            x_memory = None
            timestep_memory = None
            mouse_cond_memory = None
            keyboard_cond_memory = None
            memory_latent_idx = None
            c2ws_plucker_emb = plucker_no_mem
            if all_latents:
                # t-1, t-9, t-17, t-25, t-33
                selected_index_base = [current_end_frame_idx - offset for offset in range(1, 34, 8)]
                selected_index = select_memory_idx_fov(
                    extrinsics_all,
                    current_start_frame_idx,
                    selected_index_base,
                    height=target_h,
                    width=target_w,
                )
                if len(selected_index) > 0:
                    selected_index[-1] = 4

                memory_pluckers: list[torch.Tensor] = []
                memory_latent_idx = []
                for mem_idx, reference_idx in zip(selected_index, selected_index_base, strict=False):
                    memory_latent_idx.append(get_latent_idx(mem_idx))

                    mem_idx_aligned = align_frame_to_block(mem_idx)
                    mem_block = extrinsics_all[mem_idx_aligned:mem_idx_aligned + 4]
                    mem_src = np.linspace(mem_idx_aligned, mem_idx_aligned + mem_block.shape[0] - 1, mem_block.shape[0])
                    mem_tgt = np.array([mem_idx_aligned + 3], dtype=np.float32)
                    mem_pose = interpolate_camera_poses_handedness(
                        src_indices=mem_src,
                        src_rot_mat=mem_block[:, :3, :3].detach().cpu().numpy(),
                        src_trans_vec=mem_block[:, :3, 3].detach().cpu().numpy(),
                        tgt_indices=mem_tgt,
                    ).to(device=device, dtype=target_dtype)
                    reference_pose = extrinsics_all[reference_idx:reference_idx + 1].to(device=device,
                                                                                        dtype=target_dtype)
                    rel_pair = torch.cat([reference_pose, mem_pose], dim=0)
                    rel_pose = compute_relative_poses(rel_pair, framewise=False)[1:2]
                    memory_pluckers.append(
                        build_plucker_from_pose(
                            rel_pose,
                            target_h=target_h,
                            target_w=target_w,
                            latent_h=latent_h,
                            latent_w=latent_w,
                        ).to(device=device, dtype=target_dtype))

                if memory_pluckers:
                    c2ws_plucker_emb = torch.cat(memory_pluckers + [plucker_no_mem], dim=2)
                    history = torch.cat(all_latents, dim=2)
                    x_memory = history[:, :, memory_latent_idx]
                    mouse_cond_memory = torch.ones(
                        (clip_keyboard.shape[0], len(memory_latent_idx), clip_mouse.shape[-1]),
                        device=device,
                        dtype=target_dtype)
                    keyboard_cond_memory = -torch.ones(
                        (clip_keyboard.shape[0], len(memory_latent_idx), clip_keyboard.shape[-1]),
                        device=device,
                        dtype=target_dtype)
                    timestep_memory = x_memory.new_zeros(
                        (x_memory.shape[0], x_memory.shape[2] * x_memory.shape[3] * x_memory.shape[4] // 4))
            with self.progress_bar(total=len(timesteps)) as progress_bar:
                for timestep in timesteps:
                    # Stop if interrupted
                    if getattr(self, "interrupt", False):
                        break
                    latent_model_input = current_latents
                    if hasattr(self.scheduler, "scale_model_input"):
                        latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)

                    timestep_tokens = latent_model_input.new_full(
                        (
                            latent_model_input.shape[2],
                            latent_model_input.shape[3] * latent_model_input.shape[4] // 4,
                        ),
                        timestep,
                    )
                    timestep_tokens[:cond_frames].zero_()
                    timestep_tokens = timestep_tokens.flatten().unsqueeze(0)
                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \
                        set_forward_context(current_timestep=timestep,
                                            attn_metadata=None,
                                            forward_batch=batch):
                        noise_pred = self.transformer(
                            latent_model_input,
                            prompt_embeds,
                            timestep_tokens,
                            mouse_cond=clip_mouse,
                            keyboard_cond=clip_keyboard,
                            x_memory=x_memory,
                            timestep_memory=timestep_memory,
                            mouse_cond_memory=mouse_cond_memory,
                            keyboard_cond_memory=keyboard_cond_memory,
                            c2ws_plucker_emb=c2ws_plucker_emb,
                            memory_latent_idx=memory_latent_idx,
                            predict_latent_idx=(latent_start_idx, latent_end_idx),
                        )

                        if use_base_model and batch.do_classifier_free_guidance and negative_prompt_embeds:
                            null_mouse_cond = torch.ones_like(clip_mouse)
                            null_keyboard_cond = -torch.ones_like(clip_keyboard)
                            noise_pred_uncond = self.transformer(
                                latent_model_input,
                                negative_prompt_embeds,
                                timestep_tokens,
                                mouse_cond=null_mouse_cond,
                                keyboard_cond=null_keyboard_cond,
                                x_memory=None,
                                timestep_memory=None,
                                mouse_cond_memory=None,
                                keyboard_cond_memory=None,
                                c2ws_plucker_emb=plucker_no_mem,
                                memory_latent_idx=None,
                                predict_latent_idx=(latent_start_idx, latent_end_idx),
                            )
                            noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred - noise_pred_uncond)

                    if noise_pred.shape != current_latents.shape:
                        aligned_t = min(noise_pred.shape[2], current_latents.shape[2])
                        aligned_h = min(noise_pred.shape[3], current_latents.shape[3])
                        aligned_w = min(noise_pred.shape[4], current_latents.shape[4])
                        logger.warning(
                            "Aligning noise/sample shapes before scheduler.step: noise=%s sample=%s -> (*,*,%d,%d,%d)",
                            tuple(noise_pred.shape),
                            tuple(current_latents.shape),
                            aligned_t,
                            aligned_h,
                            aligned_w,
                        )
                        noise_pred = noise_pred[:, :, :aligned_t, :aligned_h, :aligned_w]
                        current_latents = current_latents[:, :, :aligned_t, :aligned_h, :aligned_w]

                    current_latents = self.scheduler.step(noise_pred,
                                                          timestep,
                                                          current_latents,
                                                          **extra_step_kwargs,
                                                          return_dict=False)[0]
                    current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames]
                    progress_bar.update()

            img_cond = current_latents[:, :, -4:].detach()
            denoised_pred = current_latents if first_clip else current_latents[:, :, -10:]
            all_latents.append(denoised_pred.detach())

        batch.latents = torch.cat(all_latents, dim=2)
        return batch

fastvideo.pipelines.stages.matrixgame3_denoising.MatrixGame3DenoisingStage.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

fastvideo.pipelines.stages.matrixgame3_denoising.MatrixGame3DenoisingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(
    self,
    batch: ForwardBatch,
    fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
    assert batch.latents is not None, "latents must be prepared before MatrixGame3 denoising"
    assert batch.image_latent is not None, "MatrixGame3 requires first-frame VAE latents"
    assert batch.prompt_embeds, "MatrixGame3 requires text embeddings"

    target_dtype = torch.bfloat16
    autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
    device = batch.latents.device

    extra_step_kwargs = self.prepare_extra_func_kwargs(
        self.scheduler.step,
        {
            "generator": batch.generator,
            "eta": batch.eta,
        },
    )

    prompt_embeds = batch.prompt_embeds
    negative_prompt_embeds = batch.negative_prompt_embeds or []
    use_base_model = batch.use_base_model

    latents = batch.latents
    img_cond = batch.image_latent.to(device=device, dtype=target_dtype)
    latent_h = latents.shape[-2]
    latent_w = latents.shape[-1]
    patch_h = int(getattr(self.transformer, "patch_size", (1, 2, 2))[1])
    patch_w = int(getattr(self.transformer, "patch_size", (1, 2, 2))[2])
    latent_h_aligned = (latent_h // patch_h) * patch_h
    latent_w_aligned = (latent_w // patch_w) * patch_w
    if latent_h_aligned != latent_h or latent_w_aligned != latent_w:
        logger.warning(
            "Cropping MatrixGame3 latents to patch-aligned size: (%d, %d) -> (%d, %d), patch=(%d,%d)",
            latent_h,
            latent_w,
            latent_h_aligned,
            latent_w_aligned,
            patch_h,
            patch_w,
        )
        latents = latents[:, :, :, :latent_h_aligned, :latent_w_aligned]
        img_cond = img_cond[:, :, :, :latent_h_aligned, :latent_w_aligned]
        latent_h = latent_h_aligned
        latent_w = latent_w_aligned
    spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
    target_h = latent_h * spatial_ratio
    target_w = latent_w * spatial_ratio
    num_iterations = self._infer_num_iterations(batch)
    clip_frame = 56  # hardcode for now
    first_clip_frame = clip_frame + 1
    past_frame = 16

    def align_frame_to_block(frame_idx: int) -> int:
        return (frame_idx - 1) // 4 * 4 + 1 if frame_idx > 0 else 1

    def get_latent_idx(frame_idx: int) -> int:
        return (frame_idx - 1) // 4 + 1

    total_video_frames = first_clip_frame + max(0, num_iterations - 1) * (clip_frame - past_frame)
    if batch.keyboard_cond is None or batch.mouse_cond is None:
        keyboard_cond, mouse_cond = build_matrixgame3_action_preset(total_video_frames, seed=batch.seed)
        batch.keyboard_cond = keyboard_cond.unsqueeze(0).to(device=device, dtype=target_dtype)
        batch.mouse_cond = mouse_cond.unsqueeze(0).to(device=device, dtype=target_dtype)
    else:
        batch.keyboard_cond = batch.keyboard_cond.to(device=device, dtype=target_dtype)
        batch.mouse_cond = batch.mouse_cond.to(device=device, dtype=target_dtype)

    # MG3 currently assumes batch_size=1 for camera-trajectory derivation. If a
    # future caller batches multiple action streams, build_extrinsics_from_actions
    # must be made per-batch-item; until then, fail loudly rather than silently
    # using item-0's trajectory for the whole batch.
    assert batch.keyboard_cond.shape[0] == 1, (
        "MatrixGame3 denoising stage requires batch_size=1 for action conditioning; "
        f"got batch.keyboard_cond.shape[0]={batch.keyboard_cond.shape[0]}. "
        "If batched action streams are needed, update build_extrinsics_from_actions "
        "to handle per-batch-item trajectories.")
    extrinsics_all = build_extrinsics_from_actions(batch.keyboard_cond[0], batch.mouse_cond[0]).to(device)
    all_latents: list[torch.Tensor] = []

    for clip_idx in range(num_iterations):
        try:
            self.scheduler.set_timesteps(
                batch.num_inference_steps,
                device=device,
                shift=fastvideo_args.pipeline_config.flow_shift,
            )
        except TypeError:
            self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        first_clip = clip_idx == 0
        current_end_frame_idx = first_clip_frame if first_clip else first_clip_frame + clip_idx * (clip_frame -
                                                                                                   past_frame)
        current_start_frame_idx = 0 if first_clip else current_end_frame_idx - clip_frame
        # 15 for first clip, 14 for later clips
        current_latent_frames = (first_clip_frame - 1) // 4 + 1 if first_clip else clip_frame // 4
        cond_frames = 1 if first_clip else 4
        latent_start_idx = get_latent_idx(current_start_frame_idx)
        latent_end_idx = get_latent_idx(current_end_frame_idx)

        clip_keyboard = batch.keyboard_cond[:, current_start_frame_idx:current_end_frame_idx]
        clip_mouse = batch.mouse_cond[:, current_start_frame_idx:current_end_frame_idx]

        cond_frames = min(cond_frames, img_cond.shape[2])
        clip_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
        current_latents = torch.randn(
            (latents.shape[0], latents.shape[1], latent_end_idx - latent_start_idx, latent_h, latent_w),
            generator=clip_generator if isinstance(clip_generator, torch.Generator) else None,
            dtype=target_dtype).to(device)
        current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames]

        c2ws_chunk = extrinsics_all[current_start_frame_idx:current_end_frame_idx]
        src_indices = np.linspace(current_start_frame_idx, current_end_frame_idx - 1,
                                  current_end_frame_idx - current_start_frame_idx)
        tgt_indices = np.linspace(0 if first_clip else current_start_frame_idx + 3, current_end_frame_idx - 1,
                                  current_latent_frames)

        if batch.c2ws_plucker_emb is not None:
            plucker_no_mem = batch.c2ws_plucker_emb[:, :, current_start_frame_idx:current_end_frame_idx]
            plucker_no_mem = plucker_no_mem.to(device=device, dtype=target_dtype)
        else:
            plucker_no_mem = build_plucker_from_c2ws(
                c2ws_chunk,
                src_indices,
                tgt_indices,
                target_h=target_h,
                target_w=target_w,
                latent_h=latent_h,
                latent_w=latent_w,
                framewise=True,
            ).to(device=device, dtype=target_dtype)

        x_memory = None
        timestep_memory = None
        mouse_cond_memory = None
        keyboard_cond_memory = None
        memory_latent_idx = None
        c2ws_plucker_emb = plucker_no_mem
        if all_latents:
            # t-1, t-9, t-17, t-25, t-33
            selected_index_base = [current_end_frame_idx - offset for offset in range(1, 34, 8)]
            selected_index = select_memory_idx_fov(
                extrinsics_all,
                current_start_frame_idx,
                selected_index_base,
                height=target_h,
                width=target_w,
            )
            if len(selected_index) > 0:
                selected_index[-1] = 4

            memory_pluckers: list[torch.Tensor] = []
            memory_latent_idx = []
            for mem_idx, reference_idx in zip(selected_index, selected_index_base, strict=False):
                memory_latent_idx.append(get_latent_idx(mem_idx))

                mem_idx_aligned = align_frame_to_block(mem_idx)
                mem_block = extrinsics_all[mem_idx_aligned:mem_idx_aligned + 4]
                mem_src = np.linspace(mem_idx_aligned, mem_idx_aligned + mem_block.shape[0] - 1, mem_block.shape[0])
                mem_tgt = np.array([mem_idx_aligned + 3], dtype=np.float32)
                mem_pose = interpolate_camera_poses_handedness(
                    src_indices=mem_src,
                    src_rot_mat=mem_block[:, :3, :3].detach().cpu().numpy(),
                    src_trans_vec=mem_block[:, :3, 3].detach().cpu().numpy(),
                    tgt_indices=mem_tgt,
                ).to(device=device, dtype=target_dtype)
                reference_pose = extrinsics_all[reference_idx:reference_idx + 1].to(device=device,
                                                                                    dtype=target_dtype)
                rel_pair = torch.cat([reference_pose, mem_pose], dim=0)
                rel_pose = compute_relative_poses(rel_pair, framewise=False)[1:2]
                memory_pluckers.append(
                    build_plucker_from_pose(
                        rel_pose,
                        target_h=target_h,
                        target_w=target_w,
                        latent_h=latent_h,
                        latent_w=latent_w,
                    ).to(device=device, dtype=target_dtype))

            if memory_pluckers:
                c2ws_plucker_emb = torch.cat(memory_pluckers + [plucker_no_mem], dim=2)
                history = torch.cat(all_latents, dim=2)
                x_memory = history[:, :, memory_latent_idx]
                mouse_cond_memory = torch.ones(
                    (clip_keyboard.shape[0], len(memory_latent_idx), clip_mouse.shape[-1]),
                    device=device,
                    dtype=target_dtype)
                keyboard_cond_memory = -torch.ones(
                    (clip_keyboard.shape[0], len(memory_latent_idx), clip_keyboard.shape[-1]),
                    device=device,
                    dtype=target_dtype)
                timestep_memory = x_memory.new_zeros(
                    (x_memory.shape[0], x_memory.shape[2] * x_memory.shape[3] * x_memory.shape[4] // 4))
        with self.progress_bar(total=len(timesteps)) as progress_bar:
            for timestep in timesteps:
                # Stop if interrupted
                if getattr(self, "interrupt", False):
                    break
                latent_model_input = current_latents
                if hasattr(self.scheduler, "scale_model_input"):
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)

                timestep_tokens = latent_model_input.new_full(
                    (
                        latent_model_input.shape[2],
                        latent_model_input.shape[3] * latent_model_input.shape[4] // 4,
                    ),
                    timestep,
                )
                timestep_tokens[:cond_frames].zero_()
                timestep_tokens = timestep_tokens.flatten().unsqueeze(0)
                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \
                    set_forward_context(current_timestep=timestep,
                                        attn_metadata=None,
                                        forward_batch=batch):
                    noise_pred = self.transformer(
                        latent_model_input,
                        prompt_embeds,
                        timestep_tokens,
                        mouse_cond=clip_mouse,
                        keyboard_cond=clip_keyboard,
                        x_memory=x_memory,
                        timestep_memory=timestep_memory,
                        mouse_cond_memory=mouse_cond_memory,
                        keyboard_cond_memory=keyboard_cond_memory,
                        c2ws_plucker_emb=c2ws_plucker_emb,
                        memory_latent_idx=memory_latent_idx,
                        predict_latent_idx=(latent_start_idx, latent_end_idx),
                    )

                    if use_base_model and batch.do_classifier_free_guidance and negative_prompt_embeds:
                        null_mouse_cond = torch.ones_like(clip_mouse)
                        null_keyboard_cond = -torch.ones_like(clip_keyboard)
                        noise_pred_uncond = self.transformer(
                            latent_model_input,
                            negative_prompt_embeds,
                            timestep_tokens,
                            mouse_cond=null_mouse_cond,
                            keyboard_cond=null_keyboard_cond,
                            x_memory=None,
                            timestep_memory=None,
                            mouse_cond_memory=None,
                            keyboard_cond_memory=None,
                            c2ws_plucker_emb=plucker_no_mem,
                            memory_latent_idx=None,
                            predict_latent_idx=(latent_start_idx, latent_end_idx),
                        )
                        noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred - noise_pred_uncond)

                if noise_pred.shape != current_latents.shape:
                    aligned_t = min(noise_pred.shape[2], current_latents.shape[2])
                    aligned_h = min(noise_pred.shape[3], current_latents.shape[3])
                    aligned_w = min(noise_pred.shape[4], current_latents.shape[4])
                    logger.warning(
                        "Aligning noise/sample shapes before scheduler.step: noise=%s sample=%s -> (*,*,%d,%d,%d)",
                        tuple(noise_pred.shape),
                        tuple(current_latents.shape),
                        aligned_t,
                        aligned_h,
                        aligned_w,
                    )
                    noise_pred = noise_pred[:, :, :aligned_t, :aligned_h, :aligned_w]
                    current_latents = current_latents[:, :, :aligned_t, :aligned_h, :aligned_w]

                current_latents = self.scheduler.step(noise_pred,
                                                      timestep,
                                                      current_latents,
                                                      **extra_step_kwargs,
                                                      return_dict=False)[0]
                current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames]
                progress_bar.update()

        img_cond = current_latents[:, :, -4:].detach()
        denoised_pred = current_latents if first_clip else current_latents[:, :, -10:]
        all_latents.append(denoised_pred.detach())

    batch.latents = torch.cat(all_latents, dim=2)
    return batch