vendored model

sharp/models/__init__.py

@@ -0,0 +1,79 @@
+"""Contains different Gaussian predictors.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from sharp.models.monodepth import (
+    create_monodepth_adaptor,
+    create_monodepth_dpt,
+)
+
+from .alignment import create_alignment
+from .composer import GaussianComposer
+from .gaussian_decoder import create_gaussian_decoder
+from .heads import DirectPredictionHead
+from .initializer import create_initializer
+from .params import PredictorParams
+from .predictor import RGBGaussianPredictor
+
+
+def create_predictor(params: PredictorParams) -> RGBGaussianPredictor:
+    """Create gaussian predictor model specified by name."""
+    if params.gaussian_decoder.stride < params.initializer.stride:
+        raise ValueError(
+            "We donot expected gaussian_decoder has higher resolution than initializer."
+        )
+
+    scale_factor = params.gaussian_decoder.stride // params.initializer.stride
+    gaussian_composer = GaussianComposer(
+        delta_factor=params.delta_factor,
+        min_scale=params.min_scale,
+        max_scale=params.max_scale,
+        color_activation_type=params.color_activation_type,
+        opacity_activation_type=params.opacity_activation_type,
+        color_space=params.color_space,
+        scale_factor=scale_factor,
+        base_scale_on_predicted_mean=params.base_scale_on_predicted_mean,
+    )
+    if params.num_monodepth_layers > 1 and params.initializer.num_layers != 2:
+        raise KeyError("We only support num_layers = 2 when num_monodepth_layers > 1.")
+
+    monodepth_model = create_monodepth_dpt(params.monodepth)
+    monodepth_adaptor = create_monodepth_adaptor(
+        monodepth_model,
+        params.monodepth_adaptor,
+        params.num_monodepth_layers,
+        params.sorting_monodepth,
+    )
+
+    if params.num_monodepth_layers == 2:
+        monodepth_adaptor.replicate_head(params.num_monodepth_layers)
+
+    gaussian_decoder = create_gaussian_decoder(
+        params.gaussian_decoder,
+        dims_depth_features=monodepth_adaptor.get_feature_dims(),
+    )
+    initializer = create_initializer(
+        params.initializer,
+    )
+    prediction_head = DirectPredictionHead(
+        feature_dim=gaussian_decoder.dim_out, num_layers=initializer.num_layers
+    )
+    decoder_dim = monodepth_model.decoder.dims_decoder[-1]
+    return RGBGaussianPredictor(
+        init_model=initializer,
+        feature_model=gaussian_decoder,
+        prediction_head=prediction_head,
+        monodepth_model=monodepth_adaptor,
+        gaussian_composer=gaussian_composer,
+        scale_map_estimator=create_alignment(params.depth_alignment, depth_decoder_dim=decoder_dim),
+    )
+
+
+__all__ = [
+    "PredictorParams",
+    "create_predictor",
+]

sharp/models/alignment.py

@@ -0,0 +1,126 @@
+"""Contains modules for different types of alignment.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sharp.models.decoders import UNetDecoder
+from sharp.models.encoders import UNetEncoder
+from sharp.utils import math as math_utils
+
+from .params import AlignmentParams
+
+
+def create_alignment(
+    params: AlignmentParams, depth_decoder_dim: int | None = None
+) -> nn.Module | None:
+    """Create depth alignment."""
+    if depth_decoder_dim is None:
+        raise ValueError("Requires depth_decoder_dim for LearnedAlignment.")
+    alignment = LearnedAlignment(
+        depth_decoder_features=params.depth_decoder_features,
+        depth_decoder_dim=depth_decoder_dim,
+        steps=params.steps,
+        stride=params.stride,
+        base_width=params.base_width,
+        activation_type=params.activation_type,
+    )
+
+    if params.frozen:
+        alignment.requires_grad_(False)
+
+    return alignment
+
+
+class LearnedAlignment(nn.Module):
+    """Aligns tensors using a UNet."""
+
+    def __init__(
+        self,
+        steps: int = 4,
+        stride: int = 8,
+        base_width: int = 16,
+        depth_decoder_features: bool = False,
+        depth_decoder_dim: int = 256,
+        activation_type: math_utils.ActivationType = "exp",
+    ) -> None:
+        """Initialize LearnedAlignment.
+
+        Args:
+            steps: Number of steps in the UNet.
+            stride: Effective downsampling of the alignment module.
+            base_width: Base width of the UNet.
+            depth_decoder_features: Whether to use depth decoder features.
+            depth_decoder_dim: Dimension of the depth decoder features.
+            activation_type: Activation type for the alignment output.
+        """
+        super().__init__()
+        self.activation = math_utils.create_activation_pair(activation_type)
+        bias_value = self.activation.inverse(torch.tensor(1.0))
+
+        self.depth_decoder_features = depth_decoder_features
+        if depth_decoder_features:
+            dim_in = 2 + depth_decoder_dim
+        else:
+            dim_in = 2
+
+        def is_power_of_two(n: int) -> bool:
+            """Check if a number is a power of two."""
+            if n <= 0:
+                return False
+            return (n & (n - 1)) == 0
+
+        if not is_power_of_two(stride):
+            raise ValueError(f"Stride {stride} is not a power of two.")
+
+        steps_decoder = steps - int(math.log2(stride))
+        if steps_decoder < 1:
+            raise ValueError(f"{steps_decoder} must be greater or equal to 1.")
+        widths = [min(base_width << i, 1024) for i in range(steps + 1)]
+        self.encoder = UNetEncoder(dim_in=dim_in, width=widths, steps=steps, norm_num_groups=4)
+        self.decoder = UNetDecoder(
+            dim_out=widths[0], width=widths, steps=steps_decoder, norm_num_groups=4
+        )
+        self.conv_out = nn.Conv2d(widths[0], 1, 1, bias=True)
+        nn.init.zeros_(self.conv_out.weight)
+        nn.init.constant_(self.conv_out.bias, bias_value)
+
+    def forward(
+        self,
+        tensor_src: torch.Tensor,
+        tensor_tgt: torch.Tensor,
+        depth_decoder_features: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Compute alignment map."""
+        # Since the tensors are usually given by depth which is >= 1.0, we invert
+        # the tensors to have them in a reasonable range.
+        tensor_src = 1.0 / tensor_src.clamp(min=1e-4)
+        tensor_tgt = 1.0 / tensor_tgt.clamp(min=1e-4)
+        tensor_input = torch.cat([tensor_src, tensor_tgt], dim=1)
+        if self.depth_decoder_features:
+            height, width = tensor_src.shape[-2:]
+            upsampled_encodings = F.interpolate(
+                depth_decoder_features,
+                size=(height, width),
+                mode="bilinear",
+            )
+            tensor_input = torch.cat([tensor_input, upsampled_encodings], dim=1)
+        features = self.encoder(tensor_input)
+        output = self.conv_out(self.decoder(features))
+        alignment_map_lowres = self.activation.forward(output)
+        if alignment_map_lowres.shape[-2:] != tensor_src.shape[-2]:
+            alignment_map = F.interpolate(
+                alignment_map_lowres,
+                size=tensor_src.shape[-2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+        return alignment_map

sharp/models/blocks.py

@@ -0,0 +1,210 @@
+"""Contains reusable network components.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import nn
+
+NormLayerName = Literal["noop", "batch_norm", "group_norm", "instance_norm"]
+UpsamplingMode = Literal["transposed_conv", "nearest", "bilinear"]
+
+
+def norm_layer_2d(num_features: int, norm_type: NormLayerName, num_groups: int = 8) -> nn.Module:
+    """Create normalization layer."""
+    if norm_type == "noop":
+        return nn.Identity()
+    elif norm_type == "batch_norm":
+        return nn.BatchNorm2d(num_features=num_features)
+    elif norm_type == "group_norm":
+        return nn.GroupNorm(num_channels=num_features, num_groups=num_groups)
+    elif norm_type == "instance_norm":
+        return nn.InstanceNorm2d(num_features=num_features)
+    else:
+        raise ValueError(f"Invalid normalization layer type: {norm_type}")
+
+
+def upsampling_layer(upsampling_mode: UpsamplingMode, scale_factor: int, dim_in: int) -> nn.Module:
+    """Create upsampling layer."""
+    if upsampling_mode == "transposed_conv":
+        return nn.ConvTranspose2d(
+            in_channels=dim_in,
+            out_channels=dim_in,
+            kernel_size=scale_factor,
+            stride=scale_factor,
+            padding=0,
+            bias=False,
+        )
+    elif upsampling_mode in ("nearest", "bilinear"):
+        return nn.Upsample(scale_factor=scale_factor, mode=upsampling_mode)
+    else:
+        raise ValueError(f"Invalid upsampling mode {upsampling_mode}.")
+
+
+class ResidualBlock(nn.Module):
+    """Generic implementation of residual blocks.
+
+    This implements a generic residual block from
+
+        He et al. - Identity Mappings in Deep Residual Networks (2016),
+        https://arxiv.org/abs/1603.05027
+
+    which can be further customized via factory functions.
+    """
+
+    def __init__(self, residual: nn.Module, shortcut: nn.Module | None = None) -> None:
+        """Initialize ResidualBlock."""
+        super().__init__()
+        self.residual = residual
+        self.shortcut = shortcut
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply residual block."""
+        delta_x = self.residual(x)
+
+        if self.shortcut is not None:
+            x = self.shortcut(x)
+
+        return x + delta_x
+
+
+def residual_block_2d(
+    dim_in: int,
+    dim_out: int,
+    dim_hidden: int | None = None,
+    actvn: nn.Module | None = None,
+    norm_type: NormLayerName = "noop",
+    norm_num_groups: int = 8,
+    dilation: int = 1,
+    kernel_size: int = 3,
+):
+    """Create a simple 2D residual block."""
+    if actvn is None:
+        actvn = nn.ReLU()
+
+    if dim_hidden is None:
+        dim_hidden = dim_out // 2
+
+    # Padding to maintain output size
+    # See https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
+    padding = (dilation * (kernel_size - 1)) // 2
+
+    def _create_block(dim_in: int, dim_out: int) -> list[nn.Module]:
+        layers = [
+            norm_layer_2d(dim_in, norm_type, num_groups=norm_num_groups),
+            actvn,
+        ]
+
+        layers.append(
+            nn.Conv2d(
+                dim_in,
+                dim_out,
+                kernel_size=kernel_size,
+                stride=1,
+                dilation=dilation,
+                padding=padding,
+            )
+        )
+        return layers
+
+    residual = nn.Sequential(
+        *_create_block(dim_in, dim_hidden),
+        *_create_block(dim_hidden, dim_out),
+    )
+    shortcut = None
+
+    if dim_in != dim_out:
+        shortcut = nn.Conv2d(dim_in, dim_out, 1)
+
+    return ResidualBlock(residual, shortcut)
+
+
+class FeatureFusionBlock2d(nn.Module):
+    """Feature fusion for DPT."""
+
+    # We use the name "deconv" for backward compatibility. However, "deconv" can also
+    # refer to some other upsampling layer or a no-op.
+    deconv: nn.Module
+
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int | None = None,
+        upsampling_mode: UpsamplingMode | None = None,
+        batch_norm: bool = False,
+    ):
+        """Initialize feature fusion block.
+
+        Args:
+            dim_in: Dimensions of input.
+            dim_out: Dimensions of output.
+            batch_norm: Whether to use batch normalization in resnet blocks.
+            upsampling_mode: What mode to use for upsampling. None if no upsampling
+                is required.
+        """
+        super().__init__()
+        if dim_out is None:
+            dim_out = dim_in
+        self.resnet1 = self._residual_block(dim_in, batch_norm)
+        self.resnet2 = self._residual_block(dim_in, batch_norm)
+
+        if upsampling_mode is not None:
+            self.deconv = upsampling_layer(upsampling_mode, scale_factor=2, dim_in=dim_in)
+        else:
+            self.deconv = nn.Sequential()
+
+        self.out_conv = nn.Conv2d(
+            dim_in,
+            dim_out,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor:
+        """Process and fuse input features."""
+        x = x0
+
+        if x1 is not None:
+            res = self.resnet1(x1)
+            x = self.skip_add.add(x, res)
+
+        x = self.resnet2(x)
+        x = self.deconv(x)
+        x = self.out_conv(x)
+
+        return x
+
+    @staticmethod
+    def _residual_block(num_features: int, batch_norm: bool):
+        """Create a residual block."""
+
+        def _create_block(dim: int, batch_norm: bool) -> list[nn.Module]:
+            layers = [
+                nn.ReLU(False),
+                nn.Conv2d(
+                    num_features,
+                    num_features,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    bias=not batch_norm,
+                ),
+            ]
+            if batch_norm:
+                layers.append(nn.BatchNorm2d(dim))
+            return layers
+
+        residual = nn.Sequential(
+            *_create_block(dim=num_features, batch_norm=batch_norm),
+            *_create_block(dim=num_features, batch_norm=batch_norm),
+        )
+        return ResidualBlock(residual)

sharp/models/composer.py

@@ -0,0 +1,251 @@
+"""Defines module to compose final Gaussians from base values and delta values.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sharp.models.initializer import GaussianBaseValues
+from sharp.utils import math as math_utils
+from sharp.utils.color_space import ColorSpace, sRGB2linearRGB
+from sharp.utils.gaussians import Gaussians3D
+
+from .params import DeltaFactor
+
+
+def _get_scale_activation_constant(max_scale: float, min_scale: float) -> tuple[float, float]:
+    """Return constants for scale activation function."""
+    # To ensure for delta = 0, the value of scale_factor is 1 and the gradient is 1.
+    constant_a = (max_scale - min_scale) / (1 - min_scale) / (max_scale - 1)
+    constant_b = math_utils.inverse_sigmoid(
+        torch.tensor((1.0 - min_scale) / (max_scale - min_scale))
+    ).item()
+    return constant_a, constant_b
+
+
+class GaussianComposer(nn.Module):
+    """Converts base values and deltas into Gaussians."""
+
+    color_activation_type: math_utils.ActivationType
+    opacity_activation_type: math_utils.ActivationType
+
+    def __init__(
+        self,
+        delta_factor: DeltaFactor,
+        min_scale: float,
+        max_scale: float,
+        color_activation_type: math_utils.ActivationType,
+        opacity_activation_type: math_utils.ActivationType,
+        color_space: ColorSpace,
+        base_scale_on_predicted_mean: bool,
+        scale_factor: int = 1,
+    ) -> None:
+        """Initialize GaussianComposer.
+
+        Args:
+            delta_factor: Multiply delta offsets by this factor.
+            min_scale: The minimal scale factor for gaussian scale activation.
+            max_scale: The maximal scale factor for gaussian scale activation.
+            color_activation_type: Which activation function to use for colors.
+            opacity_activation_type: Which activation function to use for opacities.
+            color_space: Which color space is used in training.
+            scale_factor: The scale factor to upsample the delta_values before composition.
+            base_scale_on_predicted_mean: Whether to account z offsets for estimating base scale.
+        """
+        super().__init__()
+        self.delta_factor = delta_factor
+        self.max_scale = max_scale
+        self.min_scale = min_scale
+        self.color_activation_type = color_activation_type
+        self.opacity_activation_type = opacity_activation_type
+        self.color_space = color_space
+        self.scale_factor = scale_factor
+        self.base_scale_on_predicted_mean = base_scale_on_predicted_mean
+
+    def upsample_delta_value(self, delta: torch.Tensor, scale_factor: int = 1):
+        """Upsample the delta value.
+
+        Args:
+            delta: The delta values predicted by gaussian predictor.
+            scale_factor: The scale factor to upsample the delta_values.
+        """
+        (
+            batch_size,
+            num_channels,
+            num_layers,
+            image_height,
+            image_width,
+        ) = delta.shape
+        new_height = image_height * scale_factor
+        new_width = image_width * scale_factor
+        upsampled_delta = F.interpolate(
+            delta.view(batch_size, num_channels * num_layers, image_height, image_width),
+            scale_factor=scale_factor,
+        ).view(batch_size, num_channels, num_layers, new_height, new_width)
+        return upsampled_delta
+
+    def forward(
+        self,
+        delta: torch.Tensor,
+        base_values: GaussianBaseValues,
+        global_scale: torch.Tensor | None = None,
+        flatten_output: bool = True,
+    ) -> Gaussians3D:
+        """Combine predicted delta values with base gaussian values and apply activation function.
+
+        Args:
+            delta: The delta values predicted by gaussian predictor.
+            base_values: The gaussian base values.
+            global_scale: Global scale of Gaussians.
+            flatten_output: Flatten the gaussian parameters.
+
+        Returns:
+            The computed 3D Gaussians.
+        """
+        # Upsample the delta if delta and base_values have different strides.
+        scale_factor = self.scale_factor
+        # For triplane head, the delta has already been upsampled.
+        actual_scale_factor = base_values.mean_x_ndc.shape[-1] // delta.shape[-1]
+        if scale_factor != 1 and actual_scale_factor != 1:
+            delta = self.upsample_delta_value(delta, scale_factor)
+
+        mean_vectors = self._forward_mean(base_values, delta)
+
+        # Account for the change in base scale due to z offsets.
+        base_scales = (
+            (base_values.scales * base_values.mean_inverse_z_ndc * mean_vectors[:, 2:3, ...])
+            if self.base_scale_on_predicted_mean
+            else base_values.scales
+        )
+        singular_values = self._scale_activation(
+            base_scales,
+            delta[:, 3:6],
+            self.min_scale,
+            self.max_scale,
+        )
+        quaternions = self._quaternion_activation(base_values.quaternions, delta[:, 6:10])
+        colors = self._color_activation(base_values.colors, delta[:, 10:13])
+        opacities = self._opacity_activation(base_values.opacities, delta[:, 13])
+
+        if flatten_output:
+            # [B, C, N, H, W] -> [B, N, H, W, C].
+            # NOTE: opacities is [B, N, H, W] so it doesn't need to permute.
+            mean_vectors = mean_vectors.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            singular_values = singular_values.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            quaternions = quaternions.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            colors = colors.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            opacities = opacities.flatten(1, 3)
+
+        # Apply global scaling to convert Gaussians to metric space.
+        if global_scale is not None:
+            mean_vectors = global_scale[:, None, None] * mean_vectors
+            singular_values = global_scale[:, None, None] * singular_values
+
+        return Gaussians3D(
+            mean_vectors=mean_vectors,
+            singular_values=singular_values,
+            quaternions=quaternions,
+            colors=colors,
+            opacities=opacities,
+        )
+
+    def _forward_mean(self, base_values: GaussianBaseValues, delta: torch.Tensor) -> torch.Tensor:
+        # Concatenate base vectors and apply mean activation.
+        delta_factor = torch.tensor(
+            [self.delta_factor.xy, self.delta_factor.xy, self.delta_factor.z],
+            device=delta.device,
+        )[None, :, None, None, None]
+
+        dtype = base_values.mean_x_ndc.dtype
+        device = base_values.mean_x_ndc.device
+        target_shape = (1, 3, 1, 1, 1)
+        mean_x_mask = torch.tensor([1.0, 0.0, 0.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+        mean_y_mask = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+        mean_z_mask = torch.tensor([0.0, 0.0, 1.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+
+        mean_vectors_ndc = (
+            base_values.mean_x_ndc.repeat(target_shape) * mean_x_mask
+            + base_values.mean_y_ndc.repeat(target_shape) * mean_y_mask
+            + base_values.mean_inverse_z_ndc.repeat(target_shape) * mean_z_mask
+        )
+
+        mean_vectors = self._mean_activation(mean_vectors_ndc, delta_factor * delta[:, :3])
+        return mean_vectors
+
+    def _mean_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        """Mean activation function.
+
+        Args:
+            base: Tensor of shape [B, 3, H, W], where first two feature dimensions
+                (x,y) are in normalized device coordinates (NDC) where (-1, -1) is
+                the top, while the third dimension is inverse depth.
+            learned_delta: Tensor of shape [B, 3, H, W] with predicted delta values.
+
+        Returns:
+            Returns: The final mean vector after combining base and delta and applying nonlinearies.
+        """
+        xx = base[:, 0:1] + learned_delta[:, 0:1]
+        yy = base[:, 1:2] + learned_delta[:, 1:2]
+
+        a = base[:, 2:3]
+        b = learned_delta[:, 2:3]
+
+        # Original formula:
+        inverse_zz = F.softplus(math_utils.inverse_softplus(a) + b)
+        zz = 1.0 / (inverse_zz + 1e-3)
+
+        mean_vectors = torch.cat([zz * xx, zz * yy, zz], dim=1)
+        return mean_vectors
+
+    def _scale_activation(
+        self,
+        base: torch.Tensor,
+        learned_delta: torch.Tensor,
+        min_scale: float,
+        max_scale: float,
+    ) -> torch.Tensor:
+        constant_a, constant_b = _get_scale_activation_constant(max_scale, min_scale)
+        scale_factor = (max_scale - min_scale) * torch.sigmoid(
+            constant_a * self.delta_factor.scale * learned_delta + constant_b
+        ) + min_scale
+        return base * scale_factor
+
+    def _quaternion_activation(
+        self, base: torch.Tensor, learned_delta: torch.Tensor
+    ) -> torch.Tensor:
+        # No need to normalize the quaternions, since this is also done in rendering.
+        return base + self.delta_factor.quaternion * learned_delta
+
+    def _color_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        # For certain activation functions we need to clamp the base value to
+        # a supported range.
+        if self.color_activation_type == "sigmoid":
+            base = torch.clamp(base, min=0.01, max=0.99)
+        elif self.color_activation_type in ("exp", "softplus"):
+            base = torch.clamp(base, min=0.01)
+
+        activation = math_utils.create_activation_pair(self.color_activation_type)
+        colors: torch.Tensor = activation.forward(
+            activation.inverse(base) + self.delta_factor.color * learned_delta
+        )
+        # Convert gaussian color to linear if linearRGB colorspace is specified.
+        if self.color_space == "linearRGB":
+            colors = sRGB2linearRGB(colors)
+        return colors
+
+    def _opacity_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        activation = math_utils.create_activation_pair(self.opacity_activation_type)
+        return activation.forward(
+            activation.inverse(base) + self.delta_factor.opacity * learned_delta
+        )

sharp/models/decoders/__init__.py

@@ -0,0 +1,22 @@
+"""Contains different decoders for Gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from .base_decoder import BaseDecoder
+from .monodepth_decoder import (
+    create_monodepth_decoder,
+)
+from .multires_conv_decoder import MultiresConvDecoder, UpsamplingMode
+from .unet_decoder import UNetDecoder
+
+__all__ = [
+    "BaseDecoder",
+    "UNetDecoder",
+    "MultiresConvDecoder",
+    "UpsamplingMode",
+    "create_monodepth_decoder",
+]

sharp/models/decoders/base_decoder.py

@@ -0,0 +1,21 @@
+"""Contains the base class for decoders.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import abc
+from typing import List
+
+import torch
+from torch import nn
+
+
+class BaseDecoder(nn.Module, abc.ABC):
+    """Base decoder class."""
+
+    dim_out: int
+
+    @abc.abstractmethod
+    def forward(self, encodings: List[torch.Tensor]) -> torch.Tensor:
+        """Decode (multi-resolution) encodings."""

sharp/models/decoders/monodepth_decoder.py

@@ -0,0 +1,37 @@
+"""Contains factory function for loading/creating monodepth decoder.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+
+from __future__ import annotations
+
+from sharp.models.presets import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    ViTPreset,
+)
+
+from .multires_conv_decoder import MultiresConvDecoder
+
+
+def create_monodepth_decoder(
+    patch_encoder_preset: ViTPreset,
+    dims_decoder=None,
+) -> MultiresConvDecoder:
+    """Create DepthDensePredictionTransformer model.
+
+    Args:
+        patch_encoder_preset: The preset patch encoder architecture in SPN.
+        dims_decoder: The decoder architecture.
+    """
+    dims_encoder = MONODEPTH_ENCODER_DIMS_MAP[patch_encoder_preset]
+    if dims_decoder is None:
+        dims_decoder = dims_encoder[0]
+    if isinstance(dims_decoder, int):
+        dims_decoder = [dims_decoder]
+    decoder = MultiresConvDecoder(
+        dims_encoder=[dims_decoder[0]] + list(dims_encoder), dims_decoder=dims_decoder
+    )
+
+    return decoder

sharp/models/decoders/multires_conv_decoder.py

@@ -0,0 +1,116 @@
+"""Contains multi-res convolutional decoder.
+
+Implements the decoder for Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+
+from sharp.models.blocks import FeatureFusionBlock2d, UpsamplingMode
+from sharp.utils.training import checkpoint_wrapper
+
+from .base_decoder import BaseDecoder
+
+
+class MultiresConvDecoder(BaseDecoder):
+    """Decoder for multi-resolution encodings."""
+
+    def __init__(
+        self,
+        dims_encoder: Iterable[int],
+        dims_decoder: Iterable[int] | int,
+        grad_checkpointing: bool = False,
+        upsampling_mode: UpsamplingMode = "transposed_conv",
+    ):
+        """Initialize multiresolution convolutional decoder.
+
+        Args:
+            dims_encoder: Expected dims at each level from the encoder.
+            dims_decoder: Dim of decoder features.
+            grad_checkpointing: Whether to checkpoint gradient during training.
+            upsampling_mode: What method to use for upsampling.
+        """
+        super().__init__()
+        self.dims_encoder = list(dims_encoder)
+
+        if isinstance(dims_decoder, int):
+            self.dims_decoder = [dims_decoder] * len(self.dims_encoder)
+        else:
+            self.dims_decoder = list(dims_decoder)
+
+        if len(self.dims_decoder) != len(self.dims_encoder):
+            raise ValueError("Received dims_encoder and dims_decoder of different sizes.")
+
+        self.dim_out = self.dims_decoder[0]
+
+        num_encoders = len(self.dims_encoder)
+
+        # At the highest resolution, i.e. level 0, we apply projection w/ 1x1 convolution
+        # when the dimensions mismatch. Otherwise we do not do anything, which is
+        # the default behavior of monodepth.
+        conv0 = (
+            nn.Conv2d(self.dims_encoder[0], self.dims_decoder[0], kernel_size=1, bias=False)
+            if self.dims_encoder[0] != self.dims_decoder[0]
+            else nn.Identity()
+        )
+
+        convs = [conv0]
+        for i in range(1, num_encoders):
+            convs.append(
+                nn.Conv2d(
+                    self.dims_encoder[i],
+                    self.dims_decoder[i],
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    bias=False,
+                )
+            )
+        self.convs = nn.ModuleList(convs)
+
+        fusions = []
+        for i in range(num_encoders):
+            fusions.append(
+                FeatureFusionBlock2d(
+                    dim_in=self.dims_decoder[i],
+                    dim_out=self.dims_decoder[i - 1] if i != 0 else self.dim_out,
+                    upsampling_mode=upsampling_mode if i != 0 else None,
+                    batch_norm=False,
+                )
+            )
+        self.fusions = nn.ModuleList(fusions)
+
+        self.grad_checkpointing = grad_checkpointing
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+
+    def forward(self, encodings: list[torch.Tensor]) -> torch.Tensor:
+        """Decode the multi-resolution encodings."""
+        num_levels = len(encodings)
+        num_encoders = len(self.dims_encoder)
+
+        if num_levels != num_encoders:
+            raise ValueError(
+                f"Encoder output levels={num_levels} at runtime "
+                f"mismatch with expected levels={num_encoders}."
+            )
+
+        # Project features of different encoder dims to the same decoder dim.
+        # Fuse features from the lowest resolution (num_levels-1)
+        # to the highest (0).
+        features = self.convs[-1](encodings[-1])
+        features = checkpoint_wrapper(self, self.fusions[-1], features)
+        for i in range(num_levels - 2, -1, -1):
+            features_i = self.convs[i](encodings[i])
+            features = checkpoint_wrapper(self, self.fusions[i], features, features_i)
+        return features

sharp/models/decoders/unet_decoder.py

@@ -0,0 +1,113 @@
+"""Contains the UNet decoder.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+from torch import nn
+
+from sharp.models.blocks import (
+    NormLayerName,
+    norm_layer_2d,
+    residual_block_2d,
+)
+
+from .base_decoder import BaseDecoder
+
+
+class UNetDecoder(BaseDecoder):
+    """Decoder of UNet model."""
+
+    def __init__(
+        self,
+        dim_out: int,
+        width: List[int] | int,
+        steps: int = 5,
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups=8,
+        blocks_per_layer=2,
+    ) -> None:
+        """Initialize UNet Decoder.
+
+        Args:
+            dim_out: The number of output channels.
+            width: Width of last input feature map from encoder
+                or the width list of all input feature maps from encoder.
+            steps: The number of upsampling steps.
+            norm_type: Which kind of normalization layer to use.
+            norm_num_groups: How many groups to use for group norm (if relevant).
+            blocks_per_layer: How many blocks per layer to use.
+        """
+        super().__init__()
+
+        if blocks_per_layer < 1:
+            raise ValueError("blocks_per_layer must be greater or equal to one.")
+
+        self.dim_out = dim_out
+
+        self.convs_up = nn.ModuleList()
+
+        self.output_dims: list[int]
+        # If only one number is specified, we assume each layer will double the channel dimension.
+        if isinstance(width, int):
+            self.input_dims = [width >> i for i in range(0, steps + 1)]
+        else:
+            self.input_dims = width[::-1][: steps + 1]
+
+        for i_step in range(steps):
+            input_width = self.input_dims[i_step]
+            current_width = self.input_dims[i_step + 1]
+            convs_up_i = nn.Sequential(
+                nn.Upsample(scale_factor=2),
+                residual_block_2d(
+                    input_width * (1 if i_step == 0 else 2),
+                    current_width,
+                    norm_type=norm_type,
+                    norm_num_groups=norm_num_groups,
+                ),
+                *[
+                    residual_block_2d(
+                        current_width,
+                        current_width,
+                        norm_type=norm_type,
+                        norm_num_groups=norm_num_groups,
+                    )
+                    for _ in range(blocks_per_layer - 1)
+                ],
+            )
+            self.convs_up.append(convs_up_i)
+            input_width = 2 * current_width
+            current_width //= 2
+
+        last_width = self.input_dims[-1]
+        self.conv_out = nn.Sequential(
+            norm_layer_2d(last_width * 2, norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+            nn.Conv2d(last_width * 2, dim_out, 1),
+            norm_layer_2d(dim_out, norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+        )
+
+    def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
+        """Apply UNet to image.
+
+        Args:
+            features: The input multi-level feature map from encoder.
+
+        Returns:
+            The output feature map.
+        """
+        i_feature_layer = len(features) - 1
+        out = self.convs_up[0](features[i_feature_layer])
+        i_feature_layer -= 1
+        for conv_up in self.convs_up[1:]:  # type: ignore
+            out = conv_up(torch.cat([out, features[i_feature_layer]], dim=1))
+            i_feature_layer -= 1
+        out = self.conv_out(torch.cat([out, features[i_feature_layer]], dim=1))
+
+        return out

sharp/models/encoders/__init__.py

@@ -0,0 +1,24 @@
+"""Contains different encoders for Gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from sharp.models.encoders.base_encoder import BaseEncoder
+
+from .monodepth_encoder import (
+    MonodepthFeatureEncoder,
+    create_monodepth_encoder,
+)
+from .spn_encoder import SlidingPyramidNetwork
+from .unet_encoder import UNetEncoder
+from .vit_encoder import create_vit
+
+__all__ = [
+    "create_vit",
+    "BaseEncoder",
+    "UNetEncoder",
+    "SlidingPyramidNetwork",
+    "MonodepthFeatureEncoder",
+    "create_monodepth_encoder",
+]

sharp/models/encoders/base_encoder.py

@@ -0,0 +1,25 @@
+"""Contains the base class for encoders.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import abc
+
+import torch
+from torch import nn
+
+
+class BaseEncoder(nn.Module, abc.ABC):
+    """Base encoder class."""
+
+    dim_in: int
+    output_dims: list[int]
+
+    @abc.abstractmethod
+    def forward(self, image: torch.Tensor) -> list[torch.Tensor]:
+        """Encode input image into multi-resolution encodings."""
+
+    def internal_resolution(self) -> int:
+        """Internal resolution of the encoder."""
+        return 1536

sharp/models/encoders/monodepth_encoder.py

@@ -0,0 +1,123 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from sharp.models.presets import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    MONODEPTH_HOOK_IDS_MAP,
+    ViTPreset,
+)
+
+from .base_encoder import BaseEncoder
+from .spn_encoder import SlidingPyramidNetwork
+from .vit_encoder import create_vit
+
+
+def create_monodepth_encoder(
+    patch_encoder_preset: ViTPreset,
+    image_encoder_preset: ViTPreset,
+    use_patch_overlap: bool = True,
+    last_encoder: int = 256,
+) -> SlidingPyramidNetwork:
+    """Creates DepthDensePredictionTransformer model.
+
+    Args:
+        patch_encoder_preset: The preset patch encoder architecture in SPN.
+        image_encoder_preset: The preset image encoder architecture in SPN.
+        use_patch_overlap: Whether to use overlap between patches in SPN.
+        last_encoder: last number of encoder features.
+    """
+    dims_encoder = [last_encoder] + MONODEPTH_ENCODER_DIMS_MAP[patch_encoder_preset]
+    patch_encoder_block_ids = MONODEPTH_HOOK_IDS_MAP[patch_encoder_preset]
+
+    patch_encoder = create_vit(
+        preset=patch_encoder_preset,
+        intermediate_features_ids=patch_encoder_block_ids,
+        # We always need to output intermediate features for assembly.
+    )
+    image_encoder = create_vit(
+        preset=image_encoder_preset,
+        intermediate_features_ids=None,
+    )
+
+    encoder = SlidingPyramidNetwork(
+        dims_encoder=dims_encoder,
+        patch_encoder=patch_encoder,
+        image_encoder=image_encoder,
+        use_patch_overlap=use_patch_overlap,
+    )
+
+    return encoder
+
+
+class ProjectionModule(nn.Module):
+    """Apply projection of features."""
+
+    def __init__(self, dims_in: list[int], dims_out: list[int]) -> None:
+        """Initialize projection module."""
+        super().__init__()
+        if len(dims_in) != len(dims_out):
+            raise ValueError("Length of dims_in must be same as length of dims_out.")
+        self.convs = nn.ModuleList(
+            [nn.Conv2d(dim_in, dim_out, 1) for dim_in, dim_out in zip(dims_in, dims_out)]
+        )
+
+    def forward(self, encodings: list[torch.Tensor]) -> list[torch.Tensor]:
+        """Apply projection module."""
+        if len(encodings) != len(self.convs):
+            raise ValueError("Number of encodings must be equal to number of projections.")
+        return [conv(encoding) for conv, encoding in zip(self.convs, encodings)]
+
+
+class MonodepthFeatureEncoder(BaseEncoder):
+    """A wrapper around monodepth network to extract features."""
+
+    def __init__(
+        self,
+        monodepth_encoder: SlidingPyramidNetwork,
+        output_dims: list[int] | None = None,
+        freeze_projection: bool = False,
+    ) -> None:
+        """Initialize MonodepthFeatureExtractor."""
+        super().__init__()
+
+        self.encoder = monodepth_encoder
+
+        # The monodepth network returns two feature maps for the first entry in
+        # backbone.encoder.dims_encoder.
+        monodepth_dims = self.encoder.dims_encoder
+        monodepth_dims = monodepth_dims
+
+        if output_dims is not None:
+            if not len(output_dims) == len(monodepth_dims):
+                raise ValueError(
+                    "When set, number of output dimensions must be equal to output "
+                    f"dimensions of monodepth model {len(monodepth_dims)}."
+                )
+
+            self.projection = ProjectionModule(monodepth_dims, output_dims)
+            self.output_dims = output_dims
+        else:
+            self.projection = nn.Identity()
+            self.output_dims = monodepth_dims
+
+        if freeze_projection:
+            self.projection.requires_grad_(False)
+
+    def forward(self, input_features: torch.Tensor) -> list[torch.Tensor]:
+        """Extract multi-resolution features."""
+        encodings = self.encoder(input_features[:, :3].contiguous())
+        return self.projection(encodings)
+
+    def internal_resolution(self) -> int:
+        """Internal resolution of the encoder."""
+        return self.encoder.internal_resolution()

sharp/models/encoders/spn_encoder.py

@@ -0,0 +1,369 @@
+"""Contains Sliding Pyramid Network architecture.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Iterable
+
+import torch
+import torch.fx
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sharp.utils.training import checkpoint_wrapper
+
+from .base_encoder import BaseEncoder
+from .vit_encoder import TimmViT
+
+# torch.fx.wrap is used here to mark functions as leaf nodes during symbolic tracing
+# ensuring they are not traced but seen as atomic operation. In short, symbolic tracing
+# struggles with native python functions and conditional flows.
+non_traceable_ops = ("len", "int")
+for op in non_traceable_ops:
+    torch.fx.wrap(op)
+
+
+class SlidingPyramidNetwork(BaseEncoder):
+    """Sliding Pyramid Network.
+
+    An encoder aimed at creating multi-resolution encodings from Vision Transformers.
+
+    Reference: Bochkovskii et al. - "Depth pro: Sharp monocular metric depth in less
+               than a second." (ICLR 2024)
+    """
+
+    def __init__(
+        self,
+        dims_encoder: Iterable[int],
+        patch_encoder: TimmViT,
+        image_encoder: TimmViT,
+        use_patch_overlap: bool = True,
+    ):
+        """Initialize Sliding Pyramid Network.
+
+        The framework
+            1. creates an image pyramid,
+            2. generates overlapping patches with a sliding window at each pyramid level,
+            3. creates batched encodings via vision transformer backbones,
+            4. produces multi-resolution encodings.
+
+        Args:
+            dims_encoder: Dimensions of the encoder at different layers.
+            patch_encoder: Backbone used for highres part of the pyramid.
+            image_encoder: Backbone used for lowres part of the pyramid.
+            use_patch_overlap: Whether to use overlap between patches in SPN.
+        """
+        super().__init__()
+
+        self.dim_in = patch_encoder.dim_in
+
+        self.dims_encoder = list(dims_encoder)
+        self.patch_encoder = patch_encoder
+        self.image_encoder = image_encoder
+
+        base_embed_dim = patch_encoder.embed_dim
+        lowres_embed_dim = image_encoder.embed_dim
+        self.patch_size = patch_encoder.internal_resolution()
+
+        self.grad_checkpointing = False
+        self.use_patch_overlap = use_patch_overlap
+
+        # Retrieve intermediate feature ids registered in create_monodepth_encoder.
+        self.patch_intermediate_features_ids = patch_encoder.intermediate_features_ids
+        if (
+            not isinstance(self.patch_intermediate_features_ids, list)
+            or not len(self.patch_intermediate_features_ids) == 4
+        ):
+            raise ValueError("Patch intermediate feature ids must be a 4-item list.")
+
+        self.image_intermediate_features_ids = image_encoder.intermediate_features_ids
+
+        def _create_project_upsample_block(
+            dim_in: int,
+            dim_out: int,
+            upsample_layers: int,
+            dim_intermediate=None,
+        ) -> nn.Module:
+            if dim_intermediate is None:
+                dim_intermediate = dim_out
+            # Projection.
+            blocks = [
+                nn.Conv2d(
+                    in_channels=dim_in,
+                    out_channels=dim_intermediate,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False,
+                )
+            ]
+
+            # Upsampling.
+            blocks += [
+                nn.ConvTranspose2d(
+                    in_channels=dim_intermediate if i == 0 else dim_out,
+                    out_channels=dim_out,
+                    kernel_size=2,
+                    stride=2,
+                    padding=0,
+                    bias=False,
+                )
+                for i in range(upsample_layers)
+            ]
+
+            return nn.Sequential(*blocks)
+
+        self.upsample_latent0 = _create_project_upsample_block(
+            dim_in=base_embed_dim,
+            dim_out=self.dims_encoder[0],
+            upsample_layers=3,
+            dim_intermediate=self.dims_encoder[1],
+        )
+        self.upsample_latent1 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[1], upsample_layers=2
+        )
+
+        self.upsample0 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[2], upsample_layers=1
+        )
+        self.upsample1 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[3], upsample_layers=1
+        )
+        self.upsample2 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[4], upsample_layers=1
+        )
+
+        self.upsample_lowres = nn.ConvTranspose2d(
+            in_channels=lowres_embed_dim,
+            out_channels=self.dims_encoder[4],
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=True,
+        )
+        self.fuse_lowres = nn.Conv2d(
+            in_channels=(self.dims_encoder[4] + self.dims_encoder[4]),
+            out_channels=self.dims_encoder[4],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+
+    def internal_resolution(self) -> int:
+        """Return the full image size of the SPN network."""
+        return self.patch_size * 4
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+        self.patch_encoder.set_grad_checkpointing(is_enabled)
+        self.image_encoder.set_grad_checkpointing(is_enabled)
+
+    @torch.jit.ignore
+    def set_requires_grad_(self, patch_encoder: bool, image_encoder: bool):
+        """Set requires grad for separate components."""
+        self.patch_encoder.requires_grad_(patch_encoder)
+        self.image_encoder.requires_grad_(image_encoder)
+
+        # Always freeze the unused TimmViT head to exclude it from the calculation of
+        # trainable parameters.
+        self.patch_encoder.head.requires_grad_(False)
+        self.image_encoder.head.requires_grad_(False)
+
+        # These upsamplers only affect patch encoder's feature maps.
+        self.upsample_latent0.requires_grad_(patch_encoder)
+        self.upsample_latent1.requires_grad_(patch_encoder)
+        self.upsample0.requires_grad_(patch_encoder)
+        self.upsample1.requires_grad_(patch_encoder)
+        self.upsample2.requires_grad_(patch_encoder)
+
+        # This upsampler affects only image encoder's feature map.
+        self.upsample_lowres.requires_grad_(image_encoder)
+
+        # This fuser affects both image and patch encoders.
+        self.fuse_lowres.requires_grad_(image_encoder or patch_encoder)
+
+    def _create_pyramid(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Creates a 3-level image pyramid."""
+        # Original resolution: 1536 by default.
+        x0 = x
+
+        # Middle resolution: 768 by default.
+        x1 = F.interpolate(x, size=None, scale_factor=0.5, mode="bilinear", align_corners=False)
+
+        # Low resolution: 384 by default, corresponding to the backbone resolution.
+        x2 = F.interpolate(x, size=None, scale_factor=0.25, mode="bilinear", align_corners=False)
+
+        return x0, x1, x2
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Encode input at multiple resolutions."""
+        batch_size = x.shape[0]
+
+        # Step 0: create a 3-level image pyramid.
+        x0, x1, x2 = self._create_pyramid(x)
+
+        if self.use_patch_overlap:
+            # Step 1: split to create batched overlapped mini-images at the ViT
+            # resolution.
+            # 5x5 @ 384x384 at the highest resolution (1536x1536).
+            x0_patches = split(x0, overlap_ratio=0.25, patch_size=self.patch_size)
+            # 3x3 @ 384x384 at the middle resolution (768x768).
+            x1_patches = split(x1, overlap_ratio=0.5, patch_size=self.patch_size)
+            # 1x1 # 384x384 at the lowest resolution (384x384).
+            x2_patches = x2
+            padding = 3
+        else:
+            # Step 1: split to create batched overlapped mini-images at the ViT
+            # resolution.
+            # 4x4 @ 384x384 at the highest resolution (1536x1536).
+            x0_patches = split(x0, overlap_ratio=0.0, patch_size=self.patch_size)
+            # 2x2 @ 384x384 at the middle resolution (768x768).
+            x1_patches = split(x1, overlap_ratio=0.0, patch_size=self.patch_size)
+            # 1x1 # 384x384 at the lowest resolution (384x384).
+            x2_patches = x2
+            padding = 0
+        x0_tile_size = x0_patches.shape[0]
+
+        # Concatenate all the sliding window patches and form a batch of size
+        # (35=5x5+3x3+1x1) or (21=4x4+2x2+1x1).
+        x_pyramid_patches = torch.cat(
+            (x0_patches, x1_patches, x2_patches),
+            dim=0,
+        )
+
+        # Run the ViT model and get the result of large batch size.
+        #
+        # For the retrieval of intermediate features forward hooks are more concise,
+        # but they are not well compatible with symbolic tracing because attributes
+        # of submodules can be lost during tracing. Therefore, forward hooks may not
+        # be preserved during graph transformation, leading to unexpected behavior.
+        # To avoid such issues it is safer not to use them because they are not
+        # essential here.
+        x_pyramid_encodings, patch_intermediate_features = self.patch_encoder(x_pyramid_patches)
+
+        # Step 3: merging.
+        # Merge highres latent encoding.
+        # NOTE: list type check has completed in init.
+        x_latent0_encodings = self.patch_encoder.reshape_feature(
+            patch_intermediate_features[self.patch_intermediate_features_ids[0]]  # type:ignore[index]
+        )
+        x_latent0_features = merge(
+            x_latent0_encodings[: batch_size * x0_tile_size],
+            batch_size=batch_size,
+            padding=padding,
+        )
+
+        x_latent1_encodings = self.patch_encoder.reshape_feature(
+            patch_intermediate_features[self.patch_intermediate_features_ids[1]]  # type:ignore[index]
+        )
+        x_latent1_features = merge(
+            x_latent1_encodings[: batch_size * x0_tile_size],
+            batch_size=batch_size,
+            padding=padding,
+        )
+
+        # Split the 35 batch size from pyramid encoding back into 5x5+3x3+1x1.
+        x0_encodings, x1_encodings, x2_encodings = torch.split(
+            x_pyramid_encodings,
+            [len(x0_patches), len(x1_patches), len(x2_patches)],
+            dim=0,
+        )
+
+        # 96x96 feature maps by merging 5x5 @ 24x24 patches with overlaps.
+        x0_features = merge(x0_encodings, batch_size=batch_size, padding=padding)
+
+        # 48x84 feature maps by merging 3x3 @ 24x24 patches with overlaps.
+        x1_features = merge(x1_encodings, batch_size=batch_size, padding=2 * padding)
+
+        # 24x24 feature maps.
+        x2_features = x2_encodings
+
+        # Apply the image encoder.
+        x_lowres_features, image_intermediate_features = self.image_encoder(x2_patches)
+
+        # Upsample feature maps.
+        x_latent0_features = checkpoint_wrapper(self, self.upsample_latent0, x_latent0_features)
+        x_latent1_features = checkpoint_wrapper(self, self.upsample_latent1, x_latent1_features)
+
+        x0_features = checkpoint_wrapper(self, self.upsample0, x0_features)
+        x1_features = checkpoint_wrapper(self, self.upsample1, x1_features)
+        x2_features = checkpoint_wrapper(self, self.upsample2, x2_features)
+
+        x_lowres_features = checkpoint_wrapper(self, self.upsample_lowres, x_lowres_features)
+        x_lowres_features = checkpoint_wrapper(
+            self, self.fuse_lowres, torch.cat((x2_features, x_lowres_features), dim=1)
+        )
+
+        output = [
+            x_latent0_features,
+            x_latent1_features,
+            x0_features,
+            x1_features,
+            x_lowres_features,
+        ]
+
+        return output
+
+
+# It seems that torch.fx.wrap can only be applied to functions, not methods.
+# Hence, split and merge were converted into functions to be marked as atomic
+# operations for symbolic tracing.
+@torch.fx.wrap
+def split(image: torch.Tensor, overlap_ratio: float = 0.25, patch_size: int = 384) -> torch.Tensor:
+    """Split the input into small patches with sliding window."""
+    patch_stride = int(patch_size * (1 - overlap_ratio))
+
+    image_size = image.shape[-1]
+    steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1
+
+    x_patch_list = []
+    for j in range(steps):
+        j0 = j * patch_stride
+        j1 = j0 + patch_size
+
+        for i in range(steps):
+            i0 = i * patch_stride
+            i1 = i0 + patch_size
+            x_patch_list.append(image[..., j0:j1, i0:i1])
+
+    return torch.cat(x_patch_list, dim=0)
+
+
+# Decorator marking function as an atomic operator for symbolic tracing.
+@torch.fx.wrap
+def merge(image_patches: torch.Tensor, batch_size: int, padding: int = 3) -> torch.Tensor:
+    """Merge the patched input into a image with sliding window."""
+    steps = int(math.sqrt(image_patches.shape[0] // batch_size))
+
+    idx = 0
+
+    output_list = []
+    for j in range(steps):
+        output_row_list = []
+        for i in range(steps):
+            output = image_patches[batch_size * idx : batch_size * (idx + 1)]
+
+            if padding != 0:
+                if j != 0:
+                    output = output[..., padding:, :]
+                if i != 0:
+                    output = output[..., :, padding:]
+                if j != steps - 1:
+                    output = output[..., :-padding, :]
+                if i != steps - 1:
+                    output = output[..., :, :-padding]
+
+            output_row_list.append(output)
+            idx += 1
+
+        output_row = torch.cat(output_row_list, dim=-1)
+        output_list.append(output_row)
+    output = torch.cat(output_list, dim=-2)
+    return output

sharp/models/encoders/unet_encoder.py

@@ -0,0 +1,117 @@
+"""Contains backbone models for feature extraction from RGBD input.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+from torch import nn
+
+from sharp.models.blocks import (
+    NormLayerName,
+    norm_layer_2d,
+    residual_block_2d,
+)
+
+from .base_encoder import BaseEncoder
+
+
+class UNetEncoder(BaseEncoder):
+    """Encoder of UNet model."""
+
+    def __init__(
+        self,
+        dim_in: int,
+        width: List[int] | int,
+        steps: int = 6,
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups=8,
+        blocks_per_layer=2,
+    ) -> None:
+        """Initialize UNet Encoder.
+
+        Args:
+            dim_in: The number of input channels.
+            width: Width multiplicator of intermediate layers or the width list of all layers.
+            steps: The number of downsampling steps.
+            norm_type: Which kind of normalization layer to use.
+            norm_num_groups: How many groups to use for group norm (if relevant).
+            blocks_per_layer: How many residual blocks per layer to use.
+        """
+        super().__init__()
+
+        if blocks_per_layer < 1:
+            raise ValueError("blocks_per_layer must be greater or equal to one.")
+
+        self.dim_in = dim_in
+        self.width = width
+        self.num_steps = steps
+
+        self.convs_down = nn.ModuleList()
+
+        self.output_dims: list[int]
+        # If only one number is specified, we assume each layer will double the channel dimension.
+        if isinstance(width, int):
+            self.output_dims = [width << i for i in range(0, steps + 1)]
+        else:
+            if len(width) != (steps + 1):
+                raise ValueError("Length of width should match the steps for UNetEncoder.")
+            self.output_dims = width
+
+        self.conv_in = nn.Sequential(
+            nn.Conv2d(self.dim_in, self.output_dims[0], 3, stride=1, padding=1),
+            norm_layer_2d(self.output_dims[0], norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+        )
+
+        for i_step in range(steps):
+            input_width = self.output_dims[i_step]
+            current_width = self.output_dims[i_step + 1]
+            convs_down_i = nn.Sequential(
+                nn.AvgPool2d(2, stride=2),
+                residual_block_2d(
+                    input_width,
+                    current_width,
+                    norm_type=norm_type,
+                    norm_num_groups=norm_num_groups,
+                ),
+                *[
+                    residual_block_2d(
+                        current_width,
+                        current_width,
+                        norm_type=norm_type,
+                        norm_num_groups=norm_num_groups,
+                    )
+                    for _ in range(blocks_per_layer - 1)
+                ],
+            )
+            self.convs_down.append(convs_down_i)
+
+    def forward(self, input: torch.Tensor) -> list[torch.Tensor]:
+        """Apply UNet Encoder to image.
+
+        Args:
+            input: The input image.
+
+        Returns:
+            The output multi-level feature map from encoder.
+        """
+        features = []
+
+        feat_i = self.conv_in(input)
+        features.append(feat_i)
+
+        for conv_down in self.convs_down:
+            feat_i = conv_down(feat_i)
+            features.append(feat_i)
+
+        return features
+
+    @property
+    def out_width(self) -> int:
+        """Compute the output width for UNet decoder."""
+        return self.output_dims[-1]

sharp/models/encoders/vit_encoder.py

@@ -0,0 +1,111 @@
+"""Contains factory functions to build and load ViT.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import timm
+import torch
+
+from sharp.models.presets.vit import VIT_CONFIG_DICT, ViTConfig, ViTPreset
+
+LOGGER = logging.getLogger(__name__)
+
+
+class TimmViT(timm.models.VisionTransformer):
+    """Contains TIMM implementation for Vanilla ViT."""
+
+    def __init__(self, config: ViTConfig):
+        """Initialize ViT from TIMM implementation."""
+        # Handle mlp layers.
+        mlp_layer = timm.layers.GluMlp if config.mlp_mode == "glu" else timm.layers.Mlp
+
+        super().__init__(
+            in_chans=config.in_chans,
+            embed_dim=config.embed_dim,
+            depth=config.depth,
+            num_heads=config.num_heads,
+            init_values=config.init_values,
+            img_size=config.img_size,
+            patch_size=config.patch_size,
+            num_classes=config.num_classes,
+            mlp_ratio=config.mlp_ratio,
+            qkv_bias=config.qkv_bias,
+            global_pool=config.global_pool,
+            mlp_layer=mlp_layer,
+        )
+
+        # Required for extracting intermediate features.
+        self.dim_in = config.in_chans
+        self.intermediate_features_ids = config.intermediate_features_ids
+
+    def reshape_feature(self, embeddings: torch.Tensor):
+        """Discard class token and reshape 1D feature map to a 2D grid."""
+        batch_size, seq_len, channel = embeddings.shape
+
+        height, width = self.patch_embed.grid_size
+
+        # Remove class token.
+        if self.num_prefix_tokens:
+            embeddings = embeddings[:, self.num_prefix_tokens :, :]
+
+        # Shape: (batch, height, width, dim) -> (batch, dim, height, width)
+        embeddings = embeddings.reshape(batch_size, height, width, channel).permute(0, 3, 1, 2)
+        return embeddings
+
+    def forward(self, input_tensor: torch.Tensor) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
+        """Override forwarding with intermediate features.
+
+        Adapted from timm ViT.
+
+        Returns:
+            Output features and list of features from intermediate layers (patch encoder only).
+        """
+        intermediate_features = {}
+
+        x = self.patch_embed(input_tensor)
+        batch_size, seq_len, _ = x.shape
+
+        x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if self.intermediate_features_ids is not None and idx in self.intermediate_features_ids:
+                intermediate_features[idx] = x
+        x = self.norm(x)
+
+        x = self.reshape_feature(x)
+        return x, intermediate_features
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        if isinstance(self.patch_embed.img_size, tuple):
+            return self.patch_embed.img_size[0]
+        else:
+            return self.patch_embed.img_size
+
+
+def create_vit(
+    config: ViTConfig | None = None,
+    preset: ViTPreset | None = "dinov2l16_384",
+    intermediate_features_ids: list[int] | None = None,
+) -> TimmViT:
+    """Factory function for creating a ViT model."""
+    if config is not None:
+        LOGGER.info("Using user-defined config.")
+    else:
+        if preset is None:
+            raise ValueError("User-defined config and preset cannot be both None.")
+        LOGGER.info("Using preset ViT %s.", preset)
+        config = VIT_CONFIG_DICT[preset]
+
+    config.intermediate_features_ids = intermediate_features_ids
+    model = TimmViT(config)
+    LOGGER.debug(model)
+    return model

sharp/models/gaussian_decoder.py

@@ -0,0 +1,267 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+
+from sharp.models.blocks import (
+    FeatureFusionBlock2d,
+    NormLayerName,
+    residual_block_2d,
+)
+from sharp.models.decoders import BaseDecoder, MultiresConvDecoder
+from sharp.models.params import DPTImageEncoderType, GaussianDecoderParams
+
+
+def create_gaussian_decoder(
+    params: GaussianDecoderParams, dims_depth_features: list[int]
+) -> GaussianDensePredictionTransformer:
+    """Create gaussian_decoder model specified by gaussian_decoder_name."""
+    decoder = MultiresConvDecoder(
+        dims_depth_features,
+        params.dims_decoder,
+        grad_checkpointing=params.grad_checkpointing,
+        upsampling_mode=params.upsampling_mode,
+    )
+
+    return GaussianDensePredictionTransformer(
+        decoder=decoder,
+        dim_in=params.dim_in,
+        dim_out=params.dim_out,
+        stride_out=params.stride,
+        norm_type=params.norm_type,
+        norm_num_groups=params.norm_num_groups,
+        use_depth_input=params.use_depth_input,
+        grad_checkpointing=params.grad_checkpointing,
+        image_encoder_type=params.image_encoder_type,
+        image_encoder_params=params,
+    )
+
+
+def _create_project_upsample_block(
+    dim_in: int,
+    dim_out: int,
+    upsample_layers: int,
+    dim_intermediate: int | None = None,
+) -> nn.Module:
+    if dim_intermediate is None:
+        dim_intermediate = dim_out
+    # Projection.
+    blocks = [
+        nn.Conv2d(
+            in_channels=dim_in,
+            out_channels=dim_intermediate,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+    ]
+
+    # Upsampling.
+    blocks += [
+        nn.ConvTranspose2d(
+            in_channels=dim_intermediate if i == 0 else dim_out,
+            out_channels=dim_out,
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=False,
+        )
+        for i in range(upsample_layers)
+    ]
+
+    return nn.Sequential(*blocks)
+
+
+class ImageFeatures(NamedTuple):
+    """Image feature extracted from decoder."""
+
+    texture_features: torch.Tensor
+    geometry_features: torch.Tensor
+
+
+class SkipConvBackbone(nn.Module):
+    """A wrapper around a conv layer that behaves like a BaseBackbone."""
+
+    def __init__(self, dim_in: int, dim_out: int, kernel_size: int, stride_out: int):
+        """Initialize SkipConvBackbone."""
+        super().__init__()
+        self.stride_out = stride_out
+        if stride_out == 1 and kernel_size != 1:
+            raise ValueError("We only support kernel_size = 1 if stride_out is 1.")
+        padding: int = (kernel_size - 1) // 2
+        self.conv = nn.Conv2d(
+            dim_in, dim_out, kernel_size=kernel_size, stride=stride_out, padding=padding
+        )
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        encodings: list[torch.Tensor] | None = None,
+    ) -> ImageFeatures:
+        """Apply SkipConvBackbone to image."""
+        output = self.conv(input_features)
+        return ImageFeatures(
+            texture_features=output,
+            geometry_features=output,
+        )
+
+    @property
+    def stride(self) -> int:
+        """Effective downsampling stride."""
+        return self.stride_out
+
+
+class GaussianDensePredictionTransformer(nn.Module):
+    """Dense Prediction Transformer for Gaussian.
+
+    Reuse monodepth decoded features for processing.
+    """
+
+    norm_type: NormLayerName
+
+    def __init__(
+        self,
+        decoder: BaseDecoder,
+        dim_in: int,
+        dim_out: int,
+        stride_out: int,
+        image_encoder_params: GaussianDecoderParams,
+        image_encoder_type: DPTImageEncoderType = "skip_conv",
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups: int = 8,
+        use_depth_input: bool = True,
+        grad_checkpointing: bool = False,
+    ):
+        """Initialize Dense Prediction Transformer for Gaussian.
+
+        Args:
+            decoder: Decoder to decode features.
+            monodepth_decoder: Optional monodepth decoder to fuse monodepth decoded features.
+            dim_in: Input dimension.
+            dim_out: Final output dimension.
+            stride_out: Stride of output feature map.
+            image_encoder_params: The backbone parameters to configurate the image encoder.
+            image_encoder_type: Type of image encoder to use.
+            encoder: Encoder to generate features using monodepth model.
+            norm_type: Type of norm layers.
+            norm_num_groups: Num groups for norm layers.
+            use_depth_input: Whether to use depth input.
+            grad_checkpointing: Whether to use gradient checkpointing.
+        """
+        super().__init__()
+
+        self.decoder = decoder
+        self.dim_in = dim_in
+        self.dim_out = dim_out
+        self.stride_out = stride_out
+        self.norm_type = norm_type
+        self.norm_num_groups = norm_num_groups
+        self.use_depth_input = use_depth_input
+        self.grad_checkpointing = grad_checkpointing
+        self.image_encoder_type = image_encoder_type
+
+        # Adopt an image encoder to lift dimension to monodepth feature and
+        # resize to be the same resolution as the decoder output.
+        dim_in = self.dim_in if use_depth_input else self.dim_in - 1
+        image_encoder_params.dim_in = dim_in
+        image_encoder_params.dim_out = decoder.dim_out
+        self.image_encoder = self._create_image_encoder(image_encoder_params, stride_out)
+
+        self.fusion = FeatureFusionBlock2d(decoder.dim_out)
+
+        if stride_out == 1:
+            self.upsample = _create_project_upsample_block(
+                decoder.dim_out,
+                decoder.dim_out,
+                upsample_layers=1,
+            )
+        elif stride_out == 2:
+            self.upsample = nn.Identity()
+        else:
+            raise ValueError("We only support stride is 1 or 2 for DPT backbone.")
+
+        self.texture_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
+        self.geometry_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
+
+    def _create_head(self, dim_decoder: int, dim_out: int) -> nn.Module:
+        return nn.Sequential(
+            residual_block_2d(
+                dim_in=dim_decoder,
+                dim_out=dim_decoder,
+                dim_hidden=dim_decoder // 2,
+                norm_type=self.norm_type,
+                norm_num_groups=self.norm_num_groups,
+            ),
+            residual_block_2d(
+                dim_in=dim_decoder,
+                dim_hidden=dim_decoder // 2,
+                dim_out=dim_decoder,
+                norm_type=self.norm_type,
+                norm_num_groups=self.norm_num_groups,
+            ),
+            nn.ReLU(),
+            nn.Conv2d(dim_decoder, dim_out, kernel_size=1, stride=1),
+            nn.ReLU(),
+        )
+
+    def _create_image_encoder(
+        self, image_encoder_params: GaussianDecoderParams, stride_out: int
+    ) -> nn.Module:
+        """Create image encoder and return based on parameters."""
+        if self.image_encoder_type == "skip_conv":
+            # Use kernel_size = 1 only if stride_out is 1.
+            return SkipConvBackbone(
+                image_encoder_params.dim_in,
+                image_encoder_params.dim_out,
+                kernel_size=3 if stride_out != 1 else 1,
+                stride_out=stride_out,
+            )
+        elif self.image_encoder_type == "skip_conv_kernel2":
+            return SkipConvBackbone(
+                image_encoder_params.dim_in,
+                image_encoder_params.dim_out,
+                kernel_size=stride_out,
+                stride_out=stride_out,
+            )
+        else:
+            raise ValueError(f"Unsupported image encoder type: {self.image_encoder_type}")
+
+    def forward(self, input_features: torch.Tensor, encodings: list[torch.Tensor]) -> ImageFeatures:
+        """Run monodepth and fuse features with input image to predict Gaussians.
+
+        Args:
+            input_features: The input features to use.
+            encodings: Feature encodings (e.g. from monodepth network).
+        """
+        features = self.decoder(encodings).contiguous()
+        features = self.upsample(features)
+
+        if self.use_depth_input:
+            skip_features = self.image_encoder(input_features).texture_features
+        else:
+            skip_features = self.image_encoder(input_features[:, :3].contiguous())
+        features = self.fusion(features, skip_features)
+
+        texture_features = self.texture_head(features)
+        geometry_features = self.geometry_head(features)
+
+        return ImageFeatures(
+            texture_features=texture_features,  # type: ignore
+            geometry_features=geometry_features,  # type: ignore
+        )
+
+    @property
+    def stride(self) -> int:
+        """Internal stride of GaussianDensePredictionTransformer."""
+        return self.stride_out

sharp/models/heads.py

@@ -0,0 +1,53 @@
+"""Contains decoder head for direct prediction of delta values.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .gaussian_decoder import ImageFeatures
+
+
+class DirectPredictionHead(nn.Module):
+    """Decodes features into delta values using convolutions."""
+
+    def __init__(self, feature_dim: int, num_layers: int) -> None:
+        """Initialize DirectGaussianPredictor.
+
+        Args:
+            feature_dim: Number of input features.
+            num_layers: The number of layers of Gaussians to predict.
+        """
+        super().__init__()
+        self.num_layers = num_layers
+
+        # 14 is 3 means, 3 scales, 4 quaternions, 3 colors and 1 opacity
+        self.geometry_prediction_head = nn.Conv2d(feature_dim, 3 * num_layers, 1)
+        self.geometry_prediction_head.weight.data.zero_()
+        assert self.geometry_prediction_head.bias is not None
+        self.geometry_prediction_head.bias.data.zero_()
+
+        self.texture_prediction_head = nn.Conv2d(feature_dim, (14 - 3) * num_layers, 1)
+        self.texture_prediction_head.weight.data.zero_()
+        assert self.texture_prediction_head.bias is not None
+        self.texture_prediction_head.bias.data.zero_()
+
+    def forward(self, image_features: ImageFeatures) -> torch.Tensor:
+        """Predict deltas for 3D Gaussians.
+
+        Args:
+            image_features: Image features from decoder.
+
+        Returns:
+            The predicted deltas for Gaussian attributes.
+        """
+        delta_values_geometry = self.geometry_prediction_head(image_features.geometry_features)
+        delta_values_texture = self.texture_prediction_head(image_features.texture_features)
+        delta_values_geometry = delta_values_geometry.unflatten(1, (3, self.num_layers))
+        delta_values_texture = delta_values_texture.unflatten(1, (14 - 3, self.num_layers))
+        delta_values = torch.cat([delta_values_geometry, delta_values_texture], dim=1)
+        return delta_values

sharp/models/initializer.py

@@ -0,0 +1,297 @@
+"""Contains modules to initialize Gaussians from RGBD.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+from torch import nn
+
+from .params import ColorInitOption, DepthInitOption, InitializerParams
+
+
+def create_initializer(params: InitializerParams) -> nn.Module:
+    """Create inpainter."""
+    return MultiLayerInitializer(
+        num_layers=params.num_layers,
+        stride=params.stride,
+        base_depth=params.base_depth,
+        scale_factor=params.scale_factor,
+        disparity_factor=params.disparity_factor,
+        color_option=params.color_option,
+        first_layer_depth_option=params.first_layer_depth_option,
+        rest_layer_depth_option=params.rest_layer_depth_option,
+        normalize_depth=params.normalize_depth,
+        feature_input_stop_grad=params.feature_input_stop_grad,
+    )
+
+
+class GaussianBaseValues(NamedTuple):
+    """Base values for gaussian predictor.
+
+    We predict x and y in normalized device coordinates (NDC) where (-1, -1) is the top
+    left corner and (1, 1) the bottom right corner. The last component of
+    mean_vectors_ndc is inverse depth.
+    """
+
+    mean_x_ndc: torch.Tensor
+    mean_y_ndc: torch.Tensor
+    mean_inverse_z_ndc: torch.Tensor
+
+    scales: torch.Tensor
+    quaternions: torch.Tensor
+    colors: torch.Tensor
+    opacities: torch.Tensor
+
+
+class InitializerOutput(NamedTuple):
+    """Output of initializer."""
+
+    # Gaussian base values.
+    gaussian_base_values: GaussianBaseValues
+
+    # Feature input to the Gaussian predictor.
+    feature_input: torch.Tensor
+
+    # Global scale to unscale output.
+    global_scale: torch.Tensor | None = None
+
+
+class MultiLayerInitializer(nn.Module):
+    """Initialize Gaussians with multilayer representation.
+
+    The returned tensors have the shape
+
+        batch_size x dim x num_layers x height x width
+
+    where dim indicates the dimensionality of the property.
+    Some of the dimensions might be set to 1 for efficiency reasons.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        stride: int,
+        base_depth: float,
+        scale_factor: float,
+        disparity_factor: float,
+        color_option: ColorInitOption = "first_layer",
+        first_layer_depth_option: DepthInitOption = "surface_min",
+        rest_layer_depth_option: DepthInitOption = "surface_min",
+        normalize_depth: bool = True,
+        feature_input_stop_grad: bool = True,
+    ) -> None:
+        """Initialize MultilayerInitializer.
+
+        Args:
+            stride: The downsample rate of output feature map.
+            base_depth: The depth of the first layer (after the foreground
+                layer if use_depth=True).
+            scale_factor: Multiply scale of Gaussians by this factor.
+            disparity_factor: Factor to convert inverse depth to disparity.
+            num_layers: How many layers of Gaussians to predict.
+            color_option: Which color option to initialize the multi-layer gaussians.
+            first_layer_depth_option: Which depth option to initialize the first layer of gaussians.
+            rest_layer_depth_option: Which depth option to initialize the rest layers of gaussians.
+            normalize_depth: # Whether to normalize depth to [DepthTransformParam.depth_min,
+                DepthTransformParam.depth_max).
+            feature_input_stop_grad: Whether to not propagate gradients through feature inputs.
+        """
+        super().__init__()
+        self.num_layers = num_layers
+        self.stride = stride
+        self.base_depth = base_depth
+        self.scale_factor = scale_factor
+        self.disparity_factor = disparity_factor
+        self.color_option = color_option
+        self.first_layer_depth_option = first_layer_depth_option
+        self.rest_layer_depth_option = rest_layer_depth_option
+        self.normalize_depth = normalize_depth
+        self.feature_input_stop_grad = feature_input_stop_grad
+
+    def prepare_feature_input(self, image: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
+        """Prepare the feature input to the Guassian predictor."""
+        if self.feature_input_stop_grad:
+            image = image.detach()
+            depth = depth.detach()
+
+        normalized_disparity = self.disparity_factor / depth
+        features_in = torch.cat([image, normalized_disparity], dim=1)
+        features_in = 2.0 * features_in - 1.0
+        return features_in
+
+    def forward(self, image: torch.Tensor, depth: torch.Tensor) -> InitializerOutput:
+        """Construct Gaussian base values and prepare feature input.
+
+        Args:
+            image: The image to process.
+            depth: The corresponding depth map from the monodepth network.
+
+        Returns:
+            The base value for Gaussians.
+        """
+        image = image.contiguous()
+        depth = depth.contiguous()
+        device = depth.device
+        batch_size, _, image_height, image_width = depth.shape
+        base_height, base_width = (
+            image_height // self.stride,
+            image_width // self.stride,
+        )
+        # global_scale is the inverse of the depth_factor, which is used to rescale
+        # the depth such that it is numerically stable for training.
+        global_scale: torch.Tensor | None = None
+        if self.normalize_depth:
+            depth, depth_factor = _rescale_depth(depth)
+            global_scale = 1.0 / depth_factor
+
+        def _create_disparity_layers(num_layers: int = 1) -> torch.Tensor:
+            """Create multiple disparity layers."""
+            disparity = torch.linspace(1.0 / self.base_depth, 0.0, num_layers + 1, device=device)
+            return disparity[None, None, :-1, None, None].repeat(
+                batch_size, 1, 1, base_height, base_width
+            )
+
+        def _create_surface_layer(
+            depth: torch.Tensor,
+            depth_pooling_mode: str,
+        ) -> torch.Tensor:
+            """Create multiple surface layers."""
+            disparity = 1.0 / depth
+            if depth_pooling_mode == "min":
+                disparity = torch.max_pool2d(disparity, self.stride, self.stride)
+            elif depth_pooling_mode == "max":
+                disparity = -torch.max_pool2d(-disparity, self.stride, self.stride)
+            else:
+                raise ValueError(f"Invalid depth pooling mode {depth_pooling_mode}.")
+
+            return disparity[:, :, None, :, :]
+
+        # Input disparity dimensions:
+        #   (batch_size, num_channels in (1, 2), height, width)
+
+        # Output disparity dimensions:
+        #   (batch_size, num_channels=1, num_layers in (1, 2), height, width)
+        if self.first_layer_depth_option == "surface_min":
+            first_disparity = _create_surface_layer(depth[:, 0:1], "min")
+        elif self.first_layer_depth_option == "surface_max":
+            first_disparity = _create_surface_layer(depth[:, 0:1], "max")
+        elif self.first_layer_depth_option in ("base_depth", "linear_disparity"):
+            first_disparity = _create_disparity_layers()
+        else:
+            raise ValueError(f"Unknown depth init option: {self.first_layer_depth_option}.")
+
+        if self.num_layers == 1:
+            disparity = first_disparity
+        else:  # Fill in the rest layers.
+            following_depth = depth if depth.shape[1] == 1 else depth[:, 1:]
+            if self.rest_layer_depth_option == "surface_min":
+                following_disparity = _create_surface_layer(following_depth, "min")
+            elif self.rest_layer_depth_option == "surface_max":
+                following_disparity = _create_surface_layer(following_depth, "max")
+            elif self.rest_layer_depth_option == "base_depth":
+                following_disparity = torch.cat(
+                    [_create_disparity_layers() for i in range(self.num_layers - 1)],
+                    dim=2,
+                )
+            elif self.rest_layer_depth_option == "linear_disparity":
+                following_disparity = _create_disparity_layers(self.num_layers - 1)
+            else:
+                raise ValueError(f"Unknown depth init option: {self.rest_layer_depth_option}.")
+
+            disparity = torch.cat([first_disparity, following_disparity], dim=2)
+
+        # Prepare base values.
+        base_x_ndc, base_y_ndc = _create_base_xy(depth, self.stride, self.num_layers)
+        disparity_scale_factor = 2 * self.scale_factor * self.stride / float(image_width)
+        base_scales = _create_base_scale(disparity, disparity_scale_factor)
+
+        base_quaternions = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device)
+        base_quaternions = base_quaternions[None, :, None, None, None]
+
+        # Initializing the opacitiy this way ensures that the initial transmittance
+        # is approximately
+        #
+        #       1 / e ~= (1 - 1 / self.num_layers)**self.num_layers
+        #
+        # and hence independent of the number of layers.
+        #
+        base_opacities = torch.tensor([min(1.0 / self.num_layers, 0.5)], device=device)
+        base_colors = torch.empty(
+            batch_size, 3, self.num_layers, base_height, base_width, device=device
+        ).fill_(0.5)
+        # Dimensions: (batch_size, num_channels, num_layers, height, width)
+        if self.color_option == "none":
+            pass
+        elif self.color_option == "first_layer":
+            base_colors[:, :, 0] = torch.nn.functional.avg_pool2d(image, self.stride, self.stride)
+        elif self.color_option == "all_layers":
+            temp = torch.nn.functional.avg_pool2d(image, self.stride, self.stride)
+            base_colors = temp[:, :, None, :, :].repeat(1, 1, self.num_layers, 1, 1)
+        else:
+            raise ValueError(f"Unknown color init option: {self.color_option}.")
+
+        features_in = self.prepare_feature_input(image, depth)
+        base_gaussians = GaussianBaseValues(
+            mean_x_ndc=base_x_ndc,
+            mean_y_ndc=base_y_ndc,
+            mean_inverse_z_ndc=disparity,
+            scales=base_scales,
+            quaternions=base_quaternions,
+            colors=base_colors,
+            opacities=base_opacities,
+        )
+
+        return InitializerOutput(
+            gaussian_base_values=base_gaussians,
+            feature_input=features_in,
+            global_scale=global_scale,
+        )
+
+
+def _create_base_xy(
+    depth: torch.Tensor, stride: int, num_layers: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Create base x and y coordinates for the gaussians in NDC space."""
+    device = depth.device
+    batch_size, _, image_height, image_width = depth.shape
+    xx = torch.arange(0.5 * stride, image_width, stride, device=device)
+    yy = torch.arange(0.5 * stride, image_height, stride, device=device)
+    xx = 2 * xx / image_width - 1.0
+    yy = 2 * yy / image_height - 1.0
+
+    xx, yy = torch.meshgrid(xx, yy, indexing="xy")
+    base_x_ndc = xx[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
+    base_y_ndc = yy[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
+
+    return base_x_ndc, base_y_ndc
+
+
+def _create_base_scale(disparity: torch.Tensor, disparity_scale_factor: float) -> torch.Tensor:
+    """Create base scale for the gaussians."""
+    inverse_disparity = torch.ones_like(disparity) / disparity
+    base_scales = inverse_disparity * disparity_scale_factor
+    return base_scales
+
+
+def _rescale_depth(
+    depth: torch.Tensor, depth_min: float = 1.0, depth_max: float = 1e2
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Rescale a depth image tensor.
+
+    Args:
+        depth: The depth tensor to transform.
+        depth_min: The min depth to scale depth to.
+        depth_max: The max clamp depth after scaling.
+
+    Returns:
+        The rescaled depth and rescale factor.
+    """
+    current_depth_min = depth.flatten(depth.ndim - 3).min(dim=-1).values
+    depth_factor = depth_min / (current_depth_min + 1e-6)
+    depth = (depth * depth_factor[..., None, None, None]).clamp(max=depth_max)
+    return depth, depth_factor

sharp/models/monodepth.py

@@ -0,0 +1,268 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import NamedTuple, Tuple
+
+import torch
+import torch.nn as nn
+
+from sharp.models import normalizers
+from sharp.models.decoders import MultiresConvDecoder, create_monodepth_decoder
+from sharp.models.encoders import (
+    SlidingPyramidNetwork,
+    create_monodepth_encoder,
+)
+from sharp.utils import module_surgery
+
+from .params import MonodepthAdaptorParams, MonodepthParams
+
+DimsDecoder = Tuple[int, int, int, int, int]
+
+
+class MonodepthDensePredictionTransformer(nn.Module):
+    """Dense Prediction Transformer for monodepth.
+
+    Attach the disparity prediction head for monodepth prediction.
+    """
+
+    def __init__(
+        self,
+        encoder: SlidingPyramidNetwork,
+        decoder: MultiresConvDecoder,
+        last_dims: tuple[int, int],
+    ):
+        """Initialize Dense Prediction Transformer.
+
+        Args:
+            encoder: The SlidingPyramidTransformer backbone.
+            decoder: The MultiresConvDecoder decoder.
+            last_dims: The dimension for the last convolution layers.
+        """
+        super().__init__()
+
+        self.normalizer = normalizers.AffineRangeNormalizer(
+            input_range=(0, 1), output_range=(-1, 1)
+        )
+        self.encoder = encoder
+        self.decoder = decoder
+
+        dim_decoder = decoder.dim_out
+        self.head = nn.Sequential(
+            nn.Conv2d(dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1),
+            nn.ConvTranspose2d(
+                in_channels=dim_decoder // 2,
+                out_channels=dim_decoder // 2,
+                kernel_size=2,
+                stride=2,
+                padding=0,
+                bias=True,
+            ),
+            nn.Conv2d(
+                dim_decoder // 2,
+                last_dims[0],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            ),
+            nn.ReLU(True),
+            nn.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0),
+            nn.ReLU(),
+        )
+
+        # Set the final convoultion layer's bias to be 0.
+        self.head[4].bias.data.fill_(0)
+
+        self.grad_checkpointing = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+        self.encoder.set_grad_checkpointing(self.grad_checkpointing)
+        self.decoder.set_grad_checkpointing(self.grad_checkpointing)
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """Decode by projection and fusion of multi-resolution encodings."""
+        encodings = self.encoder(self.normalizer(image))
+        num_encoder_features = len(self.encoder.dims_encoder)
+        features = self.decoder(encodings[:num_encoder_features])
+        disparity = self.head(features)
+        return disparity
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        return self.encoder.internal_resolution()
+
+
+def create_monodepth_dpt(
+    params: MonodepthParams | None = None,
+) -> MonodepthDensePredictionTransformer:
+    """Creates DepthDensePredictionTransformer model.
+
+    Args:
+        params: Parameters of monodepth network.
+
+    Returns:
+        The configured monodepth DPT.
+    """
+    if params is None:
+        params = MonodepthParams()
+    encoder: SlidingPyramidNetwork = create_monodepth_encoder(
+        params.patch_encoder_preset,
+        params.image_encoder_preset,
+        use_patch_overlap=params.use_patch_overlap,
+        last_encoder=params.dims_decoder[0],
+    )
+
+    decoder: MultiresConvDecoder = create_monodepth_decoder(
+        params.patch_encoder_preset, params.dims_decoder
+    )
+
+    monodepth_model = MonodepthDensePredictionTransformer(
+        encoder=encoder, decoder=decoder, last_dims=(32, 1)
+    )
+
+    # By default, we don't train the monodepth model.
+    # However, we allow to selectively unfreeze parts of the network.
+    monodepth_model.requires_grad_(False)
+
+    monodepth_model.encoder.set_requires_grad_(
+        patch_encoder=params.unfreeze_patch_encoder,
+        image_encoder=params.unfreeze_image_encoder,
+    )
+    monodepth_model.decoder.requires_grad_(params.unfreeze_decoder)
+    monodepth_model.head.requires_grad_(params.unfreeze_head)
+
+    if not params.unfreeze_norm_layers:
+        module_surgery.freeze_norm_layer(monodepth_model)
+
+    monodepth_model.set_grad_checkpointing(params.grad_checkpointing)
+
+    return monodepth_model
+
+
+class MonodepthOutput(NamedTuple):
+    """Output of the monodepth model."""
+
+    # Disparity output from the monodepth model.
+    disparity: torch.Tensor
+    # Multi-level features from monodepth encoder.
+    encoder_features: list[torch.Tensor]
+    # Single-level feature from monodepth decoder.
+    decoder_features: torch.Tensor
+    # List of monodepth features to be used in gaussian predictor.
+    output_features: list[torch.Tensor]
+    # List of intermediate encoder features to be used in distillation.
+    intermediate_features: list[torch.Tensor] = []
+
+
+class MonodepthWithEncodingAdaptor(nn.Module):
+    """Monodepth model with feature maps."""
+
+    def __init__(
+        self,
+        monodepth_predictor: MonodepthDensePredictionTransformer,
+        return_encoder_features: bool,
+        return_decoder_features: bool,
+        num_monodepth_layers: int,
+        sorting_monodepth: bool,
+    ):
+        """Initialize MonodepthWithEncodingAdaptor.
+
+        Args:
+            monodepth_predictor: The monodepth model.
+            return_encoder_features: Whether to return encoder features from monodepth model.
+            return_decoder_features: Whether to return decoder features from monodepth model.
+            num_monodepth_layers: How many layers the monodepth model predicts.
+            sorting_monodepth: Whether to sort the monodepth output (for two layer monodepth).
+        """
+        super().__init__()
+        self.monodepth_predictor = monodepth_predictor
+        self.return_encoder_features = return_encoder_features
+        self.return_decoder_features = return_decoder_features
+        self.num_monodepth_layers = num_monodepth_layers
+        self.sorting_monodepth = sorting_monodepth
+
+    def forward(self, image: torch.Tensor) -> MonodepthOutput:
+        """Process image and return disparity and feature maps."""
+        inputs = self.monodepth_predictor.normalizer(image)
+        encoder_output = self.monodepth_predictor.encoder(inputs)
+
+        num_encoder_features = len(self.monodepth_predictor.encoder.dims_encoder)
+
+        # NOTE: whether intermediate features are empty have already been decided
+        # in monodepth_predictor during create_monodepth_dpt.
+        encoder_features = encoder_output[:num_encoder_features]
+        intermediate_features = encoder_output[num_encoder_features:]
+        decoder_features = self.monodepth_predictor.decoder(encoder_features)
+        disparity = self.monodepth_predictor.head(decoder_features)
+
+        # We cannot use disparity.shape[1], otherwise the tracer will fail.
+        if self.num_monodepth_layers == 2 and self.sorting_monodepth:
+            first_layer_disparity = disparity.max(dim=1, keepdims=True).values
+            second_layer_disparity = disparity.min(dim=1, keepdims=True).values
+            disparity = torch.cat([first_layer_disparity, second_layer_disparity], dim=1)
+
+        output_features = []
+        if self.return_encoder_features:
+            output_features.extend(encoder_features)
+
+        if self.return_decoder_features:
+            output_features.append(decoder_features)
+
+        return MonodepthOutput(
+            disparity=disparity,
+            encoder_features=encoder_features,
+            decoder_features=decoder_features,
+            output_features=output_features,
+            intermediate_features=intermediate_features,
+        )
+
+    def get_feature_dims(self) -> list[int]:
+        """Return dimensions of output feature maps."""
+        dims = []
+        if self.return_encoder_features:
+            dims.extend(self.monodepth_predictor.encoder.dims_encoder)
+
+        if self.return_decoder_features:
+            dims.append(self.monodepth_predictor.decoder.dim_out)
+
+        return dims
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        return self.monodepth_predictor.internal_resolution()
+
+    def replicate_head(self, num_repeat: int):
+        """Replicate the last convolution layer (head[4] in DPT) for multi layer depth."""
+        conv_last = copy.deepcopy(self.monodepth_predictor.head[4])
+        self.monodepth_predictor.head[4].out_channels = num_repeat
+        self.monodepth_predictor.head[4].weight = nn.Parameter(
+            conv_last.weight.repeat(num_repeat, 1, 1, 1)
+        )
+        self.monodepth_predictor.head[4].bias = nn.Parameter(conv_last.bias.repeat(num_repeat))
+
+
+def create_monodepth_adaptor(
+    monodepth_predictor: MonodepthDensePredictionTransformer,
+    params: MonodepthAdaptorParams,
+    num_monodepth_layers: int,
+    sorting_monodepth: bool,
+) -> MonodepthWithEncodingAdaptor:
+    """Create an adaptor that returns both disparity and features."""
+    adaptor = MonodepthWithEncodingAdaptor(
+        monodepth_predictor=monodepth_predictor,
+        return_encoder_features=params.encoder_features,
+        return_decoder_features=params.decoder_features,
+        num_monodepth_layers=num_monodepth_layers,
+        sorting_monodepth=sorting_monodepth,
+    )
+    return adaptor

sharp/models/normalizers.py

@@ -0,0 +1,80 @@
+"""Contains an implementation of image normalizers for perceptual loss.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+import torch
+from torch import nn
+
+
+class MeanStdNormalizer(nn.Module):
+    """Normalizing image input by mean and std."""
+
+    mean: torch.Tensor
+    std_inv: torch.Tensor
+
+    def __init__(
+        self,
+        mean: Union[Sequence[float], torch.Tensor],
+        std: Union[Sequence[float], torch.Tensor],
+    ):
+        """Initialize MeanStdNormalizer."""
+        super(MeanStdNormalizer, self).__init__()
+        if not isinstance(mean, torch.Tensor):
+            mean = torch.as_tensor(mean).view(-1, 1, 1)
+        if not isinstance(std, torch.Tensor):
+            std = torch.as_tensor(std).view(-1, 1, 1)
+        self.register_buffer("mean", mean)
+        # We use inverse std to use a multiplication which is better supported by the hardware
+        self.register_buffer("std_inv", 1.0 / std)
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """Apply mean and std normalization over input image."""
+        return (image - self.mean) * self.std_inv
+
+
+class AffineRangeNormalizer(nn.Module):
+    """Perform linear mapping to map input_range to output_range.
+
+    Output_range defaults to (0, 1).
+    """
+
+    def __init__(
+        self,
+        input_range: tuple[float, float],
+        output_range: tuple[float, float] = (0, 1),
+    ):
+        """Initialize AffineRangeNormalizer."""
+        super().__init__()
+        input_min, input_max = input_range
+        output_min, output_max = output_range
+        if input_max <= input_min:
+            raise ValueError(f"Invalid input_range: {input_range}")
+        if output_max <= output_min:
+            raise ValueError(f"Invalid output_range: {output_range}")
+
+        self.scale = (output_max - output_min) / (input_max - input_min)
+        self.bias = output_min - input_min * self.scale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply affine range normalization over input image."""
+        if self.scale != 1.0:
+            x = x * self.scale
+
+        if self.bias != 0.0:
+            x = x + self.bias
+
+        return x
+
+
+class MobileNetNormalizer(AffineRangeNormalizer):
+    """Image normalization in mobilenet."""
+
+    def __init__(self, input_range: tuple[float, float] = (0, 1)):
+        """Initialize MobileNetNormalizer."""
+        super().__init__(input_range=input_range, output_range=(-1, 1))

sharp/models/params.py

@@ -0,0 +1,203 @@
+"""Contains params for backbone.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import dataclasses
+from typing import Literal
+
+import sharp.utils.math as math_utils
+from sharp.models.blocks import NormLayerName, UpsamplingMode
+from sharp.models.presets import ViTPreset
+from sharp.utils.color_space import ColorSpace
+
+DimsDecoder = tuple[int, int, int, int, int]
+DPTImageEncoderType = Literal["skip_conv", "skip_conv_kernel2"]
+
+ColorInitOption = Literal[
+    "none",  # Initialize as gray.
+    "first_layer",  # Initialize the first layer with input image, other layers with gray.
+    "all_layers",  # Initialize all layers with input image.
+]
+DepthInitOption = Literal[
+    # Initialize the layer of gaussian on surface using min pooling of input depth.
+    "surface_min",
+    # Initialize the layer of gaussian on surface using max pooling of input depth
+    "surface_max",
+    # Initialize the layer of gaussian on plane using base_depth depth.
+    "base_depth",
+    # Initialize the layer of gaussian on plane based on base_depth and index of layer.
+    "linear_disparity",
+]
+
+
+@dataclasses.dataclass
+class AlignmentParams:
+    """Parameters for depth alignment."""
+
+    kernel_size: int = 16
+    stride: int = 1
+    frozen: bool = False
+
+    # The following parameters are only used for LearnedAlignment.
+    # Number of steps in the UNet for LearnedAlignment.
+    steps: int = 4
+    # Activation type for LearnedAlignment.
+    activation_type: math_utils.ActivationType = "exp"
+    # Whether to use depth decoder features for LearnedAlignment.
+    depth_decoder_features: bool = False
+    # Base width of the UNet for LearnedAlignment.
+    base_width: int = 16
+
+
+@dataclasses.dataclass
+class DeltaFactor:
+    """Factors to multiply deltas with before activation.
+
+    These factors effectively selectively reduce the learning rate.
+    """
+
+    xy: float = 0.001
+    z: float = 0.001
+    color: float = 0.1  # We recommend 0.1 for linearRGB and 1.0 for sRGB.
+    opacity: float = 1.0
+    scale: float = 1.0
+    quaternion: float = 1.0
+
+
+@dataclasses.dataclass
+class InitializerParams:
+    """Parameters for initializer."""
+
+    # Common parameters.
+    # Multiply scale of Gaussians by this factor.
+    scale_factor: float = 1.0
+    # Factor to convert inverse depth to disparity.
+    disparity_factor: float = 1.0
+    # Stride of the initializer.
+    stride: int = 2
+
+    # Parameters that only affect MultiLayerInitializer.
+    # How many layers of Gaussians to predict (only available for MultiLayerInitializer).
+    num_layers: int = 2
+    # Which option to use for depth initialization.
+    first_layer_depth_option: DepthInitOption = "surface_min"
+    rest_layer_depth_option: DepthInitOption = "surface_min"
+    # Which option to use for color initialization.
+    color_option: ColorInitOption = "all_layers"
+    # Which depth value to use for depth layers.
+    base_depth: float = 10.0
+    # Deactivate gradient for feature inputs.
+    feature_input_stop_grad: bool = False
+    # Whether to normalize depth to [DepthTransformParam.depth_min,
+    # DepthTransformParam.depth_max).
+    normalize_depth: bool = True
+
+    # Output only the inpainted layer. In this case, num_layers = 1.
+    output_inpainted_layer_only: bool = False
+    # Whether to set the uninpainted region to zero opacities.
+    set_uninpainted_opacity_to_zero: bool = False
+    # Whether to concatenate the inpainting mask to the feature input.
+    concat_inpainting_mask: bool = False
+
+
+@dataclasses.dataclass
+class MonodepthParams:
+    """Parameters for monodepth network."""
+
+    patch_encoder_preset: ViTPreset = "dinov2l16_384"
+    image_encoder_preset: ViTPreset = "dinov2l16_384"
+
+    checkpoint_uri: str | None = None
+    unfreeze_patch_encoder: bool = False
+    unfreeze_image_encoder: bool = False
+    unfreeze_decoder: bool = False
+    unfreeze_head: bool = False
+    unfreeze_norm_layers: bool = False
+    grad_checkpointing: bool = False
+    use_patch_overlap: bool = True
+    dims_decoder: DimsDecoder = (256, 256, 256, 256, 256)
+
+
+@dataclasses.dataclass
+class MonodepthAdaptorParams:
+    """Parameters for monodepth network feature adaptor."""
+
+    encoder_features: bool = True
+    decoder_features: bool = False
+
+
+@dataclasses.dataclass
+class GaussianDecoderParams:
+    """Parameters for backbone with default values."""
+
+    dim_in: int = 5
+    dim_out: int = 32
+    # Which normalization to use in backbone.
+    norm_type: NormLayerName = "group_norm"
+    # How many groups to use for group normalization.
+    norm_num_groups: int = 8
+    # Stride of backbone.
+    stride: int = 2
+
+    patch_encoder_preset: ViTPreset = "dinov2l16_384"
+    image_encoder_preset: ViTPreset = "dinov2l16_384"
+
+    # Dimensionality of feature maps for DPT decoder.
+    dims_decoder: DimsDecoder = (128, 128, 128, 128, 128)
+
+    # Whether to use depth as input.
+    use_depth_input: bool = True
+
+    # Whether to enable gradient checkpointing for the backbone
+    grad_checkpointing: bool = False
+
+    # What mode to use for upsampling in decoder.
+    upsampling_mode: UpsamplingMode = "transposed_conv"
+
+    # The type of image encoder.
+    image_encoder_type: DPTImageEncoderType = "skip_conv_kernel2"
+
+
+@dataclasses.dataclass
+class PredictorParams:
+    """Parameters for predictors with default values."""
+
+    # Parameters for submodules.
+    initializer: InitializerParams = dataclasses.field(default_factory=InitializerParams)
+    monodepth: MonodepthParams = dataclasses.field(default_factory=MonodepthParams)
+    monodepth_adaptor: MonodepthAdaptorParams = dataclasses.field(
+        default_factory=MonodepthAdaptorParams
+    )
+    gaussian_decoder: GaussianDecoderParams = dataclasses.field(
+        default_factory=GaussianDecoderParams
+    )
+    # How to align depth map (only relevant for RGBGaussianPredictor).
+    depth_alignment: AlignmentParams = dataclasses.field(default_factory=AlignmentParams)
+
+    # Selectively reduce learning rate for different properties.
+    delta_factor: DeltaFactor = dataclasses.field(default_factory=DeltaFactor)
+    # The maximum scale of Gaussians relative to initial scale.
+    max_scale: float = 10.0
+    # The minimum scale of Gaussians relative to initial scale.
+    min_scale: float = 0.0
+    # Which normalization to use in prediction head.
+    norm_type: NormLayerName = "group_norm"
+    # How many groups to use for group normalization.
+    norm_num_groups: int = 8
+    # Whether to use predicted mean to sample triplane features.
+    use_predicted_mean: bool = False
+    # Which activation function to use for colors / opacities.
+    color_activation_type: math_utils.ActivationType = "sigmoid"
+    opacity_activation_type: math_utils.ActivationType = "sigmoid"
+    # Colorspace of the renderer ("linearRGB" or "sRGB").
+    color_space: ColorSpace = "linearRGB"
+    # A small value to avoid ill-conditioned splats
+    low_pass_filter_eps: float = 1e-2
+    # How many layer of depth does monodepth model predict.
+    num_monodepth_layers: int = 2
+    # Whether to sort the monodepth output (for two layer monodepth).
+    sorting_monodepth: bool = False
+    # Whether to account the z offsets for estimating base scale.
+    base_scale_on_predicted_mean: bool = True

sharp/models/predictor.py

@@ -0,0 +1,201 @@
+"""Contains definition of RGB-only gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from torch import nn
+
+from sharp.models.monodepth import MonodepthWithEncodingAdaptor
+from sharp.utils.gaussians import Gaussians3D
+
+from .composer import GaussianComposer
+
+LOGGER = logging.getLogger(__name__)
+
+
+class DepthAlignment(nn.Module):
+    """Depth alignment in a dedicated nn.Module.
+
+    Wrap scale_map_estimator to perform the conditional logic in a separated torch
+    module outside the forward of RGBGaussianPredictor. This module can be then
+    excluded during symbolic tracing.
+    """
+
+    def __init__(self, scale_map_estimator: nn.Module | None):
+        """Initialize DepthAlignmentWrapper.
+
+        Args:
+            scale_map_estimator: Module to align monodepth to ground truth depth.
+        """
+        super().__init__()
+        self.scale_map_estimator = scale_map_estimator
+
+    def forward(
+        self,
+        monodepth: torch.Tensor,
+        depth: torch.Tensor,
+        depth_decoder_features: torch.Tensor | None = None,
+    ):
+        """Optionally align monodepth to ground truth with a local scale map.
+
+        Args:
+            monodepth: The monodepth model with intermediate features to use.
+            depth: Ground truth depth to align predicted depth to.
+            depth_decoder_features: The (optional) monodepth decoder features.
+        """
+        if depth is not None and self.scale_map_estimator is not None:
+            depth_alignment_map = self.scale_map_estimator(
+                monodepth[:, 0:1], depth, depth_decoder_features
+            )
+            monodepth = depth_alignment_map * monodepth
+        else:
+            # Some losses rely on the presence of an alignment map.
+            # We ensure that they can be computed by creating a fake alignment map.
+            depth_alignment_map = torch.ones_like(monodepth)
+        return monodepth, depth_alignment_map
+
+
+class RGBGaussianPredictor(nn.Module):
+    """Predicts 3D Gaussians from images."""
+
+    feature_model: nn.Module
+
+    def __init__(
+        self,
+        init_model: nn.Module,
+        monodepth_model: MonodepthWithEncodingAdaptor,
+        feature_model: nn.Module,
+        prediction_head: nn.Module,
+        gaussian_composer: GaussianComposer,
+        scale_map_estimator: nn.Module | None,
+    ) -> None:
+        """Initialize RGBGaussianPredictor.
+
+        Args:
+            init_model: A model mapping image and depth to base values.
+            monodepth_model: The monodepth model with intermediate features to use.
+            feature_model: The image2image model to predict Gaussians from.
+            prediction_head: Head to decode image features.
+            gaussian_composer: Module to compose final prediction from deltas and
+                base values.
+            scale_map_estimator: Module to align monodepth to ground truth depth.
+
+        Note:
+        ----
+            when monodepth_model is trainable, using local depth alignment can
+            result in the monodepth model losing its ability to predict shapes. It is
+            hence recommend to deactivate the corresponding flag.
+        """
+        super().__init__()
+        self.init_model = init_model
+        self.feature_model = feature_model
+        self.monodepth_model = monodepth_model
+        self.prediction_head = prediction_head
+        self.gaussian_composer = gaussian_composer
+        self.depth_alignment = DepthAlignment(scale_map_estimator)
+
+    def forward(
+        self,
+        image: torch.Tensor,
+        disparity_factor: torch.Tensor,
+        depth: torch.Tensor | None = None,
+    ) -> Gaussians3D:
+        """Predict 3D Gaussians.
+
+        Args:
+            image: The image to process.
+            disparity_factor: Factor to convert depth to disparities.
+            depth: Ground truth depth to align predicted depth to.
+
+        Returns:
+            The predicted 3D Gaussians.
+
+        Note:
+        ----
+        During training, it is recommended to feed an additional ground truth depth
+        map to the network to align the predicted depth to. During inference, it is
+        recommended to use depth_gt=None and use monodepth_disparity output from the
+        model instead to compute depth.
+        """
+        # Estimate depth and align to ground truth (if available).
+        monodepth_output = self.monodepth_model(image)
+        monodepth_disparity = monodepth_output.disparity
+
+        disparity_factor = disparity_factor[:, None, None, None]
+        monodepth = disparity_factor / monodepth_disparity.clamp(min=1e-4, max=1e4)
+
+        # In the model we apply additional alignment to provided ground truth depth
+        # as well as additional normalization.
+        #
+        # The overall graph looks as follows:
+        #
+        #     monodepth        depth    # Both monodepth and depth are metric here.
+        #         |              |
+        #         +------+-------+
+        #                |
+        #        +-------+--------+     # Optionally align monodepth to ground truth
+        #        |depth_alignement|     # with a local scale map.
+        #        +-------+--------+
+        #                |
+        #                v
+        #       monodepth (aligned)     # Monodepth is now aligned to ground truth.
+        #                |
+        #          +-----+----+         # Normalize depth and compute base gaussians.
+        #          |init_model|         # in these normalized coordinates.
+        #          +-----+----+
+        #                |
+        #                v
+        #   +------ init_output         # Init_output consists of features, base
+        #   |            |              # gaussians and a global scale.
+        #   |     +------+-----+
+        #   |     |main network|        # Compute delta values to base gaussians.
+        #   |     +------+-----+
+        #   |            |
+        #   |            V
+        #   |        delta_values       # The delta values are computed with normalized depth.
+        #   |            |
+        #   |    +-------+---------+
+        #   +--> |gaussian_composer|    # Add delta to base values and unscale gaussians.
+        #        +-------+---------+
+        #                |
+        #                v
+        #            gaussians          # The final Gaussians are metric again.
+        #
+
+        # The logic to decide whether to align monodepth to the ground truth is wrapped
+        # in a submodule 'DepthAlignement' to facilitate the symbolic tracing of the
+        # predictor. This way, the depth alignment submodule containing the conditional
+        # logic can be excluded during the tracing and the graph of the predictors is
+        # static.
+        monodepth, _ = self.depth_alignment(
+            monodepth,
+            depth,
+            monodepth_output.decoder_features,
+        )
+
+        init_output = self.init_model(image, monodepth)
+        image_features = self.feature_model(
+            init_output.feature_input, encodings=monodepth_output.output_features
+        )
+        delta_values = self.prediction_head(image_features)
+        gaussians = self.gaussian_composer(
+            delta=delta_values,
+            base_values=init_output.gaussian_base_values,
+            global_scale=init_output.global_scale,
+        )
+        return gaussians
+
+    def internal_resolution(self) -> int:
+        """Internal resolution."""
+        return self.monodepth_model.internal_resolution()
+
+    @property
+    def output_resolution(self) -> int:
+        """Output resolution of Gaussians."""
+        return self.internal_resolution() // 2

sharp/models/presets/__init__.py

@@ -0,0 +1,23 @@
+"""Contains presets for pretrained neural networks.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from .monodepth import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    MONODEPTH_HOOK_IDS_MAP,
+)
+from .vit import (
+    VIT_CONFIG_DICT,
+    ViTConfig,
+    ViTPreset,
+)
+
+__all__ = [
+    "ViTConfig",
+    "ViTPreset",
+    "VIT_CONFIG_DICT",
+    "MONODEPTH_ENCODER_DIMS_MAP",
+    "MONODEPTH_HOOK_IDS_MAP",
+]

sharp/models/presets/monodepth.py

@@ -0,0 +1,21 @@
+"""Contains preset for monodepth modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from .vit import ViTPreset
+
+# Map the decoder configuration with the number of output channels
+# for each tensor from the decoder output.
+MONODEPTH_ENCODER_DIMS_MAP: dict[ViTPreset, list[int]] = {
+    # For publication
+    "dinov2l16_384": [256, 512, 1024, 1024],
+}
+
+MONODEPTH_HOOK_IDS_MAP: dict[ViTPreset, list[int]] = {
+    # For publication
+    "dinov2l16_384": [5, 11, 17, 23],
+}

sharp/models/presets/vit.py

@@ -0,0 +1,58 @@
+"""Contains preset for ViT modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Literal
+
+ViTPreset = Literal["dinov2l16_384",]
+
+MLPMode = Literal["vanilla", "glu"]
+
+
+@dataclasses.dataclass
+class ViTConfig:
+    """Configuration for ViT."""
+
+    in_chans: int
+    embed_dim: int
+    depth: int
+    num_heads: int
+    init_values: float
+
+    img_size: int = 384
+    patch_size: int = 16
+
+    num_classes: int = 21841
+    mlp_ratio: float = 4.0
+    drop_rate: float = 0.0
+    attn_drop_rate: float = 0.0
+    drop_path_rate: float = 0.0
+    qkv_bias: bool = True
+    global_pool: str = "avg"
+
+    # Properties for timm_vit.
+    mlp_mode: MLPMode = "vanilla"
+
+    # Properties for SPN.
+    intermediate_features_ids: list[int] | None = None
+
+    def asdict(self):
+        """Convenience method to convert the class to a dict."""
+        return dataclasses.asdict(self)
+
+
+VIT_CONFIG_DICT: dict[ViTPreset, ViTConfig] = {
+    "dinov2l16_384": ViTConfig(
+        in_chans=3,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        init_values=1e-5,
+        global_pool="",
+    ),
+}