模型下载与检查点选择

这个模块只做两件事：把权重放到本地，再选对检查点

LingBot-Map 的推理代码不应该关心权重来自哪里。这个模块建议把“下载”和“选择检查点”拆开：

1. 下载模型权重

   • HuggingFace：robbyant/lingbot-map
   • ModelScope：Robbyant/lingbot-map

2. 从本地权重目录里选择一个检查点

   • balanced：默认推荐，适合大多数 demo、短视频、交互式预览。
   • long：长序列推荐，适合几千到上万帧的流式重建。
   • stage1：训练阶段产物或中间模型，主要用于复现实验、消融、继续训练，不建议作为普通推理默认值。

本模块不负责网络推理、不负责构建模型结构，也不修改 demo.py 的推理逻辑。它只输出一个本地 checkpoint 路径，例如：

text📋 复制
checkpoints/lingbot-map/xxx_balanced_xxx.pth

然后把这个路径传给项目里的 demo、benchmark 或自己的加载脚本。

---

权重仓库：HuggingFace 和 ModelScope 指向同一类模型文件

LingBot-Map README 中提供了两个模型入口：

平台	模型 ID	适合场景
HuggingFace	robbyant/lingbot-map	海外网络、已有 HuggingFace 缓存体系
ModelScope	Robbyant/lingbot-map	国内网络、需要更稳定下载时

推荐在项目根目录下统一放到：

text📋 复制
checkpoints/
└── lingbot-map/
    ├── ...
    ├── *balanced*.pth
    ├── *long*.pth
    └── *stage1*.pth

实际文件名可能随项目发布版本变化，不要在业务代码里写死完整文件名，建议按关键字筛选。

---

检查点怎么选：balanced、long、stage1

检查点	推荐用途	特点	不适合
balanced	默认推理、交互 demo、普通视频	速度、显存、重建质量较均衡	极长视频的最终重建
long	长视频、长轨迹、超过数千帧的流式场景	更偏向长程上下文和轨迹稳定性	显存紧张、只跑几十帧预览
stage1	训练复现、消融实验、继续训练初始化	更接近中间阶段模型	普通用户直接推理

一般选择规则：

text📋 复制
短视频 / 快速看效果       -> balanced
长视频 / 漫游 / 上万帧    -> long
训练调试 / 对齐论文实验   -> stage1

如果只是第一次运行 LingBot-Map，先用 balanced。

---

下载接口：HuggingFace 版本

依赖安装：

bash📋 复制
python -m pip install -U huggingface_hub

下载脚本：

python📋 复制
# tools/download_lingbot_map_hf.py
from huggingface_hub import snapshot_download

repo_dir = snapshot_download(
    repo_id="robbyant/lingbot-map",
    local_dir="checkpoints/lingbot-map",
    local_dir_use_symlinks=False,
    allow_patterns=[
        "*.pth",
        "*.pt",
        "*.ckpt",
        "*.safetensors",
        "*.json",
        "*.yaml",
        "*.yml",
        "*.txt",
    ],
)

print(repo_dir)

运行：

bash📋 复制
python tools/download_lingbot_map_hf.py

下载完成后检查文件：

bash📋 复制
find checkpoints/lingbot-map -maxdepth 3 -type f

如果网络环境需要 HuggingFace 镜像，可以在运行前配置：

bash📋 复制
export HF_ENDPOINT=https://hf-mirror.com
python tools/download_lingbot_map_hf.py

---

下载接口：ModelScope 版本

依赖安装：

bash📋 复制
python -m pip install -U modelscope

下载脚本：

python📋 复制
# tools/download_lingbot_map_ms.py
from modelscope.hub.snapshot_download import snapshot_download

repo_dir = snapshot_download(
    model_id="Robbyant/lingbot-map",
    cache_dir="checkpoints",
)

print(repo_dir)

运行：

bash📋 复制
python tools/download_lingbot_map_ms.py

ModelScope 默认会在 cache_dir 下建立自己的缓存目录。下载完成后可以这样查看：

bash📋 复制
find checkpoints -maxdepth 5 -type f | grep -E '\.(pth|pt|ckpt|safetensors)$'

如果希望项目里固定使用 checkpoints/lingbot-map，可以把 ModelScope 返回目录软链接过去：

bash📋 复制
ln -s "$(python tools/download_lingbot_map_ms.py)" checkpoints/lingbot-map

如果目录已存在，不要重复软链接，直接使用脚本输出的真实路径即可。

---

本地检查点选择：按关键字找文件，不写死文件名

下面脚本会在权重目录里查找包含 balanced、long 或 stage1 的检查点文件，并输出最匹配的路径。

python📋 复制
# tools/select_lingbot_map_ckpt.py
import argparse
from pathlib import Path

CKPT_EXTS = {".pth", ".pt", ".ckpt", ".safetensors"}


def list_checkpoints(root: Path):
    if not root.exists():
        raise FileNotFoundError(f"权重目录不存在: {root}")

    files = []
    for path in root.rglob("*"):
        if path.is_file() and path.suffix.lower() in CKPT_EXTS:
            files.append(path)

    if not files:
        raise FileNotFoundError(f"没有在 {root} 下找到检查点文件: {sorted(CKPT_EXTS)}")

    return files


def score_checkpoint(path: Path, variant: str) -> int:
    name = path.name.lower()
    score = 0

    if variant in name:
        score += 100

    # 常见命名里会带 latest、final、ema 等标记，略微提高优先级。
    if "final" in name:
        score += 10
    if "latest" in name:
        score += 8
    if "ema" in name:
        score += 5

    # 避免 stage1 被误选成普通推理模型。
    if variant != "stage1" and "stage1" in name:
        score -= 50

    return score


def select_checkpoint(root: Path, variant: str) -> Path:
    candidates = list_checkpoints(root)

    matched = [
        path for path in candidates
        if variant in path.name.lower()
    ]

    if not matched:
        all_names = "\n".join(f"  - {p}" for p in candidates)
        raise FileNotFoundError(
            f"没有找到包含关键字 '{variant}' 的检查点。\n"
            f"当前可用检查点:\n{all_names}"
        )

    matched.sort(key=lambda p: score_checkpoint(p, variant), reverse=True)
    return matched[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--root",
        type=Path,
        default=Path("checkpoints/lingbot-map"),
        help="LingBot-Map 权重目录",
    )
    parser.add_argument(
        "--variant",
        choices=["balanced", "long", "stage1"],
        default="balanced",
        help="要选择的检查点类型",
    )
    args = parser.parse_args()

    ckpt = select_checkpoint(args.root, args.variant)
    print(ckpt)


if __name__ == "__main__":
    main()

运行：

bash📋 复制
python tools/select_lingbot_map_ckpt.py \
  --root checkpoints/lingbot-map \
  --variant balanced

输出类似：

text📋 复制
checkpoints/lingbot-map/lingbot_map_balanced.pth

切到长序列检查点：

bash📋 复制
python tools/select_lingbot_map_ckpt.py \
  --root checkpoints/lingbot-map \
  --variant long

---

可运行示例：下载 HuggingFace 权重并选择 balanced

在项目根目录执行：

bash📋 复制
mkdir -p tools checkpoints

cat > tools/download_lingbot_map_hf.py <<'PY'
from huggingface_hub import snapshot_download

repo_dir = snapshot_download(
    repo_id="robbyant/lingbot-map",
    local_dir="checkpoints/lingbot-map",
    local_dir_use_symlinks=False,
    allow_patterns=[
        "*.pth",
        "*.pt",
        "*.ckpt",
        "*.safetensors",
        "*.json",
        "*.yaml",
        "*.yml",
        "*.txt",
    ],
)

print(repo_dir)
PY

cat > tools/select_lingbot_map_ckpt.py <<'PY'
import argparse
from pathlib import Path

CKPT_EXTS = {".pth", ".pt", ".ckpt", ".safetensors"}


def list_checkpoints(root: Path):
    if not root.exists():
        raise FileNotFoundError(f"权重目录不存在: {root}")

    files = [
        path for path in root.rglob("*")
        if path.is_file() and path.suffix.lower() in CKPT_EXTS
    ]

    if not files:
        raise FileNotFoundError(f"没有在 {root} 下找到检查点文件")

    return files


def select_checkpoint(root: Path, variant: str) -> Path:
    candidates = list_checkpoints(root)
    matched = [p for p in candidates if variant in p.name.lower()]

    if not matched:
        print("已找到的检查点:")
        for p in candidates:
            print(f"  - {p}")
        raise FileNotFoundError(f"没有找到包含关键字 '{variant}' 的检查点")

    def score(path: Path):
        name = path.name.lower()
        value = 0
        if variant in name:
            value += 100
        if "final" in name:
            value += 10
        if "latest" in name:
            value += 8
        if "ema" in name:
            value += 5
        if variant != "stage1" and "stage1" in name:
            value -= 50
        return value

    return sorted(matched, key=score, reverse=True)[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path("checkpoints/lingbot-map"))
    parser.add_argument("--variant", choices=["balanced", "long", "stage1"], default="balanced")
    args = parser.parse_args()

    print(select_checkpoint(args.root, args.variant))


if __name__ == "__main__":
    main()
PY

python -m pip install -U huggingface_hub

python tools/download_lingbot_map_hf.py

CKPT=$(python tools/select_lingbot_map_ckpt.py --variant balanced)

echo "selected checkpoint: ${CKPT}"

如果输出了本地 .pth、.pt、.ckpt 或 .safetensors 路径，说明下载和选择都正常。

---

可运行示例：列出所有可用检查点

有时候不确定远端仓库当前发布了哪些文件，可以先列出来：

python📋 复制
# tools/list_lingbot_map_ckpts.py
from pathlib import Path

root = Path("checkpoints/lingbot-map")
exts = {".pth", ".pt", ".ckpt", ".safetensors"}

for path in sorted(root.rglob("*")):
    if path.is_file() and path.suffix.lower() in exts:
        print(path)

运行：

bash📋 复制
python tools/list_lingbot_map_ckpts.py

输出中重点看文件名是否包含：

text📋 复制
balanced
long
stage1

---

在 demo 或自定义推理里使用选中的检查点

如果项目脚本支持传入 checkpoint 路径，建议不要手写固定路径，而是先选再传：

bash📋 复制
CKPT=$(python tools/select_lingbot_map_ckpt.py \
  --root checkpoints/lingbot-map \
  --variant balanced)

echo "using checkpoint: ${CKPT}"

# 下面命令的参数名以当前仓库 demo.py 为准。
# 常见写法是 --ckpt、--checkpoint 或 --model_path。
python demo.py --ckpt "${CKPT}"

长视频场景：

bash📋 复制
CKPT=$(python tools/select_lingbot_map_ckpt.py \
  --root checkpoints/lingbot-map \
  --variant long)

echo "using checkpoint: ${CKPT}"

python demo.py --ckpt "${CKPT}"

训练复现或消融场景：

bash📋 复制
CKPT=$(python tools/select_lingbot_map_ckpt.py \
  --root checkpoints/lingbot-map \
  --variant stage1)

echo "using checkpoint: ${CKPT}"

---

自定义 Python 加载时只依赖路径

如果你在自己的脚本里加载 LingBot-Map，建议把“选择 checkpoint”作为独立步骤：

python📋 复制
from pathlib import Path
import torch

ckpt_path = Path("checkpoints/lingbot-map/lingbot_map_balanced.pth")

checkpoint = torch.load(ckpt_path, map_location="cpu")

print(type(checkpoint))
if isinstance(checkpoint, dict):
    print(checkpoint.keys())

对于 .safetensors 文件：

bash📋 复制
python -m pip install -U safetensors

python📋 复制
from safetensors.torch import load_file

ckpt_path = "checkpoints/lingbot-map/lingbot_map_balanced.safetensors"
state_dict = load_file(ckpt_path)

print(len(state_dict))
print(next(iter(state_dict)))

加载后如何 load_state_dict，以项目内模型构建函数为准。这个模块只保证能拿到正确的本地文件路径。

---

版本和缓存的几个边界

不要混用旧代码和新权重

LingBot-Map README 里提到过 KV cache、SDPA、FlashInfer 等长序列相关修复。权重、模型结构、推理缓存逻辑最好来自同一时期的代码。

推荐做法：

bash📋 复制
git pull
python tools/download_lingbot_map_hf.py

如果出现这些报错，优先检查代码和权重是否匹配：

text📋 复制
missing keys
unexpected keys
size mismatch

---

stage1 不是默认推理模型

stage1 一般用于训练阶段、消融或初始化。普通用户看重建效果时，不建议直接选它。

推荐优先级：

text📋 复制
balanced -> long -> stage1

只有明确知道自己在复现实验或继续训练时，再使用 stage1。

---

long 不一定更快，也不一定更省显存

long 更适合长序列稳定性，不代表它在短视频上一定更好。短视频、交互 demo、显存有限的机器，先用 balanced。

典型选择：

text📋 复制
100 ~ 1000 帧       balanced
3000 帧以上         long
10000 帧以上        long，并结合窗口化/缓存后端配置

---

ModelScope 的目录结构可能和 HuggingFace 不一样

HuggingFace 的 local_dir 通常就是目标目录：

text📋 复制
checkpoints/lingbot-map/

ModelScope 可能返回类似缓存目录：

text📋 复制
checkpoints/Robbyant/lingbot-map/...

所以不要假设 ModelScope 下载后一定在 checkpoints/lingbot-map。建议使用 snapshot_download 的返回值，或者把返回目录传给选择脚本：

bash📋 复制
MS_DIR=$(python tools/download_lingbot_map_ms.py)

python tools/select_lingbot_map_ckpt.py \
  --root "${MS_DIR}" \
  --variant balanced

---

权重文件很大，避免重复拷贝

如果同一台机器上多个实验都用 LingBot-Map，建议使用统一缓存目录：

bash📋 复制
export HF_HOME=/data/cache/huggingface
export MODELSCOPE_CACHE=/data/cache/modelscope

项目里只保存软链接或 checkpoint 路径，避免每个实验目录都复制一份大文件。

---

Windows 下软链接可能失败

Windows 创建软链接可能需要管理员权限。若软链接失败，直接使用真实路径即可：

powershell📋 复制
python tools/select_lingbot_map_ckpt.py --root checkpoints/lingbot-map --variant balanced

不强制要求存在 selected.ckpt 这种固定链接文件。

---

自动选择失败时先看文件名

如果选择脚本报错：

text📋 复制
没有找到包含关键字 'balanced' 的检查点

先列出本地文件：

bash📋 复制
find checkpoints -type f | grep -E '\.(pth|pt|ckpt|safetensors)$'

如果远端实际文件名没有 balanced、long、stage1 关键字，需要根据当前发布说明调整筛选规则。不要盲目把任意 .pth 当成最终推理权重。