| """模型下载管理（镜像不带权重，运行时从 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) |
|
|
