llm/scripts/plan_hf_downloads.py

from __future__ import annotations

import argparse
import fnmatch
import json
import os
import sys
from pathlib import Path

from huggingface_hub import HfApi

from common import ROOT, iter_model_card_paths, localize_workspace_path, read_yaml_mapping

DEFAULT_IGNORE_PATTERNS = [
    "*.msgpack",
    "*.h5",
    "*.ot",
    "*.tflite",
    "*.onnx",
    "*.pb",
    "*.ckpt",
    "flax_model*",
    "tf_model*",
    "rust_model*",
]


def matches_any(name: str, patterns: list[str]) -> bool:
    return any(fnmatch.fnmatch(name, pattern) for pattern in patterns)


def desired_patterns(card: dict) -> list[str]:
    if card.get("format") == "gguf" and card.get("filename"):
        return [str(card["filename"]), "README.md", "*.json", "*.txt", "*.model"]
    if card.get("upstream_id") == "Qwen/Qwen-Image-Edit":
        return [
            "README.md",
            "LICENSE.md",
            "model_index.json",
            "processor/*",
            "scheduler/*.json",
            "tokenizer/*",
            "text_encoder/*.json",
            "text_encoder/*.safetensors",
            "text_encoder/*.safetensors.index.json",
            "transformer/*.json",
            "transformer/*.safetensors",
            "transformer/*.safetensors.index.json",
            "vae/*.json",
            "vae/*.safetensors",
        ]
    if card.get("format") == "diffusers" or card.get("type") == "image-diffusion-model":
        upstream_id = str(card.get("upstream_id") or "")
        # Many SDXL community repos publish standard .safetensors names instead of *fp16.safetensors.
        # Keep the narrow fp16 defaults for repos that provide them, and widen only for known SDXL variants
        # we want to support in the current runtime.
        if upstream_id in {
            "cagliostrolab/animagine-xl-4.0",
            "glides/illustriousxl",
        }:
            return [
                "README.md",
                "LICENSE.md",
                "model_index.json",
                "scheduler/*.json",
                "tokenizer/*",
                "tokenizer_2/*",
                "text_encoder/*.json",
                "text_encoder/*.safetensors",
                "text_encoder_2/*.json",
                "text_encoder_2/*.safetensors",
                "unet/*.json",
                "unet/*.safetensors",
                "vae/*.json",
                "vae/*.safetensors",
            ]
        return [
            "README.md",
            "LICENSE.md",
            "model_index.json",
            "scheduler/*.json",
            "tokenizer/*",
            "tokenizer_2/*",
            "text_encoder/*.json",
            "text_encoder/*fp16.safetensors",
            "text_encoder_2/*.json",
            "text_encoder_2/*fp16.safetensors",
            "unet/*.json",
            "unet/*fp16.safetensors",
            "vae/*.json",
            "vae/*fp16.safetensors",
        ]
    model_type = str(card.get("type") or "")
    if model_type == "speech-model":
        return ["*.safetensors", "*.bin", "*.json", "*.txt", "*.model", "*.md"]
    if model_type in {"translation-model", "vision-language-model", "base-model"}:
        return ["*.safetensors", "*.bin", "*.json", "*.txt", "*.model", "*.jinja", "*.md"]
    return ["*.safetensors", "*.bin", "*.gguf", "*.json", "*.txt", "*.model", "*.md"]


def remote_plan(card_path: Path, *, include_ready: bool = False) -> dict | None:
    card = read_yaml_mapping(card_path)
    repo_id = card.get("upstream_id")
    if not repo_id:
        return None

    raw_storage_path = str(card.get("storage_path") or "")
    if os.environ.get("LLM_USE_RAW_MODEL_PATH") == "1" and raw_storage_path.startswith("/"):
        storage_path = Path(raw_storage_path)
    else:
        storage_path = localize_workspace_path(raw_storage_path)
    api = HfApi()
    info = api.model_info(repo_id=repo_id, files_metadata=True)
    allow = desired_patterns(card)
    files = []
    for sibling in info.siblings:
        name = sibling.rfilename
        if name == ".gitattributes":
            continue
        if matches_any(name, DEFAULT_IGNORE_PATTERNS):
            continue
        if not matches_any(name, allow):
            continue
        size = int(sibling.size or 0)
        local_path = storage_path / name
        local_size = local_path.stat().st_size if local_path.exists() else 0
        files.append(
            {
                "name": name,
                "size_bytes": size,
                "local_size_bytes": local_size,
                "missing_bytes": max(size - local_size, 0),
                "complete": size > 0 and local_size == size,
            }
        )

    missing_bytes = sum(item["missing_bytes"] for item in files if not item["complete"])
    local_ready = bool(files) and missing_bytes == 0
    if local_ready and not include_ready:
        return None

    return {
        "id": card.get("id"),
        "repo_id": repo_id,
        "storage_path": str(storage_path),
        "local_ready": local_ready,
        "allow_patterns": allow,
        "missing_bytes": missing_bytes,
        "missing_gb": round(missing_bytes / 1024**3, 2),
        "files": files,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Plan Hugging Face downloads for incomplete model cards.")
    parser.add_argument("--include-ready", action="store_true")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    plans = []
    for card_path in iter_model_card_paths():
        try:
            plan = remote_plan(card_path, include_ready=args.include_ready)
        except Exception as exc:  # noqa: BLE001 - report per model and continue.
            print(f"WARN {card_path.relative_to(ROOT)}: {exc}", file=sys.stderr)
            continue
        if plan:
            plans.append(plan)

    if args.json:
        print(json.dumps({"plans": plans}, ensure_ascii=False, indent=2))
        return 0

    total = sum(plan["missing_bytes"] for plan in plans)
    print(f"Models needing downloads: {len(plans)} | missing: {round(total / 1024**3, 2)} GB")
    for plan in plans:
        missing_files = [item for item in plan["files"] if not item["complete"]]
        print(f"{plan['id']}: {plan['missing_gb']} GB missing -> {plan['storage_path']}")
        for item in missing_files[:12]:
            print(f"  - {item['name']} {round(item['missing_bytes'] / 1024**3, 2)} GB")
        if len(missing_files) > 12:
            print(f"  ... {len(missing_files) - 12} more file(s)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())