load_model.py

import json

import equinox as eqx
import jax
from huggingface_hub import hf_hub_download
from z1t.components import Config
from z1t.model import create_model


def load_model(repo_id="Extropic-AI/Z1T-0"):
    config_path = hf_hub_download(repo_id, "config.json")
    weights_path = hf_hub_download(repo_id, "model.eqx")
    with open(config_path) as f:
        config = Config(**json.load(f))
    model = create_model(config, jax.random.key(0))
    return eqx.tree_deserialise_leaves(weights_path, model)