File size: 5,719 Bytes
0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """模型下载管理（镜像不带权重，运行时从 Hugging Face 拉取）。
- 模型源：`x6nux/laya-ad`（可通过 LAYA_HF_REPO 覆盖），默认 revision `v2`
- 精度目录：fp32 / fp16 / bf16 / fp8 / int8 / int4（HF 仓库根目录同名子目录）
- 镜像站：LAYA_HF_ENDPOINT（默认 https://huggingface.co，可设 https://hf-mirror.com）
- 下载落盘到 LAYA_MODEL_DIR/laya-<版本>-<精度>/，写 PRECISION.txt，保留 HF 的 MODEL_ID.txt/VERSION.txt
- 状态查询：status()；下载在后台线程执行，同一时间一个任务
"""
from __future__ import annotations
import os
import shutil
import threading
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
CATALOG = ["fp32", "fp16", "bf16", "fp8", "int8", "int4"]
# (文件名, 是否必需)
FILES = [
("model.safetensors", True),
("config.json", True),
("tokenizer.json", True),
("tokenizer_config.json", False),
("dual_config.json", False),
("VERSION.txt", False),
("MODEL_ID.txt", False),
]
class ModelManager:
def __init__(self, model_dir: Path, device: str | None):
self.root = Path(model_dir)
self.device = device or "cpu"
self.root.mkdir(parents=True, exist_ok=True)
self.endpoint = (os.environ.get("LAYA_HF_ENDPOINT")
or os.environ.get("HF_ENDPOINT")
or "https://huggingface.co").rstrip("/")
self.repo = os.environ.get("LAYA_HF_REPO", "x6nux/laya-ad").strip("/")
self.version = os.environ.get("LAYA_MODEL_VERSION", "v2.2").strip() or "v2.2"
self._lock = threading.Lock()
self._status = {"state": "idle", "precision": None,
"started_at": None, "error": None}
# ---------- 查询 ----------
def status(self) -> dict:
return dict(self._status)
def installed(self) -> list[str]:
out = []
for d in sorted(self.root.iterdir()):
if d.is_dir() and (d / "MODEL_ID.txt").is_file() and (d / "model.safetensors").is_file():
out.append((d / "MODEL_ID.txt").read_text(encoding="utf-8").strip())
return out
def catalog(self) -> dict:
return {
"catalog": CATALOG,
"installed": self.installed(),
"status": self.status(),
"source": {"repo": self.repo, "version": self.version, "endpoint": self.endpoint},
}
# ---------- 下载 ----------
def download(self, precision: str) -> dict:
if precision not in CATALOG:
return {"accepted": False, "error": f"unknown precision {precision!r}"}
with self._lock:
if self._status["state"] == "downloading":
return {"accepted": False, "error": "another download is running"}
self._set("downloading", precision)
threading.Thread(target=self._run, args=(precision,), daemon=True).start()
return {"accepted": True, "precision": precision}
def _set(self, state: str, precision: str | None = None, error: str | None = None):
with self._lock:
self._status = {"state": state, "precision": precision,
"started_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
"error": error}
def _url(self, precision: str, name: str) -> str:
return f"{self.endpoint}/{self.repo}/resolve/{self.version}/{precision}/{name}"
def _headers(self) -> dict[str, str]:
headers = {"User-Agent": "laya-server/2.0"}
token = os.environ.get("HF_TOKEN")
if token:
headers["Authorization"] = f"Bearer {token}"
return headers
def _fetch(self, precision: str, name: str, dst: Path) -> None:
req = urllib.request.Request(self._url(precision, name), headers=self._headers())
with urllib.request.urlopen(req, timeout=300) as resp, open(dst, "wb") as fh:
shutil.copyfileobj(resp, fh)
def _run(self, precision: str) -> None:
try:
self._log(f"downloading {self.repo}@{self.version}/{precision} from {self.endpoint}")
tmp = self.root / f".download-{precision}"
shutil.rmtree(tmp, ignore_errors=True)
tmp.mkdir(parents=True)
for name, required in FILES:
try:
self._fetch(precision, name, tmp / name)
except urllib.error.HTTPError as exc:
if exc.code == 404 and not required:
continue
raise
if not (tmp / "model.safetensors").is_file() or not (tmp / "config.json").is_file():
raise RuntimeError("模型文件不完整（缺少 model.safetensors 或 config.json）")
(tmp / "PRECISION.txt").write_text(precision, encoding="utf-8")
if not (tmp / "VERSION.txt").is_file():
(tmp / "VERSION.txt").write_text(self.version, encoding="utf-8")
if not (tmp / "MODEL_ID.txt").is_file():
(tmp / "MODEL_ID.txt").write_text(
f"laya-{self.version}-{precision}", encoding="utf-8")
variant_dir = self.root / f"laya-{self.version}-{precision}"
shutil.rmtree(variant_dir, ignore_errors=True)
shutil.move(str(tmp), str(variant_dir))
self._set("done", precision)
self._log(f"installed {variant_dir}")
except Exception as exc: # 后台任务失败必须可见
self._set("error", precision, f"{type(exc).__name__}: {exc}")
@staticmethod
def _log(message: str) -> None:
print(f"[model-manager] {message}", flush=True)
|