xsl
446d7a9391
T-P0-16/12/10 可行性验证：CPU 基准 + 代理质检指标 + 合成图体型剖面法 原型 + REPORT.md 三问三答（Sprint 1 门槛的决策依据）。 偏差说明：模型齐全未触发降级（rtmpose-s/m + birefnet_lite fp32 全部首选 候选成功，hf-mirror 实测 4~12MB/s）；huggingface_hub 不在 spike extra 且 禁新增依赖，改用 stdlib urllib 直连 resolve 端点（同 HTTP 语义）；PINTO model zoo 无法脚本化定位小模型、dhkim0115 仓不存在，姿态模型改用 bukuroo/RTMPose-ONNX；BiRefNet 官方仓无 ONNX，用 onnx-community 导出； GPU 档未测（CPU 版 ORT 无 CUDA EP，按文档不折腾）；BiRefNet 单张>3s 正式 次数自适应降为 10 并在输出注明；RTMPose 在合成剪影上置信度低（域差异）， 体型断言按文档允许走真值定高路径，关键点路径仅统计参考一致率。
270 lines
10 KiB
Python
270 lines
10 KiB
Python
"""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())
|