| """统一模型读取入口。 |
| |
| 目录约定（raw 下一个模型一个文件夹）： |
| |
| raw/ |
| <模型名>/ |
| <模型名>.safetensors 解混淆后的权重 |
| blobs/ 容器里内嵌的 ONNX / GGUF / CoreML 产物 |
| |
| 所有脚本都通过这里取权重和 blob 路径。 |
| safetensors 的读取优先用官方库；没装库时用内置的最简解析器（memmap，零拷贝）。 |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import struct |
| import warnings |
|
|
| import numpy as np |
| import torch |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| RAW = os.path.join(HERE, "..", "raw") |
|
|
| _DTYPE = { |
| "F64": "<f8", "F32": "<f4", "F16": "<f2", "BF16": None, |
| "I64": "<i8", "I32": "<i4", "I16": "<i2", "I8": "i1", |
| "U64": "<u8", "U32": "<u4", "U16": "<u2", "U8": "u1", |
| "BOOL": "?", |
| } |
|
|
|
|
| def model_dir(model: str) -> str: |
| return os.path.join(RAW, model) |
|
|
|
|
| def safetensors_path(model: str) -> str: |
| return os.path.join(RAW, model, model + ".safetensors") |
|
|
|
|
| def blobs_dir(model: str) -> str: |
| return os.path.join(RAW, model, "blobs") |
|
|
|
|
| def ckpt_path(model: str) -> str: |
| """兼容旧布局：raw/<模型>/<模型>.ckpt""" |
| return os.path.join(RAW, model, model + ".ckpt") |
|
|
|
|
| def _read_safetensors_manual(path: str) -> dict: |
| """不依赖 safetensors 库：直接解析 [u64 头长][JSON 头][数据]。memmap 零拷贝。""" |
| mm = np.memmap(path, dtype=np.uint8, mode="r") |
| with open(path, "rb") as f: |
| n = struct.unpack("<Q", f.read(8))[0] |
| hdr = json.loads(f.read(n)) |
| base = 8 + n |
| out = {} |
| for k, v in hdr.items(): |
| if k == "__metadata__": |
| continue |
| b, e = v["data_offsets"] |
| dt = _DTYPE.get(v["dtype"]) |
| if dt is None: |
| continue |
| arr = np.frombuffer(mm[base + b:base + e], dtype=np.dtype(dt)) |
| arr = arr.reshape(v["shape"]) if v["shape"] else arr.reshape(()) |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| out[k] = torch.from_numpy(arr) |
| return out |
|
|
|
|
| def load_any(path: str) -> dict: |
| """从 .safetensors 或 .ckpt 载入 {参数名: torch.Tensor}。""" |
| if path.endswith(".safetensors"): |
| try: |
| from safetensors import safe_open |
| except ImportError: |
| return _read_safetensors_manual(path) |
| out = {} |
| with safe_open(path, framework="numpy") as f: |
| for k in f.keys(): |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| out[k] = torch.from_numpy(f.get_tensor(k)) |
| return out |
|
|
| sd = torch.load(path, map_location="cpu", weights_only=True) |
| if isinstance(sd, dict) and "state_dict" in sd: |
| sd = sd["state_dict"] |
| return sd |
|
|
|
|
| def state_dict_for(model: str) -> dict: |
| """按模型名取权重：safetensors 优先，回退 ckpt。""" |
| p = safetensors_path(model) |
| if os.path.exists(p): |
| return load_any(p) |
| p = ckpt_path(model) |
| if os.path.exists(p): |
| return load_any(p) |
| raise FileNotFoundError(f"{model}: 找不到 {safetensors_path(model)}") |
|
|
|
|
| def model_path_for(model: str) -> str | None: |
| for p in (safetensors_path(model), ckpt_path(model)): |
| if os.path.exists(p): |
| return p |
| return None |
|
|
|
|
| def list_models() -> list[str]: |
| """列出 raw/ 下所有带 .safetensors 的模型文件夹。""" |
| if not os.path.isdir(RAW): |
| return [] |
| return [n for n in sorted(os.listdir(RAW)) |
| if os.path.isfile(os.path.join(RAW, n, n + ".safetensors"))] |
|
|
