matrixgame2_preprocess_pipeline
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn 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: boolWhether 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) -> FlexibleArgumentParserfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.from_cli_args
from_cli_args(args: argparse.Namespace) -> FastVideoArgsfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.from_kwargs
from_kwargs(**kwargs: Any) -> FastVideoArgsfastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.check_fastvideo_args
check_fastvideo_args() -> NoneValidate inference arguments for consistency
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.FastVideoArgs.finalize_device_offload_policy
finalize_device_offload_policy(device_id: int = 0) -> boolApply 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) -> boolDisable 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.SchemaReturn 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: PreprocessBatchThe batch containing the data to extract
schema: pa.SchemaPyArrow schema defining the expected fields
strict: boolIf True, raises an exception when required fields are missing or unfilled
Returns
dict[str, Any]Dictionary record matching the schema
Raises
ValueErrorIf 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_loggingWhether to enable logging for this stage.
is_secondaryWhether this is a secondary image encoder.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchEncode the prompt into image encoder hidden states.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe batch with encoded prompt embeddings.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify image encoding stage inputs.
fastvideo.pipelines.preprocess.matrixgame2.matrixgame2_preprocess_pipeline.ImageEncodingStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify 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</code></pre></details>
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</code></pre></details>
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</code></pre></details>