File size: 1,153 Bytes
23c283a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """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-")]
|