ffgs 0.1.0.dev0__py3-none-any.whl

ffgs/models/anysplat/gs_head.py

@@ -0,0 +1,166 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# dpt head implementation for DUST3R
+# Downstream heads assume inputs of size B x N x C (where N is the number of tokens) ;
+# or if it takes as input the output at every layer, the attribute return_all_layers should be set to True
+# the forward function also takes as input a dictionnary img_info with key "height" and "width"
+# for PixelwiseTask, the output will be of dimension B x num_channels x H x W
+# --------------------------------------------------------
+#
+# Modified by the libgs contributors, 2026, from
+# src/model/encoder/heads/vggt_dpt_gs_head.py of
+# https://github.com/InternRobotics/AnySplat (5f5e208): only `VGGT_DPT_GS_Head` kept
+# (`PixelwiseTaskWithDPT` / `create_gs_dpt_head` and the commented-out signature
+# removed, with the imports only they used: einops, dpt_block, head_modules (xFormers),
+# postprocess); `DPTHead` imported package-relative; formatted with black.
+# CC BY-NC-SA 4.0 (non-commercial): see LICENSE (this directory).
+
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .vggt.heads.dpt_head import DPTHead
+
+
+class VGGT_DPT_GS_Head(DPTHead):
+    def __init__(
+        self,
+        dim_in: int,
+        patch_size: int = 14,
+        output_dim: int = 83,
+        activation: str = "inv_log",
+        conf_activation: str = "expp1",
+        features: int = 256,
+        out_channels: List[int] = [256, 512, 1024, 1024],
+        intermediate_layer_idx: List[int] = [4, 11, 17, 23],
+        pos_embed: bool = True,
+        feature_only: bool = False,
+        down_ratio: int = 1,
+    ):
+        super().__init__(
+            dim_in,
+            patch_size,
+            output_dim,
+            activation,
+            conf_activation,
+            features,
+            out_channels,
+            intermediate_layer_idx,
+            pos_embed,
+            feature_only,
+            down_ratio,
+        )
+
+        head_features_1 = 128
+        head_features_2 = (
+            128 if output_dim > 50 else 32
+        )  # sh=0, head_features_2 = 32; sh=4, head_features_2 = 128
+        self.input_merger = nn.Sequential(
+            nn.Conv2d(3, head_features_2, 7, 1, 3),
+            nn.ReLU(),
+        )
+
+        self.scratch.output_conv2 = nn.Sequential(
+            nn.Conv2d(
+                head_features_1, head_features_2, kernel_size=3, stride=1, padding=1
+            ),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0),
+        )
+
+    def forward(
+        self,
+        encoder_tokens: List[torch.Tensor],
+        depths,
+        imgs,
+        patch_start_idx: int = 5,
+        image_size=None,
+        conf=None,
+        frames_chunk_size: int = 8,
+    ):
+        # H, W = input_info['image_size']
+        B, S, _, H, W = imgs.shape
+        image_size = self.image_size if image_size is None else image_size
+
+        # If frames_chunk_size is not specified or greater than S, process all frames at once
+        if frames_chunk_size is None or frames_chunk_size >= S:
+            return self._forward_impl(encoder_tokens, imgs, patch_start_idx)
+
+        # Otherwise, process frames in chunks to manage memory usage
+        assert frames_chunk_size > 0
+
+        # Process frames in batches
+        all_preds = []
+
+        for frames_start_idx in range(0, S, frames_chunk_size):
+            frames_end_idx = min(frames_start_idx + frames_chunk_size, S)
+
+            # Process batch of frames
+            chunk_output = self._forward_impl(
+                encoder_tokens, imgs, patch_start_idx, frames_start_idx, frames_end_idx
+            )
+            all_preds.append(chunk_output)
+
+        # Concatenate results along the sequence dimension
+        return torch.cat(all_preds, dim=1)
+
+    def _forward_impl(
+        self,
+        encoder_tokens: List[torch.Tensor],
+        imgs,
+        patch_start_idx: int = 5,
+        frames_start_idx: int = None,
+        frames_end_idx: int = None,
+    ):
+
+        if frames_start_idx is not None and frames_end_idx is not None:
+            imgs = imgs[:, frames_start_idx:frames_end_idx]
+
+        B, S, _, H, W = imgs.shape
+
+        patch_h, patch_w = H // self.patch_size[0], W // self.patch_size[1]
+
+        out = []
+        dpt_idx = 0
+        for layer_idx in self.intermediate_layer_idx:
+            # x = encoder_tokens[layer_idx][:, :, patch_start_idx:]
+            if len(encoder_tokens) > 10:
+                x = encoder_tokens[layer_idx][:, :, patch_start_idx:]
+            else:
+                list_idx = self.intermediate_layer_idx.index(layer_idx)
+                x = encoder_tokens[list_idx][:, :, patch_start_idx:]
+
+            # Select frames if processing a chunk
+            if frames_start_idx is not None and frames_end_idx is not None:
+                x = x[:, frames_start_idx:frames_end_idx].contiguous()
+
+            x = x.view(B * S, -1, x.shape[-1])
+
+            x = self.norm(x)
+
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+
+            x = self.projects[dpt_idx](x)
+            if self.pos_embed:
+                x = self._apply_pos_embed(x, W, H)
+            x = self.resize_layers[dpt_idx](x)
+
+            out.append(x)
+            dpt_idx += 1
+
+        # Fuse features from multiple layers.
+        out = self.scratch_forward(out)
+        direct_img_feat = self.input_merger(imgs.flatten(0, 1))
+        out = F.interpolate(out, size=(H, W), mode="bilinear", align_corners=True)
+        out = out + direct_img_feat
+
+        if self.pos_embed:
+            out = self._apply_pos_embed(out, W, H)
+
+        out = self.scratch.output_conv2(out)
+        out = out.view(B, S, *out.shape[1:])
+        return out

ffgs/models/anysplat/modeling.py

@@ -0,0 +1,95 @@
+"""The AnySplat network: images -> Gaussians and the input cameras, in its frame."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+from .encoder import (
+    EncoderAnySplat,
+    EncoderAnySplatCfg,
+    EncoderOutput,
+    OpacityMappingCfg,
+)
+from .gaussian_adapter import GaussianAdapterCfg
+
+
+class AnySplat(nn.Module, PyTorchModelHubMixin):
+    """
+    AnySplat (Jiang et al., TOG 2025): VGGT-based feed-forward Gaussians from
+    unposed images, with the cameras it predicts for them.
+
+    Input (from `AnySplatProcessor`): images [B, V, 3, H, W] in [0, 1], H and W
+    multiples of the patch size. Output: upstream's `EncoderOutput` -- Gaussians
+    (xyzw quaternions, SH of `sh_degree`) and `pred_context_pose` (c2w, normalised
+    K) in the frame of the first predicted camera, at the model's own scale.
+
+    The arguments are the inference fields of upstream's `encoder_cfg` (defaults:
+    the released `lhjiang/anysplat` config) and the VGGT architecture (defaults:
+    VGGT-1B). The aggregator is kept in bfloat16, as upstream; the parameter names
+    are upstream's (`encoder.`...), so its checkpoint loads as it is.
+    """
+
+    def __init__(
+        self,
+        voxel_size: float = 0.002,
+        voxelize: bool = True,
+        sh_degree: int = 4,
+        gaussian_scale_min: float = 0.5,
+        gaussian_scale_max: float = 15.0,
+        opacity_mapping: dict[str, float] | None = None,
+        intermediate_layer_idx: list[int] | None = (4, 11, 17, 23),
+        gs_prune: bool = False,
+        opacity_threshold: float = 0.001,
+        gs_keep_ratio: float = 1.0,
+        render_conf: bool = False,
+        opacity_conf: bool = False,
+        conf_threshold: float = 0.1,
+        img_size: int = 518,
+        patch_size: int = 14,
+        embed_dim: int = 1024,
+        aggregator: dict[str, Any] | None = None,
+        camera_head: dict[str, Any] | None = None,
+        depth_head: dict[str, Any] | None = None,
+        gs_head: dict[str, Any] | None = None,
+    ):
+        super().__init__()
+        mapping = {"initial": 0.0, "final": 0.0, "warm_up": 1}
+        cfg = EncoderAnySplatCfg(
+            voxel_size=voxel_size,
+            voxelize=voxelize,
+            gaussian_adapter=GaussianAdapterCfg(
+                gaussian_scale_min=gaussian_scale_min,
+                gaussian_scale_max=gaussian_scale_max,
+                sh_degree=sh_degree,
+            ),
+            opacity_mapping=OpacityMappingCfg(**(opacity_mapping or mapping)),
+            intermediate_layer_idx=(
+                None if intermediate_layer_idx is None else list(intermediate_layer_idx)
+            ),
+            gs_prune=gs_prune,
+            opacity_threshold=opacity_threshold,
+            gs_keep_ratio=gs_keep_ratio,
+            render_conf=render_conf,
+            opacity_conf=opacity_conf,
+            conf_threshold=conf_threshold,
+            img_size=img_size,
+            patch_size=patch_size,
+            embed_dim=embed_dim,
+            aggregator=dict(aggregator or {}),
+            camera_head=dict(camera_head or {}),
+            depth_head=dict(depth_head or {}),
+            gs_head=dict(gs_head or {}),
+        )
+        self.encoder = EncoderAnySplat(cfg)
+
+    @property
+    def sh_degree(self) -> int:
+        return self.encoder.cfg.gaussian_adapter.sh_degree
+
+    def forward(self, images: torch.Tensor) -> EncoderOutput:
+        # Upstream's `AnySplat.inference`: the encoder at global step 0.
+        return self.encoder(images, global_step=0)

ffgs/models/anysplat/processing.py

@@ -0,0 +1,249 @@
+"""AnySplat pre / post-processing: the author's inference data path.
+
+Per sample, as `src/utils/image.py:process_image` and the callers in
+`src/eval_nvs.py` / `inference.py` of https://github.com/InternRobotics/AnySplat
+do:
+
+1. the image is read as 8 bit (inputs are quantised to 8 bit here), resized with
+   PIL's bicubic filter so that it covers `image_shape`, the covering side by the
+   upstream arithmetic (`int`, not `round`), then centre-cropped to `image_shape`
+   ([448, 448] for the released model);
+2. `ToTensor() * 2 - 1`, then `(x + 1) * 0.5`: the [0, 1] images the model takes,
+   with the float32 rounding of that round trip.
+
+The model needs no cameras: it predicts them (`predicted_cameras`), in the frame of
+its first input camera at its own scale. With poses, the pipeline fits that frame
+to them (`ModelFrame.fit`). Gaussians come out with SH of degree 4.
+
+Rendering uses upstream's settings (its `DecoderSplattingCUDA`): near plane 1e-10
+and gsplat's default far plane, in model units, `radius_clip=0.1`, classic
+rasterisation, white background, colours clamped to [0, 1].
+
+`crop_mode="pad"` / `"none"` replace step 1's crop with the ffgs rules; the
+`resize_mode` `"bicubic"` (the default) is step 1's PIL kernel.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ...geometry import ModelFrame
+from ...image import CROP_MODES, RESIZE_MODES, center_crop, fit_images, rescale
+from ...processor import ImageFitConfig, Prepared, Processor
+from ...types import Cameras, Gaussians, Views
+from .encoder import EncoderOutput
+
+# The ffgs kernels, and upstream's (PIL bicubic on 8-bit images).
+ANYSPLAT_RESIZE_MODES = ("bicubic", *RESIZE_MODES)
+RASTERIZE_MODES = ("classic", "antialiased")
+PATCH_SIZE = 14
+
+
+@dataclass
+class AnySplatProcessorConfig(ImageFitConfig):
+    # [h, w], multiples of the 14-pixel patches; upstream: 448 x 448.
+    image_shape: list[int] = field(default_factory=lambda: [448, 448])
+    resize_mode: str = "bicubic"
+    # Render settings of the upstream renderer; planes in model units.
+    znear: float = 1e-10
+    zfar: float = 1e10
+    background: list[float] = field(default_factory=lambda: [1.0, 1.0, 1.0])
+    radius_clip: float = 0.1
+    rasterize_mode: str = "classic"
+    clamp: bool = True
+
+    def __post_init__(self) -> None:
+        # ImageFitConfig's checks, with upstream's kernel among the resize modes.
+        self.image_shape = [int(x) for x in self.image_shape]
+        if len(self.image_shape) != 2:
+            raise ValueError(f"image_shape must be [h, w], got {self.image_shape}")
+        if any(x <= 0 or x % PATCH_SIZE for x in self.image_shape):
+            raise ValueError(
+                f"image_shape must be positive multiples of {PATCH_SIZE}, got "
+                f"{self.image_shape}"
+            )
+        if self.resize_mode not in ANYSPLAT_RESIZE_MODES:
+            raise ValueError(f"resize_mode must be one of {ANYSPLAT_RESIZE_MODES}")
+        if self.crop_mode not in CROP_MODES:
+            raise ValueError(f"crop_mode must be one of {CROP_MODES}")
+        if self.rasterize_mode not in RASTERIZE_MODES:
+            raise ValueError(f"rasterize_mode must be one of {RASTERIZE_MODES}")
+
+
+class AnySplatProcessor(Processor):
+    config_cls = AnySplatProcessorConfig
+
+    def preprocess(self, views: Views, device: torch.device) -> Prepared:
+        images = self.fit_views(views).images
+        # process_image's `ToTensor() * 2 - 1`, undone by its callers.
+        images = (images * 2.0 - 1.0 + 1) * 0.5
+        # No frame yet: it is fitted to the predicted cameras (`output_frame`).
+        return Prepared(images.to(device), None)
+
+    def postprocess(self, model_output: EncoderOutput) -> Gaussians:
+        g = model_output.gaussians
+        d_sh = g.harmonics.shape[-1]
+        gaussians = Gaussians(
+            means=g.means.float(),
+            scales=g.scales.float(),
+            quats=g.rotations[..., [3, 0, 1, 2]].float(),  # xyzw -> wxyz
+            opacities=g.opacities.float(),
+            colors=g.harmonics.transpose(-1, -2).float(),  # [B, N, d_sh, 3]
+            sh_degree=round(d_sh**0.5) - 1,
+        )
+        counts = (model_output.infos or {}).get("num_gaussians")
+        return gaussians if counts is None else neutralize_padding(gaussians, counts)
+
+    def predicted_cameras(self, model_output: EncoderOutput) -> Cameras:
+        pose = model_output.pred_context_pose
+        h, w = model_output.depth_dict["depth"].shape[2:4]
+        return Cameras(
+            pose["extrinsic"].float(),
+            pose["intrinsic"].float(),
+            (h, w),
+            normalized_intrinsics=True,
+        )
+
+    def render_planes(self, frame: ModelFrame | None) -> tuple[float, float]:
+        cfg = self.config
+        if frame is None:
+            return cfg.znear, cfg.zfar
+        return cfg.znear / frame.scale, cfg.zfar / frame.scale
+
+    @property
+    def background(self) -> tuple[float, float, float]:
+        return tuple(self.config.background)
+
+    @property
+    def render_settings(self) -> dict[str, Any]:
+        cfg = self.config
+        return {
+            "radius_clip": cfg.radius_clip,
+            "rasterize_mode": cfg.rasterize_mode,
+            "clamp": cfg.clamp,
+        }
+
+    # --- fitting inputs ---------------------------------------------------------
+
+    def fit_views(self, views: Views, image_shape: list[int] | None = None) -> Views:
+        """`process_image` for `crop_mode="crop"`, the ffgs rules otherwise; the
+        images are quantised to 8 bit first (upstream reads 8-bit files)."""
+        cfg = self.config
+        shape = tuple(image_shape or cfg.image_shape)
+        scaled = upstream_scaled_shape(views.image_shape, shape)
+        k = None if views.intrinsics is None else views.normalized_k
+        images, ks = [], []
+        for i, img in enumerate(views.images):
+            ki = torch.eye(3).expand(len(img), 3, 3) if k is None else k[i]
+            img = quantize(img)
+            if cfg.crop_mode == "crop":
+                img, ki = center_crop(self._resize(img, scaled), ki, shape)
+            else:
+                img, ki = fit_images(
+                    img,
+                    ki,
+                    shape,
+                    crop_mode=cfg.crop_mode,
+                    resize=self._resize,
+                    pad_value=cfg.pad_value,
+                )
+            images.append(img)
+            ks.append(ki)
+        return Views(
+            torch.stack(images),
+            None if k is None else torch.stack(ks),
+            views.c2w,
+            normalized_intrinsics=True,
+        )
+
+    def fit_cameras(
+        self, cameras: Cameras, image_shape: list[int] | None = None
+    ) -> Cameras:
+        cfg = self.config
+        if cfg.crop_mode != "crop":
+            return super().fit_cameras(cameras, image_shape)
+        shape = tuple(image_shape or cfg.image_shape)
+        h, w = upstream_scaled_shape(cameras.image_shape, shape)
+        k = cameras.normalized_k.clone()
+        k[..., 0, 0] *= w / shape[1]
+        k[..., 1, 1] *= h / shape[0]
+        return Cameras(cameras.c2w, k, shape, normalized_intrinsics=True)
+
+    def _resize(self, images: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor:
+        if self.resize is not None:
+            return self.resize(images, shape)
+        if self.config.resize_mode == "bicubic":
+            return pil_bicubic(images, shape)
+        return rescale(images, shape, self.config.resize_mode)
+
+
+def upstream_scaled_shape(
+    image_shape: tuple[int, int], shape: tuple[int, int]
+) -> tuple[int, int]:
+    """The size `process_image` resizes an `image_shape` image to before cropping
+    it to `shape`: the covering side exactly, the other by `int` (truncation)."""
+    h, w = image_shape
+    h_out, w_out = shape
+    if h_out / h > w_out / w:  # upstream (square shape): width > height
+        return h_out, int(w * (h_out / h))
+    return int(h * (w_out / w)), w_out
+
+
+def quantize(images: torch.Tensor) -> torch.Tensor:
+    """[0, 1] images -> the nearest 8-bit values, as `ToTensor` gives them."""
+    return (images * 255).round().clamp(0, 255).to(torch.uint8).float() / 255
+
+
+def pil_bicubic(images: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor:
+    """[..., 3, h, w] -> [..., 3, *shape]: PIL's bicubic resize (`Image.resize`'s
+    default) of the 8-bit images, then `ToTensor`'s / 255."""
+    *batch, c, h, w = images.shape
+    device, dtype = images.device, images.dtype
+    h_out, w_out = shape
+    arrays = (images.reshape(-1, c, h, w) * 255).round().clamp(0, 255)
+    arrays = arrays.to(torch.uint8).permute(0, 2, 3, 1).contiguous().cpu().numpy()
+    resized = [
+        np.asarray(Image.fromarray(a).resize((w_out, h_out), Image.BICUBIC))
+        for a in arrays
+    ]
+    out = torch.from_numpy(np.stack(resized)).permute(0, 3, 1, 2).float() / 255
+    return out.to(device=device, dtype=dtype).reshape(*batch, c, h_out, w_out)
+
+
+def neutralize_padding(gaussians: Gaussians, counts: torch.Tensor) -> Gaussians:
+    """Make the padding of a batch inert: Gaussians past each sample's count.
+
+    Upstream pads the samples of a batch to the largest Gaussian count with
+    placeholder values (means -1e4, scales 0, colours -1e10, opacity 0) that a
+    ply cannot hold (log 0). Each placeholder becomes a transparent copy of its
+    sample's first Gaussian: opacity 0, colour 0, same mean, scale and rotation.
+    It renders nothing, stays within the sample's extent and is left out of
+    plys. Samples without padding (any batch of one) are returned unchanged.
+    """
+    n = gaussians.num_gaussians
+    counts = counts.to(gaussians.means.device)
+    if bool((counts >= n).all()):
+        return gaussians
+    valid = torch.arange(n, device=counts.device)[None] < counts[:, None]  # [B, N]
+
+    def first(x: torch.Tensor) -> torch.Tensor:
+        mask = valid.view(*valid.shape, *([1] * (x.dim() - 2)))
+        return torch.where(mask, x, x[:, :1])
+
+    return Gaussians(
+        means=first(gaussians.means),
+        scales=first(gaussians.scales),
+        quats=first(gaussians.quats),
+        opacities=torch.where(valid, gaussians.opacities, 0.0),
+        colors=torch.where(
+            valid.view(*valid.shape, *([1] * (gaussians.colors.dim() - 2))),
+            gaussians.colors,
+            0.0,
+        ),
+        sh_degree=gaussians.sh_degree,
+    )

ffgs/models/anysplat/vggt/heads/camera_head.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Modified by the libgs contributors, 2026, from
+# src/model/encoder/vggt/heads/camera_head.py of
+# https://github.com/InternRobotics/AnySplat (5f5e208): imports made package-relative,
+# formatted with black, unused imports removed.
+# "The license found in the LICENSE file" above is the one VGGT's code
+# had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
+# (facebookresearch/vggt c6bf698); see LICENSE (this directory).
+
+
+import torch
+import torch.nn as nn
+
+from ..layers import Mlp
+from ..layers.block import Block
+from .head_act import activate_pose
+
+
+class CameraHead(nn.Module):
+    """
+    CameraHead predicts camera parameters from token representations using iterative refinement.
+
+    It applies a series of transformer blocks (the "trunk") to dedicated camera tokens.
+    """
+
+    def __init__(
+        self,
+        dim_in: int = 2048,
+        trunk_depth: int = 4,
+        pose_encoding_type: str = "absT_quaR_FoV",
+        num_heads: int = 16,
+        mlp_ratio: int = 4,
+        init_values: float = 0.01,
+        trans_act: str = "linear",
+        quat_act: str = "linear",
+        fl_act: str = "relu",  # Field of view activations: ensures FOV values are positive.
+    ):
+        super().__init__()
+
+        if pose_encoding_type == "absT_quaR_FoV":
+            self.target_dim = 9
+        else:
+            raise ValueError(f"Unsupported camera encoding type: {pose_encoding_type}")
+
+        self.trans_act = trans_act
+        self.quat_act = quat_act
+        self.fl_act = fl_act
+        self.trunk_depth = trunk_depth
+
+        # Build the trunk using a sequence of transformer blocks.
+        self.trunk = nn.Sequential(
+            *[
+                Block(
+                    dim=dim_in,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    init_values=init_values,
+                )
+                for _ in range(trunk_depth)
+            ]
+        )
+
+        # Normalizations for camera token and trunk output.
+        self.token_norm = nn.LayerNorm(dim_in)
+        self.trunk_norm = nn.LayerNorm(dim_in)
+
+        # Learnable empty camera pose token.
+        self.empty_pose_tokens = nn.Parameter(torch.zeros(1, 1, self.target_dim))
+        self.embed_pose = nn.Linear(self.target_dim, dim_in)
+
+        # Module for producing modulation parameters: shift, scale, and a gate.
+        self.poseLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True)
+        )
+
+        # Adaptive layer normalization without affine parameters.
+        self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6)
+        self.pose_branch = Mlp(
+            in_features=dim_in,
+            hidden_features=dim_in // 2,
+            out_features=self.target_dim,
+            drop=0,
+        )
+
+    def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list:
+        """
+        Forward pass to predict camera parameters.
+
+        Args:
+            aggregated_tokens_list (list): List of token tensors from the network;
+                the last tensor is used for prediction.
+            num_iterations (int, optional): Number of iterative refinement steps. Defaults to 4.
+
+        Returns:
+            list: A list of predicted camera encodings (post-activation) from each iteration.
+        """
+        # Use tokens from the last block for camera prediction.
+        tokens = aggregated_tokens_list[-1]
+
+        # Extract the camera tokens
+        pose_tokens = tokens[:, :, 0]
+
+        pose_tokens = self.token_norm(pose_tokens)
+
+        pred_pose_enc_list = self.trunk_fn(pose_tokens, num_iterations)
+        return pred_pose_enc_list
+
+    def trunk_fn(self, pose_tokens: torch.Tensor, num_iterations: int) -> list:
+        """
+        Iteratively refine camera pose predictions.
+
+        Args:
+            pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1, C].
+            num_iterations (int): Number of refinement iterations.
+
+        Returns:
+            list: List of activated camera encodings from each iteration.
+        """
+        B, S, C = pose_tokens.shape  # S is expected to be 1.
+        pred_pose_enc = None
+        pred_pose_enc_list = []
+
+        for _ in range(num_iterations):
+            # Use a learned empty pose for the first iteration.
+            if pred_pose_enc is None:
+                module_input = self.embed_pose(self.empty_pose_tokens.expand(B, S, -1))
+            else:
+                # Detach the previous prediction to avoid backprop through time.
+                pred_pose_enc = pred_pose_enc.detach()
+                module_input = self.embed_pose(pred_pose_enc)
+
+            # Generate modulation parameters and split them into shift, scale, and gate components.
+            shift_msa, scale_msa, gate_msa = self.poseLN_modulation(module_input).chunk(
+                3, dim=-1
+            )
+
+            # Adaptive layer normalization and modulation.
+            pose_tokens_modulated = gate_msa * modulate(
+                self.adaln_norm(pose_tokens), shift_msa, scale_msa
+            )
+            pose_tokens_modulated = pose_tokens_modulated + pose_tokens
+
+            pose_tokens_modulated = torch.utils.checkpoint.checkpoint(
+                self.trunk,
+                pose_tokens_modulated,
+                use_reentrant=False,
+            )
+
+            # Compute the delta update for the pose encoding.
+            pred_pose_enc_delta = self.pose_branch(
+                self.trunk_norm(pose_tokens_modulated)
+            )
+
+            if pred_pose_enc is None:
+                pred_pose_enc = pred_pose_enc_delta
+            else:
+                pred_pose_enc = pred_pose_enc + pred_pose_enc_delta
+
+            # Apply final activation functions for translation, quaternion, and field-of-view.
+            activated_pose = activate_pose(
+                pred_pose_enc,
+                trans_act=self.trans_act,
+                quat_act=self.quat_act,
+                fl_act=self.fl_act,
+            )
+            pred_pose_enc_list.append(activated_pose)
+
+        return pred_pose_enc_list
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """
+    Modulate the input tensor using scaling and shifting parameters.
+    """
+    # modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19
+    return x * (1 + scale) + shift