#!/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()