Duplicate from google/tipsv2-l14-dpt

Browse files

Co-authored-by: Gabriele Berton <gberton@users.noreply.huggingface.co>

Files changed (8) hide show

.gitattributes +35 -0
README.md +80 -0
config.json +397 -0
configuration_dpt.py +44 -0
dpt_head.py +702 -0
model.safetensors +3 -0
modeling_dpt.py +115 -0
preprocessor_config.json +13 -0

.gitattributes ADDED Viewed

	@@ -0,0 +1,35 @@

+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md ADDED Viewed

	@@ -0,0 +1,80 @@

+---
+license: apache-2.0
+tags:
+- vision
+- depth-estimation
+- surface-normals
+- semantic-segmentation
+- dense-prediction
+- arxiv:2604.12012
+library_name: transformers
+pipeline_tag: depth-estimation
+---
+# TIPSv2 — L/14 DPT Heads
+DPT (Dense Prediction Transformer) heads for depth estimation, surface normal prediction, and semantic segmentation on top of the frozen [TIPSv2 L/14](https://huggingface.co/google/tipsv2-l14) backbone. The backbone is loaded automatically. The depth and normals heads are trained on the NYU Depth V2 dataset and segmentation is trained on the ADE20K dataset (150 classes).
+| Variant | Vision params | Text params | Embed dim | DPT Heads |
+|---------|--------------|-------------|-----------|-----------|
+| [B/14](https://huggingface.co/google/tipsv2-b14) | 86M | 110M | 768 | [B/14-dpt](https://huggingface.co/google/tipsv2-b14-dpt) |
+| [L/14](https://huggingface.co/google/tipsv2-l14) | 303M | 184M | 1024 | [L/14-dpt](https://huggingface.co/google/tipsv2-l14-dpt) |
+| [SO400m/14](https://huggingface.co/google/tipsv2-so400m14) | 412M | 448M | 1152 | [SO400m/14-dpt](https://huggingface.co/google/tipsv2-so400m14-dpt) |
+| [g/14](https://huggingface.co/google/tipsv2-g14) | 1.1B | 389M | 1536 | [g/14-dpt](https://huggingface.co/google/tipsv2-g14-dpt) |
+## Usage
+```bash
+pip install transformers torch torchvision sentencepiece
+```
+```python
+from transformers import AutoModel
+from torchvision import transforms
+from PIL import Image
+import requests
+model = AutoModel.from_pretrained("google/tipsv2-l14-dpt", trust_remote_code=True)
+model.eval().cuda()
+url = "https://huggingface.co/spaces/google/tipsv2/resolve/main/examples/depth/ade20k_00014.png"
+image = Image.open(requests.get(url, stream=True).raw)
+transform = transforms.Compose([transforms.Resize((448, 448)), transforms.ToTensor()])
+pixel_values = transform(image).unsqueeze(0).cuda()
+# All tasks at once
+outputs = model(pixel_values)
+print(outputs.depth.shape)         # (1, 1, 448, 448) — depth map
+print(outputs.normals.shape)       # (1, 3, 448, 448) — surface normals
+print(outputs.segmentation.shape)  # (1, 150, 448, 448) — segmentation logits
+# Or individual tasks (only runs the requested head)
+depth = model.predict_depth(pixel_values)
+normals = model.predict_normals(pixel_values)
+seg = model.predict_segmentation(pixel_values)
+print(seg.argmax(dim=1).shape)     # (1, 448, 448) — per-pixel class prediction
+```
+## Model details
+- **Backbone**: [TIPSv2 L/14](https://huggingface.co/google/tipsv2-l14) (loaded automatically)
+- **Heads**: ~102M total params (depth + normals + segmentation)
+- **Depth & normals**: NYU Depth V2
+- **Segmentation**: ADE20K, 150 classes
+- **Input**: images in `[0, 1]` range, any resolution (multiples of 14 recommended)
+## License
+Apache 2.0
+## Citation
+```bibtex
+@inproceedings{cao2026tipsv2,
+  title     = {{TIPSv2: Advancing Vision-Language Pretraining with Enhanced Patch-Text Alignment}},
+  author    = {Cao, Bingyi and Chen, Koert and Maninis, Kevis-Kokitsi and Chen, Kaifeng and Karpur, Arjun and Xia, Ye and Dua, Sahil and Dabral, Tanmaya and Han, Guangxing and Han, Bohyung and Ainslie, Joshua and Bewley, Alex and Jacob, Mithun and Wagner, Rene and Ramos, Washington and Choromanski, Krzysztof and Seyedhosseini, Mojtaba and Zhou, Howard and Araujo, Andre},
+  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year      = {2026},
+  url       = {https://arxiv.org/abs/2604.12012}
+}
+```

config.json ADDED Viewed

	@@ -0,0 +1,397 @@

+{
+  "architectures": [
+    "TIPSv2DPTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dpt.TIPSv2DPTConfig",
+    "AutoModel": "modeling_dpt.TIPSv2DPTModel"
+  },
+  "backbone_config": {
+    "apply_layernorm": true,
+    "attention_probs_dropout_prob": 0.0,
+    "drop_path_rate": 0.0,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1024,
+    "image_size": 448,
+    "initializer_range": 0.02,
+    "interpolate_antialias": true,
+    "interpolate_offset": 0.0,
+    "layer_norm_eps": 1e-06,
+    "layerscale_value": 1.0,
+    "mlp_ratio": 4,
+    "model_type": "tipsv2_vision_model",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "num_register_tokens": 1,
+    "out_features": [
+      "stage6",
+      "stage12",
+      "stage18",
+      "stage24"
+    ],
+    "out_indices": [
+      6,
+      12,
+      18,
+      24
+    ],
+    "patch_size": 14,
+    "qkv_bias": true,
+    "reshape_hidden_states": false,
+    "stage_names": [
+      "stem",
+      "stage1",
+      "stage2",
+      "stage3",
+      "stage4",
+      "stage5",
+      "stage6",
+      "stage7",
+      "stage8",
+      "stage9",
+      "stage10",
+      "stage11",
+      "stage12",
+      "stage13",
+      "stage14",
+      "stage15",
+      "stage16",
+      "stage17",
+      "stage18",
+      "stage19",
+      "stage20",
+      "stage21",
+      "stage22",
+      "stage23",
+      "stage24"
+    ],
+    "use_swiglu_ffn": false
+  },
+  "depth_decoder_activation": "relu",
+  "fusion_hidden_size": 256,
+  "id2label": {
+    "0": "wall",
+    "1": "building",
+    "10": "cabinet",
+    "100": "poster",
+    "101": "stage",
+    "102": "van",
+    "103": "ship",
+    "104": "fountain",
+    "105": "conveyer belt",
+    "106": "canopy",
+    "107": "washer",
+    "108": "plaything",
+    "109": "swimming pool",
+    "11": "sidewalk",
+    "110": "stool",
+    "111": "barrel",
+    "112": "basket",
+    "113": "waterfall",
+    "114": "tent",
+    "115": "bag",
+    "116": "minibike",
+    "117": "cradle",
+    "118": "oven",
+    "119": "ball",
+    "12": "person",
+    "120": "food",
+    "121": "step",
+    "122": "tank",
+    "123": "trade name",
+    "124": "microwave",
+    "125": "pot",
+    "126": "animal",
+    "127": "bicycle",
+    "128": "lake",
+    "129": "dishwasher",
+    "13": "earth",
+    "130": "screen",
+    "131": "blanket",
+    "132": "sculpture",
+    "133": "hood",
+    "134": "sconce",
+    "135": "vase",
+    "136": "traffic light",
+    "137": "tray",
+    "138": "ashcan",
+    "139": "fan",
+    "14": "door",
+    "140": "pier",
+    "141": "crt screen",
+    "142": "plate",
+    "143": "monitor",
+    "144": "bulletin board",
+    "145": "shower",
+    "146": "radiator",
+    "147": "glass",
+    "148": "clock",
+    "149": "flag",
+    "15": "table",
+    "16": "mountain",
+    "17": "plant",
+    "18": "curtain",
+    "19": "chair",
+    "2": "sky",
+    "20": "car",
+    "21": "water",
+    "22": "painting",
+    "23": "sofa",
+    "24": "shelf",
+    "25": "house",
+    "26": "sea",
+    "27": "mirror",
+    "28": "rug",
+    "29": "field",
+    "3": "floor",
+    "30": "armchair",
+    "31": "seat",
+    "32": "fence",
+    "33": "desk",
+    "34": "rock",
+    "35": "wardrobe",
+    "36": "lamp",
+    "37": "bathtub",
+    "38": "railing",
+    "39": "cushion",
+    "4": "tree",
+    "40": "base",
+    "41": "box",
+    "42": "column",
+    "43": "signboard",
+    "44": "chest of drawers",
+    "45": "counter",
+    "46": "sand",
+    "47": "sink",
+    "48": "skyscraper",
+    "49": "fireplace",
+    "5": "ceiling",
+    "50": "refrigerator",
+    "51": "grandstand",
+    "52": "path",
+    "53": "stairs",
+    "54": "runway",
+    "55": "case",
+    "56": "pool table",
+    "57": "pillow",
+    "58": "screen door",
+    "59": "stairway",
+    "6": "road",
+    "60": "river",
+    "61": "bridge",
+    "62": "bookcase",
+    "63": "blind",
+    "64": "coffee table",
+    "65": "toilet",
+    "66": "flower",
+    "67": "book",
+    "68": "hill",
+    "69": "bench",
+    "7": "bed ",
+    "70": "countertop",
+    "71": "stove",
+    "72": "palm",
+    "73": "kitchen island",
+    "74": "computer",
+    "75": "swivel chair",
+    "76": "boat",
+    "77": "bar",
+    "78": "arcade machine",
+    "79": "hovel",
+    "8": "windowpane",
+    "80": "bus",
+    "81": "towel",
+    "82": "light",
+    "83": "truck",
+    "84": "tower",
+    "85": "chandelier",
+    "86": "awning",
+    "87": "streetlight",
+    "88": "booth",
+    "89": "television receiver",
+    "9": "grass",
+    "90": "airplane",
+    "91": "dirt track",
+    "92": "apparel",
+    "93": "pole",
+    "94": "land",
+    "95": "bannister",
+    "96": "escalator",
+    "97": "ottoman",
+    "98": "bottle",
+    "99": "buffet"
+  },
+  "label2id": {
+    "airplane": 90,
+    "animal": 126,
+    "apparel": 92,
+    "arcade machine": 78,
+    "armchair": 30,
+    "ashcan": 138,
+    "awning": 86,
+    "bag": 115,
+    "ball": 119,
+    "bannister": 95,
+    "bar": 77,
+    "barrel": 111,
+    "base": 40,
+    "basket": 112,
+    "bathtub": 37,
+    "bed ": 7,
+    "bench": 69,
+    "bicycle": 127,
+    "blanket": 131,
+    "blind": 63,
+    "boat": 76,
+    "book": 67,
+    "bookcase": 62,
+    "booth": 88,
+    "bottle": 98,
+    "box": 41,
+    "bridge": 61,
+    "buffet": 99,
+    "building": 1,
+    "bulletin board": 144,
+    "bus": 80,
+    "cabinet": 10,
+    "canopy": 106,
+    "car": 20,
+    "case": 55,
+    "ceiling": 5,
+    "chair": 19,
+    "chandelier": 85,
+    "chest of drawers": 44,
+    "clock": 148,
+    "coffee table": 64,
+    "column": 42,
+    "computer": 74,
+    "conveyer belt": 105,
+    "counter": 45,
+    "countertop": 70,
+    "cradle": 117,
+    "crt screen": 141,
+    "curtain": 18,
+    "cushion": 39,
+    "desk": 33,
+    "dirt track": 91,
+    "dishwasher": 129,
+    "door": 14,
+    "earth": 13,
+    "escalator": 96,
+    "fan": 139,
+    "fence": 32,
+    "field": 29,
+    "fireplace": 49,
+    "flag": 149,
+    "floor": 3,
+    "flower": 66,
+    "food": 120,
+    "fountain": 104,
+    "glass": 147,
+    "grandstand": 51,
+    "grass": 9,
+    "hill": 68,
+    "hood": 133,
+    "house": 25,
+    "hovel": 79,
+    "kitchen island": 73,
+    "lake": 128,
+    "lamp": 36,
+    "land": 94,
+    "light": 82,
+    "microwave": 124,
+    "minibike": 116,
+    "mirror": 27,
+    "monitor": 143,
+    "mountain": 16,
+    "ottoman": 97,
+    "oven": 118,
+    "painting": 22,
+    "palm": 72,
+    "path": 52,
+    "person": 12,
+    "pier": 140,
+    "pillow": 57,
+    "plant": 17,
+    "plate": 142,
+    "plaything": 108,
+    "pole": 93,
+    "pool table": 56,
+    "poster": 100,
+    "pot": 125,
+    "radiator": 146,
+    "railing": 38,
+    "refrigerator": 50,
+    "river": 60,
+    "road": 6,
+    "rock": 34,
+    "rug": 28,
+    "runway": 54,
+    "sand": 46,
+    "sconce": 134,
+    "screen": 130,
+    "screen door": 58,
+    "sculpture": 132,
+    "sea": 26,
+    "seat": 31,
+    "shelf": 24,
+    "ship": 103,
+    "shower": 145,
+    "sidewalk": 11,
+    "signboard": 43,
+    "sink": 47,
+    "sky": 2,
+    "skyscraper": 48,
+    "sofa": 23,
+    "stage": 101,
+    "stairs": 53,
+    "stairway": 59,
+    "step": 121,
+    "stool": 110,
+    "stove": 71,
+    "streetlight": 87,
+    "swimming pool": 109,
+    "swivel chair": 75,
+    "table": 15,
+    "tank": 122,
+    "television receiver": 89,
+    "tent": 114,
+    "toilet": 65,
+    "towel": 81,
+    "tower": 84,
+    "trade name": 123,
+    "traffic light": 136,
+    "tray": 137,
+    "tree": 4,
+    "truck": 83,
+    "van": 102,
+    "vase": 135,
+    "wall": 0,
+    "wardrobe": 35,
+    "washer": 107,
+    "water": 21,
+    "waterfall": 113,
+    "windowpane": 8
+  },
+  "max_depth": 10.0,
+  "min_depth": 0.001,
+  "model_type": "tipsv2_dpt",
+  "neck_hidden_sizes": [
+    128,
+    256,
+    512,
+    1024
+  ],
+  "num_depth_bins": 256,
+  "readout_act": "gelu_pytorch_tanh",
+  "reassemble_factors": [
+    4,
+    2,
+    1,
+    0.5
+  ],
+  "semantic_loss_ignore_index": 255,
+  "transformers_version": "5.10.0.dev0"
+}

configuration_dpt.py ADDED Viewed

	@@ -0,0 +1,44 @@

+"""TIPSv2 DPT head configuration."""
+from transformers import PretrainedConfig
+_BACKBONE_REPO_BY_HIDDEN_SIZE = {
+    768: "google/tipsv2-b14",
+    1024: "google/tipsv2-l14",
+    1152: "google/tipsv2-so400m14",
+    1536: "google/tipsv2-g14",
+}
+class TIPSv2DPTConfig(PretrainedConfig):
+    """Configuration for TIPSv2 DPT dense prediction heads."""
+    model_type = "tipsv2_dpt"
+    def __init__(
+        self,
+        backbone_config=None,
+        neck_hidden_sizes=(96, 192, 384, 768),
+        fusion_hidden_size=256,
+        reassemble_factors=(4, 2, 1, 0.5),
+        readout_type="project",
+        num_depth_bins=256,
+        min_depth=1e-3,
+        max_depth=10.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        backbone_config = backbone_config or {}
+        hidden_size = backbone_config.get("hidden_size", 768)
+        out_indices = backbone_config.get("out_indices", (3, 6, 9, 12))
+        self.backbone_repo = _BACKBONE_REPO_BY_HIDDEN_SIZE[hidden_size]
+        self.embed_dim = hidden_size
+        self.channels = fusion_hidden_size
+        self.post_process_channels = list(neck_hidden_sizes)
+        self.block_indices = [out_index - 1 for out_index in out_indices]
+        self.readout_type = readout_type
+        self.num_depth_bins = num_depth_bins
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+        self.num_seg_classes = self.num_labels

dpt_head.py ADDED Viewed

	@@ -0,0 +1,702 @@

+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""DPT (Dense Prediction Transformer) depth head in PyTorch.
+Ported from the Scenic/Flax implementation at:
+  research/vision/scene_understanding/imsight/modules/dpt.py
+  scenic/projects/dense_features/models/decoders.py
+Architecture:
+  ReassembleBlocks → 4×Conv3x3 → 4×FeatureFusionBlock → project → DepthHead
+"""
+import io
+import os
+import urllib.request
+import zipfile
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+# ── Building blocks ─────────────────────────────────────────────────────────
+class PreActResidualConvUnit(nn.Module):
+    """Pre-activation residual convolution unit."""
+    def __init__(self, features: int):
+        super().__init__()
+        self.conv1 = nn.Conv2d(features, features, 3, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(features, features, 3, padding=1, bias=False)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = F.relu(x)
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        return x + residual
+class FeatureFusionBlock(nn.Module):
+    """Fuses features with optional residual input, then upsamples 2×."""
+    def __init__(self, features: int, has_residual: bool = False,
+                 expand: bool = False):
+        super().__init__()
+        self.has_residual = has_residual
+        if has_residual:
+            self.residual_unit = PreActResidualConvUnit(features)
+        self.main_unit = PreActResidualConvUnit(features)
+        out_features = features // 2 if expand else features
+        self.out_conv = nn.Conv2d(features, out_features, 1, bias=True)
+    def forward(self, x: torch.Tensor,
+                residual: torch.Tensor = None) -> torch.Tensor:
+        if self.has_residual and residual is not None:
+            if residual.shape != x.shape:
+                residual = F.interpolate(
+                    residual, size=x.shape[2:], mode="bilinear",
+                    align_corners=False)
+            residual = self.residual_unit(residual)
+            x = x + residual
+        x = self.main_unit(x)
+        # Upsample 2× with align_corners=True (matches Scenic reference)
+        x = F.interpolate(x, scale_factor=2, mode="bilinear",
+                          align_corners=True)
+        x = self.out_conv(x)
+        return x
+class ReassembleBlocks(nn.Module):
+    """Projects and resizes intermediate ViT features to different scales."""
+    def __init__(self, input_embed_dim: int = 1024,
+                 out_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project"):
+        super().__init__()
+        self.readout_type = readout_type
+        # 1×1 conv to project to per-level channels
+        self.out_projections = nn.ModuleList([
+            nn.Conv2d(input_embed_dim, ch, 1) for ch in out_channels
+        ])
+        # Spatial resize layers: 4× up, 2× up, identity, 2× down
+        self.resize_layers = nn.ModuleList([
+            nn.ConvTranspose2d(out_channels[0], out_channels[0],
+                               kernel_size=4, stride=4, padding=0),
+            nn.ConvTranspose2d(out_channels[1], out_channels[1],
+                               kernel_size=2, stride=2, padding=0),
+            nn.Identity(),
+            nn.Conv2d(out_channels[3], out_channels[3], 3, stride=2,
+                      padding=1),
+        ])
+        # Readout projection (concatenate cls_token with patch features)
+        if readout_type == "project":
+            self.readout_projects = nn.ModuleList([
+                nn.Linear(2 * input_embed_dim, input_embed_dim)
+                for _ in out_channels
+            ])
+    def forward(self, features):
+        """Process list of (cls_token, spatial_features) tuples.
+        Args:
+            features: list of (cls_token [B,D], patch_feats [B,D,H,W])
+        Returns:
+            list of tensors at different scales.
+        """
+        out = []
+        for i, (cls_token, x) in enumerate(features):
+            B, D, H, W = x.shape
+            if self.readout_type == "project":
+                # Flatten spatial → (B, HW, D)
+                x_flat = x.flatten(2).transpose(1, 2)
+                # Expand cls_token → (B, HW, D)
+                readout = cls_token.unsqueeze(1).expand(-1, x_flat.shape[1], -1)
+                # Concat + project + GELU
+                x_cat = torch.cat([x_flat, readout], dim=-1)
+                x_proj = F.gelu(self.readout_projects[i](x_cat))
+                # Reshape back to spatial
+                x = x_proj.transpose(1, 2).reshape(B, D, H, W)
+            # 1×1 projection
+            x = self.out_projections[i](x)
+            # Spatial resize
+            x = self.resize_layers[i](x)
+            out.append(x)
+        return out
+class DPTDepthHead(nn.Module):
+    """Full DPT head + depth classification decoder.
+    Takes 4 intermediate ViT features and produces a depth map.
+    """
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project",
+                 num_depth_bins: int = 256,
+                 min_depth: float = 1e-3,
+                 max_depth: float = 10.0):
+        super().__init__()
+        self.num_depth_bins = num_depth_bins
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+        # Depth classification head (Dense layer)
+        self.depth_head = nn.Linear(channels, num_depth_bins)
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT depth prediction.
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+        Returns:
+            depth map tensor (B, 1, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+        # Project
+        out = self.project(out)
+        out = F.relu(out)
+        # Depth classification
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.depth_head(out)  # (B, H, W, num_bins)
+        # Classification-based depth prediction
+        bin_centers = torch.linspace(
+            self.min_depth, self.max_depth, self.num_depth_bins,
+            device=out.device)
+        out = F.relu(out) + self.min_depth
+        out_norm = out / out.sum(dim=-1, keepdim=True)
+        depth = torch.einsum("bhwn,n->bhw", out_norm, bin_centers)
+        depth = depth.unsqueeze(1)  # (B, 1, H, W)
+        # Resize to original image size
+        if image_size is not None:
+            depth = F.interpolate(depth, size=image_size, mode="bilinear",
+                                  align_corners=False)
+        return depth
+class DPTNormalsHead(nn.Module):
+    """Full DPT head + surface normals decoder.
+    Takes 4 intermediate ViT features and produces a normal map.
+    """
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project"):
+        super().__init__()
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+        # Normals head (Dense layer)
+        self.normals_head = nn.Linear(channels, 3)
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT normals prediction.
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+        Returns:
+            normal map tensor (B, 3, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+        # Project
+        out = self.project(out)
+        # Normals head
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.normals_head(out)  # (B, H, W, 3)
+        # Normalize to unit length
+        out = F.normalize(out, p=2, dim=-1)
+        # Resize to original image size
+        if image_size is not None:
+            # PyTorch interpolate expects (B, C, H, W)
+            out = out.permute(0, 3, 1, 2)
+            out = F.interpolate(out, size=image_size, mode="bilinear",
+                                  align_corners=False)
+        else:
+            out = out.permute(0, 3, 1, 2)
+        return out
+class DPTSegmentationHead(nn.Module):
+    """Full DPT head + segmentation decoder.
+    Takes 4 intermediate ViT features and produces a segmentation map.
+    """
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project",
+                 num_classes: int = 150):
+        super().__init__()
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+        # Segmentation head (Dense layer)
+        self.segmentation_head = nn.Linear(channels, num_classes)
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT segmentation prediction.
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+        Returns:
+            segmentation map tensor (B, num_classes, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+        # Project
+        out = self.project(out)
+        # Segmentation head
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.segmentation_head(out)  # (B, H, W, num_classes)
+        # Resize to original image size
+        if image_size is not None:
+            # PyTorch interpolate expects (B, C, H, W)
+            out = out.permute(0, 3, 1, 2)
+            out = F.interpolate(out, size=image_size, mode="bilinear",
+                                  align_corners=False)
+        else:
+            out = out.permute(0, 3, 1, 2)
+        return out
+# ── Weight loading from Scenic/Flax checkpoint ─────────────────────────────
+def _load_npy_from_zip(zf, name):
+    """Load a single .npy array from a zipfile."""
+    with zf.open(name) as f:
+        return np.load(io.BytesIO(f.read()))
+def _conv_kernel_flax_to_torch(w):
+    """Convert Flax conv kernel (H,W,Cin,Cout) → PyTorch (Cout,Cin,H,W)."""
+    return torch.from_numpy(w.transpose(3, 2, 0, 1).copy())
+def _conv_transpose_kernel_flax_to_torch(w):
+    """Convert Flax ConvTranspose kernel (H,W,Cin,Cout) → PyTorch (Cin,Cout,H,W)."""
+    return torch.from_numpy(w.transpose(2, 3, 0, 1).copy())
+def _linear_kernel_flax_to_torch(w):
+    """Convert Flax Dense kernel (in,out) → PyTorch Linear (out,in)."""
+    return torch.from_numpy(w.T.copy())
+def _bias(w):
+    return torch.from_numpy(w.copy())
+def load_dpt_weights(model: DPTDepthHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+    # --- Depth classification head ---
+    sd["depth_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_depth_classif/kernel.npy"))
+    sd["depth_head.bias"] = _bias(
+        npy("decoder/pixel_depth_classif/bias.npy"))
+    zf.close()
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT depth head weights ({len(sd)} tensors)")
+    return model
+def load_normals_weights(model: DPTNormalsHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+    # --- Normals head ---
+    sd["normals_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_normals/kernel.npy"))
+    sd["normals_head.bias"] = _bias(
+        npy("decoder/pixel_normals/bias.npy"))
+    zf.close()
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT normals head weights ({len(sd)} tensors)")
+    return model
+def load_segmentation_weights(model: DPTSegmentationHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+    # --- Segmentation head ---
+    sd["segmentation_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_segmentation/kernel.npy"))
+    sd["segmentation_head.bias"] = _bias(
+        npy("decoder/pixel_segmentation/bias.npy"))
+    zf.close()
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT segmentation head weights ({len(sd)} tensors)")
+    return model

model.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:38362aae91092ca0a6f2fe7f7ed29ec2e7042e3b920b7217d9317ad4c13abdfa
+size 1622839292

modeling_dpt.py ADDED Viewed

	@@ -0,0 +1,115 @@

+"""TIPSv2 DPT dense prediction model for HuggingFace."""
+from dataclasses import dataclass
+from typing import Optional
+import torch
+from transformers import AutoConfig, AutoModel, PreTrainedModel
+from .configuration_dpt import TIPSv2DPTConfig
+from .dpt_head import DPTDepthHead, DPTNormalsHead, DPTSegmentationHead
+@dataclass
+class TIPSv2DPTOutput:
+    depth: Optional[torch.Tensor] = None
+    normals: Optional[torch.Tensor] = None
+    segmentation: Optional[torch.Tensor] = None
+class TIPSv2DPTModel(PreTrainedModel):
+    """TIPSv2 DPT dense prediction model (depth, normals, segmentation).
+    The backbone is loaded automatically from the base TIPSv2 model repo.
+    Usage::
+        model = AutoModel.from_pretrained("google/tipsv2-l14-dpt", trust_remote_code=True)
+        model.eval().cuda()
+        outputs = model(pixel_values)
+        outputs.depth          # (B, 1, H, W)
+        outputs.normals        # (B, 3, H, W)
+        outputs.segmentation   # (B, 150, H, W)
+        # Individual tasks
+        depth = model.predict_depth(pixel_values)
+        normals = model.predict_normals(pixel_values)
+        seg = model.predict_segmentation(pixel_values)
+    """
+    config_class = TIPSv2DPTConfig
+    _no_split_modules = []
+    _supports_cache_class = False
+    _tied_weights_keys = []
+    @property
+    def all_tied_weights_keys(self):
+        return {}
+    def __init__(self, config: TIPSv2DPTConfig):
+        super().__init__(config)
+        ppc = tuple(config.post_process_channels)
+        backbone_config = AutoConfig.from_pretrained(config.backbone_repo, trust_remote_code=True)
+        backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)
+        self.vision_encoder = backbone.vision_encoder
+        self.depth_head = DPTDepthHead(
+            input_embed_dim=config.embed_dim, channels=config.channels,
+            post_process_channels=ppc, readout_type=config.readout_type,
+            num_depth_bins=config.num_depth_bins,
+            min_depth=config.min_depth, max_depth=config.max_depth,
+        )
+        self.normals_head = DPTNormalsHead(
+            input_embed_dim=config.embed_dim, channels=config.channels,
+            post_process_channels=ppc, readout_type=config.readout_type,
+        )
+        self.segmentation_head = DPTSegmentationHead(
+            input_embed_dim=config.embed_dim, channels=config.channels,
+            post_process_channels=ppc, readout_type=config.readout_type,
+            num_classes=config.num_seg_classes,
+        )
+    def _extract_intermediate(self, pixel_values):
+        intermediate = self.vision_encoder.get_intermediate_layers(
+            pixel_values, n=self.config.block_indices,
+            reshape=True, return_class_token=True, norm=True,
+        )
+        return [(cls_tok, patch_feat) for patch_feat, cls_tok in intermediate]
+    @torch.no_grad()
+    def predict_depth(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Predict depth map. Returns (B, 1, H, W)."""
+        pixel_values = pixel_values.to(self.device)
+        h, w = pixel_values.shape[2:]
+        dpt_inputs = self._extract_intermediate(pixel_values)
+        return self.depth_head(dpt_inputs, image_size=(h, w))
+    @torch.no_grad()
+    def predict_normals(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Predict surface normals. Returns (B, 3, H, W)."""
+        pixel_values = pixel_values.to(self.device)
+        h, w = pixel_values.shape[2:]
+        dpt_inputs = self._extract_intermediate(pixel_values)
+        return self.normals_head(dpt_inputs, image_size=(h, w))
+    @torch.no_grad()
+    def predict_segmentation(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Predict semantic segmentation (ADE20K). Returns (B, 150, H, W)."""
+        pixel_values = pixel_values.to(self.device)
+        h, w = pixel_values.shape[2:]
+        dpt_inputs = self._extract_intermediate(pixel_values)
+        return self.segmentation_head(dpt_inputs, image_size=(h, w))
+    def forward(self, pixel_values: torch.Tensor) -> TIPSv2DPTOutput:
+        """Run all three tasks. Returns TIPSv2DPTOutput."""
+        pixel_values = pixel_values.to(self.device)
+        h, w = pixel_values.shape[2:]
+        dpt_inputs = self._extract_intermediate(pixel_values)
+        return TIPSv2DPTOutput(
+            depth=self.depth_head(dpt_inputs, image_size=(h, w)),
+            normals=self.normals_head(dpt_inputs, image_size=(h, w)),
+            segmentation=self.segmentation_head(dpt_inputs, image_size=(h, w)),
+        )

preprocessor_config.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "do_convert_rgb": true,
+  "do_normalize": false,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_processor_type": "Tipsv2DptImageProcessor",
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 448,
+    "width": 448
+  }
+}