"""Core PXDepth architecture and stable public model API. The module connects the Global Context Encoder to the Pixel-Space Depth Predictor and defines raw forward computation. Checkpoint translation and metric-scale inference live in focused helper modules, while their familiar ``from_pretrained`` and ``infer`` entry points remain methods on this class. """ from pathlib import Path from typing import Any, Dict, IO, Optional, Union import torch import torch.nn as nn import torch.nn.functional as F from .Global_Context_Encoder import GlobalContextEncoder from .Pixel_Space_Depth_Predictor import PixelSpaceDepthPredictor from .checkpoint import load_pretrained from .inference import infer as infer_model from .precision import full_precision, inference_dtype, reduced_precision from ..registry import ENCODERS, MODELS, PREDICTORS @MODELS.register() class PXDepth(nn.Module): """Complete PXDepth monocular depth model. A Global Context Encoder extracts semantic patch features and a Pixel-Space Depth Predictor estimates full-resolution normalized log-depth together with a finite-depth probability. ``forward`` exposes raw network outputs, while ``infer`` aligns them to a GT or MoGe-2 reference for metric-scale visualization and point-cloud reconstruction. """ def __init__( self, encoder: Union[nn.Module, Dict[str, Any]], predictor: Union[nn.Module, Dict[str, Any]], remap_output: str = "linear", mask_threshold: float = 0.5, ) -> None: """Construct the encoder and CM-PiT pixel predictor. Args: encoder: Encoder module or registry config. predictor: Pixel predictor module or registry config. Its context patch size and channel width default to the encoder contract. remap_output: Output remapping applied to normalized log-depth. The released model uses ``'linear'``. mask_threshold: Probability threshold used by :meth:`infer`. Returns: ``None``. Model modules and ImageNet normalization buffers are registered on the instance. """ super().__init__() if remap_output not in {"linear", "elu"}: raise ValueError(f"Unsupported remap_output: {remap_output}") self.remap_output = remap_output self.mask_threshold = float(mask_threshold) if isinstance(encoder, nn.Module): self.encoder = encoder else: encoder_config = dict(encoder) encoder_config.setdefault("type", "GlobalContextEncoder") self.encoder = ENCODERS.build(encoder_config) if not hasattr(self.encoder, "patch_size"): raise TypeError("The encoder must expose an integer patch_size attribute.") self.patch_size = self.encoder.patch_size self.p_enc = self.patch_size dim_ctx = getattr(self.encoder, "dim_out", None) if isinstance(predictor, nn.Module): self.predictor = predictor else: predictor_config = dict(predictor) predictor_config.setdefault("type", "PixelSpaceDepthPredictor") predictor_config.setdefault("in_channels", 3) predictor_config.setdefault("ctx_patch_size", self.patch_size) if dim_ctx is not None: predictor_config.setdefault("dim_ctx", int(dim_ctx)) self.predictor = PREDICTORS.build(predictor_config) self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) self._reference_model: Optional[nn.Module] = None @property def device(self) -> torch.device: """Return the device hosting PXDepth learnable parameters. No inputs are required. The value is inferred from the model's first parameter and is used when moving inference inputs or reference models. Returns: ``torch.device`` for the current model placement. """ return next(self.parameters()).device @property def dtype(self) -> torch.dtype: """Return the storage dtype of PXDepth learnable parameters. No inputs are required. This reports parameter storage, which is independent from local autocast contexts used inside attention. Returns: ``torch.dtype`` of the model's first parameter. """ return next(self.parameters()).dtype @classmethod def from_pretrained( cls, path_or_repo: Union[str, Path, IO[bytes]], model_kwargs: Optional[Dict[str, Any]] = None, strict: bool = True, **hf_kwargs: Any, ) -> "PXDepth": """Create a model from a local or Hugging Face ``model.pt`` checkpoint. Args: path_or_repo: Local checkpoint path, binary file object, or Hugging Face model repository identifier. model_kwargs: Optional constructor overrides applied after reading ``model_config`` from the checkpoint. strict: Forwarded to ``load_state_dict``. Published checkpoints should use the default exact matching. **hf_kwargs: Additional keyword arguments forwarded to ``huggingface_hub.hf_hub_download`` for remote repositories. Returns: Initialized :class:`PXDepth` instance on CPU. """ return load_pretrained( cls, path_or_repo, model_kwargs=model_kwargs, strict=strict, **hf_kwargs, ) def init_weights(self) -> None: """Initialize the Global Context Encoder from official DINOv2 weights. Predictor parameters retain the initialization created by their own module constructors. Returns: ``None``. Encoder parameters are updated in place. """ self.encoder.init_weights() def enable_gradient_checkpointing(self) -> None: """Enable activation checkpointing in both encoder and predictor. This reduces saved activation memory during backward at the cost of recomputing transformer blocks. Returns: ``None``. Child module runtime behavior is updated in place. """ self.encoder.enable_gradient_checkpointing() self.predictor.enable_gradient_checkpointing() def enable_pytorch_native_sdpa(self) -> None: """Enable the optimized SDPA attention path in the DINOv2 backbone. Decoder CM-PiT attention already uses PyTorch SDPA directly and is not modified by this method. Returns: ``None``. Encoder attention modules are wrapped in place. """ self.encoder.enable_pytorch_native_sdpa() def _remap(self, depth: torch.Tensor) -> torch.Tensor: """Apply the configured output activation to raw depth predictions. Args: depth: Raw normalized-depth tensor with arbitrary batch/spatial shape, normally ``[B, H, W]``. Returns: Tensor with the same shape. The released ``linear`` setting returns the input unchanged. """ return F.elu(depth) if self.remap_output == "elu" else depth def forward( self, image: torch.Tensor, use_fp16: bool = False, use_fp32: bool = False, ) -> Dict[str, torch.Tensor]: """Run the network without metric-scale alignment. Args: image: RGB tensor ``[B, 3, H, W]`` with values in ``[0, 1]``. ``H`` and ``W`` must be divisible by the encoder patch size. use_fp16: Run attention-heavy encoder and predictor regions under FP16 autocast. use_fp32: Disable reduced-precision autocast. It is mutually exclusive with ``use_fp16``. Returns: Dictionary with normalized log-depth ``depth`` and finite-depth probability ``mask``, both FP32 tensors ``[B, H, W]``. """ height, width = image.shape[-2:] if height % self.patch_size or width % self.patch_size: raise ValueError( f"Input resolution ({height}, {width}) must be divisible by patch size {self.patch_size}" ) dtype = inference_dtype(use_fp16=use_fp16, use_fp32=use_fp32) with full_precision(image.device): image_norm = (image.float() - self.image_mean.float()) / self.image_std.float() with reduced_precision(image.device, dtype): context = self.encoder(image, height // self.patch_size, width // self.patch_size) context = context.flatten(2).permute(0, 2, 1).contiguous() depth, mask = self.predictor(image_norm, context, autocast_dtype=dtype) with full_precision(image.device): depth = self._remap(depth.float().squeeze(1)) mask = mask.float().squeeze(1).sigmoid() return {"depth": depth, "mask": mask} def infer( self, image: torch.Tensor, gt_depth: Optional[torch.Tensor] = None, intrinsics: Optional[torch.Tensor] = None, fov_x: Optional[Union[float, torch.Tensor]] = None, ref_image: Optional[torch.Tensor] = None, apply_mask: bool = True, use_fp16: bool = True, use_fp32: bool = False, ) -> Dict[str, torch.Tensor]: """Recover metric-scale depth, validity, intrinsics, and 3D points. Raw normalized log-depth is affine-aligned in log space to ``gt_depth`` when supplied, otherwise to a lazily loaded MoGe-2 reference. Alignment parameters are estimated on a 64x64 nearest-resized valid subset. The aligned depth is exponentiated and back-projected with normalized camera intrinsics. Args: image: RGB tensor ``[3,H,W]`` or batch ``[B,3,H,W]`` in ``[0,1]``. gt_depth: Optional reference depth ``[H,W]`` or ``[B,H,W]``. Finite positive pixels define log-space alignment. intrinsics: Optional normalized camera matrices ``[3,3]`` or ``[B,3,3]`` corresponding to ``gt_depth``. fov_x: Optional horizontal field of view in degrees, scalar or tensor ``[B]``, used when intrinsics are unavailable. ref_image: Optional original-resolution RGB tensor used only by the reference model; PXDepth still consumes ``image``. apply_mask: Replace invalid predicted depth/points with infinity. use_fp16: Use FP16 for attention-heavy model regions. use_fp32: Force those regions to FP32 and override the BF16 default. Returns: Dictionary containing aligned ``depth`` ``[B,H,W]``, boolean ``mask`` ``[B,H,W]``, point map ``points`` ``[B,H,W,3]``, normalized ``intrinsics`` ``[B,3,3]``, and horizontal ``fov_x`` ``[B]``. For an unbatched input, the leading batch dimension is removed. """ return infer_model( self, image, gt_depth=gt_depth, intrinsics=intrinsics, fov_x=fov_x, ref_image=ref_image, apply_mask=apply_mask, use_fp16=use_fp16, use_fp32=use_fp32, )