utils
fastvideo.models.dits.matrixgame2.utils.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn the torch device for the current rank.
fastvideo.models.dits.matrixgame2.utils.retain_kv_with_sink
retain_kv_with_sink(k: torch.Tensor, v: torch.Tensor, target_len: int, sink_size: int) -> tuple[torch.Tensor, torch.Tensor]Source
def retain_kv_with_sink(
k: torch.Tensor,
v: torch.Tensor,
target_len: int,
sink_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
if k.shape[1] <= target_len or sink_size <= 0:
return k[:, -target_len:], v[:, -target_len:]
sink_len = min(int(sink_size), target_len, k.shape[1])
tail_len = target_len - sink_len
sink_k = k[:, :sink_len]
sink_v = v[:, :sink_len]
if tail_len <= 0:
return sink_k, sink_v
tail_k = k[:, sink_len:][:, -tail_len:]
tail_v = v[:, sink_len:][:, -tail_len:]
return torch.cat([sink_k, tail_k], dim=1), torch.cat([sink_v, tail_v], dim=1)</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.draw_rounded_rectangle
draw_rounded_rectangle(image: np.ndarray, top_left: tuple[int, int], bottom_right: tuple[int, int], color: tuple[int, int, int], radius: int = 10, alpha: float = 0.5) -> None
Source
def draw_rounded_rectangle(
image: np.ndarray,
top_left: tuple[int, int],
bottom_right: tuple[int, int],
color: tuple[int, int, int],
radius: int = 10,
alpha: float = 0.5,
) -> None:
if not _require_cv2():
return
overlay = image.copy()
x1, y1 = top_left
x2, y2 = bottom_right
cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1)
cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1)
cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)
cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)
cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)
cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)
cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.draw_keys_on_frame
draw_keys_on_frame(frame: np.ndarray, keys: dict[str, bool], key_size: tuple[int, int] = (30, 30), top_margin: int = 15) -> None
Source
def draw_keys_on_frame(
frame: np.ndarray,
keys: dict[str, bool],
key_size: tuple[int, int] = (30, 30),
top_margin: int = 15,
) -> None:
if not _require_cv2():
return
left_margin = 15
gap = 3
key_positions = {
"W": (left_margin + key_size[0] + gap, top_margin),
"A": (left_margin, top_margin + key_size[1] + gap),
"S": (left_margin + key_size[0] + gap, top_margin + key_size[1] + gap),
"D": (
left_margin + (key_size[0] + gap) * 2,
top_margin + key_size[1] + gap,
),
}
for key, (x, y) in key_positions.items():
pressed = keys.get(key, False)
color = (0, 255, 0) if pressed else (200, 200, 200)
alpha = 0.8 if pressed else 0.5
draw_rounded_rectangle(
frame,
(x, y),
(x + key_size[0], y + key_size[1]),
color,
radius=5,
alpha=alpha,
)
text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
text_x = x + (key_size[0] - text_size[0]) // 2
text_y = y + (key_size[1] + text_size[1]) // 2
cv2.putText(
frame,
key,
(text_x, text_y),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 0, 0),
1,
)</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.draw_mouse_on_frame
draw_mouse_on_frame(frame: np.ndarray, yaw: float, pitch: float, top_margin: int = 15) -> None
Source
def draw_mouse_on_frame(
frame: np.ndarray,
yaw: float,
pitch: float,
top_margin: int = 15,
) -> None:
if not _require_cv2():
return
_, width, _ = frame.shape
right_margin = 15
crosshair_radius = 25
crosshair_x = width - right_margin - crosshair_radius
crosshair_y = top_margin + crosshair_radius
dx = int(yaw * crosshair_radius * 8)
dy = int(-pitch * crosshair_radius * 8)
max_arrow = crosshair_radius - 5
dx = max(-max_arrow, min(max_arrow, dx))
dy = max(-max_arrow, min(max_arrow, dy))
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (50, 50, 50), -1)
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (200, 200, 200), 1)
cv2.line(
frame,
(crosshair_x - crosshair_radius + 5, crosshair_y),
(crosshair_x + crosshair_radius - 5, crosshair_y),
(100, 100, 100),
1,
)
cv2.line(
frame,
(crosshair_x, crosshair_y - crosshair_radius + 5),
(crosshair_x, crosshair_y + crosshair_radius - 5),
(100, 100, 100),
1,
)
if abs(dx) > 1 or abs(dy) > 1:
cv2.arrowedLine(
frame,
(crosshair_x, crosshair_y),
(crosshair_x + dx, crosshair_y + dy),
(0, 255, 0),
2,
tipLength=0.3,
)</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.overlay_validation_actions_on_frames
overlay_validation_actions_on_frames(frames: list[np.ndarray], keyboard_cond: np.ndarray | torch.Tensor | None = None, mouse_cond: np.ndarray | torch.Tensor | None = None) -> list[np.ndarray]
Source
def overlay_validation_actions_on_frames(
frames: list[np.ndarray],
keyboard_cond: np.ndarray | torch.Tensor | None = None,
mouse_cond: np.ndarray | torch.Tensor | None = None,
) -> list[np.ndarray]:
if (keyboard_cond is None and mouse_cond is None) or not _require_cv2():
return frames
if keyboard_cond is not None and torch.is_tensor(keyboard_cond):
keyboard_cond = keyboard_cond.detach().cpu().float().numpy()
if mouse_cond is not None and torch.is_tensor(mouse_cond):
mouse_cond = mouse_cond.detach().cpu().float().numpy()
if keyboard_cond is not None:
keyboard_cond = np.asarray(keyboard_cond, dtype=np.float32)
if keyboard_cond.ndim == 3 and keyboard_cond.shape[0] == 1:
keyboard_cond = keyboard_cond[0]
if mouse_cond is not None:
mouse_cond = np.asarray(mouse_cond, dtype=np.float32)
if mouse_cond.ndim == 3 and mouse_cond.shape[0] == 1:
mouse_cond = mouse_cond[0]
processed_frames: list[np.ndarray] = []
for frame_idx, frame in enumerate(frames):
frame = np.ascontiguousarray(frame.copy())
if keyboard_cond is not None and frame_idx < len(keyboard_cond):
draw_keys_on_frame(
frame,
_matrixgame2_overlay_keys_from_keyboard(keyboard_cond[frame_idx]),
)
if mouse_cond is not None and frame_idx < len(mouse_cond):
# MG action convention: mouse[0]=pitch, mouse[1]=yaw.
pitch = float(mouse_cond[frame_idx, 0])
yaw = float(mouse_cond[frame_idx, 1])
draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
processed_frames.append(frame)
return processed_frames</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.scale_keyboard_condition
scale_keyboard_condition(keyboard_condition: torch.Tensor, scale: float) -> torch.Tensor
Source
def scale_keyboard_condition(
keyboard_condition: torch.Tensor,
scale: float,
) -> torch.Tensor:
if float(scale) == 1.0:
return keyboard_condition
return keyboard_condition * float(scale)
fastvideo.models.dits.matrixgame2.utils.expand_action_to_frames
expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]
Source
def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]:
result = {}
for key, tensor in action.items():
if tensor is not None:
# Expand to [num_frames, D] then unsqueeze to [1, num_frames, D]
result[key] = tensor.unsqueeze(0).repeat(num_frames, 1).unsqueeze(0)
else:
result[key] = None
if "mouse" not in result or result["mouse"] is None:
# keyboard device if available, otherwise default
device = result.get("keyboard", torch.tensor(
[])).device if result.get("keyboard") is not None else get_local_torch_device()
result["mouse"] = torch.zeros(1, num_frames, 2, device=device)
return result["keyboard"], result["mouse"]</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.get_current_action
get_current_action(mode = 'universal')
Source
def get_current_action(mode="universal"):
CAM_VALUE = 0.1
if mode == 'universal':
logger.info("")
logger.info('-' * 30)
logger.info("PRESS [I, K, J, L, U] FOR CAMERA TRANSFORM\n (I: up, K: down, J: left, L: right, U: no move)")
logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)")
logger.info('-' * 30)
CAMERA_VALUE_MAP = {
"i": [CAM_VALUE, 0],
"k": [-CAM_VALUE, 0],
"j": [0, -CAM_VALUE],
"l": [0, CAM_VALUE],
"u": [0, 0]
}
KEYBOARD_IDX = {"w": [1, 0, 0, 0], "s": [0, 1, 0, 0], "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], "q": [0, 0, 0, 0]}
flag = 0
while flag != 1:
try:
idx_mouse = input('Please input the mouse action (e.g. `U`):\n').strip().lower()
idx_keyboard = input('Please input the keyboard action (e.g. `W`):\n').strip().lower()
if idx_mouse in CAMERA_VALUE_MAP and idx_keyboard in KEYBOARD_IDX:
flag = 1
except Exception:
pass
mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse]).cuda()
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda()
elif mode == 'gta_drive':
logger.info("")
logger.info('-' * 30)
logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)")
logger.info('-' * 30)
CAMERA_VALUE_MAP = {"a": [0, -CAM_VALUE], "d": [0, CAM_VALUE], "q": [0, 0]}
KEYBOARD_IDX = {"w": [1, 0], "s": [0, 1], "q": [0, 0]}
flag = 0
while flag != 1:
try:
indexes = input(
'Please input the actions (split with ` `):\n(e.g. `W` for forward, `W A` for forward and left)\n'
).strip().lower().split(' ')
idx_mouse = []
idx_keyboard = []
for i in indexes:
if i in CAMERA_VALUE_MAP.keys():
idx_mouse += [i]
elif i in KEYBOARD_IDX.keys():
idx_keyboard += [i]
if len(idx_mouse) == 0:
idx_mouse += ['q']
if len(idx_keyboard) == 0:
idx_keyboard += ['q']
assert idx_mouse in [['a'], ['d'], ['q']] and idx_keyboard in [['q'], ['w'], ['s']]
flag = 1
except Exception:
pass
mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse[0]]).cuda()
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard[0]]).cuda()
elif mode == 'templerun':
logger.info("")
logger.info('-' * 30)
logger.info(
"PRESS [W, S, A, D, Z, C, Q] FOR ACTIONS\n (W: jump, S: slide, A: left side, D: right side, Z: turn left, C: turn right, Q: no move)"
)
logger.info('-' * 30)
KEYBOARD_IDX = {
"w": [0, 1, 0, 0, 0, 0, 0],
"s": [0, 0, 1, 0, 0, 0, 0],
"a": [0, 0, 0, 0, 0, 1, 0],
"d": [0, 0, 0, 0, 0, 0, 1],
"z": [0, 0, 0, 1, 0, 0, 0],
"c": [0, 0, 0, 0, 1, 0, 0],
"q": [1, 0, 0, 0, 0, 0, 0]
}
flag = 0
while flag != 1:
try:
idx_keyboard = input(
'Please input the action: \n(e.g. `W` for forward, `Z` for turning left)\n').strip().lower()
if idx_keyboard in KEYBOARD_IDX.keys():
flag = 1
except Exception:
pass
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda()
if mode != 'templerun':
return {"mouse": mouse_cond, "keyboard": keyboard_cond}
return {"keyboard": keyboard_cond}</code></pre></details>
fastvideo.models.dits.matrixgame2.utils.get_current_action_async
get_current_action_async(mode = 'universal')
Source
async def get_current_action_async(mode="universal"):
return await asyncio.to_thread(get_current_action, mode)
fastvideo.models.dits.matrixgame2.utils.load_initial_image
load_initial_image(image_path: str = None) -> Image.Image
Source
def load_initial_image(image_path: str = None) -> Image.Image:
if image_path and os.path.exists(image_path):
return Image.open(image_path).convert("RGB")
logger.warning("No image provided, creating placeholder...")
return Image.new("RGB", (640, 352), (128, 128, 128))
fastvideo.models.dits.matrixgame2.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.")
# Set seed for reproducibility if provided
if seed is not None:
random.seed(seed)
num_samples_per_action = 4
# Define actions based on keyboard_dim
if keyboard_dim == 4:
# Universal model: W, S, A, D
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:
# GTA model: W, S only (steering via mouse)
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
# Temple Run model: still, w, s, left, right, a, d (no mouse)
actions_single_action = ["forward", "back", "left", "right"]
actions_double_action = []
actions_single_camera = [] # No mouse for Temple Run
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}")
# Ensure we have at least some actions
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>
fastvideo.models.dits.matrixgame2.utils.parse_config
parse_config(config, mode = 'universal')
Source
def parse_config(config, mode="universal"):
assert mode in ['universal', 'gta_drive', 'templerun']
key_data = {}
mouse_data = {}
if mode != 'templerun':
key, mouse = config
else:
key = config
for i in range(len(key)):
if mode == 'templerun':
still, w, s, left, right, a, d = key[i]
elif mode == 'universal':
w, s, a, d = key[i]
else:
w, s, a, d = key[i][0], key[i][1], mouse[i][1] < 0, mouse[i][1] > 0
if mode == 'universal':
mouse_y, mouse_x = mouse[i]
mouse_y = -1 * mouse_y
key_data[i] = {"W": bool(w), "A": bool(a), "S": bool(s), "D": bool(d)}
if mode == 'templerun':
key_data[i].update({"left": bool(left), "right": bool(right)})
if mode == 'universal':
if i == 0:
mouse_data[i] = (320, 352 // 2)
else:
global_scale_factor = 0.1
mouse_scale_x = 15 * global_scale_factor
mouse_scale_y = 15 * 4 * global_scale_factor
mouse_data[i] = (
mouse_data[i - 1][0] + mouse_x * mouse_scale_x,
mouse_data[i - 1][1] + mouse_y * mouse_scale_y,
)
return key_data, mouse_data</code></pre></details>