scripts/setup_space.py

#!/usr/bin/env python3
"""
HF Space 一次性环境搭建：下载权重 + 生成 BEV 地图缓存
运行一次后文件持久保存在挂载存储上，之后不再需要重复运行。

用法:
    python scripts/setup_space.py          # 全部下载+生成
    python scripts/setup_space.py --check  # 只检查状态
"""

import argparse
import os
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

# ---- 下载清单 ----
DOWNLOAD_TASKS = [
    {
        "name": "t5",
        "repo": "google/t5-v1_1-xxl",
        "local": ROOT / "pretrained" / "t5-v1_1-xxl",
        "desc": "T5-XXL 文本编码器 (~40 GB)",
    },
    {
        "name": "ckpt",
        "repo": "flymin/MagicDriveDiT-stage3-40k-ft",
        "local": ROOT / "ckpts" / "MagicDriveDiT-stage3-40k-ft",
        "desc": "Stage-3 DiT 扩散模型 (~8 GB)",
    },
    {
        "name": "vae",
        "repo": "THUDM/CogVideoX-2b",
        "local": ROOT / "pretrained" / "CogVideoX-2b",
        "allow_patterns": ["vae/**"],
        "desc": "CogVideoX-2b VAE 编码/解码器 (~2 GB)",
    },
    {
        "name": "data_pkl",
        "repo": "flymin/MagicDriveDiT-nuScenes-metadata",
        "repo_type": "dataset",
        "local": ROOT / "data" / "nuscenes_mmdet3d-12Hz",
        "allow_patterns": [
            "nuscenes_mmdet3d-12Hz/nuscenes_interp_12Hz_infos_val_with_bid.pkl",
        ],
        "desc": "nuScenes 标注 PKL (~430 MB)",
    },
    {
        "name": "ghost_peek",
        "repo": "doradream/magicdrive-ghost-peek",
        "repo_type": "dataset",
        "local": ROOT / "data" / "nuscenes_mmdet3d-12Hz",
        "allow_patterns": [
            "ghost_peek_s8.pkl",
            "ghost_peek_s25.pkl",
        ],
        "desc": "鬼探头场景 PKL (~660 MB, 需先上传到 HF Dataset)",
    },
    {
        "name": "nuscenes_maps",
        "repo": "doradream/magicdrive-data",
        "repo_type": "dataset",
        "local": ROOT / "data" / "nuscenes" / "maps",
        "allow_patterns": ["expansion/*.json"],
        "desc": "nuScenes 官方地图文件 (~10 MB)",
    },
    {
        "name": "map_json",
        "repo": "flymin/MagicDriveDiT-nuScenes-metadata",
        "repo_type": "dataset",
        "local": ROOT / "data" / "nuscenes" / "interp_12Hz_trainval",
        "allow_patterns": ["nuscenes/interp_12Hz_trainval/*.json"],
        "desc": "nuScenes 地图 JSON (~50 MB)",
    },
]

# ---- 地图缓存生成配置 ----
MAP_CACHE_DIR = ROOT / "data" / "nuscenes_mmdet3d-12Hz" / "nuscenes_map_aux_12Hz"
MAP_CACHE_CONFIG = ROOT / "configs" / "cache_gen" / "map_cache_gen_interp.yaml"


def check_path(path: Path) -> tuple[bool, str]:
    """检查目录/文件是否存在且非空"""
    if not path.exists():
        return False, "不存在"
    files = list(path.rglob("*")) if path.is_dir() else [path]
    real_files = [f for f in files if f.is_file()]
    if not real_files:
        return False, "空目录"
    size_gb = sum(f.stat().st_size for f in real_files) / 1e9
    return True, f"{len(real_files)} 个文件, {size_gb:.1f} GB"


def download_from_hf(task: dict):
    """从 HuggingFace 下载单个任务"""
    from huggingface_hub import snapshot_download

    kwargs = dict(
        repo_id=task["repo"],
        local_dir=str(task["local"]),
        repo_type=task.get("repo_type", "model"),
        resume_download=True,
        max_workers=4,
    )
    if "allow_patterns" in task:
        kwargs["allow_patterns"] = task["allow_patterns"]

    print(f"  下载中... ({task['desc']})")
    t0 = time.time()
    snapshot_download(**kwargs)
    elapsed = time.time() - t0
    print(f"  完成 ({elapsed / 60:.1f} min)")


def generate_map_cache():
    """调用 prepare_map_aux.py 生成 BEV 地图 h5 缓存"""
    print("\n--- 生成 BEV 地图缓存 ---")

    # 安装 hydra（如果还没装）
    try:
        import hydra  # noqa
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "hydra-core==1.3.0"])

    # 生成 train 缓存
    for split in ["val"]:  # 推理只需要 val
        h5_file = ROOT / f"{split}_map_cache.h5"
        if h5_file.exists():
            print(f"  {h5_file} 已存在，跳过")
            continue

        print(f"  生成 {split} 地图缓存...")
        result = subprocess.run(
            [
                sys.executable, "-u",
                str(ROOT / "tools" / "prepare_data" / "prepare_map_aux.py"),
                f"+process={split}",
                f"subfix=map_cache",
            ],
            cwd=str(ROOT),
            capture_output=False,
            env={**os.environ, "PYTHONPATH": str(ROOT)},
        )
        if result.returncode != 0:
            print(f"  [WARN] {split} 缓存生成失败 (exit {result.returncode})，将回退到实时生成")
        else:
            # 移动到正确位置
            MAP_CACHE_DIR.mkdir(parents=True, exist_ok=True)
            h5_file.rename(MAP_CACHE_DIR / h5_file.name)
            print(f"  [OK] 已保存到 {MAP_CACHE_DIR / h5_file.name}")


def main():
    parser = argparse.ArgumentParser(description="HF Space 一次性环境搭建")
    parser.add_argument("--check", action="store_true", help="只检查状态，不执行")
    parser.add_argument("--skip-download", action="store_true", help="跳过下载")
    parser.add_argument("--skip-cache", action="store_true", help="跳过地图缓存生成")
    args = parser.parse_args()

    os.chdir(str(ROOT))
    sys.path.insert(0, str(ROOT))

    print("=" * 60)
    print("MagicDrive-V2 HF Space 环境检查")
    print(f"Root: {ROOT}")
    print("=" * 60)

    # ---- 检查状态 ----
    all_ok = True
    print("\n【权重/数据文件】")
    for task in DOWNLOAD_TASKS:
        ok, msg = check_path(task["local"])
        status = "[OK]" if ok else "[MISS]"
        if not ok:
            all_ok = False
        print(f"  {status} {task['name']:12s} {msg}")

    print("\n【BEV 地图缓存】")
    cache_ok, cache_msg = check_path(MAP_CACHE_DIR)
    if not cache_ok:
        all_ok = False
    print(f"  {'[OK]' if cache_ok else '[MISS]'} map_cache   {cache_msg}")

    if args.check:
        if all_ok:
            print("\n所有文件就绪，可以直接推理！")
        else:
            print("\n部分文件缺失，运行: python scripts/setup_space.py")
        return

    # ---- 下载 ----
    if not args.skip_download:
        print("\n" + "=" * 60)
        print("下载权重与数据")
        print("=" * 60)
        for task in DOWNLOAD_TASKS:
            ok, msg = check_path(task["local"])
            if ok:
                print(f"\n[{task['name']}] 已存在 -> 跳过")
                continue
            print(f"\n[{task['name']}] {task['desc']}")
            try:
                download_from_hf(task)
            except Exception as e:
                print(f"  [ERR] 下载失败: {e}")
                print(f"  你可以稍后重试: python scripts/setup_space.py")

    # ---- 生成地图缓存 ----
    if not args.skip_cache:
        cache_ok, _ = check_path(MAP_CACHE_DIR)
        if cache_ok:
            print("\n[BEV 地图缓存] 已存在 -> 跳过")
        else:
            generate_map_cache()

    # ---- 最终状态 ----
    print("\n" + "=" * 60)
    print("最终状态")
    print("=" * 60)
    all_ok = True
    for task in DOWNLOAD_TASKS:
        ok, msg = check_path(task["local"])
        if not ok:
            all_ok = False
        print(f"  {'[OK]' if ok else '[MISS]'} {task['name']}")
    cache_ok, _ = check_path(MAP_CACHE_DIR)
    if not cache_ok:
        all_ok = False
    print(f"  {'[OK]' if cache_ok else '[MISS]'} map_cache")

    if all_ok:
        print("\n全部就绪！可以运行推理了。")
    else:
        print("\n有缺失项，请重新运行或检查网络。")


if __name__ == "__main__":
    main()