"""Download only official assets required by the four evaluation adapters."""

import argparse
from pathlib import Path

from huggingface_hub import snapshot_download

parser = argparse.ArgumentParser()
parser.add_argument("--root", type=Path, required=True)
args = parser.parse_args()
root = args.root

snapshot_download("KlingTeam/DecMem", revision="0ca8acbb5e32237bfcb01dad0edf450acc2d7fba", allow_patterns=["decmem.pt"], local_dir=root / "checkpoints/decmem")
snapshot_download(
    "Wan-AI/Wan2.1-T2V-1.3B",
    revision="37ec512624d61f7aa208f7ea8140a131f93afc9a",
    allow_patterns=["config.json", "diffusion_pytorch_model.safetensors", "Wan2.1_VAE.pth"],
    local_dir=root / "DecMem/wan_models/Wan2.1-T2V-1.3B",
)
snapshot_download(
    "Skywork/Matrix-Game-2.0",
    revision="f1729d99a80e0f07993a77d7dad4a3190e23c2c8",
    allow_patterns=["base_distilled_model/*", "Wan2.1_VAE.pth", "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", "xlm-roberta-large/*"],
    local_dir=root / "checkpoints/matrix-game2",
)
snapshot_download(
    "junchaoh-cs/LIVE-Re10K",
    revision="28b8396ba1e569a3ae7f06cd64009b22ad74204b",
    allow_patterns=["live-re10k.ckpt", "vae-kl16.ckpt"],
    local_dir=root / "checkpoints/live",
)
snapshot_download(
    "Haoyuwu/GeometryForcing",
    revision="3bd52c727ebb4762cb0829cb3c26fe677ad354d1",
    allow_patterns=["geometry_forcing_state_dict.ckpt"],
    local_dir=root / "checkpoints/geometry-forcing",
)
snapshot_download(
    "zeqixiao/worldmem_checkpoints",
    revision="a3ae4d1c3a89260f800ddeab01ca56b7e63b10e2",
    allow_patterns=["vae_only.ckpt"],
    local_dir=root / "checkpoints/worldmem",
)