# SPDX-License-Identifier: Apache-2.0
"""Device-profiler driver: 2 eager forwards (JIT + warm, profiled) then trace capture and
3 replays, real checkpoint, sample1 crop.  Run under `python -m tracy -r -p -v`."""
import os, sys
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
os.environ.setdefault("TT_FUSED_HASH", "0")
import torch, ttnn
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
from hamer.preprocess import preprocess
hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
ref = Hamer().eval(); load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
x, _ = preprocess(_REPO.parent / "media/sample1.jpg", bbox_xyxy=(20, 170, 295, 260), is_right=True)
tt = build_from_reference(ref, device_id=0, strict=True)
with torch.no_grad():
    tt(x)                      # eager, JIT
    ttnn.ReadDeviceProfiler(tt.device)
    tt(x)                      # eager warm -> trace captured on this call? (warmup_done) -- see model
    ttnn.ReadDeviceProfiler(tt.device)
    for _ in range(3):
        tt(x)
    ttnn.ReadDeviceProfiler(tt.device)
tt.close()