decoder.py

"""D4RT decoder: lightweight cross-attention transformer (paper Section 2.2)."""

from __future__ import annotations

import torch
import torch.nn as nn

from attention import LightweightDecoder
from query_encoder import QueryEncoder

try:  # VGGT is optional; only needed to build the real encoder.
    from encoder import VideoEncoder
except ImportError:
    VideoEncoder = None


class D4RTDecoder(nn.Module):
    """Cross-attention decoder + one linear head per predicted attribute."""

    _HEAD_SIZES = {
        "points_3d": 3,
        "points_2d": 2,
        "visibility": 1,
        "surface_normal": 3,
        "motion_vector": 3,
        "confidence": 1,
    }

    def __init__(
        self,
        d_model: int = 768,
        num_layers: int = 4,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers

        self.stack = LightweightDecoder(
            dim=d_model,
            depth=num_layers,
            heads=num_heads,
            mlp_ratio=mlp_ratio,
        )

        self.heads = nn.ModuleDict({
            name: nn.Linear(d_model, out_dim)
            for name, out_dim in self._HEAD_SIZES.items()
        })

    def forward(
        self,
        queries: torch.Tensor,   # (B, Q, D)
        F: torch.Tensor,         # (B, N, D)
    ) -> dict[str, torch.Tensor]:
        feat = self.stack(queries, F)   # (B, Q, D)
        return {
            name: head(feat)            # (B, Q, out_dim)
            for name, head in self.heads.items()
        }


class D4RTModel(nn.Module):
    """Full model: VideoEncoder -> scene_proj -> QueryEncoder -> D4RTDecoder."""

    def __init__(
        self,
        encoder: nn.Module | None = None,
        embed_dim: int = 768,
        scene_token_dim: int = 2048,
        num_layers: int = 4,
        num_heads: int = 8,
        max_timesteps: int = 32,
        num_frequencies: int = 64,
        patch_size: int = 9,
        device: str = "cpu",
    ) -> None:
        super().__init__()
        self.embed_dim = embed_dim

        if encoder is not None:
            self.encoder = encoder
        else:
            if VideoEncoder is None:
                raise ImportError(
                    "No encoder provided and VGGT is not installed. Run:\n"
                    "  pip install git+https://github.com/facebookresearch/vggt.git"
                )
            self.encoder = VideoEncoder(device=device)

        # VGGT is a frozen pretrained backbone.
        for p in self.encoder.parameters():
            p.requires_grad_(False)

        # Map 2048-d VGGT tokens down to the decoder's D-d space.
        self.scene_proj = nn.Linear(scene_token_dim, embed_dim)

        self.query_encoder = QueryEncoder(
            embed_dim=embed_dim,
            num_frequencies=num_frequencies,
            max_timesteps=max_timesteps,
            patch_size=patch_size,
        )

        self.decoder = D4RTDecoder(
            d_model=embed_dim,
            num_layers=num_layers,
            num_heads=num_heads,
        )

    def forward(
        self,
        video: torch.Tensor,             # (B, T, 3, H, W) in [0, 1]
        query_params: dict[str, torch.Tensor],
    ) -> dict[str, torch.Tensor]:
        """video + queries -> predictions. query_params has u, v, t_src, t_tgt, t_cam."""
        B, T, C, H, W = video.shape

        F_raw = self.encoder(video)                    # (B, T, P, D_in)
        B, T, P, D_in = F_raw.shape
        F_raw = F_raw.reshape(B, T * P, D_in)          # (B, N, D_in)
        F = self.scene_proj(F_raw)                     # (B, N, D)

        queries = self.query_encoder(
            query_params["u"],
            query_params["v"],
            query_params["t_src"],
            query_params["t_tgt"],
            query_params["t_cam"],
            video,
        )                                              # (B, Q, D)

        return self.decoder(queries, F)


class _MockEncoder(nn.Module):
    """Stand-in for VideoEncoder when VGGT is unavailable. Frozen conv patch embed."""

    def __init__(self, scene_token_dim: int = 2048, patch_size: int = 32) -> None:
        super().__init__()
        self.patch_embed = nn.Conv2d(
            3, scene_token_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, video: torch.Tensor) -> torch.Tensor:
        B, T, C, H, W = video.shape
        x = self.patch_embed(video.reshape(B * T, C, H, W))  # (BT, Din, h, w)
        x = x.flatten(2).transpose(1, 2)                     # (BT, P, Din)
        return x.reshape(B, T, -1, x.shape[-1])              # (B, T, P, Din)


if __name__ == "__main__":
    torch.manual_seed(0)

    B, T, H, W = 2, 4, 128, 128
    Q = 32
    D = 768

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

    enc = _MockEncoder(scene_token_dim=2048, patch_size=32)
    for p in enc.parameters():
        p.requires_grad_(False)

    model = D4RTModel(encoder=enc, device=device)
    model.train()

    video = torch.rand(B, T, 3, H, W, device=device)
    query_params = {
        "u": torch.rand(B, Q, device=device),
        "v": torch.rand(B, Q, device=device),
        "t_src": torch.randint(0, T, (B, Q), device=device),
        "t_tgt": torch.randint(0, T, (B, Q), device=device),
        "t_cam": torch.randint(0, T, (B, Q), device=device),
    }

    preds = model(video, query_params)

    print("\n--- Output shapes ---")
    for name, tensor in preds.items():
        assert tensor.shape[0] == B and tensor.shape[1] == Q, (
            f"{name}: expected batch {B} x queries {Q}, got {tensor.shape}"
        )
        print(f"  {name:<15} {tuple(tensor.shape)}")

    # Decoder isolated shape check.
    F_iso = torch.randn(B, T * 4, D, device=device)
    queries_iso = torch.randn(B, Q, D, device=device)
    iso_preds = model.decoder(queries_iso, F_iso)
    assert iso_preds["points_3d"].shape == (B, Q, 3)
    assert iso_preds["visibility"].shape == (B, Q, 1)
    print(f"\nDecoder (isolated) F: {tuple(F_iso.shape)} -> points_3d "
          f"{tuple(iso_preds['points_3d'].shape)}")

    # Query independence: no self-attention between queries.
    q2 = queries_iso.clone()
    q2[:, 0, :] = 0.0
    out_full = model.decoder(queries_iso, F_iso)
    out_masked = model.decoder(q2, F_iso)
    diff = (out_full["points_3d"][:, 1:, :] - out_masked["points_3d"][:, 1:, :]).abs().max().item()
    print(f"\nQuery-independence check (other queries): max diff = {diff:.2e}")
    assert diff < 1e-6, "Queries are NOT independent — self-attention leaked in!"

    # Gradient flow.
    loss = sum(t.square().mean() for t in preds.values())
    loss.backward()

    print("\n--- Gradient flow ---")
    enc_grads = [p.grad for p in model.encoder.parameters() if p.requires_grad]
    assert len(enc_grads) == 0, "Encoder has trainable params — it must be frozen!"
    print(f"  encoder            frozen OK: 0 trainable params")

    trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
    for n, p in trainable:
        assert p.grad is not None, f"No gradient for {n}"
        print(f"  {n:<28} grad {tuple(p.grad.shape)}")

    assert all(p.grad is not None for _, p in trainable), "Missing gradients!"
    print("\nAll decoder tests passed.")