Add ExecuTorch WebGPU program for depth-anything-v2-small

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+depth_v2s_webgpu.pte filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

NOTICE

@@ -0,0 +1,6 @@
+This repository contains a modified version of depth-anything/Depth-Anything-V2-Small-hf.
+
+Modifications by Meta Platforms, Inc.: conversion to the ExecuTorch program format (.pte) for the
+ExecuTorch WebGPU backend, and quantization as described in README.md.
+
+The original work and these modifications are distributed under the license in LICENSE.

README.md

@@ -0,0 +1,54 @@
+---
+license: apache-2.0
+base_model:
+- depth-anything/Depth-Anything-V2-Small-hf
+base_model_relation: quantized
+library_name: executorch
+pipeline_tag: depth-estimation
+tags:
+- executorch
+- webgpu
+- on-device
+- browser
+---
+# Depth Anything V2 Small for ExecuTorch WebGPU
+
+This repository contains an [ExecuTorch](https://github.com/pytorch/executorch) conversion of
+[depth-anything/Depth-Anything-V2-Small-hf](https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf) for the ExecuTorch **WebGPU** backend, so the model runs entirely in a web browser on
+the visitor's GPU. It powers the corresponding demo in the
+[ExecuTorch Playground](https://executorch.ai/playground/webgpu/).
+
+- **Base model:** [depth-anything/Depth-Anything-V2-Small-hf](https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf)
+- **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend
+- **Quantization:** fp32 (as exported)
+- **Total size:** 92.8 MB
+
+> **Note:** Only the Small variant is Apache-2.0; Base and Large are CC BY-NC 4.0 and must not be published.
+
+## Files
+
+| File | Size | SHA-256 |
+|---|---|---|
+| `depth_v2s_webgpu.pte` | 92.8 MB | `285602ed3a34a26a…` |
+
+Full checksums are in `manifest.json`. The playground verifies them after download.
+
+## Running it
+
+These programs need the ExecuTorch WebGPU runtime compiled to WebAssembly
+([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)). The easiest way to
+try the model is the [ExecuTorch Playground](https://executorch.ai/playground/webgpu/), which downloads these
+files once, caches them in the browser, and runs everything locally.
+
+## Provenance
+
+- **Weights:** `depth-anything/Depth-Anything-V2-Small-hf` @ `5426e4f0f36572d16453bbda7a8389317b1bef99`, file `None`, sha256 `None` (official)
+- **Exported with:** executorch 1.5.1+cpu, torch 2.14.1+cpu, torchao 0.18.0+cpu, transformers 5.19.0 (executorch git `3b6068392324`)
+- **Reproduce:** `python recipes/depth-anything-v2-small.py --out out/depth-anything-v2-small` (recipe in `export/`)
+- Re-running the recipe produces byte-identical files.
+
+## License
+
+This is a derivative of the upstream model and is distributed under the same terms
+(`apache-2.0`); see `LICENSE`.
+Modifications (conversion to the ExecuTorch format and quantization) by Meta Platforms, Inc.

depth_v2s_webgpu.pte

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:285602ed3a34a26a4d9d36355ce978587ed555e68b6b5fa8a567c1155a3b5160
+size 97307268

export/_common.py

@@ -0,0 +1,330 @@
+"""Shared helpers for the ExecuTorch WebGPU export recipes in this directory.
+
+Each recipe builds an eager model from a pinned upstream revision and calls
+`export_vision` (or its own LLM path), which writes:
+
+  <out>/<name>_webgpu.pte    program lowered with WebGPUPartitioner (delegate id "VulkanBackend")
+  <out>/<name>_portable.pte  same exported graph without a delegate, for CPU validation only
+  <out>/recipe.json          upstream id + revision, versions, shapes, dtype, command line
+  <out>/validation.json      eager vs exported-program vs portable .pte (and goldens, if given)
+"""
+
+import hashlib
+import io
+import json
+import platform
+import shlex
+import sys
+import urllib.request
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+PROXY_HINT = (
+    "set https_proxy if the machine needs a proxy for huggingface.co / pypi.org"
+)
+
+
+def versions():
+    import executorch
+    import torchao
+    from executorch.version import git_version
+
+    out = {
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "torchao": torchao.__version__,
+        "executorch": getattr(executorch, "__version__", None)
+        or __import__("executorch.version").version.__version__,
+        "executorch_git": git_version,
+    }
+    try:
+        import transformers
+
+        out["transformers"] = transformers.__version__
+    except ImportError:
+        pass
+    try:
+        import huggingface_hub
+
+        out["huggingface_hub"] = huggingface_hub.__version__
+    except ImportError:
+        pass
+    return out
+
+
+def sha256_file(path):
+    h = hashlib.sha256()
+    with open(path, "rb") as fh:
+        for chunk in iter(lambda: fh.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def hf_file(repo_id, filename, revision, expected_sha256=None, cache_dir=None):
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(
+        repo_id=repo_id, filename=filename, revision=revision, cache_dir=cache_dir
+    )
+    digest = sha256_file(path)
+    if expected_sha256 and digest != expected_sha256:
+        sys.exit(
+            f"{repo_id}@{revision}:{filename} sha256 {digest} != {expected_sha256}"
+        )
+    return Path(path), digest
+
+
+def pypi_wheel_member(url, expected_sha256, member):
+    """Read one file out of a PyPI wheel, verifying the wheel's sha256."""
+    with urllib.request.urlopen(url, timeout=120) as resp:
+        blob = resp.read()
+    digest = hashlib.sha256(blob).hexdigest()
+    if digest != expected_sha256:
+        sys.exit(f"{url}: sha256 {digest} != {expected_sha256} ({PROXY_HINT})")
+    with zipfile.ZipFile(io.BytesIO(blob)) as zf:
+        return zf.read(member).decode()
+
+
+def scrub(text, out_dir):
+    """Replace machine-specific paths so recipe.json can ship with the model."""
+    text = text.replace(str(Path(out_dir).resolve()), "<out>")
+    text = text.replace(str(Path(__file__).resolve().parent), "recipes")
+    return text.replace(str(Path.home()), "~")
+
+
+def command_line(out_dir):
+    argv = [f"recipes/{Path(sys.argv[0]).name}", *sys.argv[1:]]
+    args = []
+    for a in argv:
+        if a.startswith("/") or a.startswith("."):
+            a = scrub(str(Path(a).resolve()), out_dir)
+        args.append(a)
+    return scrub(" ".join(["python", *(shlex.quote(a) for a in args)]), out_dir)
+
+
+def lower_webgpu(ep):
+    from executorch.backends.webgpu.partitioner import WebGPUPartitioner
+    from executorch.exir import to_edge_transform_and_lower
+
+    return to_edge_transform_and_lower(
+        ep, partitioner=[WebGPUPartitioner()]
+    ).to_executorch()
+
+
+def lower_portable(ep):
+    from executorch.exir import to_edge_transform_and_lower
+
+    return to_edge_transform_and_lower(ep).to_executorch()
+
+
+def delegate_summary(pte_bytes):
+    """Non-delegated kernels and per-delegate op histograms of a serialized program."""
+    import collections
+
+    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
+        extract_vk_flatbuffer,
+        flatbuffer_to_vk_graph,
+    )
+    from executorch.exir._serialize._program import deserialize_pte_binary
+
+    pte = deserialize_pte_binary(pte_bytes)
+    prog = pte.program if hasattr(pte, "program") else pte
+    seg_base = int.from_bytes(pte_bytes[24:32], "little")
+    plan = prog.execution_plan[0]
+    ops = [f"{o.name}.{o.overload}" for o in plan.operators]
+    kernels = collections.Counter()
+    for chain in plan.chains:
+        for ins in chain.instructions:
+            if type(ins.instr_args).__name__ == "KernelCall":
+                kernels[ops[ins.instr_args.op_index]] += 1
+    delegates = []
+    for d in plan.delegates:
+        proc = d.processed
+        if proc.location.name == "INLINE":
+            blob = prog.backend_delegate_data[proc.index].data
+        else:
+            seg = prog.segments[proc.index]
+            blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
+        g = flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))
+        delegates.append(
+            {
+                "id": d.id,
+                "ops": dict(collections.Counter(c.name for c in g.chain)),
+                "inputs": len(g.input_ids),
+                "outputs": len(g.output_ids),
+            }
+        )
+    return {"non_delegated_kernels": dict(kernels), "delegates": delegates}
+
+
+def compare(a, b):
+    a = np.asarray(a, dtype=np.float64).ravel()
+    b = np.asarray(b, dtype=np.float64).ravel()
+    n = min(a.size, b.size)
+    a, b = a[:n], b[:n]
+    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
+    return {
+        "max_abs": float(np.max(np.abs(a - b))),
+        "mean_abs": float(np.mean(np.abs(a - b))),
+        "cosine": cos,
+    }
+
+
+def run_pte(pte_bytes, inputs):
+    from executorch.extension.pybindings.portable_lib import (
+        _load_for_executorch_from_buffer,
+    )
+
+    m = _load_for_executorch_from_buffer(pte_bytes)
+    return [
+        o.detach().numpy() if isinstance(o, torch.Tensor) else o
+        for o in m.forward(inputs)
+    ]
+
+
+def try_run_webgpu_pte_on_cpu(pte_bytes, inputs):
+    try:
+        return run_pte(pte_bytes, inputs), None
+    except Exception as e:  # noqa: BLE001 - the error text is the result
+        return None, f"{type(e).__name__}: {str(e).splitlines()[0][:300]}"
+
+
+def export_vision(
+    *,
+    name,
+    model,
+    example_input,
+    out_dir,
+    upstream,
+    extra_recipe,
+    golden=None,
+    output_index=None,
+):
+    """Export, lower, write both .pte files, recipe.json, and validation.json.
+
+    `golden` is an optional (input ndarray, output ndarray) pair; when its input
+    shape matches `example_input`, all paths are also compared against it.
+    `output_index` picks one tensor when the eager model returns a tuple.
+    """
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    torch.manual_seed(0)
+    model = model.eval()
+
+    class Wrapped(torch.nn.Module):
+        def __init__(self, m):
+            super().__init__()
+            self.m = m
+
+        def forward(self, x):
+            y = self.m(x)
+            return y[output_index] if output_index is not None else y
+
+    wrapped = Wrapped(model).eval()
+    with torch.no_grad():
+        ep = torch.export.export(wrapped, (example_input,), strict=False)
+        webgpu = lower_webgpu(ep)
+        portable = lower_portable(ep)
+
+    webgpu_path = out_dir / f"{name}_webgpu.pte"
+    portable_path = out_dir / f"{name}_portable.pte"
+    webgpu_path.write_bytes(webgpu.buffer)
+    portable_path.write_bytes(portable.buffer)
+
+    summary = delegate_summary(webgpu.buffer)
+    out_shape = list(ep.module()(example_input).shape)
+    recipe = {
+        "model": name,
+        "upstream": upstream,
+        "versions": versions(),
+        "export": {
+            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
+            "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
+            "inputs": [
+                {
+                    "shape": list(example_input.shape),
+                    "dtype": "float32",
+                    "layout": "NCHW",
+                }
+            ],
+            "outputs": [{"shape": out_shape, "dtype": "float32"}],
+            "dtype": "float32",
+            "quantization": "none (fp32)",
+            **extra_recipe,
+        },
+        "files": {
+            webgpu_path.name: {
+                "bytes": webgpu_path.stat().st_size,
+                "sha256": sha256_file(webgpu_path),
+            },
+            portable_path.name: {
+                "bytes": portable_path.stat().st_size,
+                "sha256": sha256_file(portable_path),
+                "note": "CPU-only validation artifact (portable kernels, no delegate); not for release",
+            },
+        },
+        "graph": summary,
+        "command": command_line(out_dir),
+    }
+    (out_dir / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
+
+    # Validation on CPU.
+    gen = torch.Generator().manual_seed(1234)
+    rand_input = torch.rand(example_input.shape, generator=gen)
+    cases = {"random": rand_input}
+    if golden is not None and tuple(golden[0].shape) == tuple(example_input.shape):
+        cases["golden"] = torch.from_numpy(golden[0].copy())
+    report = {
+        "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel",
+        "cases": {},
+    }
+    webgpu_cpu_error = None
+    for label, x in cases.items():
+        with torch.no_grad():
+            eager = wrapped(x).numpy()
+            exported = ep.module()(x).numpy()
+        port = run_pte(portable.buffer, (x,))[0]
+        res = {
+            "exported_vs_eager": compare(exported, eager),
+            "portable_pte_vs_eager": compare(port, eager),
+        }
+        wg, err = try_run_webgpu_pte_on_cpu(webgpu.buffer, (x,))
+        if wg is not None:
+            res["webgpu_pte_vs_eager"] = compare(wg[0], eager)
+        else:
+            webgpu_cpu_error = err
+        if label == "golden":
+            ref = golden[1]
+            res["internal_golden_vs_eager"] = compare(eager, ref)
+            res["internal_golden_vs_portable_pte"] = compare(port, ref)
+            res["golden_note"] = (
+                "reference output from the internal goldens.manifest.json"
+            )
+        report["cases"][label] = res
+    report["webgpu_pte_on_cpu"] = (
+        "ran"
+        if webgpu_cpu_error is None
+        else f"not runnable on CPU pybindings: {webgpu_cpu_error}"
+    )
+    (out_dir / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
+    print(
+        json.dumps(
+            {"recipe": recipe["files"], "graph": summary, "validation": report},
+            indent=2,
+        )
+    )
+    return ep, webgpu, portable
+
+
+def load_golden(manifest_path, key):
+    man = json.loads(Path(manifest_path).read_text())[key]["tensors"]
+    base = Path(manifest_path).parent
+
+    def ld(t):
+        return np.fromfile(base / t["file"], dtype=np.float32).reshape(t["shape"])
+
+    inp = man.get("input")
+    return ld(inp), ld(man["output"])

export/depth-anything-v2-small.py

@@ -0,0 +1,120 @@
+#!/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()

export/recipe.json

@@ -0,0 +1,99 @@
+{
+  "model": "depth_v2s",
+  "upstream": {
+    "model": "depth-anything/Depth-Anything-V2-Small-hf",
+    "weights": {
+      "repo": "depth-anything/Depth-Anything-V2-Small-hf",
+      "revision": "5426e4f0f36572d16453bbda7a8389317b1bef99",
+      "official": true
+    },
+    "license": "Apache-2.0"
+  },
+  "versions": {
+    "python": "3.12.15+meta",
+    "torch": "2.14.1+cpu",
+    "torchao": "0.18.0+cpu",
+    "executorch": "1.5.1+cpu",
+    "executorch_git": "3b60683923245cf472b7323426920e15623ba361",
+    "transformers": "5.19.0",
+    "huggingface_hub": "2.2.0"
+  },
+  "export": {
+    "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
+    "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
+    "inputs": [
+      {
+        "shape": [
+          1,
+          3,
+          224,
+          224
+        ],
+        "dtype": "float32",
+        "layout": "NCHW"
+      }
+    ],
+    "outputs": [
+      {
+        "shape": [
+          1,
+          224,
+          224
+        ],
+        "dtype": "float32"
+      }
+    ],
+    "dtype": "float32",
+    "quantization": "none (fp32)",
+    "attn_implementation": "eager",
+    "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"
+  },
+  "files": {
+    "depth_v2s_webgpu.pte": {
+      "bytes": 97307268,
+      "sha256": "285602ed3a34a26a4d9d36355ce978587ed555e68b6b5fa8a567c1155a3b5160"
+    },
+    "depth_v2s_portable.pte": {
+      "bytes": 97290408,
+      "sha256": "ce840e1c17ea65fcc52b7d8912dfe2d4d598350e3bdf6b8150f3d8c05d18423c",
+      "note": "CPU-only validation artifact (portable kernels, no delegate); not for release"
+    }
+  },
+  "graph": {
+    "non_delegated_kernels": {},
+    "delegates": [
+      {
+        "id": "VulkanBackend",
+        "ops": {
+          "et_vk.prepack.default": 39,
+          "aten.convolution.default": 24,
+          "aten.view_copy.default": 113,
+          "aten.permute_copy.default": 65,
+          "aten.clone.default": 7,
+          "aten.cat.default": 1,
+          "aten.add.Tensor": 35,
+          "aten.native_layer_norm.default": 28,
+          "aten.squeeze_copy.dims": 97,
+          "aten.linear.default": 72,
+          "aten.bmm.default": 24,
+          "aten.mul.Tensor": 37,
+          "aten._softmax.default": 12,
+          "aten.gelu.default": 12,
+          "aten.slice_copy.Tensor": 4,
+          "aten.relu.default": 7,
+          "et_vk.conv_with_clamp.default": 9,
+          "aten.upsample_bilinear2d.vec": 5,
+          "dim_order_ops._to_dim_order_copy.default": 1
+        },
+        "inputs": 1,
+        "outputs": 1
+      }
+    ]
+  },
+  "command": "python recipes/depth-anything-v2-small.py --out out/depth-anything-v2-small"
+}

export/requirements.txt

@@ -0,0 +1,15 @@
+# Environment the re-exports in REEXPORT.md were produced with (Python 3.12, Linux x86_64, CPU only).
+#   python -m venv venv
+#   venv/bin/pip install -r recipes/requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
+# executorch 1.5.1 = pytorch/executorch git 3b60683923245cf472b7323426920e15623ba361; its wheel ships
+# backends/webgpu (WebGPUPartitioner) and examples/models/llama (export_llm).
+executorch==1.5.1
+torch==2.14.1+cpu
+torchvision==0.29.1+cpu
+torchao==0.18.0+cpu
+transformers==5.19.0
+huggingface_hub==2.2.0
+safetensors==0.8.0
+tokenizers==0.23.3
+numpy==2.5.4
+pillow==12.3.0

manifest.json

@@ -0,0 +1,11 @@
+{
+  "model": "depth-anything-v2-small",
+  "repo": "executorch-community/Depth-Anything-V2-Small-ExecuTorch-WebGPU",
+  "files": [
+    {
+      "path": "depth_v2s_webgpu.pte",
+      "bytes": 97307268,
+      "sha256": "285602ed3a34a26a4d9d36355ce978587ed555e68b6b5fa8a567c1155a3b5160"
+    }
+  ]
+}