| """Load the exported weights (models/maya2, written by export/export_weights.py).""" | |
| import json | |
| import os | |
| import mlx.core as mx | |
| MODELS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models") | |
| # the dynamic 8-bit weights (the published model); the 16-bit export (models/maya2, made by export_weights.py) when | |
| # only that one is present | |
| DEFAULT_DIR = next((os.path.join(MODELS, d) for d in ("maya2-8bit", "maya2") if os.path.isdir(os.path.join(MODELS, d))), | |
| os.path.join(MODELS, "maya2-8bit")) | |
| RELEASE_16BIT = os.path.join(MODELS, "maya2") | |
| def load_meta(path=DEFAULT_DIR): | |
| return json.load(open(os.path.join(path, "meta.json"))) | |
| def load(names, path=DEFAULT_DIR): | |
| """Merge the given safetensors files into one dict of mx arrays (dtypes as stored).""" | |
| out = {} | |
| for n in names: | |
| part = mx.load(os.path.join(path, n)) | |
| dup = set(part) & set(out) | |
| if dup: | |
| raise RuntimeError(f"{n}: duplicate tensors {sorted(dup)[:4]}") | |
| out.update(part) | |
| return out | |
| def backbone_files(meta): | |
| return [f for f in meta["files"] if f.startswith("backbone-")] | |