nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

read action from npy and assert action length

Browse Source
This commit is contained in:
H1yori233
2026-02-27 15:46:19 -08:00
parent 69dd3c68f6
commit 306ade1ecf
4 changed files with 63 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/dataset/validation_dataset.py
+32
View File
@@ -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
fastvideo/models/dits/matrixgame/causal_model.py
+6
View File
@@ -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
fastvideo/models/dits/matrixgame/model.py
+6
View File
@@ -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
fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py
+19 -59
View File
@@ -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: