356 lines
14 KiB
Python
356 lines
14 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
|||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|||
|
|
"""Camera-control helpers for SANA-WM.
|
|||
|
|
|
|||
|
|
This module mirrors the public NVlabs/Sana SANA-WM camera preparation path:
|
|||
|
|
WASD/IJKL action rollout or explicit camera-to-world poses are converted into
|
|||
|
|
relative poses, per-latent-frame ray metadata, and per-VAE-chunk Plucker maps.
|
|||
|
|
Model-local projection layers consume these tensors in the Stage-1 DiT.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import math
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
SANA_WM_DEFAULT_TRANSLATION_SPEED = 0.05
|
|||
|
|
SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2
|
|||
|
|
SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0
|
|||
|
|
SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl")
|
|||
|
|
SANA_WM_DEFAULT_VAE_STRIDE = (8, 32, 32)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass(frozen=True)
|
|||
|
|
class SanaWmCameraCondition:
|
|||
|
|
poses: Any | None = None
|
|||
|
|
intrinsics: Any | None = None
|
|||
|
|
action: str | None = None
|
|||
|
|
num_frames: int | None = None
|
|||
|
|
height: int = 704
|
|||
|
|
width: int = 1280
|
|||
|
|
translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED
|
|||
|
|
rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG
|
|||
|
|
pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _rot_x(angle_rad: float) -> np.ndarray:
|
|||
|
|
c, s = np.cos(angle_rad), np.sin(angle_rad)
|
|||
|
|
return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _rot_y(angle_rad: float) -> np.ndarray:
|
|||
|
|
c, s = np.cos(angle_rad), np.sin(angle_rad)
|
|||
|
|
return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _parse_action_segments(action: str) -> list[tuple[list[str], int]]:
|
|||
|
|
"""Parse an action string into ``(keys, duration)`` pairs without expanding it.
|
|||
|
|
|
|||
|
|
A single segment names an unbounded rollout (``"w-1000000000"`` is valid
|
|||
|
|
syntax), so expanding here would allocate one list entry per frame before
|
|||
|
|
anything checks the frame count. Callers total the durations and validate
|
|||
|
|
against ``num_frames`` first; only :func:`_parse_action_string` expands.
|
|||
|
|
"""
|
|||
|
|
cleaned = "".join(action.replace("，", ",").split())
|
|||
|
|
if not cleaned:
|
|||
|
|
raise ValueError("Sana-WM action string is empty.")
|
|||
|
|
|
|||
|
|
segments: list[tuple[list[str], int]] = []
|
|||
|
|
for segment in cleaned.split(","):
|
|||
|
|
if not segment or "-" not in segment:
|
|||
|
|
raise ValueError(f"Invalid Sana-WM action segment {segment!r}: expected '<keys>-<duration>'.")
|
|||
|
|
keys_part, duration = segment.rsplit("-", 1)
|
|||
|
|
if not duration.isdigit() or int(duration) <= 0:
|
|||
|
|
raise ValueError(f"Sana-WM action segment {segment!r} has a non-positive duration {duration!r}.")
|
|||
|
|
keys_lower = keys_part.lower()
|
|||
|
|
if keys_lower == "none":
|
|||
|
|
keys: list[str] = []
|
|||
|
|
else:
|
|||
|
|
bad = sorted({key for key in keys_lower if key not in SANA_WM_ALLOWED_ACTION_KEYS})
|
|||
|
|
if bad:
|
|||
|
|
allowed = "".join(sorted(SANA_WM_ALLOWED_ACTION_KEYS))
|
|||
|
|
raise ValueError(f"Sana-WM action segment {segment!r} contains unknown keys {bad}; allowed: {allowed}.")
|
|||
|
|
keys = sorted(set(keys_lower))
|
|||
|
|
segments.append((keys, int(duration)))
|
|||
|
|
return segments
|
|||
|
|
|
|||
|
|
|
|||
|
|
def action_rollout_num_frames(action: str) -> int:
|
|||
|
|
"""Frames ``action`` rolls out: the identity start pose plus one per step.
|
|||
|
|
|
|||
|
|
Totals the durations arithmetically, so an oversized request is rejected
|
|||
|
|
before anything allocates per-frame state.
|
|||
|
|
"""
|
|||
|
|
return sum(duration for _, duration in _parse_action_segments(action)) + 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _parse_action_string(action: str) -> list[list[str]]:
|
|||
|
|
per_frame: list[list[str]] = []
|
|||
|
|
for keys, duration in _parse_action_segments(action):
|
|||
|
|
per_frame.extend([keys] * duration)
|
|||
|
|
return per_frame
|
|||
|
|
|
|||
|
|
|
|||
|
|
def action_string_to_c2w(
|
|||
|
|
action: str,
|
|||
|
|
*,
|
|||
|
|
translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED,
|
|||
|
|
rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG,
|
|||
|
|
pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
"""Roll out an OpenCV-convention ``(N + 1, 4, 4)`` c2w trajectory."""
|
|||
|
|
|
|||
|
|
per_frame = _parse_action_string(action)
|
|||
|
|
rotate_rad = math.radians(rotation_speed_deg)
|
|||
|
|
pitch_limit_rad = math.radians(pitch_limit_deg)
|
|||
|
|
current = np.eye(4, dtype=np.float64)
|
|||
|
|
current_pitch = 0.0
|
|||
|
|
poses = [current.copy()]
|
|||
|
|
|
|||
|
|
for keys in per_frame:
|
|||
|
|
held = set(keys)
|
|||
|
|
rotation = current[:3, :3]
|
|||
|
|
translation = current[:3, 3]
|
|||
|
|
|
|||
|
|
pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0)
|
|||
|
|
new_pitch = current_pitch + pitch_delta
|
|||
|
|
if -pitch_limit_rad <= new_pitch <= pitch_limit_rad:
|
|||
|
|
current_pitch = new_pitch
|
|||
|
|
else:
|
|||
|
|
pitch_delta = 0.0
|
|||
|
|
|
|||
|
|
yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0)
|
|||
|
|
new_rotation = _rot_y(yaw_delta) @ rotation @ _rot_x(pitch_delta)
|
|||
|
|
|
|||
|
|
forward = new_rotation[:, 2].copy()
|
|||
|
|
forward[1] = 0.0
|
|||
|
|
right = new_rotation[:, 0].copy()
|
|||
|
|
right[1] = 0.0
|
|||
|
|
forward_norm = float(np.linalg.norm(forward))
|
|||
|
|
right_norm = float(np.linalg.norm(right))
|
|||
|
|
if forward_norm > 0:
|
|||
|
|
forward /= forward_norm + 1e-6
|
|||
|
|
if right_norm > 0:
|
|||
|
|
right /= right_norm + 1e-6
|
|||
|
|
|
|||
|
|
move = np.zeros(3, dtype=np.float64)
|
|||
|
|
if "w" in held:
|
|||
|
|
move += forward * translation_speed
|
|||
|
|
if "s" in held:
|
|||
|
|
move -= forward * translation_speed
|
|||
|
|
if "d" in held:
|
|||
|
|
move += right * translation_speed
|
|||
|
|
if "a" in held:
|
|||
|
|
move -= right * translation_speed
|
|||
|
|
|
|||
|
|
current = np.eye(4, dtype=np.float64)
|
|||
|
|
current[:3, :3] = new_rotation
|
|||
|
|
current[:3, 3] = translation + move
|
|||
|
|
poses.append(current.copy())
|
|||
|
|
|
|||
|
|
return np.stack(poses, axis=0).astype(np.float32)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Invert homogeneous rigid transforms with shape ``(..., 4, 4)``."""
|
|||
|
|
|
|||
|
|
rotation = transform[..., :3, :3]
|
|||
|
|
translation = transform[..., :3, 3]
|
|||
|
|
rotation_inv = rotation.transpose(-1, -2)
|
|||
|
|
translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1)
|
|||
|
|
output = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat(transform.shape[:-2] + (1, 1))
|
|||
|
|
output[..., :3, :3] = rotation_inv
|
|||
|
|
output[..., :3, 3] = translation_inv
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
|
|||
|
|
def compute_raymap(
|
|||
|
|
intrinsics: torch.Tensor,
|
|||
|
|
poses: torch.Tensor,
|
|||
|
|
height: int,
|
|||
|
|
width: int,
|
|||
|
|
*,
|
|||
|
|
use_plucker: bool = True,
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
"""Compute SANA-WM geometry ray maps.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
intrinsics: ``(T, 4)`` tensor in ``[fx, fy, cx, cy]`` order.
|
|||
|
|
poses: ``(T, 4, 4)`` camera-to-world matrices.
|
|||
|
|
height: raymap height in latent pixels.
|
|||
|
|
width: raymap width in latent pixels.
|
|||
|
|
use_plucker: return ``[direction, moment]`` if true, otherwise
|
|||
|
|
``[origin, direction]``.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
if intrinsics.ndim != 2 or intrinsics.shape[-1] != 4:
|
|||
|
|
raise ValueError(f"Sana-WM intrinsics must have shape (T, 4), got {tuple(intrinsics.shape)}.")
|
|||
|
|
if poses.ndim != 3 or poses.shape[-2:] != (4, 4):
|
|||
|
|
raise ValueError(f"Sana-WM poses must have shape (T, 4, 4), got {tuple(poses.shape)}.")
|
|||
|
|
if intrinsics.shape[0] != poses.shape[0]:
|
|||
|
|
raise ValueError("Sana-WM intrinsics and poses must have the same temporal length.")
|
|||
|
|
|
|||
|
|
frames = intrinsics.shape[0]
|
|||
|
|
device = intrinsics.device
|
|||
|
|
dtype = intrinsics.dtype
|
|||
|
|
y_grid, x_grid = torch.meshgrid(
|
|||
|
|
torch.arange(height, device=device, dtype=dtype),
|
|||
|
|
torch.arange(width, device=device, dtype=dtype),
|
|||
|
|
indexing="ij",
|
|||
|
|
)
|
|||
|
|
x_grid = x_grid[None].expand(frames, -1, -1)
|
|||
|
|
y_grid = y_grid[None].expand(frames, -1, -1)
|
|||
|
|
|
|||
|
|
fx = intrinsics[:, 0].view(frames, 1, 1)
|
|||
|
|
fy = intrinsics[:, 1].view(frames, 1, 1)
|
|||
|
|
cx = intrinsics[:, 2].view(frames, 1, 1)
|
|||
|
|
cy = intrinsics[:, 3].view(frames, 1, 1)
|
|||
|
|
dirs_cam = torch.stack([(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], dim=-1)
|
|||
|
|
|
|||
|
|
rotation = poses[:, :3, :3]
|
|||
|
|
translation = poses[:, :3, 3]
|
|||
|
|
dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam)
|
|||
|
|
dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True).clamp_min(1e-12)
|
|||
|
|
origins = translation.view(frames, 1, 1, 3).expand_as(dirs_world)
|
|||
|
|
|
|||
|
|
if use_plucker:
|
|||
|
|
moments = torch.cross(origins, dirs_world, dim=-1)
|
|||
|
|
return torch.cat([dirs_world, moments], dim=-1)
|
|||
|
|
return torch.cat([origins, dirs_world], dim=-1)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def intrinsics_to_vec4_array(intrinsics: Any, *, num_frames: int, height: int, width: int) -> np.ndarray:
|
|||
|
|
"""Normalize the intrinsics payload to an ``(F, 4)`` ``[fx, fy, cx, cy]`` array.
|
|||
|
|
|
|||
|
|
Accepts either ``None`` (derive from the output resolution) or the
|
|||
|
|
``{fx, fy, cx, cy}`` mapping — the only form the request contract exposes.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
if intrinsics is None:
|
|||
|
|
focal = float(max(height, width))
|
|||
|
|
vec = np.array([focal, focal, width / 2.0, height / 2.0], dtype=np.float32)
|
|||
|
|
return np.broadcast_to(vec, (num_frames, 4)).copy()
|
|||
|
|
|
|||
|
|
if isinstance(intrinsics, dict):
|
|||
|
|
vec = np.array(
|
|||
|
|
[intrinsics["fx"], intrinsics["fy"], intrinsics["cx"], intrinsics["cy"]],
|
|||
|
|
dtype=np.float32,
|
|||
|
|
)
|
|||
|
|
return np.broadcast_to(vec, (num_frames, 4)).copy()
|
|||
|
|
|
|||
|
|
raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _pack_camera_conditions(
|
|||
|
|
poses: torch.Tensor,
|
|||
|
|
intrinsics_latent: torch.Tensor,
|
|||
|
|
*,
|
|||
|
|
num_frames: int,
|
|||
|
|
latent_height: int,
|
|||
|
|
latent_width: int,
|
|||
|
|
vae_time_stride: int,
|
|||
|
|
) -> dict[str, torch.Tensor]:
|
|||
|
|
# This yields exactly the latent frame count the VAE produces:
|
|||
|
|
# ``len(arange(0, n, s)) == ceil(n / s) == (n - 1) // s + 1``.
|
|||
|
|
time_indices = torch.arange(0, num_frames, vae_time_stride)
|
|||
|
|
|
|||
|
|
raymap = torch.cat([poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], dim=-1)
|
|||
|
|
|
|||
|
|
chunks: list[torch.Tensor] = []
|
|||
|
|
spatial_raymap_frames: list[torch.Tensor] = []
|
|||
|
|
chunk_starts = time_indices - (vae_time_stride - 1)
|
|||
|
|
for start, t_idx in zip(chunk_starts, time_indices):
|
|||
|
|
start_idx = max(0, int(start))
|
|||
|
|
end_idx = start_idx + vae_time_stride
|
|||
|
|
chunk_poses = poses[start_idx:end_idx]
|
|||
|
|
chunk_intrinsics = intrinsics_latent[start_idx:end_idx]
|
|||
|
|
if chunk_poses.shape[0] < vae_time_stride:
|
|||
|
|
pad = vae_time_stride - chunk_poses.shape[0]
|
|||
|
|
chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0)
|
|||
|
|
chunk_intrinsics = torch.cat([chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)], dim=0)
|
|||
|
|
plucker = compute_raymap(chunk_intrinsics, chunk_poses, latent_height, latent_width, use_plucker=True)
|
|||
|
|
chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_height, latent_width))
|
|||
|
|
|
|||
|
|
# Per-frame spatial ray-direction map for raymap_embedder: take the
|
|||
|
|
# representative frame in each chunk (the anchor at t_idx) and use only
|
|||
|
|
# the direction channels (last 3 of 6 in the origin+direction layout).
|
|||
|
|
t = int(t_idx)
|
|||
|
|
od = compute_raymap(
|
|||
|
|
intrinsics_latent[t : t + 1],
|
|||
|
|
poses[t : t + 1],
|
|||
|
|
latent_height,
|
|||
|
|
latent_width,
|
|||
|
|
use_plucker=False,
|
|||
|
|
) # [1, H, W, 6]
|
|||
|
|
spatial_raymap_frames.append(od[0, :, :, 3:]) # [H, W, 3] — direction only
|
|||
|
|
|
|||
|
|
chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3)
|
|||
|
|
# spatial_raymap: [3, F_latent, H, W] — consumed by raymap_embedder
|
|||
|
|
spatial_raymap = torch.stack(spatial_raymap_frames).permute(3, 0, 1, 2)
|
|||
|
|
return {"raymap": raymap, "chunk_plucker": chunk_plucker, "spatial_raymap": spatial_raymap}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_plucker_condition(
|
|||
|
|
condition: SanaWmCameraCondition,
|
|||
|
|
*,
|
|||
|
|
vae_stride: tuple[int, int, int] | list[int] = SANA_WM_DEFAULT_VAE_STRIDE,
|
|||
|
|
) -> dict[str, torch.Tensor]:
|
|||
|
|
"""Build native Stage-1 camera tensors from normalized request metadata."""
|
|||
|
|
|
|||
|
|
if bool(condition.action) == (condition.poses is not None):
|
|||
|
|
raise ValueError("Sana-WM camera condition requires exactly one of action or poses.")
|
|||
|
|
|
|||
|
|
if condition.action:
|
|||
|
|
poses_c2w = action_string_to_c2w(
|
|||
|
|
condition.action,
|
|||
|
|
translation_speed=condition.translation_speed,
|
|||
|
|
rotation_speed_deg=condition.rotation_speed_deg,
|
|||
|
|
pitch_limit_deg=condition.pitch_limit_deg,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
poses_c2w = np.asarray(condition.poses, dtype=np.float32)
|
|||
|
|
|
|||
|
|
if poses_c2w.ndim != 3 or poses_c2w.shape[-2:] != (4, 4):
|
|||
|
|
raise ValueError(f"Sana-WM c2w poses must have shape (F, 4, 4), got {poses_c2w.shape}.")
|
|||
|
|
|
|||
|
|
num_frames = condition.num_frames or int(poses_c2w.shape[0])
|
|||
|
|
if poses_c2w.shape[0] < num_frames:
|
|||
|
|
if condition.action:
|
|||
|
|
num_frames = int(poses_c2w.shape[0])
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"Sana-WM c2w poses length {poses_c2w.shape[0]} is shorter than num_frames {num_frames}.")
|
|||
|
|
poses_c2w = poses_c2w[:num_frames]
|
|||
|
|
|
|||
|
|
intrinsics = intrinsics_to_vec4_array(
|
|||
|
|
condition.intrinsics,
|
|||
|
|
num_frames=num_frames,
|
|||
|
|
height=condition.height,
|
|||
|
|
width=condition.width,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
vae_time_stride = int(vae_stride[0])
|
|||
|
|
vae_spatial_stride = int(vae_stride[-1])
|
|||
|
|
latent_height = condition.height // vae_spatial_stride
|
|||
|
|
latent_width = condition.width // vae_spatial_stride
|
|||
|
|
|
|||
|
|
poses = torch.from_numpy(poses_c2w).float()
|
|||
|
|
first_inv = get_pose_inverse(poses[0:1]).squeeze(0)
|
|||
|
|
poses_rel = torch.matmul(first_inv, poses[1:])
|
|||
|
|
poses = torch.cat([torch.eye(4, dtype=poses.dtype).unsqueeze(0), poses_rel], dim=0)
|
|||
|
|
|
|||
|
|
intrinsics_latent = torch.from_numpy(intrinsics).float()
|
|||
|
|
intrinsics_latent[:, [0, 2]] *= latent_width / float(condition.width)
|
|||
|
|
intrinsics_latent[:, [1, 3]] *= latent_height / float(condition.height)
|
|||
|
|
|
|||
|
|
return _pack_camera_conditions(
|
|||
|
|
poses,
|
|||
|
|
intrinsics_latent,
|
|||
|
|
num_frames=num_frames,
|
|||
|
|
latent_height=latent_height,
|
|||
|
|
latent_width=latent_width,
|
|||
|
|
vae_time_stride=vae_time_stride,
|
|||
|
|
)
|