| |
| """Depth Anything V2 Small (Hugging Face transformers port) -> ExecuTorch WebGPU. |
| |
| Upstream: depth-anything/Depth-Anything-V2-Small-hf (Apache-2.0; the Base and Large |
| variants are CC BY-NC 4.0 and must not be used here). |
| |
| Input: [1, 3, 224, 224] float32 RGB, ImageNet-normalized outside the graph |
| ((x / 255 - mean) / std, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)). |
| Output: [1, 224, 224] relative inverse depth (`predicted_depth`). |
| """ |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import _common |
|
|
| UPSTREAM = { |
| "repo": "depth-anything/Depth-Anything-V2-Small-hf", |
| "revision": "5426e4f0f36572d16453bbda7a8389317b1bef99", |
| } |
|
|
|
|
| class PredictedDepth(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m |
|
|
| def forward(self, pixel_values): |
| return self.m(pixel_values=pixel_values).predicted_depth |
|
|
|
|
| def freeze_position_embeddings(model, height, width): |
| """Bake the bicubic-resized DINOv2 position embeddings for one input size into a buffer. |
| |
| Traced live, the resize leaves lifted tensor constants and aten.arange nodes whose fake |
| tensors belong to a different FakeTensorMode, and the Vulkan preprocess passes fail on |
| them. The values are what the eager module computes for this size. |
| """ |
| import types |
|
|
| emb = model.m.backbone.embeddings |
| with torch.no_grad(): |
| dummy = torch.zeros(1, (height // 14) * (width // 14) + 1, emb.position_embeddings.shape[-1]) |
| frozen = emb.interpolate_pos_encoding(dummy, height, width).detach().clone() |
| emb.register_buffer("frozen_position_embeddings", frozen, persistent=False) |
| emb.interpolate_pos_encoding = types.MethodType(lambda self, e, h, w: self.frozen_position_embeddings, emb) |
| return model |
|
|
|
|
| def build_model(attn): |
| from transformers import AutoModelForDepthEstimation |
| from transformers.models.dinov2 import modeling_dinov2 |
|
|
| |
| |
| |
| create_mask = modeling_dinov2.create_bidirectional_mask |
| modeling_dinov2.create_bidirectional_mask = lambda **kw: ( |
| None if kw.get("attention_mask") is None else create_mask(**kw) |
| ) |
|
|
| model = AutoModelForDepthEstimation.from_pretrained( |
| UPSTREAM["repo"], |
| revision=UPSTREAM["revision"], |
| attn_implementation=attn, |
| torch_dtype=torch.float32, |
| ) |
| return PredictedDepth(model.eval()).eval() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--out", required=True, help="output directory") |
| ap.add_argument( |
| "--size", |
| type=int, |
| nargs=2, |
| default=[224, 224], |
| metavar=("H", "W"), |
| help="multiples of 14", |
| ) |
| ap.add_argument( |
| "--attn", |
| default="eager", |
| choices=["eager", "sdpa"], |
| help="attention implementation traced into the graph", |
| ) |
| ap.add_argument( |
| "--golden", |
| help="optional goldens.manifest.json to compare against (key: depth)", |
| ) |
| args = ap.parse_args() |
|
|
| golden = _common.load_golden(args.golden, "depth") if args.golden else None |
| _common.export_vision( |
| name="depth_v2s", |
| model=freeze_position_embeddings(build_model(args.attn), *args.size), |
| example_input=torch.randn(1, 3, *args.size), |
| out_dir=args.out, |
| upstream={ |
| "model": UPSTREAM["repo"], |
| "weights": {**UPSTREAM, "official": True}, |
| "license": "Apache-2.0", |
| }, |
| extra_recipe={ |
| "attn_implementation": args.attn, |
| "graph_rewrites": ["transformers 5.x all-zero bidirectional attention mask skipped (attention_mask=None; no padding in images)", "DINOv2 position embeddings resized (bicubic, as upstream) for the fixed input size in eager and stored as a constant"], |
| "preprocessing_outside_graph": "RGB / 255, ImageNet mean/std normalization, resize to input size", |
| "postprocessing_outside_graph": "min-max normalize predicted_depth for display", |
| }, |
| golden=golden, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
