src/mavt/model/mavt.py

"""MAVT: Memory-Augmented Vision Tokenizer.

Unified 7-stage pipeline:
  1. Patchify (Conv3d, modality-specific)
  2. Hybrid Transformer-RGAT Backbone (12 blocks)
  3. Content-Detail Split (slot attention)
  4. Dual Latent Projection (VAE + Semantic)
  5. Modality-Specific Decoder
  6. Losses (handled by LightningModule)
  7. Outputs
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional, Tuple

import torch
import torch.nn as nn

from mavt.model.patchify import PatchifyEncoder
from mavt.model.backbone import HybridBackbone
from mavt.model.content_detail_split import ContentDetailSplit
from mavt.model.latent_heads import VAEHead
from mavt.model.decoder import AsymmetricDecoder, UnderstandingDecoder


# Compression ratios per modality (content, detail)
_MODALITY_RATIOS = {
    'image':  (0.25, 0.25),
    'video':  (0.25, 0.25),
    'threed': (0.35, 0.25),
}


@dataclass
class MAVTOutput:
    reconstruction: torch.Tensor      # pixel-space reconstruction
    z: torch.Tensor                    # VAE latent
    mu: torch.Tensor
    logvar: torch.Tensor
    latent_positions: torch.Tensor     # (N_z, 4), content zeros + local detail centers
    latent_token_types: torch.Tensor    # (N_z,), 0=content, 1=detail
    semantic: torch.Tensor             # (B, semantic_dim)
    loss_kl: torch.Tensor
    cd_metrics: Dict[str, torch.Tensor]  # slot_diversity, residual_ratio


class MAVT(nn.Module):
    """Full MAVT model.

    All hyper-parameters are configurable via YAML (Lightning CLI).
    """

    def __init__(
        self,
        embed_dim: int = 1152,
        num_heads: int = 16,
        num_blocks: int = 12,
        patch_size: int = 16,
        t_patch: int = 2,
        # C-D Split
        num_slot_heads: int = 8,
        num_slot_layers: int = 2,
        local_detail_window_size: int = 1,
        local_detail_temporal_window_size: int = 1,
        # VAE
        latent_dim: int = 32,
        kl_weight: float = 1e-4,
        # Semantic
        semantic_dim: int = 768,
        # Decoder
        dec_dim: int = 768,
        num_dec_attn_blocks: int = 4,
        # RGAT
        r_s: int = 2,
        r_t: int = 1,
        # Training
        use_gradient_checkpointing: bool = False,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.latent_dim = latent_dim
        self.patch_size = patch_size

        # Stage 1
        self.patchify = PatchifyEncoder(embed_dim, patch_size, t_patch)

        # Stage 2
        self.backbone = HybridBackbone(
            dim=embed_dim, num_heads=num_heads, num_blocks=num_blocks,
            mlp_ratio=mlp_ratio, dropout=dropout,
            r_s=r_s, r_t=r_t,
            use_gradient_checkpointing=use_gradient_checkpointing,
        )

        # Stage 3
        self.cd_split = ContentDetailSplit(
            dim=embed_dim, num_heads=num_slot_heads, num_slot_layers=num_slot_layers,
            local_detail_window_size=local_detail_window_size,
            local_detail_temporal_window_size=local_detail_temporal_window_size,
        )

        # Stage 4 — VAE bottleneck only (semantic moved downstream of z)
        self.vae_head = VAEHead(embed_dim, latent_dim, kl_weight)

        # Stage 5 — two heads decoding from the shared latent z
        # 5a. Reconstruction head: z → pixel
        self.decoder = AsymmetricDecoder(
            latent_dim=latent_dim, dec_dim=dec_dim,
            num_attn_blocks=num_dec_attn_blocks, num_heads=num_heads,
            mlp_ratio=mlp_ratio,
        )
        # 5b. Understanding head: z → semantic vector aligned with vision teacher
        self.understanding_decoder = UnderstandingDecoder(
            latent_dim=latent_dim, dec_dim=dec_dim,
            semantic_dim=semantic_dim, num_heads=8, num_layers=2,
            mlp_ratio=mlp_ratio,
        )

    # ------------------------------------------------------------------ #
    #  Helpers                                                             #
    # ------------------------------------------------------------------ #

    def _grid_shape(self, modality: str, x: torch.Tensor) -> tuple:
        """Return (H_grid, W_grid) or (Tp, Hg, Wg) based on input shape."""
        if modality == 'image':
            _, _, H, W = x.shape
            return (H // self.patch_size, W // self.patch_size)
        elif modality == 'video':
            _, _, T, H, W = x.shape
            return (T // 2, H // self.patch_size, W // self.patch_size)
        elif modality == 'threed':
            _, _, _, S, _ = x.shape   # (B, 3planes, 3ch, S, S)
            return (S // self.patch_size, S // self.patch_size)
        raise ValueError(modality)

    # ------------------------------------------------------------------ #
    #  Forward                                                             #
    # ------------------------------------------------------------------ #

    def forward(
        self,
        x: torch.Tensor,
        modality: str,
        decode: bool = True,
    ) -> MAVTOutput:
        """
        x : raw input tensor (see patchify.py for shapes per modality)
        modality : 'image' | 'video' | 'threed'
        decode : if False, skip decoder (encoder-only mode for downstream tasks)
        """
        grid_shape = self._grid_shape(modality, x)

        # Stage 1 — Patchify
        tokens, positions, plane_ids = self.patchify(x, modality)
        # tokens: (B, N, D), positions: (N, 4), plane_ids: (N,)

        # Stage 2 — Hybrid backbone
        features = self.backbone(tokens, positions, plane_ids, modality)

        # Stage 3 — Content-Detail Split
        content_ratio, detail_ratio = _MODALITY_RATIOS[modality]
        compressed, cd_metrics, latent_positions, latent_token_types = self.cd_split(
            features,
            positions=positions,
            plane_ids=plane_ids,
            content_ratio=content_ratio,
            detail_ratio=detail_ratio,
            return_metadata=True,
        )  # (B, N_c + N_d, D)

        # Stage 4 — VAE bottleneck (semantic now derives from z, not compressed)
        z, mu, logvar, loss_kl = self.vae_head(compressed)

        # Stage 5a — Understanding head: z → semantic
        # Always run (cheap, gives semantic supervision signal even when decode=False)
        semantic = self.understanding_decoder(z)

        # Stage 5b — Reconstruction head: z → pixel
        if decode:
            recon = self.decoder(
                z, positions, modality, grid_shape,
                latent_positions=latent_positions,
                latent_token_types=latent_token_types,
            )
        else:
            recon = torch.zeros(1, device=x.device)  # placeholder

        return MAVTOutput(
            reconstruction=recon,
            z=z,
            mu=mu,
            logvar=logvar,
            latent_positions=latent_positions,
            latent_token_types=latent_token_types,
            semantic=semantic,
            loss_kl=loss_kl,
            cd_metrics=cd_metrics,
        )

    def encode(self, x: torch.Tensor, modality: str) -> Tuple[torch.Tensor, torch.Tensor]:
        """Convenience: return (z, semantic) without decoding."""
        out = self.forward(x, modality, decode=False)
        return out.z, out.semantic

    def load_siglip2_weights(self, model_name: str = "google/siglip2-base-patch16-224",
                              freeze_stages: int = 10) -> None:
        self.backbone.load_siglip2_weights(model_name, freeze_stages)

    # ------------------------------------------------------------------ #
    #  Eager pre-creation of slot poolers                                 #
    # ------------------------------------------------------------------ #

    def prepare_for_modalities(self, specs: Iterable[Dict[str, Any]]) -> None:
        """Pre-create every SlotPooler the trainer will need.

        Must be called BEFORE the optimizer is built (e.g. from
        LightningModule.setup) so the pooler params are picked up by the
        optimizer's param_groups. Without this, poolers are created lazily
        in ContentDetailSplit.forward and their parameters never receive
        gradient updates.

        Each spec dict has key 'modality' plus modality-specific shape keys:
            image  : {'modality': 'image',  'resolution': H}
            video  : {'modality': 'video',  'resolution': H, 'frames': T,
                      't_patch': 2}                # t_patch optional
            threed : {'modality': 'threed', 'resolution': S}
        """
        for spec in specs:
            modality = spec['modality']
            if modality == 'image':
                H = spec['resolution']
                Hp = H // self.patch_size
                N  = Hp * Hp
            elif modality == 'video':
                H  = spec['resolution']
                T  = spec['frames']
                tp = spec.get('t_patch', 2)
                Tp = T // tp
                Hp = H // self.patch_size
                N  = Tp * Hp * Hp
            elif modality == 'threed':
                S  = spec['resolution']
                Sp = S // self.patch_size
                N  = 3 * Sp * Sp                  # 3 planes
            else:
                raise ValueError(f"Unknown modality in spec: {modality!r}")
            c_r, d_r = _MODALITY_RATIOS[modality]
            N_c = max(1, int(N * c_r))
            N_d = max(1, int(N * d_r))
            self.cd_split.prepare_poolers(N_c, N_d)