src/model/model.py

import os
import torch
import traceback
import torch.nn.functional as F
from typing import Any, Callable
from torch import nn
from easydict import EasyDict as edict
from einops.layers.torch import Rearrange
from einops import rearrange
from safetensors.torch import load_file
from dataclasses import dataclass

from src.model.sh_eval import _spherical_harmonics
from src.model.transformer import TransformerBlock
from src.model.dpt_head import DPTHead
from src.model.prope_custom import PropeDotProductAttention
from src.model.depth_anything.da_model.da3 import DepthAnything3
from src.model.gaussians import GaussianRenderer, GaussianField
from src.utils.camera_utils import (
    invert_SE3, compute_rays, compute_plucmap, fxfycxcy_to_K, mat_to_quat, quat_to_mat
)


# ---------------------------------------------------------------------------
# Module-level helpers
# ---------------------------------------------------------------------------

def _init_weights(module: nn.Module) -> None:
    """Initialise Linear and Embedding weights with N(0, 0.02) and reset norm layers."""
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            torch.nn.init.zeros_(module.bias)
    elif isinstance(module, (nn.RMSNorm, nn.LayerNorm)):
        module.reset_parameters()
    elif isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)


# ---------------------------------------------------------------------------
# Structured output types
# ---------------------------------------------------------------------------

@dataclass
class CameraBundle:
    pred_i_fxfycxcy: torch.Tensor  # (B, V, 4)
    pred_i_c2w:      torch.Tensor  # (B, V, 4, 4)
    pred_t_fxfycxcy: torch.Tensor  # (B, T, 4)
    pred_t_c2w:      torch.Tensor  # (B, T, 4, 4)
    gt_i_fxfycxcy:   torch.Tensor  # (B, V, 4)
    gt_i_c2w:        torch.Tensor  # (B, V, 4, 4)
    gt_t_fxfycxcy:   torch.Tensor  # (B, T, 4)
    gt_t_c2w:        torch.Tensor  # (B, T, 4, 4)


# ---------------------------------------------------------------------------
# Geometry Expert (DepthAnything3)
# ---------------------------------------------------------------------------

class GeometryExpert:
    """
    Predicts camera poses and intrinsics for all views using DA3.
    """

    def __init__(
        self,
        pose_regressor: DepthAnything3,
        scene_scale: float,
        inference_mode: bool,
        camera_mode: str | None,
    ):
        """
        Args:
            pose_regressor: DA3 model used to predict poses and intrinsics.
            scene_scale: Denominator for translation normalisation (from data config).
            inference_mode: If True, ``camera_mode`` controls which cameras are used.
            camera_mode: One of ``gt_pose_gt_intr``, ``pred_pose_gt_intr``,
                ``pred_pose_pred_intr``; ignored during training.
        """
        self.pose_regressor = pose_regressor
        self.scene_scale = scene_scale
        self.inference_mode = inference_mode
        self.camera_mode = camera_mode

    def predict_cameras(
        self,
        input_data_dict: dict,
        target_data_dict: dict,
    ) -> CameraBundle:
        """Run DA3 on all views, normalize poses to scene scale, and apply camera mode."""
        gt_i_fxfycxcy = input_data_dict["fxfycxcy"].float()
        gt_i_c2w = input_data_dict["c2w"].float()
        gt_t_fxfycxcy = target_data_dict["fxfycxcy"].float()
        gt_t_c2w = target_data_dict["c2w"].float()

        pred_i_fxfycxcy, pred_i_c2w, pred_t_fxfycxcy, pred_t_c2w = self._run_da3_and_normalize(
            input_data_dict["image"], target_data_dict["image"],
        )

        if self.inference_mode:
            pred_i_fxfycxcy, pred_i_c2w, pred_t_fxfycxcy, pred_t_c2w = self._select_cameras(
                pred_i_fxfycxcy, pred_i_c2w, pred_t_fxfycxcy, pred_t_c2w,
                gt_i_fxfycxcy, gt_i_c2w, gt_t_fxfycxcy, gt_t_c2w,
            )

        return CameraBundle(
            pred_i_fxfycxcy=pred_i_fxfycxcy,
            pred_i_c2w=pred_i_c2w,
            pred_t_fxfycxcy=pred_t_fxfycxcy,
            pred_t_c2w=pred_t_c2w,
            gt_i_fxfycxcy=gt_i_fxfycxcy,
            gt_i_c2w=gt_i_c2w,
            gt_t_fxfycxcy=gt_t_fxfycxcy,
            gt_t_c2w=gt_t_c2w,
        )

    def _run_da3_and_normalize(
        self,
        input_images: torch.Tensor,
        target_images: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward DA3 on all input+target images and normalize cameras to scene scale."""
        _, v, _, h, w = input_images.shape
        with torch.autocast(device_type="cuda", enabled=False):
            all_images = torch.cat([input_images, target_images], dim=1)
            output = self.pose_regressor(all_images)

            i_fxfycxcy_raw = output["fxfycxcy"][:, :v].float()
            i_c2w_raw = output["extrinsics"][:, :v].float()
            t_fxfycxcy_raw = output["fxfycxcy"][:, v:].float()
            t_c2w_raw = output["extrinsics"][:, v:].float()

            # DA3 outputs intrinsics at its own resolution; rescale to the model's input resolution.
            da_h, da_w = input_images.shape[-2:]
            intr_scale = torch.tensor(
                [w / float(da_w), h / float(da_h), w / float(da_w), h / float(da_h)],
                dtype=i_fxfycxcy_raw.dtype,
                device=i_fxfycxcy_raw.device,
            ).view(1, 1, 4)

            i_fxfycxcy = i_fxfycxcy_raw * intr_scale
            t_fxfycxcy = t_fxfycxcy_raw * intr_scale
            i_c2w, t_c2w = self._shared_scene_normalization(
                input_c2ws=i_c2w_raw,
                target_c2ws=t_c2w_raw,
                scene_scale=self.scene_scale,
            )
        return i_fxfycxcy, i_c2w, t_fxfycxcy, t_c2w

    @staticmethod
    def _shared_scene_normalization(
        input_c2ws: torch.Tensor,                 # (V, 4, 4) or (B, V, 4, 4)
        target_c2ws: torch.Tensor | None = None,  # (T, 4, 4) or (B, T, 4, 4)
        scene_scale: float = 1.0,
    ):
        """
        Match scene normalization used in dataset.py:
        1) Build canonical frame from input poses only (Gram-Schmidt).
        2) Apply same transform to both input/target poses.
        3) Scale translations by max abs translation from normalized input poses.
        """
        squeeze_input = input_c2ws.ndim == 3
        squeeze_target = target_c2ws is not None and target_c2ws.ndim == 3

        if squeeze_input:
            input_c2ws = input_c2ws.unsqueeze(0)
        if target_c2ws is not None and squeeze_target:
            target_c2ws = target_c2ws.unsqueeze(0)

        position_avg = input_c2ws[:, :, :3, 3].mean(dim=1)   # (B, 3)
        forward_avg = input_c2ws[:, :, :3, 2].mean(dim=1)    # (B, 3)
        down_avg = input_c2ws[:, :, :3, 1].mean(dim=1)       # (B, 3)

        forward_avg = F.normalize(forward_avg, dim=-1)
        down_proj = (down_avg * forward_avg).sum(dim=-1, keepdim=True) * forward_avg
        down_avg = F.normalize(down_avg - down_proj, dim=-1)
        right_avg = torch.cross(down_avg, forward_avg, dim=-1)

        pos_avg = torch.eye(
            4, dtype=input_c2ws.dtype, device=input_c2ws.device
        ).expand(input_c2ws.shape[0], 4, 4).clone()
        pos_avg[:, :3, 0] = right_avg
        pos_avg[:, :3, 1] = down_avg
        pos_avg[:, :3, 2] = forward_avg
        pos_avg[:, :3, 3] = position_avg
        pos_avg_inv = torch.linalg.inv(pos_avg)

        input_c2ws = torch.matmul(pos_avg_inv.unsqueeze(1), input_c2ws)
        if target_c2ws is not None:
            target_c2ws = torch.matmul(pos_avg_inv.unsqueeze(1), target_c2ws)

        translations = input_c2ws[:, :, :3, 3].clone().detach()
        scene_extent = translations.abs().amax(dim=(1, 2))

        scale = 1.0 / (scene_scale * scene_extent)
        # Avoid in-place writes on sliced views (e.g. [..., :3, 3]) to keep
        # autograd version tracking consistent.
        input_scaled_t = input_c2ws[:, :, :3, 3] * scale[:, None, None]
        input_c2ws = torch.cat(
            [
                torch.cat([input_c2ws[:, :, :3, :3], input_scaled_t.unsqueeze(-1)], dim=-1),
                input_c2ws[:, :, 3:, :],
            ],
            dim=-2,
        )
        if target_c2ws is not None:
            target_scaled_t = target_c2ws[:, :, :3, 3] * scale[:, None, None]
            target_c2ws = torch.cat(
                [
                    torch.cat([target_c2ws[:, :, :3, :3], target_scaled_t.unsqueeze(-1)], dim=-1),
                    target_c2ws[:, :, 3:, :],
                ],
                dim=-2,
            )

        if squeeze_input:
            input_c2ws = input_c2ws.squeeze(0)
        if target_c2ws is not None and squeeze_target:
            target_c2ws = target_c2ws.squeeze(0)

        if target_c2ws is None:
            return input_c2ws
        return input_c2ws, target_c2ws

    def _select_cameras(
        self,
        pred_i_fxfycxcy: torch.Tensor,
        pred_i_c2w: torch.Tensor,
        pred_t_fxfycxcy: torch.Tensor,
        pred_t_c2w: torch.Tensor,
        gt_i_fxfycxcy: torch.Tensor,
        gt_i_c2w: torch.Tensor,
        gt_t_fxfycxcy: torch.Tensor,
        gt_t_c2w: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Choose between predicted and GT cameras according to camera_mode."""
        if self.camera_mode == "gt_pose_gt_intr":
            return gt_i_fxfycxcy, gt_i_c2w, gt_t_fxfycxcy, gt_t_c2w
        if self.camera_mode == "pred_pose_gt_intr":
            return gt_i_fxfycxcy, pred_i_c2w, gt_t_fxfycxcy, pred_t_c2w
        if self.camera_mode == "pred_pose_pred_intr":
            return pred_i_fxfycxcy, pred_i_c2w, pred_t_fxfycxcy, pred_t_c2w
        raise ValueError(
            f"Unsupported camera_mode: {self.camera_mode!r}. "
            "Use one of: gt_pose_gt_intr, pred_pose_gt_intr, pred_pose_pred_intr."
        )


# ---------------------------------------------------------------------------
# Appearance Expert (MVP)
# ---------------------------------------------------------------------------

class AppearanceExpert:
    """
    Three-stage transformer (MVP) that maps ray-encoded images to a 3D Gaussian field.
    """

    def __init__(
        self,
        # nn.Module references (owned by TwoExpertModel)
        image_tokenizer: nn.Sequential,
        stage1: nn.ModuleList,
        stage2: nn.ModuleList,
        stage3: nn.ModuleList,
        merge_block1: nn.Conv2d,
        resize_block1: nn.Linear,
        merge_block2: nn.Conv2d,
        resize_block2: nn.Linear,
        dpt_head: DPTHead,
        gaussian_decoder: nn.Sequential,
        register_token_init: nn.Parameter,
        # resolution -> (stage2, stage3) PRoPE attention modules
        get_prope: Callable[[int, int], tuple[PropeDotProductAttention, PropeDotProductAttention]],
        # scalar hyper-parameters
        patch_size: int,
        num_register_tokens: int,
        group_size: int,
        pos_dim: int,
        color_dim: int,
        opacity_dim: int,
        sh_degree: int,
        opacity_degree: int,
        scale_bias: float,
        scale_max: float,
        opacity_bias: float,
        max_dist: float,
        inference_mode: bool,
    ):
        """Store references to all nn.Module components and scalar hyper-parameters.

        All nn.Module arguments are owned and registered by ``TwoExpertModel``;
        this class holds non-owning references and orchestrates the forward pass.
        """
        self.image_tokenizer = image_tokenizer
        self.stage1 = stage1
        self.stage2 = stage2
        self.stage3 = stage3
        self.merge_block1 = merge_block1
        self.resize_block1 = resize_block1
        self.merge_block2 = merge_block2
        self.resize_block2 = resize_block2
        self.dpt_head = dpt_head
        self.gaussian_decoder = gaussian_decoder
        self.register_token_init = register_token_init
        self.get_prope = get_prope
        self.patch_size = patch_size
        self.num_register_tokens = num_register_tokens
        self.group_size = group_size
        self.pos_dim = pos_dim
        self.color_dim = color_dim
        self.opacity_dim = opacity_dim
        self.sh_degree = sh_degree
        self.opacity_degree = opacity_degree
        self.scale_bias = scale_bias
        self.scale_max = scale_max
        self.opacity_bias = opacity_bias
        self.max_dist = max_dist
        self.inference_mode = inference_mode

    def predict_gaussians(
        self,
        raymap_images: torch.Tensor,   # (B, V, C, H, W)
        i_w2c: torch.Tensor,           # (B, V, 4, 4)
        Ks: torch.Tensor,              # (B, V, 3, 3)
        i_fxfycxcy: torch.Tensor,      # (B, V, 4)
        i_c2w: torch.Tensor,           # (B, V, 4, 4)
        t_c2w: torch.Tensor,           # (B, T, 4, 4)
    ) -> GaussianField:
        """Run the full three-stage transformer pipeline and decode Gaussian parameters."""
        _, _, _, h, w = raymap_images.shape
        attn2, attn3 = self.get_prope(h, w)
        x, s1_patch_tokens = self._stage1_tokenize_and_encode(raymap_images)
        x, s2_patch_tokens = self._stage2_cross_view_and_downsample(x, i_w2c, Ks, h, w, attn2)
        s3_patch_tokens = self._stage3_global_cross_view(x, i_w2c, Ks, attn3)
        return self._decode_gaussians(
            s1_patch_tokens, s2_patch_tokens, s3_patch_tokens,
            i_fxfycxcy, i_c2w, t_c2w, h, w,
        )

    def _stage1_tokenize_and_encode(
        self,
        raymap_images: torch.Tensor,
    ):
        """Patch-embed ray-map images with register tokens, then run stage-1 transformer."""
        b, v, _, h, w = raymap_images.shape
        register_tokens = self.register_token_init.repeat(b, v, 1, 1)
        x = self.image_tokenizer(raymap_images)
        x = rearrange(x, "b (v l) d -> b v l d", v=v)
        x = torch.cat([register_tokens, x], dim=2)
        x = rearrange(x, "b v l d -> (b v) l d")
        x = self._run_stage1_blocks(x, None)

        r_tokens, s1_patch_tokens = x[:, :self.num_register_tokens], x[:, self.num_register_tokens:]
        r_tokens = self.resize_block1(r_tokens)

        h_patches = h // self.patch_size
        w_patches = w // self.patch_size
        i_tokens = rearrange(s1_patch_tokens, "b (hh ww) d -> b d hh ww", hh=h_patches, ww=w_patches)
        i_tokens = self.merge_block1(i_tokens)
        i_tokens = rearrange(i_tokens, "b d hh ww -> b (hh ww) d", hh=h_patches // 2, ww=w_patches // 2)

        x = torch.cat([r_tokens, i_tokens], dim=1)
        x = rearrange(x, "(b g v) l d -> (b g) (v l) d", g=v // self.group_size, v=self.group_size)
        return x, s1_patch_tokens

    def _stage2_cross_view_and_downsample(
        self,
        x: torch.Tensor,
        i_w2c: torch.Tensor,
        Ks: torch.Tensor,
        h: int,
        w: int,
        attn2: PropeDotProductAttention,
    ):
        """Run grouped cross-view attention (stage 2) and spatially downsample patch tokens."""
        v = i_w2c.shape[1]
        info = {
            "num_input_views": v,
            "w2c": rearrange(i_w2c, "b (g v) ... -> (b g) v ...", g=v // self.group_size, v=self.group_size),
            "Ks": rearrange(Ks, "b (g v) ... -> (b g) v ...", g=v // self.group_size, v=self.group_size),
            "attn2": attn2,
        }
        x = self._run_stage2_blocks(x, info)

        r_tokens, s2_patch_tokens = x[:, :self.num_register_tokens], x[:, self.num_register_tokens:]
        r_tokens = self.resize_block2(r_tokens)

        h_patches = (h // self.patch_size) // 2
        w_patches = (w // self.patch_size) // 2
        i_tokens = rearrange(s2_patch_tokens, "b (hh ww) d -> b d hh ww", hh=h_patches, ww=w_patches)
        i_tokens = self.merge_block2(i_tokens)
        i_tokens = rearrange(i_tokens, "b d hh ww -> b (hh ww) d", hh=h_patches // 2, ww=w_patches // 2)

        x = torch.cat([r_tokens, i_tokens], dim=1)
        x = rearrange(x, "(b v) l d -> b (v l) d", v=v)
        return x, s2_patch_tokens

    def _stage3_global_cross_view(
        self,
        x: torch.Tensor,
        i_w2c: torch.Tensor,
        Ks: torch.Tensor,
        attn3: PropeDotProductAttention,
    ) -> torch.Tensor:
        """Run global cross-view attention (stage 3) across all input views."""
        v = i_w2c.shape[1]
        info = {
            "num_input_views": v,
            "attn3": attn3,
            "w2c": i_w2c,
            "Ks": Ks,
        }
        x = self._run_stage3_blocks(x, info)
        return x[:, self.num_register_tokens:]

    def _decode_gaussians(
        self,
        s1_patch_tokens: torch.Tensor,
        s2_patch_tokens: torch.Tensor,
        s3_patch_tokens: torch.Tensor,
        i_fxfycxcy: torch.Tensor,
        i_c2w: torch.Tensor,
        t_c2w: torch.Tensor,
        h: int,
        w: int,
    ) -> GaussianField:
        """Fuse multi-scale tokens with DPT head and decode into 3D Gaussian parameters."""
        b, v = i_c2w.shape[:2]
        t = t_c2w.shape[1] if t_c2w is not None else 0
        output_tokens = self.dpt_head(
            [s1_patch_tokens, s2_patch_tokens, s3_patch_tokens], [h, w], self.patch_size,
        )
        output_tokens = rearrange(output_tokens, "(b v) l d -> b (v l) d", v=v)
        gaussians = self.gaussian_decoder(output_tokens)
        gaussians = rearrange(
            gaussians, "b (v hh ww) (ph pw d) -> b (v hh ph ww pw) d",
            v=v,
            hh=h // self.patch_size,
            ww=w // self.patch_size,
            ph=self.patch_size,
            pw=self.patch_size,
        )
        pos, feature, scale, rotation, opacity = torch.split(
            gaussians, [self.pos_dim, self.color_dim, 3, 4, self.opacity_dim], dim=-1,
        )
        pos = pos.float()
        feature = feature.float()
        scale = scale.float()
        rotation = rotation.float()
        opacity = opacity.float()

        with torch.autocast(device_type="cuda", enabled=False):
            rayo_gs, rayd_gs = compute_rays(i_fxfycxcy, i_c2w, h, w)
            scale = self._activate_scale(scale)
            # Bias only the DC (sh0) component; higher-order terms are bias-free.
            opacity[..., 0] = opacity[..., 0] + self.opacity_bias
            feature = rearrange(feature, "b n (c d) -> b n d c", c=3).contiguous()
            opacity = rearrange(opacity, "b n (c d) -> b n d c", c=1).contiguous()
            dist = self._ray_distance(pos)
            xyz = dist * rayd_gs + rayo_gs

            if not self.inference_mode:
                dirs = xyz[:, None, :, :] - t_c2w[..., :3, 3][..., None, :]  # (B, T, N, 3)
                opacity_broad = torch.broadcast_to(
                    opacity[..., None, :, :, :], (b, t, opacity.shape[1], -1, 1),
                )
                opacity_precompute = _spherical_harmonics(self.opacity_degree, dirs, opacity_broad)
            else:
                opacity_precompute = None

        return GaussianField(
            xyz=xyz,
            feature=feature,
            scale=scale,
            rotation=rotation,
            opacity=opacity,
            opacity_precompute=opacity_precompute,
        )

    def _activate_scale(self, scale: torch.Tensor) -> torch.Tensor:
        """Map raw scale channels to log-space scales (GaussianRenderer exponentiates)."""
        return (scale + self.scale_bias).clamp(max=self.scale_max)

    def _ray_distance(self, pos: torch.Tensor) -> torch.Tensor:
        """Map the head's position channels to a per-pixel distance along the pixel ray."""
        return pos.mean(dim=-1, keepdim=True).sigmoid() * self.max_dist

    def _run_stage1_blocks(self, x: torch.Tensor, info: dict | None) -> torch.Tensor:
        """Run all stage-1 transformer blocks with per-view self-attention."""
        for block in self.stage1:
            x = block(x, False, 1, info)
        return x

    def _run_stage2_blocks(self, x: torch.Tensor, info: dict) -> torch.Tensor:
        """Run stage-2 blocks, alternating between per-view and grouped cross-view attention."""
        g = self.group_size
        v = info["num_input_views"]
        for i, block in enumerate(self.stage2):
            if i % 2 == 0:
                x = rearrange(x, "(b g) (v l) d -> (b g v) l d", g=v // g, v=g)
                x = block(x, False, 2, info)
                x = rearrange(x, "(b g v) l d -> (b g) (v l) d", g=v // g, v=g)
            else:
                x = block(x, True, 2, info)
        return rearrange(x, "(b g) (v l) d -> (b g v) l d", g=v // g, v=g)

    def _run_stage3_blocks(self, x: torch.Tensor, info: dict) -> torch.Tensor:
        """Run stage-3 blocks, alternating between per-view and global cross-view attention."""
        v = info["num_input_views"]
        for i, block in enumerate(self.stage3):
            if i % 2 == 0:
                x = rearrange(x, "b (v l) d -> (b v) l d", v=v)
                x = block(x, False, 3, info)
                x = rearrange(x, "(b v) l d -> b (v l) d", v=v)
            else:
                x = block(x, True, 3, info)
        return rearrange(x, "b (v l) d -> (b v) l d", v=v)


# ---------------------------------------------------------------------------
# TwoExpertModel — compositor
# ---------------------------------------------------------------------------

class TwoExpertModel(nn.Module):
    """Compositor that wires the GeometryExpert and AppearanceExpert together."""

    # Number of position channels the Gaussian head emits: 3 = xyz, 1 = ray depth.
    POS_DIM = 3
    APPEARANCE_EXPERT_CLS = AppearanceExpert

    def __init__(self, config: Any) -> None:
        """Initialise all sub-modules and compose experts from the OmegaConf config.

        Args:
            config: OmegaConf config object. Presence of a ``config.inference`` key
                selects inference mode; absence means training mode.
        """
        super().__init__()

        # Extract all config values (no self.config stored)
        self.dim1 = config.model.dim1
        self.dim2 = config.model.dim2
        self.dim3 = config.model.dim3
        self.patch_size = config.model.patch_size
        self.num_register_tokens = config.model.num_register_tokens
        self.group_size = config.model.group_size
        self.head_dim = config.model.head_dim
        self.inter_multi = config.model.inter_multi
        self.qk_norm = config.model.qk_norm
        self.in_channels = config.model.in_channels
        self.stage1_nlayer = config.model.stage1_nlayer
        self.stage2_nlayer = config.model.stage2_nlayer
        self.stage3_nlayer = config.model.stage3_nlayer
        self.sh_degree = config.model.gaussians.sh_degree
        self.opacity_degree = config.model.gaussians.opacity_degree
        self.near_plane = config.model.gaussians.near_plane
        self.far_plane = config.model.gaussians.far_plane
        self.scale_bias = config.model.gaussians.scale_bias
        self.scale_max = config.model.gaussians.scale_max
        self.opacity_bias = config.model.gaussians.opacity_bias
        self.max_dist = config.model.gaussians.max_dist
        self.da_model_name = config.model.da_model_name
        self.da_weights_path = getattr(config.model, "da_model_weights_path", None)
        self.mvp_weights_path = getattr(config.model, "mvp_weights_path", None)
        self.scene_scale = config.data.scene_scale
        self.inference_mode = hasattr(config, "inference")
        self.camera_mode = (
            getattr(config.inference, "camera_mode", "pred_pose_pred_intr")
            if self.inference_mode else None
        )
        self.use_pose_optimization = bool(
            getattr(config.inference, "pose_optimization", False)
            if self.inference_mode else False
        )

        # Derived from SH degree config; computed once here, not on every forward.
        self.color_dim = 3 * (self.sh_degree + 1) ** 2
        self.opacity_dim = 1 * (self.opacity_degree + 1) ** 2

        # Build all nn.Module components (names must match checkpoint keys)
        self._build_geometry_modules()
        self._build_appearance_modules()

        # Compose experts
        self.geometry_expert = GeometryExpert(
            pose_regressor=self.pose_regressor,
            scene_scale=self.scene_scale,
            inference_mode=self.inference_mode,
            camera_mode=self.camera_mode,
        )
        self.appearance_expert = self.APPEARANCE_EXPERT_CLS(
            image_tokenizer=self.image_tokenizer,
            stage1=self.stage1,
            stage2=self.stage2,
            stage3=self.stage3,
            merge_block1=self.merge_block1,
            resize_block1=self.resize_block1,
            merge_block2=self.merge_block2,
            resize_block2=self.resize_block2,
            dpt_head=self.dpt_head,
            gaussian_decoder=self.gaussian_decoder,
            register_token_init=self.register_token_init,
            get_prope=self._get_prope_attention,
            patch_size=self.patch_size,
            num_register_tokens=self.num_register_tokens,
            group_size=self.group_size,
            pos_dim=self.POS_DIM,
            color_dim=self.color_dim,
            opacity_dim=self.opacity_dim,
            sh_degree=self.sh_degree,
            opacity_degree=self.opacity_degree,
            scale_bias=self.scale_bias,
            scale_max=self.scale_max,
            opacity_bias=self.opacity_bias,
            max_dist=self.max_dist,
            inference_mode=self.inference_mode,
        )

        if not self.inference_mode:
            from src.model.loss import LossComputer
            self.loss_computer = LossComputer(config)

    # --- Module builders ---

    def _build_appearance_modules(self):
        """Build all nn.Module components for the appearance expert."""
        self.image_tokenizer = self._create_patch_tokenizer(
            self.in_channels, self.patch_size, self.dim1,
        )
        self.gaussian_decoder = nn.Sequential(
            nn.LayerNorm(self.dim3, bias=False),
            nn.Linear(
                self.dim3,
                (self.patch_size ** 2) * (self.POS_DIM + self.color_dim + 3 + 4 + self.opacity_dim),
                bias=False,
            ),
        )
        self.stage1 = self._build_transformer_stage(self.dim1, self.stage1_nlayer)
        self.stage2 = self._build_transformer_stage(self.dim2, self.stage2_nlayer)
        self.stage3 = self._build_transformer_stage(self.dim3, self.stage3_nlayer)

        self.register_token_init = nn.Parameter(
            torch.randn(1, 1, self.num_register_tokens, self.dim1),
        )
        nn.init.normal_(self.register_token_init, mean=0.0, std=0.02)

        # PRoPE grids depend on the input resolution, which is chosen per request, so cache
        # one module pair per resolution (see _get_prope_attention). These modules hold only
        # non-persistent buffers, so they never appear in or expect anything from the state dict.
        self.prope_s2 = nn.ModuleDict()
        self.prope_s3 = nn.ModuleDict()
        self.merge_block1 = nn.Conv2d(
            self.dim1, self.dim2, kernel_size=2, stride=2,
            padding=0, bias=True, groups=self.dim1,
        )
        self.resize_block1 = nn.Linear(self.dim1, self.dim2)
        self.merge_block2 = nn.Conv2d(
            self.dim2, self.dim3, kernel_size=2, stride=2,
            padding=0, bias=True, groups=self.dim2,
        )
        self.resize_block2 = nn.Linear(self.dim2, self.dim3)
        self.dpt_head = DPTHead(
            dim_in=[self.dim1, self.dim2, self.dim3],
            features=self.dim3,
            out_channels=[self.dim1, self.dim2, self.dim3],
        )

        if self.mvp_weights_path is not None:
            checkpoint = torch.load(self.mvp_weights_path, map_location="cpu", weights_only=True)
            state_dict = checkpoint["ema"] if "ema" in checkpoint else checkpoint
            result = self.load_state_dict(state_dict, strict=False)
            # print(f"{result.missing_keys} missing keys")
            # print(f"{result.unexpected_keys} unexpected keys")
            print(f"Loaded MVP appearance weights from {self.mvp_weights_path}")

    def _get_prope_attention(self, h: int, w: int):
        """Return the (stage2, stage3) PRoPE modules for this resolution, building them once.

        The stage-2 grid is the patch grid halved once (merge_block1), the stage-3 grid
        halved twice (merge_block1 + merge_block2).
        """
        key = f"{w}x{h}"
        if key not in self.prope_s2:
            device = next(self.parameters()).device
            for cache, factor in ((self.prope_s2, 2), (self.prope_s3, 4)):
                cache[key] = PropeDotProductAttention(
                    head_dim=self.head_dim,
                    patches_x=w // (self.patch_size * factor),
                    patches_y=h // (self.patch_size * factor),
                    image_width=w,
                    image_height=h,
                    num_register_tokens=self.num_register_tokens,
                ).to(device)
        return self.prope_s2[key], self.prope_s3[key]

    def _build_geometry_modules(self):
        """Build the DA3 pose regressor and optionally load pretrained weights."""
        self.pose_regressor = DepthAnything3(model_name=self.da_model_name)
        if self.da_weights_path is not None:
            state_dict = load_file(self.da_weights_path)
            results = self.pose_regressor.load_state_dict(state_dict, strict=False)
            # print(f"{results.missing_keys} missing keys")
            # print(f"{results.unexpected_keys} unexpected keys")
            print(f"Loaded DA3 pose regressor weights from {self.da_weights_path}")
        self.pose_regressor.prune_layers()

    def _build_transformer_stage(self, dim: int, nlayer: int) -> nn.ModuleList:
        """Create a list of identical TransformerBlocks for one transformer stage.

        Args:
            dim: Hidden dimension for all blocks in this stage.
            nlayer: Number of transformer blocks to create.

        Returns:
            An ``nn.ModuleList`` of ``nlayer`` TransformerBlock instances.
        """
        return nn.ModuleList([
            TransformerBlock(dim, False, self.head_dim, self.inter_multi, self.qk_norm)
            for _ in range(nlayer)
        ])

    @staticmethod
    def _create_patch_tokenizer(
        in_channels: int, patch_size: int, d_model: int,
    ) -> nn.Sequential:
        """Build a patch-embedding tokenizer: rearrange → linear projection → LayerNorm.

        Args:
            in_channels: Number of input image channels (e.g. 12 for ray-map images).
            patch_size: Side length of each square patch in pixels.
            d_model: Output embedding dimension.

        Returns:
            An ``nn.Sequential`` that maps ``(B, V, C, H, W)`` to ``(B, V*L, d_model)``.
        """
        return nn.Sequential(
            Rearrange(
                "b v c (hh ph) (ww pw) -> b (v hh ww) (ph pw c)",
                ph=patch_size, pw=patch_size,
            ),
            nn.Linear(in_channels * (patch_size ** 2), d_model, bias=False),
            nn.LayerNorm(d_model, bias=False),
        )

    # --- Training mode override ---

    def train(self, mode: bool = True):
        """Override train() to keep loss modules permanently in eval mode."""
        super().train(mode)
        if not self.inference_mode:
            self.loss_computer.eval()

    # --- Forward helpers ---

    def _build_raymap_input(
        self,
        input_data_dict: dict,
        i_fxfycxcy: torch.Tensor,
        i_c2w: torch.Tensor,
    ):
        """Construct plucker ray-map images: [ray_origin, ray_dir, origin×dir, image]."""
        h, w = input_data_dict["image"].shape[-2:]
        with torch.autocast(device_type="cuda", enabled=False):
            ray_o, ray_d = compute_plucmap(i_fxfycxcy, i_c2w, h, w)
            o_cross_d = torch.cross(ray_o, ray_d, dim=2)
            i_normalized_image = input_data_dict["image"] * 2.0 - 1.0
            i_raymap_images = torch.concat([ray_o, ray_d, o_cross_d, i_normalized_image], dim=2)
            Ks = fxfycxcy_to_K(i_fxfycxcy)
            i_w2c = invert_SE3(i_c2w)
        return i_raymap_images, Ks, i_w2c

    def _render_and_compute_loss(
        self,
        gaussians: GaussianField,
        cameras: CameraBundle,
        target_images: torch.Tensor,
    ) -> tuple[torch.Tensor, dict]:
        """Render Gaussians to all target views and compute photometric + pose losses."""
        h, w = target_images.shape[-2:]
        with torch.autocast(device_type="cuda", enabled=False):
            renderings = GaussianRenderer.apply(
                gaussians.xyz, gaussians.feature, gaussians.scale,
                gaussians.rotation, gaussians.opacity_precompute,
                cameras.pred_t_c2w, cameras.pred_t_fxfycxcy, w, h,
                self.sh_degree, self.near_plane, self.far_plane,
            )
        renderings = renderings.permute(0, 1, 4, 2, 3).contiguous()  # (B, V, 3, H, W)

        cam_info = {
            "pred_fxfycxcy": torch.cat([cameras.pred_i_fxfycxcy, cameras.pred_t_fxfycxcy], dim=1),
            "pred_c2w": torch.cat([cameras.pred_i_c2w, cameras.pred_t_c2w], dim=1),
            "gt_fxfycxcy": torch.cat([cameras.gt_i_fxfycxcy, cameras.gt_t_fxfycxcy], dim=1),
            "gt_c2w": torch.cat([cameras.gt_i_c2w, cameras.gt_t_c2w], dim=1),
        }
        loss_metrics = self.loss_computer(renderings, target_images, cam_info)

        with torch.autocast(device_type="cuda", enabled=False):
            rand_dirs = F.normalize(torch.randn_like(gaussians.xyz), p=2, dim=-1)
            opacity_random = _spherical_harmonics(
                self.opacity_degree, rand_dirs, gaussians.opacity,
            )
            opacity_random = opacity_random.sigmoid().mean()

        loss_metrics["opacity_loss"] = opacity_random * 0.001
        loss_metrics["loss"] = loss_metrics["loss"] + loss_metrics["opacity_loss"]
        return renderings, loss_metrics

    def _render_target_views(
        self,
        gaussians: GaussianField,
        t_c2w: torch.Tensor,
        t_fxfycxcy: torch.Tensor,
        w: int,
        h: int,
    ) -> torch.Tensor:
        """Render each target view sequentially without gradient tracking."""
        t = t_c2w.shape[1]
        xyz = gaussians.xyz[0]
        feature = gaussians.feature[0]
        scale = gaussians.scale[0]
        rotation = gaussians.rotation[0]
        opacity = gaussians.opacity[0]
        renderings = []
        with torch.no_grad(), torch.autocast(device_type="cuda", enabled=False):
            for i in range(t):
                dir = xyz - t_c2w[0, i:i+1, :3, 3][None, ...]
                opacity_i = _spherical_harmonics(
                    self.opacity_degree, dir, opacity[None, ...],
                )[0]
                rendering = GaussianRenderer.render(
                    xyz, feature, scale, rotation, opacity_i,
                    t_c2w[0, i], t_fxfycxcy[0, i], w, h,
                    self.sh_degree, self.near_plane, self.far_plane,
                )
                renderings.append(rendering)
        renderings = torch.cat(renderings, dim=0)[None, ...]   # (1, T, H, W, 3)
        return renderings.permute(0, 1, 4, 2, 3).contiguous()  # (1, T, 3, H, W)

    def _training_forward(
        self,
        gaussians: GaussianField,
        cameras: CameraBundle,
        target_data_dict: dict,
        input_data_dict: dict,
    ) -> edict:
        """Render Gaussians and compute all training losses.

        Returns:
            An edict with keys ``input``, ``target``, ``loss_metrics``, and ``render``.
        """
        renderings, loss_metrics = self._render_and_compute_loss(
            gaussians, cameras, target_data_dict["image"],
        )
        return edict(
            input=input_data_dict,
            target=target_data_dict,
            loss_metrics=loss_metrics,
            render=renderings,
        )

    def _inference_forward(
        self,
        gaussians: GaussianField,
        cameras: CameraBundle,
        target_data_dict: dict,
        input_data_dict: dict,
    ) -> edict:
        """Render target views at inference time, optionally with pose optimization.

        Returns:
            An edict with keys ``input``, ``target``, and ``render``.
        """
        h, w = target_data_dict["image"].shape[-2:]

        if self.use_pose_optimization:
            prev_grad_state = torch.is_grad_enabled()
            torch.set_grad_enabled(True)
            renderings, _ = self.pose_optimization(
                gaussians=gaussians,
                t_c2w=cameras.pred_t_c2w,
                t_fxfycxcy=cameras.pred_t_fxfycxcy,
                w=w, h=h,
                target_images=target_data_dict["image"],
            )
            torch.set_grad_enabled(prev_grad_state)
        else:
            renderings = self._render_target_views(
                gaussians, cameras.pred_t_c2w, cameras.pred_t_fxfycxcy, w, h,
            )

        return edict(
            input=input_data_dict,
            target=target_data_dict,
            render=renderings,
        )

    # --- Main forward ---

    def forward(self, input_data_dict: dict, target_data_dict: dict) -> edict:
        """Full forward pass: predict cameras → build ray maps → predict Gaussians → render.

        Args:
            input_data_dict: Batch dict for context views, must contain ``image`` and
                camera ground-truth tensors.
            target_data_dict: Batch dict for novel views to render.

        Returns:
            An edict whose contents depend on the mode (training vs. inference).
        """
        cameras = self.geometry_expert.predict_cameras(input_data_dict, target_data_dict)
        raymap_images, Ks, i_w2c = self._build_raymap_input(
            input_data_dict, cameras.pred_i_fxfycxcy, cameras.pred_i_c2w,
        )
        gaussians = self.appearance_expert.predict_gaussians(
            raymap_images, i_w2c, Ks,
            cameras.pred_i_fxfycxcy, cameras.pred_i_c2w, cameras.pred_t_c2w,
        )

        if not self.inference_mode:
            return self._training_forward(gaussians, cameras, target_data_dict, input_data_dict)
        return self._inference_forward(gaussians, cameras, target_data_dict, input_data_dict)

    # --- Pose optimization (optional inference refinement) ---

    def pose_optimization(
        self,
        gaussians: GaussianField,
        t_c2w: torch.Tensor,
        t_fxfycxcy: torch.Tensor,
        w: int,
        h: int,
        target_images: torch.Tensor,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        """Refine target poses with 100 steps of Adam (EPA — evaluation-time pose alignment).

        Quaternion and translation parameters are jointly optimised against an MSE
        photometric loss while all Gaussian parameters remain frozen.

        Args:
            gaussians: Frozen Gaussian field from the appearance expert.
            t_c2w: Initial target camera-to-world matrices (1, T, 4, 4).
            t_fxfycxcy: Target camera intrinsics (1, T, 4).
            w: Image width in pixels.
            h: Image height in pixels.
            target_images: Ground-truth target images (1, T, 3, H, W).

        Returns:
            A tuple of (renderings, (refined_c2w, refined_fxfycxcy)) where
            renderings has shape (1, T, 3, H, W).
        """
        num_target_views = t_c2w.shape[1]
        quats = mat_to_quat(t_c2w[..., :3, :3]).clone().detach().requires_grad_(True)
        trans = t_c2w[..., :3, 3].clone().detach().requires_grad_(True)
        t_fxfycxcy = t_fxfycxcy.clone().detach()
        xyz = gaussians.xyz[0].detach()
        feature = gaussians.feature[0].detach()
        scale = gaussians.scale[0].detach()
        rotation = gaussians.rotation[0].detach()
        opacity = gaussians.opacity[0].detach()
        optimizer = torch.optim.Adam([quats, trans], lr=1e-4)

        for _ in range(100):
            R = quat_to_mat(quats / quats.norm(dim=-1, keepdim=True))
            t_c2w_new = torch.zeros_like(t_c2w)
            t_c2w_new[..., :3, :3] = R
            t_c2w_new[..., :3, 3] = trans
            t_c2w_new[..., 3, 3] = 1.0

            optimizer.zero_grad()
            renderings = []
            with torch.autocast(device_type="cuda", enabled=False):
                for i in range(num_target_views):
                    dir = xyz - t_c2w_new[0, i:i+1, :3, 3][None, ...]
                    opacity_i = _spherical_harmonics(
                        self.opacity_degree, dir, opacity[None, ...],
                    )[0]
                    rendering = GaussianRenderer.render(
                        xyz, feature, scale, rotation, opacity_i,
                        t_c2w_new[0, i], t_fxfycxcy[0, i], w, h,
                        self.sh_degree, self.near_plane, self.far_plane,
                    )
                    renderings.append(rendering)
            renderings = torch.cat(renderings, dim=0)[None, ...]

            loss = F.mse_loss(renderings.permute(0, 1, 4, 2, 3).contiguous(), target_images)
            loss.backward()
            optimizer.step()

        refined_c2w = t_c2w_new.detach()
        refined_fxfycxcy = t_fxfycxcy.detach()
        return renderings.permute(0, 1, 4, 2, 3).contiguous(), (refined_c2w, refined_fxfycxcy)

    # --- Checkpoint loading ---

    @torch.no_grad()
    def load_ckpt(self, load_path: str) -> int | None:
        """Load EMA weights from a checkpoint file or the latest .pt in a directory.

        The checkpoint is expected to be a dict with an ``"ema"`` key containing
        the state dict. Loading is done with ``strict=False`` so DA3 weights (loaded
        separately) do not cause missing-key errors.

        Args:
            load_path: Path to a ``.pt`` checkpoint file, or a directory containing
                one or more ``.pt`` files (the lexicographically last is used).

        Returns:
            0 on success, None on failure.
        """
        if os.path.isdir(load_path):
            ckpt_names = sorted(f for f in os.listdir(load_path) if f.endswith(".pt"))
            ckpt_path = os.path.join(load_path, ckpt_names[-1])
        else:
            ckpt_path = load_path
        try:
            checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        except:
            traceback.print_exc()
            print(f"Failed to load {ckpt_path}")
            return None
        result = self.load_state_dict(checkpoint["ema"], strict=False)
        # print(f"{result.missing_keys} missing keys")
        # print(f"{result.unexpected_keys} unexpected keys")
        print(f"Loaded 2Xplat checkpoint in {ckpt_path}")
        return 0