File size: 4,387 Bytes
f6696e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
"""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 # noqa: E402
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
# transformers 5.x materializes an all-zero additive mask for eager attention even
# without padding (arange/ge/where plus a CPU scalar_tensor in the graph). Images never
# pad, so drop it; outputs are unchanged.
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()
|