#!/usr/bin/env bash
# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
#
# Pre-download the pinned nvidia/GR00T-N1.7-3B snapshot -- exactly the files the server reads -- into the HF cache
# (HF_HOME / HF_HUB_CACHE, default ~/.cache/huggingface). `tt-model serve` does the same by itself before the
# container starts; this script is for host runs and offline preparation. The weights are distributed by NVIDIA
# under the NVIDIA License (non-commercial: research or evaluation use only) -- see LICENSE in the snapshot.
#
#   scripts/download_weights.sh            # prints the snapshot directory
#   PYTHON=/path/to/python scripts/download_weights.sh
set -euo pipefail
REPO=${GR00T_WEIGHTS_REPO:-nvidia/GR00T-N1.7-3B}
REV=${GR00T_WEIGHTS_REVISION:-2fc962b973bccdd5d8ce4f67cc63b264d6886495}
PY=${PYTHON:-python3}
"$PY" - "$REPO" "$REV" <<'PYEOF'
import sys

from huggingface_hub import snapshot_download

repo, rev = sys.argv[1], sys.argv[2]
patterns = [
    "model-*.safetensors",
    "model.safetensors.index.json",
    "config.json",
    "processor_config.json",
    "statistics.json",
    "embodiment_id.json",
    "LICENSE",
]
print(snapshot_download(repo_id=repo, revision=rev, allow_patterns=patterns))
PYEOF