"""T9 SPIKE · 模型下载脚本（本批唯一允许联网的环节，任务文档 §3.1）。 为什么不用 huggingface_hub：该包不在 spike extra 依赖中，且协作规程 §4.2 禁止新增 依赖，故用 stdlib urllib 直连 `{endpoint}/{repo}/resolve/main/{file}`——与 huggingface_hub 底层相同的 HTTP 语义（30x 跳转与 Range 断点续传均支持）。 行为要点： - 端点顺序：环境变量 HF_ENDPOINT > https://hf-mirror.com（国内镜像优先）> 官方站； - 候选链：每个目标文件按候选依次降级（BiRefNet_lite fp32 → fp16 → RMBG-1.4）； - 超时纪律：单文件超过 --budget 秒（默认 600s ≈ 10 分钟）即放弃当前候选换下一个； - 幂等：本地已存在且字节数与远端一致则跳过；.part 断点续传；每候选失败重试 1 次； - 结果（repo/revision/大小/耗时/端点）写入 models/manifest.json，供 REPORT.md 引用。 用法（在 worktree 根目录）： .venv/bin/python scripts/spikes/download_models.py [--only matting] [--budget 600] """ from __future__ import annotations import argparse import json import os import sys import time import urllib.error import urllib.request from dataclasses import dataclass from pathlib import Path from typing import Any REPO_ROOT = Path(__file__).resolve().parents[2] DEFAULT_MODELS_DIR = REPO_ROOT / "models" DEFAULT_BUDGET_SECONDS = 600 # 任务文档：单个模型 ~10 分钟即降级 CHUNK_BYTES = 1 << 20 ATTEMPTS_PER_CANDIDATE = 2 # 首次 + 失败重试一次 PROGRESS_EVERY_BYTES = 32 * CHUNK_BYTES UA = "meida-t9-spike/1.0" @dataclass(frozen=True, slots=True) class Candidate: """一个候选模型文件的下载描述。""" repo: str filename: str dest: str # models/ 下的目标文件名 note: str # 变体说明（写入 manifest，REPORT 引用） # 目标文件 → 候选链（顺序即优先级）。REQUIRED 目标缺失时退出码为 1， # 基准脚本随后走 HOG / Otsu 降级路线（见各脚本的降级分支）。 TARGETS: dict[str, tuple[Candidate, ...]] = { "pose-s": ( Candidate( repo="bukuroo/RTMPose-ONNX", filename="rtmpose-s.onnx", dest="rtmpose-s.onnx", note="RTMPose-s body(COCO-17, SimCC) 社区 ONNX 导出", ), ), "pose-m": ( Candidate( repo="bukuroo/RTMPose-ONNX", filename="rtmpose-m.onnx", dest="rtmpose-m.onnx", note="RTMPose-m body(COCO-17, SimCC)，规模-时延对比用（可选）", ), ), "matting": ( Candidate( repo="onnx-community/BiRefNet_lite-ONNX", filename="onnx/model.onnx", dest="birefnet_lite.onnx", note="BiRefNet_lite fp32 1024x1024（任务文档首选 lite 变体）", ), Candidate( repo="onnx-community/BiRefNet_lite-ONNX", filename="onnx/model_fp16.onnx", dest="birefnet_lite_fp16.onnx", note="BiRefNet_lite fp16 降级档（体积约减半）", ), Candidate( repo="briaai/RMBG-1.4", filename="onnx/model.onnx", dest="rmbg14.onnx", note="RMBG-1.4 fp32 末级降级（非 BiRefNet，需在报告注明）", ), ), } REQUIRED = ("pose-s", "matting") def _endpoints() -> list[str]: """镜像优先：HF_ENDPOINT 可覆盖，默认 hf-mirror.com，再兜底官方站。""" eps = [] for ep in ( os.environ.get("HF_ENDPOINT", "").rstrip("/"), "https://hf-mirror.com", "https://huggingface.co", ): if ep and ep not in eps: eps.append(ep) return eps def _url(endpoint: str, cand: Candidate) -> str: return f"{endpoint}/{cand.repo}/resolve/main/{cand.filename}" def _repo_revision(endpoint: str, repo: str) -> str: """通过 model API 拿当前 revision（commit sha），失败记 unknown。""" req = urllib.request.Request(f"{endpoint}/api/models/{repo}", headers={"User-Agent": UA}) try: with urllib.request.urlopen(req, timeout=20) as resp: data = json.load(resp) return str(data.get("sha", "unknown")) except (urllib.error.URLError, json.JSONDecodeError, OSError): return "unknown" def _download_once( endpoint: str, cand: Candidate, models_dir: Path, deadline: float, ) -> bool: """单次下载（支持 .part 续传）。超预算/网络错误返回 False，由外层重试或降级。""" partial = models_dir / f"{cand.dest}.part" offset = partial.stat().st_size if partial.exists() else 0 headers = {"User-Agent": UA} if offset: headers["Range"] = f"bytes={offset}-" req = urllib.request.Request(_url(endpoint, cand), headers=headers) try: with urllib.request.urlopen(req, timeout=30) as resp, open(partial, "ab") as fh: content_length = resp.headers.get("Content-Length") total = int(content_length) + offset if content_length else None last_bytes = 0 last_time = time.perf_counter() while True: if time.perf_counter() > deadline: print(" [超时] 超过下载预算，放弃当前候选", flush=True) return False chunk = resp.read(CHUNK_BYTES) if not chunk: break fh.write(chunk) offset += len(chunk) now = time.perf_counter() if offset - last_bytes >= PROGRESS_EVERY_BYTES or now - last_time > 10.0: total_str = f"/{total / 1048576:.1f}MB" if total else "" print(f" {offset / 1048576:.1f}MB{total_str} ...", flush=True) last_bytes = offset last_time = now partial.replace(models_dir / cand.dest) return True except urllib.error.HTTPError as exc: if exc.code == 416 and partial.exists(): # Range 起点已到文件尾 = 实际已下完 partial.replace(models_dir / cand.dest) return True print(f" [HTTP {exc.code}] {exc.reason}", flush=True) return False except OSError as exc: # 含 socket.timeout / 连接重置：保留 .part 供续传 print(f" [网络错误] {exc}", flush=True) return False def _mark_cached(cand: Candidate, dest: Path, manifest: dict[str, Any]) -> None: """缓存命中也补全 manifest（repo/revision 供 REPORT 引用；离线时 revision 记 unknown）。""" entry = manifest.get(cand.dest, {}) manifest[cand.dest] = { "repo": entry.get("repo", cand.repo), "revision": entry.get("revision") or _repo_revision(_endpoints()[0], cand.repo), "bytes": dest.stat().st_size, "cached": True, "note": cand.note, } def _try_candidate( cand: Candidate, models_dir: Path, budget_s: int, manifest: dict[str, Any], ) -> bool: """单候选全流程：双端点 × 重试一次，成功写 manifest。""" dest = models_dir / cand.dest for endpoint in _endpoints(): deadline = time.perf_counter() + budget_s for attempt in range(1, ATTEMPTS_PER_CANDIDATE + 1): print(f"[下载] {cand.repo}/{cand.filename} <- {endpoint}（第 {attempt} 次）") t0 = time.perf_counter() if _download_once(endpoint, cand, models_dir, deadline): manifest[cand.dest] = { "repo": cand.repo, "revision": _repo_revision(endpoint, cand.repo), "endpoint": endpoint, "bytes": dest.stat().st_size, "elapsed_s": round(time.perf_counter() - t0, 1), "cached": False, "note": cand.note, } print(f"[完成] {cand.dest} {manifest[cand.dest]['bytes'] / 1048576:.1f}MB") return True return False def ensure_target( name: str, models_dir: Path, budget_s: int, manifest: dict[str, Any], ) -> bool: """目标文件级编排：按候选链降级。返回该目标是否本地可用。""" for cand in TARGETS[name]: # 候选链内任何一个目标文件已在本地（含上次跑过其它候选）都算满足 dest = models_dir / cand.dest if dest.exists() and dest.stat().st_size > 0: print(f"[缓存] {name}: {cand.dest} 已在本地") _mark_cached(cand, dest, manifest) return True if _try_candidate(cand, models_dir, budget_s, manifest): return True print(f"[降级] {name} 的候选 {cand.repo}/{cand.filename} 失败，尝试下一候选") return False def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description="T9 SPIKE 模型下载（唯一联网环节）") parser.add_argument("--models-dir", type=Path, default=DEFAULT_MODELS_DIR) parser.add_argument( "--budget", type=int, default=DEFAULT_BUDGET_SECONDS, help="单文件下载预算（秒）" ) parser.add_argument( "--only", action="append", choices=sorted(TARGETS), help="只下载指定目标（可多次）" ) args = parser.parse_args(argv) models_dir: Path = args.models_dir models_dir.mkdir(parents=True, exist_ok=True) manifest_path = models_dir / "manifest.json" manifest: dict[str, Any] = {} if manifest_path.exists(): manifest = json.loads(manifest_path.read_text(encoding="utf-8")) selected = args.only or sorted(TARGETS) missing_required = False for name in selected: ok = ensure_target(name, models_dir, args.budget, manifest) required = name in REQUIRED if not ok: print(f"[失败] {name}：全部候选未取到（{'必需' if required else '可选'}）") missing_required |= required dump = dict(sorted(manifest.items())) manifest_path.write_text(json.dumps(dump, indent=2, ensure_ascii=False), encoding="utf-8") print(f"\nmanifest 已写入 {manifest_path}") for key, entry in dump.items(): repo = entry.get("repo", "?") revision = entry.get("revision", "?") print(f" {key}: {repo}@{revision} cached={entry.get('cached')}") if missing_required: print("\n有必需目标缺失：基准脚本将走降级路线（HOG/Otsu），详见 REPORT.md") return 1 return 0 if __name__ == "__main__": sys.exit(main())