--- title: utils ---

fastvideo.models.dits.matrixgame2.utils.get_local_torch_device

get_local_torch_device() -> torch.device

Return 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)

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)

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,
        )

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,
        )

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

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"]

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}

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}

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