download_model.py

#!/usr/bin/env python3

import os
import shutil
import tempfile
import stat
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download


def require_env(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise SystemExit(f"missing required environment variable: {name}")
    return value


def quant_rank(file_name: str) -> tuple[int, str]:
    name = file_name.lower()
    # Allow user-specified preference list via QUANT_PREFERENCE env (comma-separated)
    pref = os.environ.get("QUANT_PREFERENCE", "q4_k_m").split(",")
    pref = [p.strip().lower() for p in pref if p.strip()]
    for idx, token in enumerate(pref):
        if token in name:
            return (idx, file_name)

    # fallback groups (lower priority than explicit preference)
    groups = [
        ("q8_k_l", "q8_k_m", "q8_k_s", "q8_k", "q8_0", "q8_1"),
        ("f16", "bf16"),
        ("q6_k", "q6_0", "q6_1", "q6"),
        ("q5_k_m", "q5_k_s", "q5_k", "q5_0", "q5_1", "q5"),
        ("q4_k_m", "q4_k_s", "q4_k", "q4_0", "q4_1", "q4"),
        ("q3_k_m", "q3_k_s", "q3_k", "q3_0", "q3_1", "q3"),
        ("q2_k", "q2_0", "q2_1", "q2"),
        ("f32",),
    ]
    base = len(pref)
    for gi, tokens in enumerate(groups):
        for token in tokens:
            if token in name:
                return (base + gi, file_name)

    return (base + len(groups), file_name)


def main() -> None:
    model_name = require_env("MODEL_NAME")
    hf_token = require_env("HF_TOKEN")
    require_env("API_PASSWORD")

    models_dir = Path(os.environ.get("MODEL_DIR", "/data/models"))
    models_dir.mkdir(parents=True, exist_ok=True)

    repo_marker = models_dir / "model.repo"
    model_file_override = os.environ.get("MODEL_FILE", "").strip()
    selection_key = f"file:{model_file_override}" if model_file_override else ""

    # If a repo marker exists, check whether the previously downloaded file exists
    if repo_marker.exists():
        marker = repo_marker.read_text(encoding="utf-8").strip()
        marker_parts = marker.split("|")
        prev_repo = marker_parts[0] if len(marker_parts) > 0 else ""
        prev_file = marker_parts[1] if len(marker_parts) > 1 else ""
        prev_selection_key = marker_parts[2] if len(marker_parts) > 2 else ""
        if prev_repo == model_name and prev_file and (not selection_key or prev_selection_key == selection_key):
            cached = models_dir / prev_file
            if cached.exists():
                print(f"using cached model: {cached}")
                return
        elif prev_repo == model_name and prev_file and selection_key and prev_selection_key != selection_key:
            print("cached model does not match MODEL_FILE override; selecting requested file", flush=True)
    # Use a temporary directory (under /tmp) for downloads to avoid permission issues
    download_dir = Path(tempfile.mkdtemp(prefix="hf-download-"))

    # previous cache check moved above

    # Ensure any stale temp dir is removed with a permissive onerror handler
    def _rmtree_onerror(func, path, exc_info):
        # try to make the file writable then retry
        try:
            os.chmod(path, stat.S_IWUSR | stat.S_IRUSR)
        except Exception:
            pass
        try:
            func(path)
        except Exception:
            try:
                if os.path.isdir(path):
                    for root, dirs, files in os.walk(path):
                        for name in files:
                            fp = os.path.join(root, name)
                            try:
                                os.chmod(fp, stat.S_IWUSR | stat.S_IRUSR)
                            except Exception:
                                pass
                    func(path)
                else:
                    os.remove(path)
            except Exception as e:
                print(f"warning: failed to remove {path}: {e}", flush=True)

    api = HfApi(token=hf_token)
    print(f"listing model files: {model_name}", flush=True)
    files = api.list_repo_files(repo_id=model_name, repo_type="model")

    # If a specific filename is provided via MODEL_FILE, try to download it directly.
    if model_file_override:
        if model_file_override not in files:
            print(f"requested MODEL_FILE '{model_file_override}' not found in repo. Available files:", flush=True)
            for f in files[:200]:
                print(f"- {f}", flush=True)
            raise SystemExit("MODEL_FILE override not present in the repository")

        print(f"downloading requested file: {model_name}/{model_file_override}", flush=True)
        snapshot_download(
            repo_id=model_name,
            token=hf_token,
            allow_patterns=[model_file_override],
            local_dir=str(download_dir),
            cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
        )
        downloaded_path = download_dir / model_file_override

    else:
        candidates = [file_name for file_name in files if file_name.lower().endswith(".gguf")]
        if not candidates:
            print("no gguf files found in repository:", model_name, flush=True)
            print("available files (first 200):", flush=True)
            for f in files[:200]:
                print(f"- {f}", flush=True)
            print("")
            print("Options:")
            print(" - set MODEL_NAME to a repo that contains a .gguf file")
            print(" - if the repo contains a model file you want, set MODEL_FILE to the exact filename (env var) to download it")
            print(" - add a .gguf build/asset to the repo or upload a GGUF to a storage bucket mounted at /data/models")
            raise SystemExit(f"no gguf files found in repository: {model_name}")

        print("available gguf files:", flush=True)
        for candidate in candidates:
            print(f"- {candidate}", flush=True)

        selected = min(candidates, key=quant_rank)
        print(f"selected model: {selected}", flush=True)

        print(f"downloading model: {model_name}/{selected}", flush=True)
        snapshot_download(
            repo_id=model_name,
            token=hf_token,
            allow_patterns=[selected],
            local_dir=str(download_dir),
            cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
        )
        downloaded_path = download_dir / selected

    # Save using the original filename from the repo
    selected_basename = Path(downloaded_path).name
    target_model = models_dir / selected_basename
    if target_model.exists():
        target_model.unlink()
    shutil.move(str(downloaded_path), str(target_model))

    if repo_marker.exists():
        repo_marker.unlink()
    repo_marker.write_text(f"{model_name}|{selected_basename}|{selection_key}", encoding="utf-8")

    # --- mmproj (vision projector) support ---
    mmproj_file_override = os.environ.get("MMPROJ_FILE", "").strip()
    mmproj_marker = models_dir / "model.mmproj"

    mmproj_candidates = [f for f in files if "mmproj" in f.lower() and f.lower().endswith(".gguf")]

    if mmproj_file_override:
        if mmproj_file_override not in files:
            print(f"warning: MMPROJ_FILE '{mmproj_file_override}' not found in repo; vision support disabled", flush=True)
        else:
            print(f"downloading mmproj: {model_name}/{mmproj_file_override}", flush=True)
            snapshot_download(
                repo_id=model_name,
                token=hf_token,
                allow_patterns=[mmproj_file_override],
                local_dir=str(download_dir),
                cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
            )
            mmproj_src = download_dir / mmproj_file_override
            mmproj_target = models_dir / mmproj_file_override
            if mmproj_target.exists():
                mmproj_target.unlink()
            shutil.move(str(mmproj_src), str(mmproj_target))
            if mmproj_marker.exists():
                mmproj_marker.unlink()
            mmproj_marker.write_text(mmproj_file_override, encoding="utf-8")
            print(f"mmproj ready: {mmproj_target}")
    elif mmproj_candidates:
        # auto-detect: pick the first mmproj file
        auto_mmproj = mmproj_candidates[0]
        print(f"auto-detected mmproj: {model_name}/{auto_mmproj}", flush=True)
        snapshot_download(
            repo_id=model_name,
            token=hf_token,
            allow_patterns=[auto_mmproj],
            local_dir=str(download_dir),
            cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
        )
        mmproj_src = download_dir / auto_mmproj
        mmproj_target = models_dir / auto_mmproj
        if mmproj_target.exists():
            mmproj_target.unlink()
        shutil.move(str(mmproj_src), str(mmproj_target))
        if mmproj_marker.exists():
            mmproj_marker.unlink()
        mmproj_marker.write_text(auto_mmproj, encoding="utf-8")
        print(f"mmproj ready: {mmproj_target}")
    else:
        if mmproj_marker.exists():
            mmproj_marker.unlink()
        print("no mmproj file found; vision support disabled", flush=True)

    try:
        shutil.rmtree(download_dir, onerror=_rmtree_onerror)
    except Exception as e:
        print(f"warning: unable to fully remove temp download dir: {e}", flush=True)
    print(f"ready: {target_model}")


if __name__ == "__main__":
    main()