Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import Dict | |
| import torch | |
| from torch import nn | |
| from .depth_pro import DepthProConfig, create_model_and_transforms | |
| def _inference_precision(device: torch.device) -> torch.dtype: | |
| return torch.float16 if device.type == "cuda" else torch.float32 | |
| def _as_depth_map(depth: torch.Tensor, batch_size: int) -> torch.Tensor: | |
| if depth.ndim == 2: | |
| depth = depth.unsqueeze(0) | |
| if depth.ndim != 3: | |
| raise ValueError(f"Expected depth with 2 or 3 dims, got shape {tuple(depth.shape)}") | |
| if depth.shape[0] != batch_size: | |
| raise ValueError( | |
| f"Depth batch mismatch: expected {batch_size}, got {depth.shape[0]}" | |
| ) | |
| return depth.unsqueeze(1) | |
| def _as_normalized_intrinsics_batch( | |
| intrinsics: torch.Tensor | None, | |
| batch_size: int, | |
| device: torch.device, | |
| ) -> torch.Tensor | None: | |
| if intrinsics is None: | |
| return None | |
| intrinsics = intrinsics.to(device=device, dtype=torch.float32) | |
| if intrinsics.ndim == 2: | |
| intrinsics = intrinsics.unsqueeze(0) | |
| if intrinsics.ndim != 3 or intrinsics.shape[-2:] != (3, 3): | |
| raise ValueError( | |
| f"Expected normalized intrinsics with shape [B,3,3] or [3,3], got {tuple(intrinsics.shape)}" | |
| ) | |
| if intrinsics.shape[0] == 1 and batch_size > 1: | |
| intrinsics = intrinsics.expand(batch_size, -1, -1) | |
| if intrinsics.shape[0] != batch_size: | |
| raise ValueError( | |
| f"Intrinsics batch mismatch: expected {batch_size}, got {intrinsics.shape[0]}" | |
| ) | |
| return intrinsics | |
| class DepthPro(nn.Module): | |
| """InfiniSplat wrapper around the vendored upstream DepthPro package.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.device_hint = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.precision = _inference_precision(self.device_hint) | |
| depthpro_cfg = DepthProConfig( | |
| patch_encoder_preset="dinov2l16_384", | |
| image_encoder_preset="dinov2l16_384", | |
| decoder_features=256, | |
| checkpoint_uri=None, | |
| fov_encoder_preset="dinov2l16_384", | |
| use_fov_head=True, | |
| ) | |
| self.model, self.transform = create_model_and_transforms( | |
| config=depthpro_cfg, | |
| device=self.device_hint, | |
| precision=self.precision, | |
| ) | |
| self.model.eval() | |
| def _prepare_input(self, image: torch.Tensor) -> torch.Tensor: | |
| if image.ndim != 4: | |
| raise ValueError(f"Expected image tensor [B,3,H,W], got shape {tuple(image.shape)}") | |
| self._ensure_runtime_dtype() | |
| model_device = next(self.model.parameters()).device | |
| image = image.to(device=model_device, dtype=torch.float32) | |
| image = image.clamp(0.0, 1.0) | |
| image = 2.0 * image - 1.0 | |
| return image.to(dtype=self.precision) | |
| def _ensure_runtime_dtype(self) -> None: | |
| param = next(self.model.parameters()) | |
| desired_dtype = _inference_precision(param.device) | |
| if param.dtype != desired_dtype: | |
| self.model.to(dtype=desired_dtype) | |
| self.precision = desired_dtype | |
| def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: | |
| """Run DepthPro with the standardized InfiniSplat depth-only interface. | |
| Args: | |
| batch: Dictionary containing: | |
| - image: Tensor with shape [B, 3, H, W]. | |
| - intrinsics: Optional normalized intrinsics with shape [B, 3, 3] | |
| or [3, 3]. When provided, fx is converted to pixel-space focal | |
| length and passed to DepthPro. If missing, DepthPro falls back | |
| to its internal FoV / focal estimation path. | |
| Returns: | |
| Dense metric depth with shape [B, 1, H, W]. | |
| """ | |
| image = batch["image"] | |
| prepared = self._prepare_input(image) | |
| batch_size = prepared.shape[0] | |
| focal_px: torch.Tensor | None = None | |
| intrinsics = _as_normalized_intrinsics_batch( | |
| batch.get("intrinsics"), | |
| batch_size, | |
| prepared.device, | |
| ) | |
| if intrinsics is not None: | |
| focal_px = intrinsics[:, 0, 0] * float(image.shape[-1]) | |
| with torch.no_grad(): | |
| prediction = self.model.infer(prepared, f_px=focal_px) | |
| return _as_depth_map(prediction["depth"], batch_size).to(dtype=torch.float32) | |