"""Download and inspect BrainAxL weights without custom remote code."""

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

path = hf_hub_download(
    repo_id="hieuphamha/BrainAxL",
    filename="brainaxl-b/model.safetensors",
)
state_dict = load_file(path, device="cpu")
print(f"Loaded {len(state_dict)} tensors")

# Instantiate brainaxl_b from the linked GitHub repository, then:
# model.load_state_dict(state_dict, strict=True)