#!/usr/bin/env python3
# DPT (Dense Prediction Transformer) monocular depth estimation on Neuron
import argparse
import logging
import time

import torch
from transformers import DPTImageProcessor, DPTForDepthEstimation
from datasets import load_dataset
import torch_neuronx  # ensures Neuron backend

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(description="Run DPT depth estimation on Neuron")
    parser.add_argument(
        "--model",
        type=str,
        default="Intel/dpt-large",
        help="DPT model name on Hugging Face Hub",
    )
    args = parser.parse_args()

    torch.set_default_dtype(torch.float32)
    torch.manual_seed(42)

    # load dataset image
    dataset = load_dataset("huggingface/cats-image")
    image = dataset["test"]["image"][0]

    # load processor & DPT model
    processor = DPTImageProcessor.from_pretrained(args.model)
    model = DPTForDepthEstimation.from_pretrained(
        args.model, torch_dtype=torch.float32, attn_implementation="eager"
    ).eval()

    # preprocess
    inputs = processor(images=image, return_tensors="pt")

    # pre-run to lock shapes
    with torch.no_grad():
        _ = model(**inputs).predicted_depth

    # compile
    model.forward = torch.compile(model.forward, backend="neuron", fullgraph=False)

    # warmup
    warmup_start = time.time()
    with torch.no_grad():
        _ = model(**inputs)
    warmup_time = time.time() - warmup_start

    # benchmark run
    run_start = time.time()
    with torch.no_grad():
        depth = model(**inputs).predicted_depth
    run_time = time.time() - run_start

    logger.info("Warmup: %.2f s, Run: %.4f s", warmup_time, run_time)
    logger.info("Output depth shape: %s", depth.shape)  # [B, 1, H, W]


if __name__ == "__main__":
    main()