scripts/download_artifacts.py

#!/usr/bin/env python3
"""Download SATA demo artifacts.

Models are downloaded from HuggingFace. Demo assets are downloaded from a
Google Drive/direct URL zip and extracted into the repository root.
"""

import argparse
import os
import shutil
import sys
import tempfile
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path
from typing import List, Optional


HF_REPO_ID = "SteveZh/sata_models"
DEFAULT_ASSETS_URL = os.environ.get("SATA_ASSETS_URL", "")
DEFAULT_ASSETS_GDRIVE_ID = os.environ.get("SATA_ASSETS_GDRIVE_ID", "")

MODEL_PATTERNS = [
    "result/vae_merge/*",
    "result/vae_human/*",
    "result/rvq_human/*",
    "src/mdm/save/paper_vae_human_0125/*",
    "src/momask-preenc/checkpoints/t2m/t2m_base/*",
    "src/momask-preenc/checkpoints/t2m/t2m_base/model/*",
    "src/momask-preenc/checkpoints/t2m/t2m_rvq/*",
    "src/momask-preenc/checkpoints/t2m/t2m_rvq/model/*",
]

REQUIRED_MODEL_FILES = [
    "result/vae_merge/config.yaml",
    "result/vae_merge/ms_dict.pt",
    "result/vae_merge/last_model.pt",
    "result/vae_human/config.yaml",
    "result/vae_human/ms_dict.pt",
    "result/vae_human/model_390.pt",
    "result/rvq_human/config.yaml",
    "result/rvq_human/ms_dict.pt",
    "result/rvq_human/last_model.pt",
    "src/mdm/save/paper_vae_human_0125/args.json",
    "src/mdm/save/paper_vae_human_0125/model000100000.pt",
    "src/momask-preenc/checkpoints/t2m/t2m_base/opt.txt",
    "src/momask-preenc/checkpoints/t2m/t2m_base/model/latest.tar",
    "src/momask-preenc/checkpoints/t2m/t2m_rvq/opt.txt",
    "src/momask-preenc/checkpoints/t2m/t2m_rvq/model/latest.tar",
]

REQUIRED_ASSET_FILES = [
    "data/gradio_human/bvh/000327.bvh",
    "data/gradio_human/processed/000327.npz",
    "data/gradio_animo/bvh/arctic_wolf_female_fightflee.bvh",
    "data/gradio_skel/human/processed/guard.npz",
    "data/gradio_skel/animo/processed/wisent_male_attackfence.npz",
    "data/gradio_target_bvh/human/guard.bvh",
    "data/gradio_target_bvh/animo/wisent_male_attackfence.bvh",
    "data/gradio_z/human/000327.pt",
    "data/gradio_z/animo/arctic_wolf_female_fightflee.pt",
    "data/test/character/default.txt",
]


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def missing_files(root: Path, files: List[str]) -> List[str]:
    return [path for path in files if not (root / path).is_file()]


def verify(root: Path, *, models: bool, assets: bool) -> None:
    missing = []  # type: List[str]
    if models:
        missing.extend(missing_files(root, REQUIRED_MODEL_FILES))
    if assets:
        missing.extend(missing_files(root, REQUIRED_ASSET_FILES))

    if missing:
        print("Missing required artifact files:", file=sys.stderr)
        for path in missing:
            print(f"  {path}", file=sys.stderr)
        raise SystemExit(1)


def download_models(root: Path, repo_id: str) -> None:
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise SystemExit(
            "Missing dependency: huggingface_hub. Install it with:\n"
            "  python -m pip install huggingface-hub"
        ) from exc

    print(f"Downloading model artifacts from HuggingFace: {repo_id}")
    snapshot_download(
        repo_id=repo_id,
        repo_type="model",
        local_dir=str(root),
        allow_patterns=MODEL_PATTERNS,
    )
    verify(root, models=True, assets=False)
    print("Model artifacts are ready.")


def google_drive_url(file_id: str) -> str:
    return f"https://drive.google.com/uc?export=download&id={file_id}"


def extract_google_drive_id(value: str) -> Optional[str]:
    value = value.strip()
    if not value:
        return None
    if "/" not in value and "?" not in value:
        return value

    parsed = urllib.parse.urlparse(value)
    query = urllib.parse.parse_qs(parsed.query)
    if "id" in query and query["id"]:
        return query["id"][0]

    parts = [part for part in parsed.path.split("/") if part]
    if "d" in parts:
        idx = parts.index("d")
        if idx + 1 < len(parts):
            return parts[idx + 1]
    return None


def resolve_assets_url(url: str, gdrive_id: str) -> str:
    if gdrive_id:
        return google_drive_url(extract_google_drive_id(gdrive_id) or gdrive_id)

    if not url:
        raise SystemExit(
            "No assets URL configured. Pass --assets-url, pass --assets-gdrive-id, "
            "or set SATA_ASSETS_URL after uploading sata_demo_assets.zip."
        )

    parsed = urllib.parse.urlparse(url)
    if "drive.google.com" in parsed.netloc:
        file_id = extract_google_drive_id(url)
        if file_id:
            return google_drive_url(file_id)
    return url


def download_file(url: str, output_path: Path) -> None:
    print(f"Downloading assets zip: {url}")
    request = urllib.request.Request(url, headers={"User-Agent": "sata-artifact-downloader"})
    with urllib.request.urlopen(request) as response, output_path.open("wb") as handle:
        shutil.copyfileobj(response, handle)


def download_assets(root: Path, assets_url: str, assets_gdrive_id: str) -> None:
    url = resolve_assets_url(assets_url, assets_gdrive_id)
    with tempfile.TemporaryDirectory(prefix="sata_assets_") as tmp:
        zip_path = Path(tmp) / "sata_demo_assets.zip"
        download_file(url, zip_path)
        if not zipfile.is_zipfile(zip_path):
            raise SystemExit(
                "Downloaded assets file is not a valid zip. Check that the Google "
                "Drive file is shared publicly and that the URL points directly to "
                "sata_demo_assets.zip."
            )
        print(f"Extracting {zip_path.name} into {root}")
        with zipfile.ZipFile(zip_path) as archive:
            archive.extractall(root)

    verify(root, models=False, assets=True)
    print("Demo assets are ready.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    actions = parser.add_argument_group("actions")
    actions.add_argument("--all", action="store_true", help="Download both models and assets.")
    actions.add_argument("--models", action="store_true", help="Download model checkpoints from HuggingFace.")
    actions.add_argument("--assets", action="store_true", help="Download and extract demo assets.")
    actions.add_argument("--check", action="store_true", help="Only verify required artifact files.")

    parser.add_argument("--hf-repo", default=HF_REPO_ID, help=f"HuggingFace repo id. Default: {HF_REPO_ID}")
    parser.add_argument("--assets-url", default=DEFAULT_ASSETS_URL, help="Direct/Google Drive URL for sata_demo_assets.zip.")
    parser.add_argument("--assets-gdrive-id", default=DEFAULT_ASSETS_GDRIVE_ID, help="Google Drive file id for sata_demo_assets.zip.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    root = repo_root()

    selected = args.all or args.models or args.assets or args.check
    if not selected:
        args.all = True

    want_models = args.all or args.models
    want_assets = args.all or args.assets

    if args.check:
        verify(root, models=want_models or not args.assets, assets=want_assets or not args.models)
        print("Artifact check passed.")
        return

    if want_models:
        download_models(root, args.hf_repo)
    if want_assets:
        download_assets(root, args.assets_url, args.assets_gdrive_id)

    verify(root, models=want_models, assets=want_assets)
    print("All requested artifacts are ready.")


if __name__ == "__main__":
    main()