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