`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.get_local_torch_device`

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs`

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.is_output_rank`

is_output_rank: bool

Whether this process may materialize user-facing outputs.

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.add_cli_args`

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.from_cli_args`

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.from_kwargs`

from_kwargs(**kwargs: Any) -> FastVideoArgs

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.check_fastvideo_args`

check_fastvideo_args() -> None

Validate inference arguments for consistency

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.finalize_device_offload_policy`

finalize_device_offload_policy(device_id: int = 0) -> bool

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.disable_offload_on_unified_memory`

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

Disable host offload after a worker has selected its device.

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

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.set_forward_context`

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline`

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

Base class for preprocessing pipelines that handles common functionality.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Set up pipeline stages with proper dependency injection.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.forward`

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs, args)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.get_extra_features`

get_extra_features(valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]

Get additional features specific to the pipeline type. Override in subclasses.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.get_pyarrow_schema`

get_pyarrow_schema() -> pa.Schema

Return the PyArrow schema for this pipeline. Must be overridden.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.get_schema_fields`

get_schema_fields() -> list[str]

Get the schema fields for the pipeline type.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.create_record_for_schema`

create_record_for_schema(preprocess_batch: PreprocessBatch, schema: pa.Schema, strict: bool = False) -> dict[str, Any]

Create a record for the Parquet dataset using a generic schema-based approach.

Parameters

preprocess_batch: PreprocessBatch: The batch containing the data to extract
schema: pa.Schema: PyArrow schema defining the expected fields
strict: bool: If True, raises an exception when required fields are missing or unfilled

Returns

dict[str, Any]: Dictionary record matching the schema

Raises

ValueError: If strict=True and required fields are missing or unfilled

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.create_record`

create_record(video_name: str, vae_latent: np.ndarray, text_embedding: np.ndarray, valid_data: dict[str, Any], idx: int, extra_features: dict[str, Any] | None = None) -> dict[str, Any]

Create a record for the Parquet dataset.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.BasePreprocessPipeline.preprocess_video_and_text`

preprocess_video_and_text(fastvideo_args: FastVideoArgs, args)

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage`

ImageEncodingStage(image_encoder, image_processor)

Stage for encoding image prompts into embeddings for diffusion models.

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

Initialize the prompt encoding stage.

Parameters

enable_logging: Whether to enable logging for this stage.
is_secondary: Whether this is a secondary image encoder.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.forward`

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

Encode the prompt into image encoder hidden states.

Parameters

batch: ForwardBatch: The current batch information.
fastvideo_args: FastVideoArgs: The inference arguments.

Returns

ForwardBatch: The batch with encoded prompt embeddings.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.verify_input`

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

Verify image encoding stage inputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.verify_output`

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

Verify image encoding stage outputs.

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.PreprocessPipeline_MatrixGame2`

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

I2V preprocessing pipeline implementation.

Source

class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline):
    """I2V preprocessing pipeline implementation."""

    _required_config_modules = ["vae", "image_encoder", "image_processor"]

    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
        self.add_stage(stage_name="image_encoding_stage",
                       stage=ImageEncodingStage(
                           image_encoder=self.get_module("image_encoder"),
                           image_processor=self.get_module("image_processor"),
                       ))

    def get_pyarrow_schema(self):
        """Return the PyArrow schema for I2V pipeline."""
        return pyarrow_schema_matrixgame2

    def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:

        # TODO(will): move these to cpu at some point
        self.get_module("image_encoder").to(get_local_torch_device())
        self.get_module("vae").to(get_local_torch_device())

        features = {}
        """Get CLIP features from the first frame of each video."""
        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
        _, _, num_frames, height, width = valid_data["pixel_values"].shape
        # latent_height = height // self.get_module(
        #     "vae").spatial_compression_ratio
        # latent_width = width // self.get_module("vae").spatial_compression_ratio

        processed_images = []
        # Frame has values between -1 and 1
        for frame in first_frame:
            frame = (frame + 1) * 127.5
            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
            processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt")
            processed_images.append(processed_img)

        # Get CLIP features
        pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device())
        with torch.no_grad():
            image_inputs = {'pixel_values': pixel_values}
            with set_forward_context(current_timestep=0, attn_metadata=None):
                clip_features = self.get_module("image_encoder")(**image_inputs)
            clip_features = clip_features.last_hidden_state

        features["clip_feature"] = clip_features
        # Get CLIP features from the first frame of each video.
        video_conditions = []
        for frame in first_frame:
            processed_img = frame.to(device="cpu", dtype=torch.float32)
            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
            # (B, H, W, C) -> (B, C, 1, H, W)
            video_condition = torch.cat([
                processed_img,
                processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1], num_frames - 1, height, width)
            ],
                                        dim=2)
            video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32)
            video_conditions.append(video_condition)

        video_conditions = torch.cat(video_conditions, dim=0)

        with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
            encoder_outputs = self.get_module("vae").encode(video_conditions)

        latent_condition = encoder_outputs.mean
        if (hasattr(self.get_module("vae"), "shift_factor") and self.get_module("vae").shift_factor is not None):
            if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
                latent_condition -= self.get_module("vae").shift_factor.to(latent_condition.device,
                                                                           latent_condition.dtype)
            else:
                latent_condition -= self.get_module("vae").shift_factor

        if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
            latent_condition = latent_condition * self.get_module("vae").scaling_factor.to(
                latent_condition.device, latent_condition.dtype)
        else:
            latent_condition = latent_condition * self.get_module("vae").scaling_factor

        # mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
        #                            latent_width)
        # mask_lat_size[:, :, list(range(1, num_frames))] = 0
        # first_frame_mask = mask_lat_size[:, :, 0:1]
        # first_frame_mask = torch.repeat_interleave(
        #     first_frame_mask,
        #     dim=2,
        #     repeats=self.get_module("vae").temporal_compression_ratio)
        # mask_lat_size = torch.concat(
        #     [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
        # mask_lat_size = mask_lat_size.view(
        #     batch_size, -1,
        #     self.get_module("vae").temporal_compression_ratio, latent_height,
        #     latent_width)
        # mask_lat_size = mask_lat_size.transpose(1, 2)
        # mask_lat_size = mask_lat_size.to(latent_condition.device)

        # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)

        features["first_frame_latent"] = latent_condition

        if "action_path" in valid_data and valid_data["action_path"]:
            keyboard_cond_list = []
            mouse_cond_list = []
            for action_path in valid_data["action_path"]:
                if action_path:
                    action_data = np.load(action_path, allow_pickle=True)
                    if isinstance(action_data, np.ndarray) and action_data.dtype == np.dtype('O'):
                        action_dict = action_data.item()
                        if "keyboard" in action_dict:
                            kb_length = len(action_dict["keyboard"])
                            assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than "
                                                             f"num_frames {num_frames} for {action_path}.")
                            keyboard_cond_list.append(action_dict["keyboard"][:num_frames].astype(np.float32))
                        if "mouse" in action_dict:
                            ms_length = len(action_dict["mouse"])
                            assert ms_length >= num_frames, (f"mouse length {ms_length} is smaller than "
                                                             f"num_frames {num_frames} for {action_path}.")
                            mouse_cond_list.append(action_dict["mouse"][:num_frames])
                    else:
                        kb_length = len(action_data)
                        assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than "
                                                         f"num_frames {num_frames} for {action_path}.")
                        keyboard_cond_list.append(action_data.astype(np.float32)[:num_frames])
            if keyboard_cond_list:
                features["keyboard_cond"] = keyboard_cond_list
            if mouse_cond_list:
                features["mouse_cond"] = mouse_cond_list

        return features

    def create_record(self,
                      video_name: str,
                      vae_latent: np.ndarray,
                      text_embedding: np.ndarray,
                      valid_data: dict[str, Any],
                      idx: int,
                      extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
        """Create a record for the Parquet dataset with CLIP features."""
        record = super().create_record(video_name=video_name,
                                       vae_latent=vae_latent,
                                       text_embedding=text_embedding,
                                       valid_data=valid_data,
                                       idx=idx,
                                       extra_features=extra_features)

        if extra_features and "clip_feature" in extra_features:
            clip_feature = extra_features["clip_feature"]
            record.update({
                "clip_feature_bytes": clip_feature.tobytes(),
                "clip_feature_shape": list(clip_feature.shape),
                "clip_feature_dtype": str(clip_feature.dtype),
            })
        else:
            record.update({
                "clip_feature_bytes": b"",
                "clip_feature_shape": [],
                "clip_feature_dtype": "",
            })

        if extra_features and "first_frame_latent" in extra_features:
            first_frame_latent = extra_features["first_frame_latent"]
            record.update({
                "first_frame_latent_bytes": first_frame_latent.tobytes(),
                "first_frame_latent_shape": list(first_frame_latent.shape),
                "first_frame_latent_dtype": str(first_frame_latent.dtype),
            })
        else:
            record.update({
                "first_frame_latent_bytes": b"",
                "first_frame_latent_shape": [],
                "first_frame_latent_dtype": "",
            })

        if extra_features and "pil_image" in extra_features:
            pil_image = extra_features["pil_image"]
            record.update({
                "pil_image_bytes": pil_image.tobytes(),
                "pil_image_shape": list(pil_image.shape),
                "pil_image_dtype": str(pil_image.dtype),
            })
        else:
            record.update({
                "pil_image_bytes": b"",
                "pil_image_shape": [],
                "pil_image_dtype": "",
            })

        if extra_features and "keyboard_cond" in extra_features:
            keyboard_cond = extra_features["keyboard_cond"]
            record.update({
                "keyboard_cond_bytes": keyboard_cond.tobytes(),
                "keyboard_cond_shape": list(keyboard_cond.shape),
                "keyboard_cond_dtype": str(keyboard_cond.dtype),
            })
        else:
            record.update({
                "keyboard_cond_bytes": b"",
                "keyboard_cond_shape": [],
                "keyboard_cond_dtype": "",
            })

        if extra_features and "mouse_cond" in extra_features:
            mouse_cond = extra_features["mouse_cond"]
            record.update({
                "mouse_cond_bytes": mouse_cond.tobytes(),
                "mouse_cond_shape": list(mouse_cond.shape),
                "mouse_cond_dtype": str(mouse_cond.dtype),
            })
        else:
            record.update({
                "mouse_cond_bytes": b"",
                "mouse_cond_shape": [],
                "mouse_cond_dtype": "",
            })

        return record

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.PreprocessPipeline_MatrixGame2.create_pipeline_stages`

create_pipeline_stages(fastvideo_args: FastVideoArgs)

Source

def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
    self.add_stage(stage_name="image_encoding_stage",
                   stage=ImageEncodingStage(
                       image_encoder=self.get_module("image_encoder"),
                       image_processor=self.get_module("image_processor"),
                   ))

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.PreprocessPipeline_MatrixGame2.get_pyarrow_schema`

get_pyarrow_schema()

Return the PyArrow schema for I2V pipeline.

Source

def get_pyarrow_schema(self):
    """Return the PyArrow schema for I2V pipeline."""
    return pyarrow_schema_matrixgame2

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.PreprocessPipeline_MatrixGame2.get_extra_features`

get_extra_features(valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]

Source

def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]:

    # TODO(will): move these to cpu at some point
    self.get_module("image_encoder").to(get_local_torch_device())
    self.get_module("vae").to(get_local_torch_device())

    features = {}
    """Get CLIP features from the first frame of each video."""
    first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
    _, _, num_frames, height, width = valid_data["pixel_values"].shape
    # latent_height = height // self.get_module(
    #     "vae").spatial_compression_ratio
    # latent_width = width // self.get_module("vae").spatial_compression_ratio

    processed_images = []
    # Frame has values between -1 and 1
    for frame in first_frame:
        frame = (frame + 1) * 127.5
        frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
        processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt")
        processed_images.append(processed_img)

    # Get CLIP features
    pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device())
    with torch.no_grad():
        image_inputs = {'pixel_values': pixel_values}
        with set_forward_context(current_timestep=0, attn_metadata=None):
            clip_features = self.get_module("image_encoder")(**image_inputs)
        clip_features = clip_features.last_hidden_state

    features["clip_feature"] = clip_features
    # Get CLIP features from the first frame of each video.
    video_conditions = []
    for frame in first_frame:
        processed_img = frame.to(device="cpu", dtype=torch.float32)
        processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
        # (B, H, W, C) -> (B, C, 1, H, W)
        video_condition = torch.cat([
            processed_img,
            processed_img.new_zeros(processed_img.shape[0], processed_img.shape[1], num_frames - 1, height, width)
        ],
                                    dim=2)
        video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32)
        video_conditions.append(video_condition)

    video_conditions = torch.cat(video_conditions, dim=0)

    with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
        encoder_outputs = self.get_module("vae").encode(video_conditions)

    latent_condition = encoder_outputs.mean
    if (hasattr(self.get_module("vae"), "shift_factor") and self.get_module("vae").shift_factor is not None):
        if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
            latent_condition -= self.get_module("vae").shift_factor.to(latent_condition.device,
                                                                       latent_condition.dtype)
        else:
            latent_condition -= self.get_module("vae").shift_factor

    if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
        latent_condition = latent_condition * self.get_module("vae").scaling_factor.to(
            latent_condition.device, latent_condition.dtype)
    else:
        latent_condition = latent_condition * self.get_module("vae").scaling_factor

    # mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
    #                            latent_width)
    # mask_lat_size[:, :, list(range(1, num_frames))] = 0
    # first_frame_mask = mask_lat_size[:, :, 0:1]
    # first_frame_mask = torch.repeat_interleave(
    #     first_frame_mask,
    #     dim=2,
    #     repeats=self.get_module("vae").temporal_compression_ratio)
    # mask_lat_size = torch.concat(
    #     [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
    # mask_lat_size = mask_lat_size.view(
    #     batch_size, -1,
    #     self.get_module("vae").temporal_compression_ratio, latent_height,
    #     latent_width)
    # mask_lat_size = mask_lat_size.transpose(1, 2)
    # mask_lat_size = mask_lat_size.to(latent_condition.device)

    # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)

    features["first_frame_latent"] = latent_condition

    if "action_path" in valid_data and valid_data["action_path"]:
        keyboard_cond_list = []
        mouse_cond_list = []
        for action_path in valid_data["action_path"]:
            if action_path:
                action_data = np.load(action_path, allow_pickle=True)
                if isinstance(action_data, np.ndarray) and action_data.dtype == np.dtype('O'):
                    action_dict = action_data.item()
                    if "keyboard" in action_dict:
                        kb_length = len(action_dict["keyboard"])
                        assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than "
                                                         f"num_frames {num_frames} for {action_path}.")
                        keyboard_cond_list.append(action_dict["keyboard"][:num_frames].astype(np.float32))
                    if "mouse" in action_dict:
                        ms_length = len(action_dict["mouse"])
                        assert ms_length >= num_frames, (f"mouse length {ms_length} is smaller than "
                                                         f"num_frames {num_frames} for {action_path}.")
                        mouse_cond_list.append(action_dict["mouse"][:num_frames])
                else:
                    kb_length = len(action_data)
                    assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than "
                                                     f"num_frames {num_frames} for {action_path}.")
                    keyboard_cond_list.append(action_data.astype(np.float32)[:num_frames])
        if keyboard_cond_list:
            features["keyboard_cond"] = keyboard_cond_list
        if mouse_cond_list:
            features["mouse_cond"] = mouse_cond_list

    return features

`fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.PreprocessPipeline_MatrixGame2.create_record`

create_record(video_name: str, vae_latent: np.ndarray, text_embedding: np.ndarray, valid_data: dict[str, Any], idx: int, extra_features: dict[str, Any] | None = None) -> dict[str, Any]

Create a record for the Parquet dataset with CLIP features.

Source

def create_record(self,
                  video_name: str,
                  vae_latent: np.ndarray,
                  text_embedding: np.ndarray,
                  valid_data: dict[str, Any],
                  idx: int,
                  extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
    """Create a record for the Parquet dataset with CLIP features."""
    record = super().create_record(video_name=video_name,
                                   vae_latent=vae_latent,
                                   text_embedding=text_embedding,
                                   valid_data=valid_data,
                                   idx=idx,
                                   extra_features=extra_features)

    if extra_features and "clip_feature" in extra_features:
        clip_feature = extra_features["clip_feature"]
        record.update({
            "clip_feature_bytes": clip_feature.tobytes(),
            "clip_feature_shape": list(clip_feature.shape),
            "clip_feature_dtype": str(clip_feature.dtype),
        })
    else:
        record.update({
            "clip_feature_bytes": b"",
            "clip_feature_shape": [],
            "clip_feature_dtype": "",
        })

    if extra_features and "first_frame_latent" in extra_features:
        first_frame_latent = extra_features["first_frame_latent"]
        record.update({
            "first_frame_latent_bytes": first_frame_latent.tobytes(),
            "first_frame_latent_shape": list(first_frame_latent.shape),
            "first_frame_latent_dtype": str(first_frame_latent.dtype),
        })
    else:
        record.update({
            "first_frame_latent_bytes": b"",
            "first_frame_latent_shape": [],
            "first_frame_latent_dtype": "",
        })

    if extra_features and "pil_image" in extra_features:
        pil_image = extra_features["pil_image"]
        record.update({
            "pil_image_bytes": pil_image.tobytes(),
            "pil_image_shape": list(pil_image.shape),
            "pil_image_dtype": str(pil_image.dtype),
        })
    else:
        record.update({
            "pil_image_bytes": b"",
            "pil_image_shape": [],
            "pil_image_dtype": "",
        })

    if extra_features and "keyboard_cond" in extra_features:
        keyboard_cond = extra_features["keyboard_cond"]
        record.update({
            "keyboard_cond_bytes": keyboard_cond.tobytes(),
            "keyboard_cond_shape": list(keyboard_cond.shape),
            "keyboard_cond_dtype": str(keyboard_cond.dtype),
        })
    else:
        record.update({
            "keyboard_cond_bytes": b"",
            "keyboard_cond_shape": [],
            "keyboard_cond_dtype": "",
        })

    if extra_features and "mouse_cond" in extra_features:
        mouse_cond = extra_features["mouse_cond"]
        record.update({
            "mouse_cond_bytes": mouse_cond.tobytes(),
            "mouse_cond_shape": list(mouse_cond.shape),
            "mouse_cond_dtype": str(mouse_cond.dtype),
        })
    else:
        record.update({
            "mouse_cond_bytes": b"",
            "mouse_cond_shape": [],
            "mouse_cond_dtype": "",
        })

    return record