import os
from huggingface_hub import hf_hub_download

REPO_ID = "Deva8/GENvqa-model"

# We use the token from the environment variable (which the user must set in Settings -> Secrets)
HF_TOKEN = os.getenv("HF_TOKEN")

print("Downloading models from HuggingFace Hub...")

# Download base checkpoint
hf_hub_download(
    repo_id=REPO_ID,
    filename="vqa_checkpoint.pt",
    local_dir=".",
    token=HF_TOKEN
)
print("Base checkpoint downloaded successfully.")

# Download spatial checkpoint
hf_hub_download(
    repo_id=REPO_ID,
    filename="vqa_spatial_checkpoint.pt",
    local_dir=".",
    token=HF_TOKEN
)
print("Spatial checkpoint downloaded successfully.")