encoder.py

"""VideoEncoder: frozen VGGT-1B feature extractor for D4RT.

F = patch tokens from VGGT's aggregator final layer, (B, S, P, 2*embed_dim),
with camera/register tokens stripped. VGGT source: facebookresearch/vggt.
"""

import torch
import torch.nn as nn

try:
    from vggt.models.vggt import VGGT
except ImportError:
    VGGT = None


class VideoEncoder(nn.Module):
    """Frozen VGGT-1B encoder; outputs (B, S, num_patches, 2*embed_dim)."""

    def __init__(self, img_size: int = 518, device: str = "cuda") -> None:
        super().__init__()
        if VGGT is None:
            raise ImportError(
                "VGGT not installed. Run:\n"
                "  pip install git+https://github.com/facebookresearch/vggt.git\n"
                "and ensure 'torch', 'einops' are available."
            )
        self.vggt = VGGT()
        self.vggt.to(device)

        for p in self.vggt.parameters():
            p.requires_grad_(False)

        self.embed_dim = 1024
        self.patch_size = 14
        self.img_size = img_size
        self.last_layer_idx = 23  # aggregator default intermediate_layer_idx
        self.patch_start_idx = 5  # 1 camera + 4 register tokens

        self._num_patches = (img_size // self.patch_size) ** 2  # 1369

    @torch.no_grad()
    def forward(
        self, images: torch.Tensor, return_all_layers: bool = False
    ) -> torch.Tensor | list[torch.Tensor]:
        """images (B, S, 3, H, W) in [0, 1] -> F (B, S, P, 2048).

        With return_all_layers=True, returns the 4 cached layers instead.
        """
        B, S, C, H, W = images.shape
        assert C == 3, f"Expected 3 channels, got {C}"

        images = images.clamp(0, 1)

        aggregated_tokens_list, patch_start_idx = self.vggt.aggregator(images)

        def _strip(tokens_layer: torch.Tensor) -> torch.Tensor:
            # Remove the camera + register tokens from each layer.
            return tokens_layer[:, :, patch_start_idx:]

        if return_all_layers:
            return [_strip(aggregated_tokens_list[idx]) for idx in (4, 11, 17, 23)]

        return _strip(aggregated_tokens_list[self.last_layer_idx])

    @property
    def output_dim(self) -> int:
        return 2 * self.embed_dim  # 2048

    @property
    def num_patches(self) -> int:
        return self._num_patches

    @property
    def spatial_shape(self) -> tuple[int, int]:
        side = self.img_size // self.patch_size
        return (side, side)


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}")

    print("Loading VGGT-1B encoder...")
    enc = VideoEncoder(device=device)
    enc.eval()

    B, S = 1, 4
    fake = torch.rand(B, S, 3, 518, 518, device=device)

    with torch.no_grad():
        F_map = enc(fake)

    side = 518 // 14
    print(f"F shape:  {F_map.shape}")        # (1, 4, 1369, 2048)
    print(f"Expected: ({B}, {S}, {side*side}, {enc.output_dim})")
    print(f"F mean:   {F_map.mean().item():.4f}")
    print(f"F std:    {F_map.std().item():.4f}")
    print(f"Num patches: {enc.num_patches}  spatial: {enc.spatial_shape}")
    print("Encoder OK ✓")