vllm-omni/vllm_omni/diffusion/models/lingbot_world/camera.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Camera trajectory loading and geometry for LingBot World v2 conditioning."""

from __future__ import annotations

import io
import math
import os
import stat
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO

import numpy as np
import torch

_REFERENCE_HEIGHT = 480
_REFERENCE_WIDTH = 832
_MAX_ACTION_FRAMES = 117
_MAX_SOURCE_ACTION_FRAMES = 4096
_MAX_NPY_HEADER_BYTES = 10_000
_DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
_FILE_OPEN_FLAGS = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK


@dataclass(frozen=True)
class CameraTrajectory:
    """Camera-to-world poses and ``(fx, fy, cx, cy)`` intrinsics by frame."""

    poses: torch.Tensor
    intrinsics: torch.Tensor


@dataclass(frozen=True)
class TrustedActionDirectory:
    """Canonical action location tied to the trusted root's filesystem identity."""

    root: Path
    relative: Path
    root_device: int
    root_inode: int


def resolve_trusted_action_directory(
    action_path: str | os.PathLike[str],
    trusted_root: str | os.PathLike[str],
) -> TrustedActionDirectory:
    """Resolve an action directory beneath a trusted root without exposing paths."""

    try:
        root = Path(trusted_root).expanduser().resolve(strict=True)
    except (OSError, RuntimeError):
        raise ValueError("The configured LingBot trusted action root is unavailable.") from None
    if not root.is_dir():
        raise ValueError("The configured LingBot trusted action root must be a directory.")

    try:
        candidate = Path(action_path).expanduser()
        if not candidate.is_absolute():
            candidate = root / candidate
        candidate = candidate.resolve(strict=True)
        relative = candidate.relative_to(root)
    except (OSError, RuntimeError, ValueError):
        raise ValueError(
            "sampling_params.extra_args.action_path must be contained by the trusted action root."
        ) from None
    if not candidate.is_dir():
        raise ValueError("sampling_params.extra_args.action_path must identify a directory in the trusted action root.")

    try:
        root_stat = root.stat()
    except OSError:
        raise ValueError("The configured LingBot trusted action root is unavailable.") from None
    return TrustedActionDirectory(
        root=root,
        relative=relative,
        root_device=root_stat.st_dev,
        root_inode=root_stat.st_ino,
    )


def _validate_trajectory(
    trajectory: CameraTrajectory,
    *,
    file_suffix: str = "",
) -> None:
    poses = trajectory.poses
    intrinsics = trajectory.intrinsics
    if poses.ndim != 3 or poses.shape[1:] != (4, 4):
        raise ValueError(f"poses{file_suffix} must have shape [frames, 4, 4], got {tuple(poses.shape)}")
    if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
        raise ValueError(
            f"intrinsics{file_suffix} must have shape [frames, 4] ordered as fx, fy, cx, cy, "
            f"got {tuple(intrinsics.shape)}"
        )
    if poses.shape[0] == 0:
        raise ValueError("camera trajectory must contain at least one frame")
    if poses.shape[0] != intrinsics.shape[0]:
        raise ValueError("poses and intrinsics must contain the same number of frames")
    if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all():
        raise ValueError("camera trajectory values must all be finite")


@contextmanager
def _open_action_directory(action_directory: TrustedActionDirectory) -> Iterator[int]:
    root_fd: int | None = None
    action_fd: int | None = None
    try:
        root_fd = os.open(action_directory.root, _DIRECTORY_OPEN_FLAGS)
        root_stat = os.fstat(root_fd)
        if (root_stat.st_dev, root_stat.st_ino) != (
            action_directory.root_device,
            action_directory.root_inode,
        ):
            raise ValueError("The configured LingBot trusted action root changed before camera files were opened.")

        action_fd = os.dup(root_fd)
        for component in action_directory.relative.parts:
            next_fd = os.open(component, _DIRECTORY_OPEN_FLAGS, dir_fd=action_fd)
            os.close(action_fd)
            action_fd = next_fd
    except OSError:
        if action_fd is not None:
            os.close(action_fd)
        if root_fd is not None:
            os.close(root_fd)
        raise ValueError("Unable to securely open the LingBot action directory within the trusted root.") from None
    except ValueError:
        if action_fd is not None:
            os.close(action_fd)
        if root_fd is not None:
            os.close(root_fd)
        raise
    assert action_fd is not None
    assert root_fd is not None
    try:
        yield action_fd
    finally:
        os.close(action_fd)
        os.close(root_fd)


@contextmanager
def _open_camera_file(action_fd: int, filename: str) -> Iterator[tuple[BinaryIO, os.stat_result]]:
    file_fd = -1
    try:
        file_fd = os.open(filename, _FILE_OPEN_FLAGS, dir_fd=action_fd)
        file_stat = os.fstat(file_fd)
        if not stat.S_ISREG(file_stat.st_mode):
            raise ValueError(f"{filename} must be a regular NPY file.")
        with os.fdopen(file_fd, "rb", closefd=True) as file_handle:
            file_fd = -1
            yield file_handle, file_stat
    except FileNotFoundError:
        raise FileNotFoundError(f"camera trajectory file not found: {filename}") from None
    except OSError:
        raise ValueError(f"{filename} could not be opened securely from the trusted action directory.") from None
    finally:
        if file_fd >= 0:
            os.close(file_fd)


def _load_bounded_npy(
    action_fd: int,
    filename: str,
    *,
    trailing_shape: tuple[int, ...],
) -> np.ndarray:
    with _open_camera_file(action_fd, filename) as (file_handle, file_stat):
        max_payload_bytes = _MAX_SOURCE_ACTION_FRAMES * math.prod(trailing_shape) * 8
        max_file_bytes = _MAX_NPY_HEADER_BYTES + max_payload_bytes + 32
        if file_stat.st_size > max_file_bytes:
            raise ValueError(f"{filename} byte size exceeds the bounded camera-array limit.")
        snapshot = file_handle.read(max_file_bytes + 1)
        if len(snapshot) != file_stat.st_size:
            raise ValueError(f"{filename} changed while its bounded contents were read.")
        snapshot_file = io.BytesIO(snapshot)
        try:
            version = np.lib.format.read_magic(snapshot_file)
            if version == (1, 0):
                shape, _fortran_order, dtype = np.lib.format.read_array_header_1_0(snapshot_file)
            elif version == (2, 0):
                shape, _fortran_order, dtype = np.lib.format.read_array_header_2_0(snapshot_file)
            else:
                raise ValueError(f"{filename} must use NPY format version 1.0 or 2.0.")
        except (EOFError, TypeError, ValueError):
            raise ValueError(f"{filename} has an invalid or oversized NPY header.") from None

        if len(shape) != len(trailing_shape) + 1 or tuple(shape[1:]) != trailing_shape:
            expected = ", ".join(str(value) for value in trailing_shape)
            raise ValueError(f"{filename} must have shape [frames, {expected}], got {shape}.")
        frames = shape[0]
        if isinstance(frames, bool) or not isinstance(frames, int) or not 1 <= frames <= _MAX_SOURCE_ACTION_FRAMES:
            raise ValueError(f"{filename} must contain between 1 and {_MAX_SOURCE_ACTION_FRAMES} source frames.")
        if dtype.kind not in "fiu" or dtype.hasobject or dtype.itemsize > 8:
            raise ValueError(f"{filename} must contain real numeric values with at most 8 bytes per element.")

        data_offset = snapshot_file.tell()
        expected_bytes = math.prod(shape) * dtype.itemsize
        if len(snapshot) != data_offset + expected_bytes:
            raise ValueError(f"{filename} byte size does not match its NPY header.")
        snapshot_file.seek(0)
        try:
            array = np.load(snapshot_file, allow_pickle=False)
        except (EOFError, TypeError, ValueError):
            raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None
    # Official LingBot trajectories may be longer than one request. Only
    # materialize the bounded prefix that any supported request can consume.
    result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32)
    return result


def load_camera_trajectory(action_directory: TrustedActionDirectory) -> CameraTrajectory:
    """Load bounded numeric camera arrays through a trusted directory handle."""

    if not isinstance(action_directory, TrustedActionDirectory):
        raise TypeError("action_directory must be resolved against a trusted action root before loading.")
    with _open_action_directory(action_directory) as action_fd:
        poses_array = _load_bounded_npy(action_fd, "poses.npy", trailing_shape=(4, 4))
        intrinsics_array = _load_bounded_npy(action_fd, "intrinsics.npy", trailing_shape=(4,))

    trajectory = CameraTrajectory(
        poses=torch.from_numpy(poses_array),
        intrinsics=torch.from_numpy(intrinsics_array),
    )
    _validate_trajectory(trajectory, file_suffix=".npy")
    return trajectory


def _linear_interpolate(values: torch.Tensor, target_times: np.ndarray) -> torch.Tensor:
    source_times = np.linspace(0.0, 1.0, values.shape[0])
    values_array = values.detach().cpu().double().numpy()
    columns = [
        np.interp(target_times, source_times, values_array[:, column]) for column in range(values_array.shape[1])
    ]
    interpolated = torch.from_numpy(np.stack(columns, axis=1))
    return interpolated.to(device=values.device, dtype=values.dtype)


def _rotation_matrices_to_quaternions(matrices: np.ndarray) -> np.ndarray:
    """Convert rotation matrices to normalized ``(w, x, y, z)`` quaternions."""

    quaternions: np.ndarray = np.empty((matrices.shape[0], 4), dtype=np.float64)
    for index, matrix in enumerate(matrices):
        trace = float(np.trace(matrix))
        if trace > 0.0:
            scale = math.sqrt(trace + 1.0) * 2.0
            quaternion = np.array(
                [
                    0.25 * scale,
                    (matrix[2, 1] - matrix[1, 2]) / scale,
                    (matrix[0, 2] - matrix[2, 0]) / scale,
                    (matrix[1, 0] - matrix[0, 1]) / scale,
                ]
            )
        else:
            axis = int(np.argmax(np.diag(matrix)))
            if axis == 0:
                scale = math.sqrt(max(0.0, 1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2])) * 2.0
                quaternion = np.array(
                    [
                        (matrix[2, 1] - matrix[1, 2]) / scale,
                        0.25 * scale,
                        (matrix[0, 1] + matrix[1, 0]) / scale,
                        (matrix[0, 2] + matrix[2, 0]) / scale,
                    ]
                )
            elif axis == 1:
                scale = math.sqrt(max(0.0, 1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2])) * 2.0
                quaternion = np.array(
                    [
                        (matrix[0, 2] - matrix[2, 0]) / scale,
                        (matrix[0, 1] + matrix[1, 0]) / scale,
                        0.25 * scale,
                        (matrix[1, 2] + matrix[2, 1]) / scale,
                    ]
                )
            else:
                scale = math.sqrt(max(0.0, 1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1])) * 2.0
                quaternion = np.array(
                    [
                        (matrix[1, 0] - matrix[0, 1]) / scale,
                        (matrix[0, 2] + matrix[2, 0]) / scale,
                        (matrix[1, 2] + matrix[2, 1]) / scale,
                        0.25 * scale,
                    ]
                )
        norm = np.linalg.norm(quaternion)
        if not np.isfinite(norm) or norm <= np.finfo(np.float64).eps:
            raise ValueError("camera rotation matrices must describe valid rotations")
        quaternions[index] = quaternion / norm
    return quaternions


def _quaternions_to_rotation_matrices(quaternions: np.ndarray) -> np.ndarray:
    quaternions = quaternions / np.linalg.norm(quaternions, axis=1, keepdims=True)
    w, x, y, z = quaternions.T
    matrices: np.ndarray = np.stack(
        (
            1 - 2 * (y * y + z * z),
            2 * (x * y - z * w),
            2 * (x * z + y * w),
            2 * (x * y + z * w),
            1 - 2 * (x * x + z * z),
            2 * (y * z - x * w),
            2 * (x * z - y * w),
            2 * (y * z + x * w),
            1 - 2 * (x * x + y * y),
        ),
        axis=1,
    ).reshape(-1, 3, 3)
    return matrices


def _slerp_rotations(matrices: np.ndarray, target_times: np.ndarray) -> np.ndarray:
    quaternions = _rotation_matrices_to_quaternions(matrices)
    interpolated = np.empty((target_times.shape[0], 4), dtype=np.float64)
    for target_index, target_time in enumerate(target_times):
        scaled_time = float(target_time) * (matrices.shape[0] - 1)
        left = min(max(math.floor(scaled_time), 0), matrices.shape[0] - 2)
        fraction = scaled_time - left
        first = quaternions[left]
        second = quaternions[left + 1]
        cosine = float(np.dot(first, second))
        if cosine < 0.0:
            second = -second
            cosine = -cosine
        cosine = min(cosine, 1.0)
        if cosine > 0.9995:
            quaternion = first + fraction * (second - first)
            interpolated[target_index] = quaternion / np.linalg.norm(quaternion)
            continue
        angle = math.acos(cosine)
        denominator = math.sin(angle)
        interpolated[target_index] = (
            math.sin((1.0 - fraction) * angle) / denominator * first + math.sin(fraction * angle) / denominator * second
        )
    return _quaternions_to_rotation_matrices(interpolated)


def interpolate_camera_trajectory(
    trajectory: CameraTrajectory,
    num_frames: int,
) -> CameraTrajectory:
    """Resample translations/intrinsics linearly and rotations with quaternion SLERP."""

    _validate_trajectory(trajectory)
    if num_frames <= 0:
        raise ValueError(f"num_frames must be positive, got {num_frames}")

    source_frames = trajectory.poses.shape[0]
    if num_frames == source_frames:
        return trajectory
    if source_frames == 1:
        return CameraTrajectory(
            poses=trajectory.poses.repeat(num_frames, 1, 1),
            intrinsics=trajectory.intrinsics.repeat(num_frames, 1),
        )

    # Rotations use quaternion SLERP; translations and intrinsics are linear.
    target_times = np.linspace(0.0, 1.0, num_frames)
    rotation_matrices = trajectory.poses[:, :3, :3].detach().cpu().double().numpy()
    interpolated_rotations = _slerp_rotations(rotation_matrices, target_times)
    translations = _linear_interpolate(trajectory.poses[:, :3, 3], target_times)
    intrinsics = _linear_interpolate(trajectory.intrinsics, target_times)

    poses = torch.eye(
        4,
        device=trajectory.poses.device,
        dtype=trajectory.poses.dtype,
    ).repeat(num_frames, 1, 1)
    poses[:, :3, :3] = torch.from_numpy(interpolated_rotations).to(
        device=trajectory.poses.device,
        dtype=trajectory.poses.dtype,
    )
    poses[:, :3, 3] = translations
    return CameraTrajectory(poses=poses, intrinsics=intrinsics)


def _invert_camera_poses(poses: torch.Tensor) -> torch.Tensor:
    """Invert rigid transforms using R^-1 = R^T and t^-1 = -R^T t."""

    rotations = poses[:, :3, :3]
    translations = poses[:, :3, 3:]
    inverse_rotations = rotations.transpose(-1, -2)
    inverse = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(poses.shape[0], 1, 1)
    inverse[:, :3, :3] = inverse_rotations
    inverse[:, :3, 3:] = -torch.bmm(inverse_rotations, translations)
    return inverse


def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor:
    """Convert raw C2W poses to normalized first-relative framewise deltas."""

    # Remove global placement, then convert the sequence to framewise deltas.
    first_world_to_camera = _invert_camera_poses(raw_poses[:1])
    relative_poses = torch.matmul(first_world_to_camera, raw_poses).clone()
    relative_poses[0] = torch.eye(4, device=raw_poses.device, dtype=raw_poses.dtype)
    if relative_poses.shape[0] > 1:
        relative_poses[1:] = torch.bmm(
            _invert_camera_poses(relative_poses[:-1]),
            relative_poses[1:],
        )

    # Only relative direction/magnitude matters to the checkpoint; normalizing
    # by the largest step removes the trajectory's arbitrary translation scale.
    translations = relative_poses[:, :3, 3]
    max_translation_norm = torch.linalg.vector_norm(translations, dim=-1).max()
    if max_translation_norm > 0:
        relative_poses[:, :3, 3] = translations / max_translation_norm
    return relative_poses


def build_plucker_embedding(
    trajectory: CameraTrajectory,
    *,
    height: int,
    width: int,
    target_height: int,
    target_width: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Build checkpoint-ordered ``(ray origin, ray direction)`` channels."""

    _validate_trajectory(trajectory)
    dimensions = {
        "height": height,
        "width": width,
        "target_height": target_height,
        "target_width": target_width,
    }
    for name, value in dimensions.items():
        if value <= 0:
            raise ValueError(f"{name} must be positive, got {value}")
    if dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
        raise ValueError(f"dtype must be floating point, got {dtype}")

    # Compute geometry in at least FP32, then cast at the model boundary.
    compute_dtype = torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype
    poses = _prepare_framewise_poses(trajectory.poses.to(device=device, dtype=compute_dtype))
    intrinsics = trajectory.intrinsics.to(device=device, dtype=compute_dtype).clone()
    intrinsics[:, (0, 2)] *= width / _REFERENCE_WIDTH
    intrinsics[:, (1, 3)] *= height / _REFERENCE_HEIGHT
    if torch.any(intrinsics[:, :2] == 0):
        raise ValueError("camera focal lengths fx and fy must be non-zero")

    # Sample requested-pixel coordinates directly at the conditioning resolution.
    x_coordinates = (torch.arange(target_width, device=device, dtype=compute_dtype) + 0.5) * (width / target_width)
    y_coordinates = (torch.arange(target_height, device=device, dtype=compute_dtype) + 0.5) * (height / target_height)
    grid_y, grid_x = torch.meshgrid(y_coordinates, x_coordinates, indexing="ij")
    grid_x = grid_x.unsqueeze(0)
    grid_y = grid_y.unsqueeze(0)

    # Pinhole back-projection: pixel (x, y) becomes the camera-space direction
    # ((x-cx)/fx, (y-cy)/fy, 1), which is normalized before rotation.
    fx, fy, cx, cy = intrinsics.unbind(dim=1)
    camera_x = (grid_x - cx[:, None, None]) / fx[:, None, None]
    camera_y = (grid_y - cy[:, None, None]) / fy[:, None, None]
    camera_directions = torch.stack(
        (camera_x, camera_y, torch.ones_like(camera_x)),
        dim=-1,
    )
    camera_directions = camera_directions / torch.linalg.vector_norm(camera_directions, dim=-1, keepdim=True)

    # A ray is represented by its camera center (origin) and unit direction in
    # the prepared relative coordinate frame. Channel order is checkpoint
    # data, not an implementation preference: [origin_xyz, direction_xyz].
    world_directions = torch.einsum(
        "fij,fhwj->fhwi",
        poses[:, :3, :3],
        camera_directions,
    )
    world_directions = world_directions / torch.linalg.vector_norm(world_directions, dim=-1, keepdim=True)
    ray_origins = poses[:, None, None, :3, 3].expand_as(world_directions)

    embedding = torch.cat((ray_origins, world_directions), dim=-1)
    return embedding.permute(0, 3, 1, 2).contiguous().to(dtype=dtype)