query_encoder.py

"""Query encoder: turns a 5-tuple (u, v, t_src, t_tgt, t_cam) into a D-dim token."""

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def fourier_features(
    coords: torch.Tensor,          # (B, Q, 2) in [0, 1]
    num_frequencies: int = 64,
) -> torch.Tensor:                 # (B, Q, 4 * num_frequencies)
    """Sinusoidal (NeRF-style) encoding of 2D coords, sin/cos at 2^l * pi * x."""
    freqs = 2.0 ** torch.arange(
        num_frequencies,
        device=coords.device,
        dtype=coords.dtype,
    )

    scaled = coords.unsqueeze(-1) * freqs * math.pi   # (B, Q, 2, L)
    encoded = torch.stack(
        [torch.sin(scaled), torch.cos(scaled)], dim=-1
    )                                                  # (B, Q, 2, L, 2)

    B, Q = coords.shape[:2]
    return encoded.reshape(B, Q, 4 * num_frequencies)


class PatchEmbedder(nn.Module):
    """Crop a patch_size x patch_size RGB patch around (u, v) and embed it."""

    def __init__(self, patch_size: int = 9, in_channels: int = 3,
                 embed_dim: int = 768) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Linear(patch_size * patch_size * in_channels, embed_dim)

    def forward(
        self,
        video: torch.Tensor,      # (B, T, 3, H, W)
        u: torch.Tensor,          # (B, Q) in [0, 1]
        v: torch.Tensor,          # (B, Q) in [0, 1]
        t_src: torch.Tensor,      # (B, Q) int frame index
    ) -> torch.Tensor:            # (B, Q, D)
        B, T, C, H, W = video.shape
        Q = u.shape[1]
        ps = self.patch_size
        half = ps // 2

        px = (u * (W - 1)).round().long()
        py = (v * (H - 1)).round().long()

        x0, y0 = px - half, py - half
        x1, y1 = x0 + ps, y0 + ps

        # Bounding box corners in [-1, 1] for grid_sample.
        gx0 = (x0.float() / (W - 1)) * 2 - 1
        gx1 = (x1.float() / (W - 1)) * 2 - 1
        gy0 = (y0.float() / (H - 1)) * 2 - 1
        gy1 = (y1.float() / (H - 1)) * 2 - 1

        t_range = torch.linspace(0, 1, ps, device=video.device, dtype=video.dtype)
        grid_x = gx0.unsqueeze(-1) + t_range * (gx1 - gx0).unsqueeze(-1)  # (B,Q,ps)
        grid_y = gy0.unsqueeze(-1) + t_range * (gy1 - gy0).unsqueeze(-1)

        grid = torch.stack([
            grid_x.unsqueeze(2).expand(-1, -1, ps, -1),
            grid_y.unsqueeze(2).expand(-1, -1, ps, -1),
        ], dim=-1)                                                     # (B,Q,ps,ps,2)

        # Gather frame t_src for each query, then sample the patch.
        frames = video.gather(
            1,
            t_src[:, :, None, None, None].expand(-1, -1, C, H, W)
        )   # (B, Q, C, H, W)

        N = B * Q
        patches = F.grid_sample(
            frames.reshape(N, C, H, W),
            grid.reshape(N, ps, ps, 2),
            mode="bilinear",
            padding_mode="border",
            align_corners=False,
        )   # (BQ, 3, ps, ps)

        embeddings = self.proj(patches.reshape(N, -1))   # (BQ, D)
        return embeddings.reshape(B, Q, -1)


class TimestepEmbedding(nn.Module):
    """Learned embedding table per timestep role (src/tgt/cam)."""

    def __init__(self, max_timesteps: int = 32, embed_dim: int = 768) -> None:
        super().__init__()
        self.emb_src = nn.Embedding(max_timesteps, embed_dim)
        self.emb_tgt = nn.Embedding(max_timesteps, embed_dim)
        self.emb_cam = nn.Embedding(max_timesteps, embed_dim)

    def forward(
        self,
        t_src: torch.Tensor,  # (B, Q) int
        t_tgt: torch.Tensor,  # (B, Q) int
        t_cam: torch.Tensor,  # (B, Q) int
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        return (
            self.emb_src(t_src),
            self.emb_tgt(t_tgt),
            self.emb_cam(t_cam),
        )   # each (B, Q, D)


class QueryEncoder(nn.Module):
    """Sum of positional, temporal, and appearance embeddings -> query token."""

    def __init__(
        self,
        embed_dim: int = 768,
        num_frequencies: int = 64,
        max_timesteps: int = 32,
        patch_size: int = 9,
    ) -> None:
        super().__init__()
        self.embed_dim = embed_dim

        fourier_dim = 4 * num_frequencies
        self.fourier_proj = nn.Linear(fourier_dim, embed_dim)
        self.timestep_emb = TimestepEmbedding(max_timesteps, embed_dim)
        self.patch_emb = PatchEmbedder(patch_size, in_channels=3, embed_dim=embed_dim)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(
        self,
        u: torch.Tensor,          # (B, Q) in [0, 1]
        v: torch.Tensor,          # (B, Q) in [0, 1]
        t_src: torch.Tensor,      # (B, Q) int
        t_tgt: torch.Tensor,      # (B, Q) int
        t_cam: torch.Tensor,      # (B, Q) int
        video: torch.Tensor,      # (B, T, 3, H, W)
    ) -> torch.Tensor:            # (B, Q, D)
        coords = torch.stack([u, v], dim=-1)
        pos_emb = self.fourier_proj(fourier_features(coords))

        t_src_emb, t_tgt_emb, t_cam_emb = self.timestep_emb(t_src, t_tgt, t_cam)

        patch_emb = self.patch_emb(video, u, v, t_src)

        query = pos_emb + t_src_emb + t_tgt_emb + t_cam_emb + patch_emb
        return self.norm(query)


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

    B, T, H, W = 2, 8, 256, 256
    Q = 32
    D = 768

    video = torch.rand(B, T, 3, H, W)
    u = torch.rand(B, Q)
    v = torch.rand(B, Q)
    t_src = torch.randint(0, T, (B, Q))
    t_tgt = torch.randint(0, T, (B, Q))
    t_cam = torch.randint(0, T, (B, Q))

    encoder = QueryEncoder(embed_dim=D, num_frequencies=64, max_timesteps=T,
                           patch_size=9)
    out = encoder(u, v, t_src, t_tgt, t_cam, video)

    assert out.shape == (B, Q, D), f"Expected ({B}, {Q}, {D}), got {out.shape}"
    print(f"QueryEncoder output shape: {out.shape}")

    out.sum().backward()
    has_grad = all(p.grad is not None for p in encoder.parameters() if p.requires_grad)
    assert has_grad, "Not all parameters received gradients!"
    print("Gradient check OK")

    coords = torch.rand(2, 16, 2)
    ff = fourier_features(coords, num_frequencies=32)
    assert ff.shape == (2, 16, 4 * 32), f"Fourier shape mismatch: {ff.shape}"
    print(f"fourier_features output shape: {ff.shape}")

    print("All query_encoder tests passed.")