diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 5ab467d75..ec937520a 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -4,6 +4,7 @@ import os import pathlib import datasets +import numpy as np from torch.utils.data import IterableDataset from fastvideo.distributed import (get_sp_world_size, get_world_rank, @@ -160,5 +161,36 @@ class ValidationDataset(IterableDataset): else: sample["control_video"] = load_video(control_video_path) + if sample.get("action_path", None) is not None: + action_path = sample["action_path"] + action_path = os.path.join(self.dir, action_path) + sample["action_path"] = action_path + if not pathlib.Path(action_path).is_file(): + logger.warning("Action file %s does not exist.", + action_path) + else: + action_data = np.load(action_path, allow_pickle=True) + num_frames = sample["num_frames"] + if isinstance(action_data, + np.ndarray) and action_data.dtype == object: + action_data = action_data.item() + if isinstance(action_data, dict): + kb_length = len(action_data["keyboard"]) + ms_length = len(action_data["mouse"]) + assert kb_length >= num_frames, ( + f"keyboard length {kb_length} is smaller than " + f"num_frames {num_frames} for {action_path}.") + assert ms_length >= num_frames, ( + f"mouse length {ms_length} is smaller than " + f"num_frames {num_frames} for {action_path}.") + sample["keyboard_cond"] = action_data["keyboard"][:num_frames] + sample["mouse_cond"] = action_data["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}.") + sample["keyboard_cond"] = action_data[:num_frames] + sample = {k: v for k, v in sample.items() if v is not None} yield sample diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 99bb8e134..dd999e94b 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -487,6 +487,12 @@ class CausalMatrixGameTransformerBlock(nn.Module): assert temb.ndim == 4 num_frames = temb.shape[1] + + if mouse_cond is not None: + assert mouse_cond.shape[1] == num_frames, f"mouse_cond length ({mouse_cond.shape[1]}) must be exactly num_frames ({num_frames})" + if keyboard_cond is not None: + assert keyboard_cond.shape[1] == num_frames, f"keyboard_cond length ({keyboard_cond.shape[1]}) must be exactly num_frames ({num_frames})" + frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype diff --git a/fastvideo/models/dits/matrixgame/model.py b/fastvideo/models/dits/matrixgame/model.py index 1fa520c6b..14e840d6c 100644 --- a/fastvideo/models/dits/matrixgame/model.py +++ b/fastvideo/models/dits/matrixgame/model.py @@ -231,6 +231,12 @@ class MatrixGameTransformerBlock(nn.Module): mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, ) -> torch.Tensor: + num_frames = int(grid_sizes[0]) + if mouse_cond is not None: + assert mouse_cond.shape[1] == num_frames, f"mouse_cond length ({mouse_cond.shape[1]}) must be exactly num_frames ({num_frames})" + if keyboard_cond is not None: + assert keyboard_cond.shape[1] == num_frames, f"keyboard_cond length ({keyboard_cond.shape[1]}) must be exactly num_frames ({num_frames})" + if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) orig_dtype = hidden_states.dtype diff --git a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py index 5079ec264..c762d03d3 100644 --- a/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py +++ b/fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py @@ -132,7 +132,6 @@ class PreprocessPipeline_MatrixGame(BasePreprocessPipeline): if "action_path" in valid_data and valid_data["action_path"]: keyboard_cond_list = [] mouse_cond_list = [] - num_bits = 6 for action_path in valid_data["action_path"]: if action_path: action_data = np.load(action_path, allow_pickle=True) @@ -141,66 +140,27 @@ class PreprocessPipeline_MatrixGame(BasePreprocessPipeline): np.ndarray) and action_data.dtype == np.dtype('O'): action_dict = action_data.item() if "keyboard" in action_dict: - keyboard_raw = action_dict["keyboard"] - # Convert 1D bit-flag values to 2D multi-hot encoding - if isinstance(keyboard_raw, np.ndarray): - if keyboard_raw.ndim == 1: - # [T] -> [T, num_bits] - T = len(keyboard_raw) - multi_hot = np.zeros((T, num_bits), - dtype=np.float32) - action_values = keyboard_raw.astype(int) - for bit_idx in range(num_bits): - target_idx = ( - 2 - - (bit_idx % 3)) + 3 * (bit_idx // 3) - if target_idx < num_bits: - multi_hot[:, target_idx] = ( - (action_values >> bit_idx) - & 1).astype(np.float32) - keyboard_cond_list.append(multi_hot) - else: - # If already 2D, pad to num_bits if necessary - k_data = keyboard_raw.astype(np.float32) - if k_data.ndim == 2 and k_data.shape[ - -1] < num_bits: - padding = np.zeros( - (k_data.shape[0], - num_bits - k_data.shape[-1]), - dtype=np.float32) - k_data = np.concatenate( - [k_data, padding], axis=-1) - keyboard_cond_list.append(k_data) - else: - keyboard_cond_list.append(keyboard_raw) + 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: - mouse_cond_list.append(action_dict["mouse"]) + 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: - if isinstance(action_data, - np.ndarray) and action_data.ndim == 1: - T = len(action_data) - multi_hot = np.zeros((T, num_bits), - dtype=np.float32) - action_values = action_data.astype(int) - for bit_idx in range(num_bits): - target_idx = ( - 2 - (bit_idx % 3)) + 3 * (bit_idx // 3) - if target_idx < num_bits: - multi_hot[:, target_idx] = ( - (action_values >> bit_idx) & 1).astype( - np.float32) - keyboard_cond_list.append(multi_hot) - else: - # If already 2D, pad to num_bits if necessary - k_data = action_data.astype(np.float32) - if k_data.ndim == 2 and k_data.shape[-1] < num_bits: - padding = np.zeros( - (k_data.shape[0], - num_bits - k_data.shape[-1]), - dtype=np.float32) - k_data = np.concatenate([k_data, padding], - axis=-1) - keyboard_cond_list.append(k_data) + 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: