"""Depth Anything 3 (ByteDance-Seed/depth-anything-3 @ 3d835ec1) as a transformers depth-estimation model. ``DepthAnything3ForDepthEstimation.model`` is the package's ``DepthAnything3Net`` (or ``NestedDepthAnything3Net``: an any-view model aligned to a metric model), built from ``config.da3_config`` -- the checkpoint's own object graph -- the way ``depth_anything_3.cfg.create_object`` builds it, so the ``model.*`` tensors of the depth-anything/DA3* checkpoints load by name. ``forward`` runs the network's ``forward`` and returns its outputs as a ``DepthAnything3Output`` (``depth`` renamed ``predicted_depth``; depth_conf, extrinsics, intrinsics, sky, aux, gaussians with ``infer_gs=True``, ...). Autocast is the caller's: the package API runs it under bf16 (fp16 without bf16 support). The vendored code is in ``da3_*.py``. """ import sys from dataclasses import dataclass from typing import Any import torch from transformers import PreTrainedModel from transformers.utils import ModelOutput from .configuration_depth_anything_3 import DepthAnything3Config from .da3_backbone import DinoV2 from .da3_heads import DPT, CameraDec, CameraEnc, DualDPT from .da3_net import DepthAnything3Net, NestedDepthAnything3Net # depth_anything_3.cfg.create_object, for the classes these checkpoints name (all built with args: as_params) _CLASSES = { "depth_anything_3.model.da3.DepthAnything3Net": DepthAnything3Net, "depth_anything_3.model.da3.NestedDepthAnything3Net": NestedDepthAnything3Net, "depth_anything_3.model.dinov2.dinov2.DinoV2": DinoV2, "depth_anything_3.model.dpt.DPT": DPT, "depth_anything_3.model.dualdpt.DualDPT": DualDPT, "depth_anything_3.model.cam_enc.CameraEnc": CameraEnc, "depth_anything_3.model.cam_dec.CameraDec": CameraDec, } def build_object(spec): """Nested ``{"__object__": {path, name, args}, **arguments}`` -> module, children first (cfg.create_object).""" if not (isinstance(spec, dict) and "__object__" in spec): return spec obj = spec["__object__"] key = f"{obj['path']}.{obj['name']}" if key not in _CLASSES: raise NotImplementedError(f"{key} is not part of this package") if obj.get("args", "as_config") != "as_params": raise NotImplementedError(f"{key}: only args=as_params configs are supported") return _CLASSES[key](**{k: build_object(v) for k, v in spec.items() if k != "__object__"}) # NestedDepthAnything3Net.__init__ calls cfg.create_object on its two branches sys.modules[DepthAnything3Net.__module__].create_object = build_object @dataclass class DepthAnything3Output(ModelOutput): """The network's output dict, as fields (``depth`` -> ``predicted_depth``, the transformers depth-estimation name). predicted_depth (B, N, H, W) / depth_conf (B, N, H, W); extrinsics (B, N, 3, 4) world-to-camera / intrinsics (B, N, 3, 3) for any-view models; sky (B, N, H, W) for monocular models; ray / ray_conf when a DualDPT model has no camera decoder; aux: requested backbone features (``feat_layer_``); gaussians: 3D Gaussians (``infer_gs=True``); is_metric / scale_factor: the nested model's metric alignment. """ predicted_depth: torch.Tensor | None = None depth_conf: torch.Tensor | None = None extrinsics: torch.Tensor | None = None intrinsics: torch.Tensor | None = None sky: torch.Tensor | None = None ray: torch.Tensor | None = None ray_conf: torch.Tensor | None = None aux: Any = None gaussians: Any = None is_metric: Any = None scale_factor: Any = None class DepthAnything3PreTrainedModel(PreTrainedModel): config_class = DepthAnything3Config config: DepthAnything3Config base_model_prefix = "model" main_input_name = "pixel_values" def _init_weights(self, module): super()._init_weights(module) # GaussianAdapter.sh_mask is a non-persistent buffer: not in checkpoints, left uninitialised by meta-device # loading -> rebuild it as GaussianAdapter.__init__ does if getattr(module, "sh_mask", None) is not None and hasattr(module, "sh_degree"): mask = torch.ones((module.d_sh,), dtype=torch.float32) for degree in range(1, module.sh_degree + 1): mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree with torch.no_grad(): module.sh_mask.copy_(mask) class DepthAnything3ForDepthEstimation(DepthAnything3PreTrainedModel): def __init__(self, config: DepthAnything3Config): super().__init__(config) self.model = build_object(config.da3_config) # DualDPT builds its aux heads with one shared LayerNorm: declare those names as tied (target -> first name), # so save_pretrained stores the tensor once and from_pretrained keeps the sharing first, tied = {}, {} for name, parameter in self.named_parameters(remove_duplicate=False): if id(parameter) in first: tied[name] = first[id(parameter)] else: first[id(parameter)] = name self._tied_weights_keys = tied or None self.post_init() def forward( self, pixel_values: torch.Tensor, extrinsics: torch.Tensor | None = None, intrinsics: torch.Tensor | None = None, export_feat_layers: list[int] | None = None, infer_gs: bool = False, use_ray_pose: bool = False, ref_view_strategy: str = "saddle_balanced", return_dict: bool | None = None, **kwargs, ) -> tuple | DepthAnything3Output: """pixel_values: (B, N, 3, H, W), ImageNet-normalised (``DepthAnything3Processor``); extrinsics (B, N, 4, 4) / intrinsics (B, N, 3, 3) optional (pose conditioning, any-view models); export_feat_layers: backbone layers returned in ``aux``; infer_gs: also predict 3D Gaussians (checkpoints with a GS head).""" return_dict = return_dict if return_dict is not None else self.config.return_dict output = self.model(pixel_values, extrinsics, intrinsics, export_feat_layers or [], infer_gs, use_ray_pose, ref_view_strategy) fields = {("predicted_depth" if key == "depth" else key): value for key, value in output.items()} unknown = set(fields) - set(DepthAnything3Output.__dataclass_fields__) if unknown: raise ValueError(f"the network returned keys this output class does not hold: {sorted(unknown)}") result = DepthAnything3Output(**fields) return result if return_dict else result.to_tuple() __all__ = ["DepthAnything3ForDepthEstimation", "DepthAnything3Output", "DepthAnything3PreTrainedModel", "build_object"]