233 lines
9.7 KiB
Python
233 lines
9.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# schema.py
|
|
"""
|
|
Unified data schema and format for saving and loading image/video data after
|
|
preprocessing.
|
|
|
|
It uses apache arrow in-memory format that can be consumed by modern data
|
|
frameworks that can handle parquet or lance file.
|
|
"""
|
|
|
|
import pyarrow as pa
|
|
|
|
pyarrow_schema_i2v = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Image/Video VAE latents ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("vae_latent_bytes", pa.binary()),
|
|
# e.g., [C, T, H, W] or [C, H, W]
|
|
pa.field("vae_latent_shape", pa.list_(pa.int64())),
|
|
# e.g., 'float32'
|
|
pa.field("vae_latent_dtype", pa.string()),
|
|
# --- Text encoder output tensor ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("text_embedding_bytes", pa.binary()),
|
|
# e.g., [SeqLen, Dim]
|
|
pa.field("text_embedding_shape", pa.list_(pa.int64())),
|
|
# e.g., 'bfloat16' or 'float32'
|
|
pa.field("text_embedding_dtype", pa.string()),
|
|
#I2V
|
|
pa.field("clip_feature_bytes", pa.binary()),
|
|
pa.field("clip_feature_shape", pa.list_(pa.int64())),
|
|
pa.field("clip_feature_dtype", pa.string()),
|
|
pa.field("first_frame_latent_bytes", pa.binary()),
|
|
pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
|
|
pa.field("first_frame_latent_dtype", pa.string()),
|
|
# I2V Validation
|
|
pa.field("pil_image_bytes", pa.binary()),
|
|
pa.field("pil_image_shape", pa.list_(pa.int64())),
|
|
pa.field("pil_image_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()), # 'image' or 'video'
|
|
pa.field("width", pa.int64()),
|
|
pa.field("height", pa.int64()),
|
|
# -- Video-specific (can be null/default for images) ---
|
|
# Number of frames processed (e.g., 1 for image, N for video)
|
|
pa.field("num_frames", pa.int64()),
|
|
pa.field("duration_sec", pa.float64()),
|
|
pa.field("fps", pa.float64()),
|
|
])
|
|
|
|
pyarrow_schema_i2v_track = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
pa.field("vae_latent_bytes", pa.binary()),
|
|
pa.field("vae_latent_shape", pa.list_(pa.int64())),
|
|
pa.field("vae_latent_dtype", pa.string()),
|
|
pa.field("text_embedding_bytes", pa.binary()),
|
|
pa.field("text_embedding_shape", pa.list_(pa.int64())),
|
|
pa.field("text_embedding_dtype", pa.string()),
|
|
pa.field("first_frame_latent_bytes", pa.binary()),
|
|
pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
|
|
pa.field("first_frame_latent_dtype", pa.string()),
|
|
pa.field("clip_feature_bytes", pa.binary()),
|
|
pa.field("clip_feature_shape", pa.list_(pa.int64())),
|
|
pa.field("clip_feature_dtype", pa.string()),
|
|
# MotionStream point tracks, normalized to [0, 1].
|
|
pa.field("track_points_bytes", pa.binary()),
|
|
pa.field("track_points_shape", pa.list_(pa.int64())), # [T, N, 2]
|
|
pa.field("track_points_dtype", pa.string()),
|
|
pa.field("track_visibility_bytes", pa.binary()),
|
|
pa.field("track_visibility_shape", pa.list_(pa.int64())), # [T, N]
|
|
pa.field("track_visibility_dtype", pa.string()),
|
|
# Optional preprocessing signals used by object-covering track sampling.
|
|
# New preprocessing writes dense fallback values when segmentation or
|
|
# informativeness weights are unavailable.
|
|
pa.field("object_ids_bytes", pa.binary()),
|
|
pa.field("object_ids_shape", pa.list_(pa.int64())), # [N]
|
|
pa.field("object_ids_dtype", pa.string()),
|
|
pa.field("track_weights_bytes", pa.binary()),
|
|
pa.field("track_weights_shape", pa.list_(pa.int64())), # [N]
|
|
pa.field("track_weights_dtype", pa.string()),
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()),
|
|
pa.field("width", pa.int64()),
|
|
pa.field("height", pa.int64()),
|
|
pa.field("num_frames", pa.int64()),
|
|
pa.field("duration_sec", pa.float64()),
|
|
pa.field("fps", pa.float64()),
|
|
])
|
|
|
|
|
|
pyarrow_schema_t2v = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Image/Video VAE latents ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("vae_latent_bytes", pa.binary()),
|
|
# e.g., [C, T, H, W] or [C, H, W]
|
|
pa.field("vae_latent_shape", pa.list_(pa.int64())),
|
|
# e.g., 'float32'
|
|
pa.field("vae_latent_dtype", pa.string()),
|
|
# --- Text encoder output tensor ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("text_embedding_bytes", pa.binary()),
|
|
# e.g., [SeqLen, Dim]
|
|
pa.field("text_embedding_shape", pa.list_(pa.int64())),
|
|
# e.g., 'bfloat16' or 'float32'
|
|
pa.field("text_embedding_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()), # 'image' or 'video'
|
|
pa.field("width", pa.int64()),
|
|
pa.field("height", pa.int64()),
|
|
# -- Video-specific (can be null/default for images) ---
|
|
# Number of frames processed (e.g., 1 for image, N for video)
|
|
pa.field("num_frames", pa.int64()),
|
|
pa.field("duration_sec", pa.float64()),
|
|
pa.field("fps", pa.float64()),
|
|
])
|
|
|
|
|
|
pyarrow_schema_ode_trajectory_text_only = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Text encoder output tensor ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("text_embedding_bytes", pa.binary()),
|
|
# e.g., [SeqLen, Dim]
|
|
pa.field("text_embedding_shape", pa.list_(pa.int64())),
|
|
# e.g., 'bfloat16' or 'float32'
|
|
pa.field("text_embedding_dtype", pa.string()),
|
|
# --- ODE Trajectory ---
|
|
pa.field("trajectory_latents_bytes", pa.binary()),
|
|
pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
|
|
pa.field("trajectory_latents_dtype", pa.string()),
|
|
pa.field("trajectory_timesteps_bytes", pa.binary()),
|
|
pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
|
|
pa.field("trajectory_timesteps_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()), # Always 'text' for text-only
|
|
])
|
|
|
|
|
|
pyarrow_schema_text_only = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Text encoder output tensor ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("text_embedding_bytes", pa.binary()),
|
|
# e.g., [SeqLen, Dim]
|
|
pa.field("text_embedding_shape", pa.list_(pa.int64())),
|
|
# e.g., 'bfloat16' or 'float32'
|
|
pa.field("text_embedding_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("caption", pa.string()),
|
|
])
|
|
|
|
|
|
pyarrow_schema_matrixgame2 = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Image/Video VAE latents ---
|
|
# Tensors are stored as raw bytes with shape and dtype info for loading
|
|
pa.field("vae_latent_bytes", pa.binary()),
|
|
# e.g., [C, T, H, W] or [C, H, W]
|
|
pa.field("vae_latent_shape", pa.list_(pa.int64())),
|
|
# e.g., 'float32'
|
|
pa.field("vae_latent_dtype", pa.string()),
|
|
#I2V
|
|
pa.field("clip_feature_bytes", pa.binary()),
|
|
pa.field("clip_feature_shape", pa.list_(pa.int64())),
|
|
pa.field("clip_feature_dtype", pa.string()),
|
|
pa.field("first_frame_latent_bytes", pa.binary()),
|
|
pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
|
|
pa.field("first_frame_latent_dtype", pa.string()),
|
|
# --- Action ---
|
|
pa.field("mouse_cond_bytes", pa.binary()),
|
|
pa.field("mouse_cond_shape", pa.list_(pa.int64())), # [T, 2]
|
|
pa.field("mouse_cond_dtype", pa.string()),
|
|
pa.field("keyboard_cond_bytes", pa.binary()),
|
|
pa.field("keyboard_cond_shape", pa.list_(pa.int64())), # [T, 4]
|
|
pa.field("keyboard_cond_dtype", pa.string()),
|
|
# I2V Validation
|
|
pa.field("pil_image_bytes", pa.binary()),
|
|
pa.field("pil_image_shape", pa.list_(pa.int64())),
|
|
pa.field("pil_image_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()), # 'image' or 'video'
|
|
pa.field("width", pa.int64()),
|
|
pa.field("height", pa.int64()),
|
|
# -- Video-specific (can be null/default for images) ---
|
|
# Number of frames processed (e.g., 1 for image, N for video)
|
|
pa.field("num_frames", pa.int64()),
|
|
pa.field("duration_sec", pa.float64()),
|
|
pa.field("fps", pa.float64()),
|
|
])
|
|
|
|
pyarrow_schema_matrixgame2_ode_trajectory = pa.schema([
|
|
pa.field("id", pa.string()),
|
|
# --- Action ---
|
|
pa.field("mouse_cond_bytes", pa.binary()),
|
|
pa.field("mouse_cond_shape", pa.list_(pa.int64())), # [T, 2]
|
|
pa.field("mouse_cond_dtype", pa.string()),
|
|
pa.field("keyboard_cond_bytes", pa.binary()),
|
|
pa.field("keyboard_cond_shape", pa.list_(pa.int64())), # [T, 4]
|
|
pa.field("keyboard_cond_dtype", pa.string()),
|
|
#I2V
|
|
pa.field("clip_feature_bytes", pa.binary()),
|
|
pa.field("clip_feature_shape", pa.list_(pa.int64())),
|
|
pa.field("clip_feature_dtype", pa.string()),
|
|
pa.field("first_frame_latent_bytes", pa.binary()),
|
|
pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
|
|
pa.field("first_frame_latent_dtype", pa.string()),
|
|
# I2V Validation
|
|
pa.field("pil_image_bytes", pa.binary()),
|
|
pa.field("pil_image_shape", pa.list_(pa.int64())),
|
|
pa.field("pil_image_dtype", pa.string()),
|
|
# --- ODE Trajectory ---
|
|
pa.field("trajectory_latents_bytes", pa.binary()),
|
|
pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
|
|
pa.field("trajectory_latents_dtype", pa.string()),
|
|
pa.field("trajectory_timesteps_bytes", pa.binary()),
|
|
pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
|
|
pa.field("trajectory_timesteps_dtype", pa.string()),
|
|
# --- Metadata ---
|
|
pa.field("file_name", pa.string()),
|
|
pa.field("caption", pa.string()),
|
|
pa.field("media_type", pa.string()), # 'image' or 'video'
|
|
])
|