Source
def compute_next_pose_from_action(current_pose, keyboard_action, mouse_action):
x, y, z, pitch, yaw = current_pose
w, s, a, d = keyboard_action[:4]
mouse_x, mouse_y = mouse_action[:2]
delta_pitch = MOUSE_PITCH_SENSITIVITY * mouse_x if abs(mouse_x) >= MOUSE_THRESHOLD else 0.0
delta_yaw = MOUSE_YAW_SENSITIVITY * mouse_y if abs(mouse_y) >= MOUSE_THRESHOLD else 0.0
new_pitch = pitch + delta_pitch
new_yaw = yaw + delta_yaw
new_yaw = (new_yaw + 180) % 360 - 180
local_forward = 0.0
if w > 0.5 and s < 0.5:
local_forward = WSAD_OFFSET
elif s > 0.5 and w < 0.5:
local_forward = -WSAD_OFFSET
local_right = 0.0
if d > 0.5 and a < 0.5:
local_right = WSAD_OFFSET
elif a > 0.5 and d < 0.5:
local_right = -WSAD_OFFSET
if abs(local_forward) > 0.1 and abs(local_right) > 0.1:
local_forward = np.sign(local_forward) * DIAGONAL_OFFSET
local_right = np.sign(local_right) * DIAGONAL_OFFSET
avg_yaw = float((yaw + new_yaw) / 2.0)
yaw_rad = float(np.deg2rad(avg_yaw))
cos_yaw = np.cos(yaw_rad)
sin_yaw = np.sin(yaw_rad)
delta_x = cos_yaw * local_forward - sin_yaw * local_right
delta_y = sin_yaw * local_forward + cos_yaw * local_right
return np.array([x + delta_x, y + delta_y, z, new_pitch, new_yaw], dtype=np.float32)</code></pre></details>
fastvideo.models.dits.matrixgame3.utils.compute_all_poses_from_actions
compute_all_poses_from_actions(keyboard_conditions, mouse_conditions)
Source
def compute_all_poses_from_actions(keyboard_conditions, mouse_conditions):
poses = np.zeros((len(keyboard_conditions), 5), dtype=np.float32)
for idx in range(len(keyboard_conditions) - 1):
poses[idx + 1] = compute_next_pose_from_action(poses[idx], keyboard_conditions[idx], mouse_conditions[idx])
return poses
fastvideo.models.dits.matrixgame3.utils.build_intrinsics
build_intrinsics(height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor
Source
def build_intrinsics(height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
fov_rad = float(np.deg2rad(90.0))
fx = float(width) / (2.0 * float(np.tan(fov_rad / 2.0)))
fy = float(height) / (2.0 * float(np.tan(fov_rad / 2.0)))
return torch.tensor([[fx, fy, width / 2.0, height / 2.0]], device=device, dtype=dtype)
fastvideo.models.dits.matrixgame3.utils.build_matrixgame3_action_preset
build_matrixgame3_action_preset(num_frames: int, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def build_matrixgame3_action_preset(num_frames: int, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
presets = create_action_presets(num_frames=num_frames, keyboard_dim=6, seed=seed)
return presets["keyboard"], presets["mouse"]
fastvideo.models.dits.matrixgame3.utils.interpolate_camera_poses_handedness
interpolate_camera_poses_handedness(src_indices: np.ndarray, src_rot_mat: np.ndarray, src_trans_vec: np.ndarray, tgt_indices: np.ndarray) -> torch.Tensor
Source
def interpolate_camera_poses_handedness(
src_indices: np.ndarray,
src_rot_mat: np.ndarray,
src_trans_vec: np.ndarray,
tgt_indices: np.ndarray,
) -> torch.Tensor:
dets = np.linalg.det(src_rot_mat)
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
if flip_handedness:
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
src_rot_mat = src_rot_mat @ flip_mat
c2ws = interpolate_camera_poses(
src_indices=src_indices,
src_rot_mat=src_rot_mat,
src_trans_vec=src_trans_vec,
tgt_indices=tgt_indices,
)
if flip_handedness:
flip_mat_t = torch.from_numpy(flip_mat).to(c2ws.device, dtype=c2ws.dtype)
c2ws[:, :3, :3] = c2ws[:, :3, :3] @ flip_mat_t
return c2ws</code></pre></details>
fastvideo.models.dits.matrixgame3.utils.build_extrinsics_from_actions
build_extrinsics_from_actions(keyboard_conditions: torch.Tensor, mouse_conditions: torch.Tensor) -> torch.Tensor
Source
def build_extrinsics_from_actions(
keyboard_conditions: torch.Tensor,
mouse_conditions: torch.Tensor,
) -> torch.Tensor:
keyboard_np = keyboard_conditions.detach().to(dtype=torch.float32).cpu().numpy()
mouse_np = mouse_conditions.detach().to(dtype=torch.float32).cpu().numpy()
poses = compute_all_poses_from_actions(keyboard_np, mouse_np)
rotations = np.concatenate(
[np.zeros((poses.shape[0], 1), dtype=np.float32), poses[:, 3:5]],
axis=1,
)
positions = poses[:, :3]
return build_extrinsics(rotations, positions)
fastvideo.models.dits.matrixgame3.utils.build_extrinsics
build_extrinsics(video_rotation: np.ndarray, video_position: np.ndarray) -> torch.Tensor
Source
def build_extrinsics(
video_rotation: np.ndarray,
video_position: np.ndarray,
) -> torch.Tensor:
extrinsics = []
for frame_rotation, frame_position in zip(video_rotation, video_position, strict=False):
roll_deg, pitch_deg, yaw_deg = frame_rotation
roll, pitch, yaw = np.radians([roll_deg, pitch_deg, yaw_deg])
rot_z = np.array(
[[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]],
dtype=np.float32,
)
rot_y = np.array(
[[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]],
dtype=np.float32,
)
rot_x = np.array(
[[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]],
dtype=np.float32,
)
rot = rot_z @ rot_y @ rot_x
ext = np.eye(4, dtype=np.float32)
ext[:3, :3] = rot
ext[:3, 3] = np.asarray(frame_position, dtype=np.float32)
extrinsics.append(ext)
r_init = np.array([[0, 0, 1], [1, 0, 0], [0, -1, 0]], dtype=np.float32)
extrinsics_tensor = torch.from_numpy(np.stack(extrinsics, axis=0)).float()
extrinsics_tensor[:, :3, :3] = extrinsics_tensor[:, :3, :3] @ torch.from_numpy(r_init)
extrinsics_tensor[:, :3, 3] = extrinsics_tensor[:, :3, 3] * 0.01
return extrinsics_tensor</code></pre></details>
fastvideo.models.dits.matrixgame3.utils.build_plucker_from_c2ws
build_plucker_from_c2ws(c2ws_seq: torch.Tensor, src_indices: np.ndarray, tgt_indices: np.ndarray, *, target_h: int, target_w: int, latent_h: int, latent_w: int, framewise: bool = True) -> torch.Tensor
Source
def build_plucker_from_c2ws(
c2ws_seq: torch.Tensor,
src_indices: np.ndarray,
tgt_indices: np.ndarray,
*,
target_h: int,
target_w: int,
latent_h: int,
latent_w: int,
framewise: bool = True,
) -> torch.Tensor:
c2ws_np = c2ws_seq.detach().to(dtype=torch.float32).cpu().numpy()
c2ws_infer = interpolate_camera_poses_handedness(
src_indices=src_indices,
src_rot_mat=c2ws_np[:, :3, :3],
src_trans_vec=c2ws_np[:, :3, 3],
tgt_indices=tgt_indices,
).to(device=c2ws_seq.device, dtype=c2ws_seq.dtype)
c2ws_infer = compute_relative_poses(c2ws_infer, framewise=framewise)
return build_plucker_from_pose(
c2ws_infer,
target_h=target_h,
target_w=target_w,
latent_h=latent_h,
latent_w=latent_w,
)
fastvideo.models.dits.matrixgame3.utils.build_plucker_from_pose
build_plucker_from_pose(c2ws_pose: torch.Tensor, *, target_h: int, target_w: int, latent_h: int, latent_w: int) -> torch.Tensor
Source
def build_plucker_from_pose(
c2ws_pose: torch.Tensor,
*,
target_h: int,
target_w: int,
latent_h: int,
latent_w: int,
) -> torch.Tensor:
ks = build_intrinsics(target_h, target_w, c2ws_pose.device, c2ws_pose.dtype).repeat(c2ws_pose.shape[0], 1)
plucker = get_plucker_embeddings(c2ws_pose, ks, target_h, target_w)
c1 = target_h // latent_h
c2 = target_w // latent_w
plucker = plucker.view(c2ws_pose.shape[0], latent_h, c1, latent_w, c2, 6)
plucker = plucker.permute(0, 1, 3, 5, 2, 4).reshape(c2ws_pose.shape[0], latent_h, latent_w, 6 * c1 * c2)
plucker = plucker.permute(3, 0, 1, 2).unsqueeze(0)
return plucker
fastvideo.models.dits.matrixgame3.utils.select_memory_idx_fov
select_memory_idx_fov(extrinsics_all: torch.Tensor | np.ndarray, current_start_frame_idx: int, selected_index_base: list[int], *, height: int = 720, width: int = 1280, return_confidence: bool = False) -> list[int] | tuple[list[int], list[float]]
Source
def select_memory_idx_fov(
extrinsics_all: torch.Tensor | np.ndarray,
current_start_frame_idx: int,
selected_index_base: list[int],
*,
height: int = 720,
width: int = 1280,
return_confidence: bool = False,
) -> list[int] | tuple[list[int], list[float]]:
if isinstance(extrinsics_all, np.ndarray):
extrinsics_tensor = torch.from_numpy(extrinsics_all).float()
else:
extrinsics_tensor = extrinsics_all.float()
device = extrinsics_tensor.device
if current_start_frame_idx <= 1:
selected = [0] * len(selected_index_base)
if return_confidence:
return selected, [0.0] * len(selected)
return selected
fov_rad = np.deg2rad(90.0)
fx = width / (2 * np.tan(fov_rad / 2))
fy = height / (2 * np.tan(fov_rad / 2))
near, far = 0.1, 30.0
candidate_indices = torch.arange(1, current_start_frame_idx, device=device)
r_cand = extrinsics_tensor[candidate_indices, :3, :3]
t_cand = extrinsics_tensor[candidate_indices, :3, 3:4]
r_cand_inv = r_cand.transpose(1, 2)
t_cand_inv = -torch.bmm(r_cand_inv, t_cand)
num_side = 10
z_samples = torch.linspace(near, far, num_side, device=device)
x_samples = torch.linspace(-1, 1, num_side, device=device)
y_samples = torch.linspace(-1, 1, num_side, device=device)
grid_x, grid_y, grid_z = torch.meshgrid(x_samples, y_samples, z_samples, indexing="ij")
points_cam_base = torch.stack([
grid_x.reshape(-1) * grid_z.reshape(-1) * (width / (2 * fx)),
grid_y.reshape(-1) * grid_z.reshape(-1) * (height / (2 * fy)),
grid_z.reshape(-1),
],
dim=0)
selected_index: list[int] = []
selected_confidence: list[float] = []
for frame_idx in selected_index_base:
base_pose = extrinsics_tensor[frame_idx]
points_world = base_pose[:3, :3] @ points_cam_base + base_pose[:3, 3:4]
points_world_batched = points_world.unsqueeze(0)
points_in_cands = torch.bmm(r_cand_inv, points_world_batched.expand(len(candidate_indices), -1,
-1)) + t_cand_inv
x = points_in_cands[:, 0, :]
y = points_in_cands[:, 1, :]
z = points_in_cands[:, 2, :]
u = (x * fx / torch.clamp(z, min=1e-6)) + width / 2
v = (y * fy / torch.clamp(z, min=1e-6)) + height / 2
in_view = (z > near) & (z < far) & (u >= 0) & (u <= width) & (v >= 0) & (v <= height)
ratios = in_view.float().mean(dim=1)
best_idx = torch.argmax(ratios)
selected_index.append(candidate_indices[best_idx].item())
selected_confidence.append(ratios[best_idx].item())
if return_confidence:
return selected_index, selected_confidence
return selected_index</code></pre></details>
fastvideo.models.dits.matrixgame3.utils.create_action_presets
create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None)
Source
def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None):
if keyboard_dim not in (2, 4, 6, 7):
raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}")
if num_frames % 4 != 1:
raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.")
if seed is not None:
random.seed(seed)
num_samples_per_action = 4
if keyboard_dim == 4:
actions_single_action = ["forward", "left", "right"]
actions_double_action = ["forward_left", "forward_right"]
actions_single_camera = ["camera_l", "camera_r"]
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3}
elif keyboard_dim == 2:
actions_single_action = ["forward", "back"]
actions_double_action = []
actions_single_camera = ["camera_l", "camera_r"]
keyboard_idx = {"forward": 0, "back": 1}
elif keyboard_dim == 6:
actions_single_action = ["forward", "back", "left", "right"]
actions_double_action = ["forward_left", "forward_right"]
actions_single_camera = ["camera_l", "camera_r"]
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5}
else: # keyboard_dim == 7
actions_single_action = ["forward", "back", "left", "right"]
actions_double_action = []
actions_single_camera = []
keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6}
actions_to_test = (actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5)
for action in (actions_single_action + actions_double_action):
for camera in actions_single_camera:
actions_to_test.append(f"{action}_{camera}")
if not actions_to_test:
actions_to_test = actions_single_action * 5
base_action = actions_single_action + actions_single_camera
cam_value = 0.1
camera_value_map = {
"camera_up": [cam_value, 0],
"camera_down": [-cam_value, 0],
"camera_l": [0, -cam_value],
"camera_r": [0, cam_value],
"camera_ur": [cam_value, cam_value],
"camera_ul": [cam_value, -cam_value],
"camera_dr": [-cam_value, cam_value],
"camera_dl": [-cam_value, -cam_value],
}
data = []
for action_name in actions_to_test:
keyboard_condition = torch.zeros((num_samples_per_action, keyboard_dim))
mouse_condition = torch.zeros((num_samples_per_action, 2))
for sub_act in base_action:
if sub_act not in action_name:
continue
if sub_act in camera_value_map:
mouse_condition = torch.tensor(
[camera_value_map[sub_act] for _ in range(num_samples_per_action)],
dtype=mouse_condition.dtype,
)
elif sub_act in keyboard_idx:
keyboard_condition[:, keyboard_idx[sub_act]] = 1
data.append({
"keyboard_condition": keyboard_condition,
"mouse_condition": mouse_condition,
})
keyboard_condition = torch.zeros((num_frames, keyboard_dim))
mouse_condition = torch.zeros((num_frames, 2))
current_frame = 0
selections = [12]
while current_frame < num_frames:
rd_frame = selections[random.randint(0, len(selections) - 1)]
entry = data[random.randint(0, len(data) - 1)]
key_seq = entry["keyboard_condition"]
mouse_seq = entry["mouse_condition"]
if current_frame == 0:
keyboard_condition[:1] = key_seq[:1]
mouse_condition[:1] = mouse_seq[:1]
current_frame = 1
else:
rd_frame = min(rd_frame, num_frames - current_frame)
repeat_time = rd_frame // 4
keyboard_condition[current_frame:current_frame + rd_frame] = key_seq.repeat(repeat_time, 1)
mouse_condition[current_frame:current_frame + rd_frame] = mouse_seq.repeat(repeat_time, 1)
current_frame += rd_frame
return {"keyboard": keyboard_condition, "mouse": mouse_condition}</code></pre></details>