#!/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()